diff --git a/.gitattributes b/.gitattributes
index daacfe9b3fb9cb81ff1c9f85ee021080f57a7f1b..604ffa217ded3a6b645774342eb32d3e53b12936 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -226,3 +226,9 @@ workspace/outputs/audit_venv/lib/python3.11/site-packages/pip/_vendor/distlib/t6
 workspace/outputs/audit_venv/lib/python3.11/site-packages/pip/_vendor/distlib/w64-arm.exe filter=lfs diff=lfs merge=lfs -text
 workspace/outputs/audit_venv/lib/python3.11/site-packages/pip/_vendor/distlib/w64.exe filter=lfs diff=lfs merge=lfs -text
 workspace/outputs/audit_venv/lib/python3.11/site-packages/pydantic_core/_pydantic_core.cpython-311-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10.so filter=lfs diff=lfs merge=lfs -text
+workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10_cuda.so filter=lfs diff=lfs merge=lfs -text
+workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so filter=lfs diff=lfs merge=lfs -text
+workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so filter=lfs diff=lfs merge=lfs -text
+workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda_linalg.so filter=lfs diff=lfs merge=lfs -text
+workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_python.so filter=lfs diff=lfs merge=lfs -text
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.h
new file mode 100644
index 0000000000000000000000000000000000000000..fb9db4ccb84e2f961a3829d2ec28dfec7fcb2136
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.h
@@ -0,0 +1,63 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/autograd/profiler.h>
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_agent.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+
+namespace torch::distributed::autograd {
+class TORCH_API RpcWithProfilingResp : public rpc::RpcCommandBase {
+ public:
+  // For sending RPCs over the wire
+  RpcWithProfilingResp(
+      rpc::MessageType messageType,
+      c10::intrusive_ptr<rpc::Message> wrappedMessage,
+      std::vector<torch::autograd::profiler::LegacyEvent> profiledEvents,
+      rpc::ProfilingId profilingId);
+
+  // For receiving RPCs. Used in from message when converting a message received
+  // over the wire.
+  RpcWithProfilingResp(
+      rpc::MessageType messageType,
+      std::unique_ptr<rpc::RpcCommandBase> wrappedRpc,
+      rpc::MessageType wrappedMessageType,
+      std::vector<torch::Tensor> tensors,
+      std::vector<torch::autograd::profiler::LegacyEvent> profiledEvents,
+      rpc::ProfilingId profilingId);
+  c10::intrusive_ptr<rpc::Message> toMessageImpl() && override;
+  static std::unique_ptr<RpcWithProfilingResp> fromMessage(
+      const rpc::Message& message);
+  // Retrieve remote Events
+  std::vector<torch::autograd::profiler::LegacyEvent> getProfiledEvents() const;
+  // Retrieve the globally unique profiling ID corresponding to this command.
+  const rpc::ProfilingId& getProfilingId() const;
+  // Retrieve the original RPC which this ProfilingRPC wraps.
+  RpcCommandBase& wrappedRpc();
+  // Destructively move the wrapped RPC.
+  std::unique_ptr<RpcCommandBase> moveWrappedRpc() &&;
+  // Message type of the wrapped RPC
+  rpc::MessageType wrappedMessageType() const;
+  // Set the wrapped RPC for this RPC.
+  void setWrappedRpc(std::unique_ptr<RpcCommandBase> wrappedRpc);
+
+ private:
+  // message type
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const rpc::MessageType messageType_;
+  // wrapped message
+  c10::intrusive_ptr<rpc::Message> wrappedMessage_;
+  std::unique_ptr<RpcCommandBase> wrappedRpc_;
+  rpc::MessageType wrappedMessageType_;
+  std::vector<torch::Tensor> tensors_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::vector<torch::autograd::profiler::LegacyEvent> profiledEvents_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const rpc::ProfilingId profilingId_;
+};
+} // namespace torch::distributed::autograd
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_req.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_req.h
new file mode 100644
index 0000000000000000000000000000000000000000..1cb78980aa6deeef29a2e918eaf87ed970b7a127
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_req.h
@@ -0,0 +1,43 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+
+namespace torch::distributed::autograd {
+
+// Internal system RPC to invoke distributed backward pass on remote nodes when
+// 'rref.backward()' is invoked.
+class TORCH_API RRefBackwardReq : public rpc::RpcCommandBase {
+ public:
+  RRefBackwardReq(
+      const rpc::RRefId& rrefId,
+      int64_t autogradContextId,
+      bool retainGraph = false);
+
+  const rpc::RRefId& getRRefId() const;
+
+  int64_t getAutogradContextId() const;
+
+  bool retainGraph() const;
+
+  // Serialization and deserialization methods.
+  c10::intrusive_ptr<rpc::Message> toMessageImpl() && override;
+  static std::unique_ptr<RRefBackwardReq> fromMessage(
+      const rpc::Message& message);
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const rpc::RRefId rrefId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const int64_t autogradContextId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const bool retainGraph_;
+};
+
+} // namespace torch::distributed::autograd
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_resp.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_resp.h
new file mode 100644
index 0000000000000000000000000000000000000000..c4bf412302ced62d07a5ca8a24675376f1ed2b68
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_resp.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+
+namespace torch::distributed::autograd {
+
+// Response for the RRefBackwardReq.
+class TORCH_API RRefBackwardResp : public rpc::RpcCommandBase {
+ public:
+  RRefBackwardResp() = default;
+  c10::intrusive_ptr<rpc::Message> toMessageImpl() && override;
+  static std::unique_ptr<RRefBackwardResp> fromMessage(
+      const rpc::Message& message);
+};
+
+} // namespace torch::distributed::autograd
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..47903ceb5e9bf2491e8896cdc2fed920b03e9448
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/autograd/utils.h
@@ -0,0 +1,61 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/autograd/context/context.h>
+#include <torch/csrc/distributed/autograd/rpc_messages/rpc_with_autograd.h>
+#include <torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_req.h>
+#include <torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.h>
+
+namespace torch::distributed::autograd {
+
+// This method is used to attach the 'send' autograd function to the autograd
+// graph when we use RPC. This method creates a new 'send' autograd function
+// and attaches the provided tensors as next_edges to the 'send' function. In
+// addition to this, it also registers the send function in the provided
+// autograd context. Finally, the RPC message is updated with appropriate
+// autograd information for the recipient.
+TORCH_API void addSendRpcBackward(
+    const ContextPtr& autogradContext,
+    const AutogradMetadata& autogradMetadata,
+    std::vector<torch::Tensor>& tensors);
+
+// This method is used to attach the 'recv' autograd function to the autograd
+// graph when we use RPC. This method creates a new 'recv' autograd function
+// and attaches the provided tensors as inputs to the 'recv' function. It
+// creates a new autograd context if needed and registers the 'recv' function
+// with this context.
+//
+// Returns a pointer to the autograd context created.
+TORCH_API ContextPtr addRecvRpcBackward(
+    const AutogradMetadata& autogradMetadata,
+    std::vector<torch::Tensor>& tensors,
+    rpc::worker_id_t fromWorkerId,
+    const rpc::DeviceMap& deviceMap);
+
+// This method is a wrapper utility used internally to wrap autograd info
+// and attach autograd function for each type of rpc call if it has valid
+// context and tensors require grads or forceGradRecording is true, in this
+// case, return RpcWithAutograd message; otherwise return original rpc message.
+// NB: forceGradRecording is useful when the request does not contain any tensor
+// but the corresponding response does.
+TORCH_API c10::intrusive_ptr<rpc::Message> getMessageWithAutograd(
+    const rpc::worker_id_t dstId,
+    c10::intrusive_ptr<rpc::Message> wrappedRpcMsg,
+    rpc::MessageType msgType,
+    bool forceGradRecording = false,
+    const rpc::DeviceMap& deviceMap = {});
+
+// Send message after autograd checking
+TORCH_API c10::intrusive_ptr<c10::ivalue::Future> sendMessageWithAutograd(
+    rpc::RpcAgent& agent,
+    const rpc::WorkerInfo& dst,
+    c10::intrusive_ptr<rpc::Message> wrappedRpcMsg,
+    bool forceGradRecording = false,
+    const float rpcTimeoutSeconds = torch::distributed::rpc::kUnsetRpcTimeout,
+    bool forceDisableProfiling = false);
+
+} // namespace torch::distributed::autograd
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Backend.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Backend.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..bb018312224b7b5a008373a25e1e822ef3084061
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Backend.hpp
@@ -0,0 +1,559 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <c10/core/Allocator.h>
+#include <c10/macros/Macros.h>
+
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+#include <torch/csrc/distributed/c10d/Work.hpp>
+#include <torch/csrc/distributed/c10d/debug.h>
+
+constexpr auto kBackendDefaultTimeout =
+    std::chrono::milliseconds(30 * 60 * 1000);
+
+namespace c10d {
+
+enum class ErrorType {
+  SUCCESS = 0,
+  TIMEOUT = 1,
+  // e.g., NCCL error, etc
+  COMM_ERROR = 2,
+  // TODO, do we need to distinguish between remote timeout or remote COMM
+  // errors?
+  REMOTE_ERROR = 3
+};
+
+class TORCH_API Backend : public torch::CustomClassHolder {
+ public:
+  // Backend Options is a base struct that defines the basic options
+  // when constructing a Backend. Each Backend subclass should
+  // extend this struct and define its options if it wants to provide more
+  // config options (beyond basic ones defined here) to end user.
+  struct TORCH_API Options : torch::CustomClassHolder {
+    explicit Options(
+        std::string backend,
+        std::chrono::milliseconds timeout = kBackendDefaultTimeout)
+        : timeout(timeout), backend(std::move(backend)) {}
+    ~Options() override = default;
+    Options(const Options&) = default;
+
+    std::chrono::milliseconds timeout;
+
+    // backend name
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+    const std::string backend;
+    std::string group_name;
+    std::string group_desc;
+    std::vector<uint64_t> global_ranks_in_group;
+  };
+
+  explicit Backend(int rank, int size);
+  ~Backend() override = 0;
+
+  int getRank() const {
+    return rank_;
+  }
+
+  int getSize() const {
+    return size_;
+  }
+
+  // Returns an unique opaque ID of this backend that can be used to correlate
+  // with its collectives.
+  int64_t getID() const {
+    return reinterpret_cast<std::intptr_t>(this);
+  }
+
+  virtual bool supportsSplitting() const {
+    return false;
+  }
+
+  virtual bool supportsCoalescing() const {
+    return false;
+  }
+
+  virtual bool supportsTimeEstimation() const {
+    return false;
+  }
+
+  virtual bool supportsShrinking() const {
+    return false;
+  }
+
+  // Shrink the backend by excluding specified ranks. Backends that support
+  // communicator shrinking should override this and return a new backend
+  // instance representing the shrunken group. Backends may use opts_override
+  // to supply backend-specific options for the new group.
+  virtual c10::intrusive_ptr<Backend> shrink(
+      const std::vector<int64_t>& /*ranks_to_exclude*/,
+      int /*shrink_flags*/ = 0,
+      const c10::intrusive_ptr<Options>& /*opts_override*/ = nullptr) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support shrink"));
+  }
+
+  virtual void setTimeout(std::chrono::milliseconds timeout) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support setting timeout"));
+  }
+
+  virtual void startCoalescing() {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not implement startCoalescing"));
+  }
+
+  virtual c10::intrusive_ptr<Work> endCoalescing() {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not implement endCoalescing"));
+  }
+
+  // Subclasses must override this method to return the backend name
+  virtual const std::string getBackendName() const {
+    TORCH_INTERNAL_ASSERT(false, "getBackendName is not implemented.");
+  }
+
+  // Subclasses must override this method to return the backend name
+  virtual c10::intrusive_ptr<Options> getBackendOptions() {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not implement getBackendOptions."));
+  }
+
+  virtual c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& /* tensors */,
+      const BroadcastOptions& /* opts */ = BroadcastOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support broadcast"));
+  }
+
+  virtual c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& /* tensors */,
+      const AllreduceOptions& /* opts */ = AllreduceOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support allreduce"));
+  }
+
+  virtual c10::intrusive_ptr<Work> allreduce_sparse(
+      std::vector<at::Tensor>& /* tensors */,
+      const AllreduceOptions& /* opts */ = AllreduceOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not support allreduce sparse"));
+  }
+
+  virtual c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& /* tensors */,
+      const AllreduceCoalescedOptions& /* opts */ =
+          AllreduceCoalescedOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not support allreduce_coalesced"));
+  }
+
+  virtual c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& /* tensors */,
+      const ReduceOptions& /* opts */ = ReduceOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support reduce"));
+  }
+
+  virtual c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& /* outputTensors */,
+      std::vector<at::Tensor>& /* inputTensors */,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support allgather"));
+  }
+
+  // Gathers a single tensor inputBuffer into a single buffer outputBuffer that
+  // is interpreted as a contiguous collection of size inputBuffer * WORLD_SIZE.
+  // For implementers of ProcessGroup API and advanced users only.
+  // Note: this function will be deprecated in near future.
+  virtual c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& /* outputBuffer */,
+      at::Tensor& /* inputBuffer */,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support _allgather_base"));
+  }
+
+  // This function is deprecated and will be moved out of Backend to comms:
+  // * do not add dependencies on this function,
+  // * do not implement it in your Backend, implement _allgather_base
+  //   instead.
+  virtual c10::intrusive_ptr<Work> allgather_coalesced(
+      std::vector<std::vector<at::Tensor>>& /* outputTensorLists */,
+      std::vector<at::Tensor>& /* inputTensors */,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not support allgather_coalesced"));
+  }
+
+  // This function is a coalesced version of `allgather_into_tensor` (currently
+  // still named as `_allgather_base`). Each tensor in the vector corresponds to
+  // an input/output of one `allgather_into_tensor` operation.
+  virtual c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& /* outputs */,
+      std::vector<at::Tensor>& /* inputs */,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not support allgather_into_tensor_coalesced"));
+  }
+
+  virtual c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& /* outputTensors */,
+      std::vector<at::Tensor>& /* inputTensors */,
+      const GatherOptions& /* opts */ = GatherOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support gather"));
+  }
+
+  virtual c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& /* outputTensors */,
+      std::vector<std::vector<at::Tensor>>& /* inputTensors */,
+      const ScatterOptions& /* opts */ = ScatterOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support scatter"));
+  }
+
+  virtual c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& /* outputTensors */,
+      std::vector<std::vector<at::Tensor>>& /* inputTensors */,
+      const ReduceScatterOptions& /* opts */ = ReduceScatterOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support reduce_scatter"));
+  }
+
+  virtual c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& /* outputBuffer */,
+      at::Tensor& /* inputBuffer */,
+      const ReduceScatterOptions& /* opts */ = ReduceScatterOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not support _reduce_scatter_base"));
+  }
+
+  // This function is a coalesced version of `reduce_scatter_tensor` (currently
+  // still named as `_reduce_scatter_base`). Each tensor in the vector
+  // corresponds to an input/output of one `reduce_scatter_tensor` operation.
+  virtual c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& /* outputs */,
+      std::vector<at::Tensor>& /* inputs */,
+      const ReduceScatterOptions& /* opts */ = ReduceScatterOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            getBackendName(),
+            " does not support reduce_scatter_tensor_coalesced"));
+  }
+
+  virtual c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& /* outputBuffer */,
+      at::Tensor& /* inputBuffer */,
+      std::vector<int64_t>& /* outputSplitSizes */,
+      std::vector<int64_t>& /* inputSplitSizes */,
+      const AllToAllOptions& /* opts */ = AllToAllOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support alltoall_base"));
+  }
+
+  virtual c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& /* outputTensors */,
+      std::vector<at::Tensor>& /* inputTensors */,
+      const AllToAllOptions& opts = AllToAllOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support alltoall"));
+  }
+
+  virtual void monitoredBarrier(
+      const BarrierOptions& /* unused */,
+      bool /* unused */ = false) {
+    auto backendName = getBackendName();
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            backendName,
+            " does not support monitoredBarrier, only GLOO supports monitored barrier."));
+  }
+
+  // Agrees on an initial sequence number for the whole group by having rank 0
+  // create it and broadcast it to other ranks using the store. Only implemented
+  // for GLOO and NCCL backends currently.
+  virtual void setSequenceNumberForGroup() {
+    auto backendName = getBackendName();
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            backendName,
+            " does not yet support sequence numbers."));
+  }
+
+  // Retrieves the current sequence number for the whole group, which should be
+  // in sync. If the returned number is not consistent across the group, it
+  // may indicate that there is some sort of collective desynchronization.
+  virtual uint64_t getSequenceNumberForGroup() {
+    auto backendName = getBackendName();
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ",
+            backendName,
+            " does not yet support sequence numbers."));
+  }
+
+  virtual c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& /* tensors */,
+      int /* dstRank */,
+      int /* tag */) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support send"));
+  }
+
+  virtual c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& /* tensors */,
+      int /* srcRank */,
+      int /* tag */) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support recv"));
+  }
+
+  virtual c10::intrusive_ptr<Work> recvAnysource(
+      std::vector<at::Tensor>& /* tensors */,
+      int /* tag */) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support recvAnysource"));
+  }
+
+  virtual c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& /* opts */ = BarrierOptions()) {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support barrier"));
+  }
+
+  virtual void registerOnCompletionHook(
+      std::function<void(std::shared_ptr<WorkInfo>)>&& hook) {
+    TORCH_CHECK(
+        false,
+        "Only ProcessGrouppNCCL supports onCompletion hook, but got ",
+        getBackendName(),
+        " backend.");
+  }
+
+  virtual void waitForPendingWorks() {
+    TORCH_CHECK(
+        false,
+        "Only ProcessGrouppNCCL supports waitForPendingWorks, but got ",
+        getBackendName(),
+        " backend.");
+  }
+
+  virtual void enableCollectivesTiming() {
+    TORCH_CHECK(
+        false,
+        "Backend ",
+        getBackendName(),
+        " is missing implementation of enableCollectivesTiming.");
+  }
+
+  virtual c10::intrusive_ptr<Backend> split(
+      const c10::intrusive_ptr<Store>& store,
+      const std::vector<int>& ranks,
+      const c10::intrusive_ptr<Options>& opts) {
+    TORCH_CHECK(
+        false,
+        "Backend ",
+        getBackendName(),
+        " is missing implementation of split.");
+  }
+
+  virtual c10::intrusive_ptr<Backend> merge(
+      const c10::intrusive_ptr<Store>& store,
+      const c10::intrusive_ptr<Options>& opts,
+      const int& rank,
+      const int& size) {
+    TORCH_CHECK(
+        false,
+        "Backend ",
+        getBackendName(),
+        " is missing implementation of merge.");
+  }
+
+  bool hasHooks() const {
+    return onCompletionHook_ != nullptr;
+  }
+
+  // Do not call this directly, use ProcessGroup::setGroupName instead.
+  virtual void setGroupUid(const std::string& pg_uid) {
+    pg_uid_ = pg_uid;
+  }
+
+  const std::string& getGroupUid() const {
+    return pg_uid_;
+  }
+
+  void setGroupDesc(const std::string& desc) {
+    pg_desc_ = desc;
+  }
+
+  const std::string& getGroupDesc() const {
+    return pg_desc_;
+  }
+
+  // See similar functions in ProcessGroup.hpp for context.
+  std::optional<at::Device> getBoundDeviceId() const {
+    return bound_device_id_;
+  }
+
+  // Perform an eager connect to the specified device if the backend supports
+  // it.
+  virtual void eagerConnectSingleDevice(at::Device device) {
+    // no-op in the default case; this is an optimization some
+    // backends may perform
+  }
+
+  void setBoundDeviceId(std::optional<at::Device> device) {
+    if (device) {
+      TORCH_CHECK(device->has_index(), "setBoundDeviceId must have an index");
+    }
+    bound_device_id_ = device;
+  }
+
+  virtual ErrorType getError() {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support getError"));
+  }
+
+  virtual std::shared_ptr<c10::Allocator> getMemAllocator() {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support getMemAllocator"));
+  }
+
+  // Allocate tensor (aten::empty) from backend's communication-optimized memory
+  // pool
+  virtual at::Tensor allocateTensor(long size, at::TensorOptions options = {}) {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support allocateTensor"));
+  }
+
+  // Returns true if backend supports tensor allocation
+  virtual bool supportsTensorAlloc(c10::DeviceIndex deviceIdx) {
+    // Change to true in concrete backend if supported
+    return false;
+  }
+
+  // Aborts all pending operations and connections in the backend if the backend
+  // supports it.
+  virtual void abort() {}
+
+  // Shutdown the backend if the backend supports it. This should be used for
+  // normal shutdown.
+  virtual void shutdown() {}
+
+  // APIs related to memory offload
+  virtual void suspend() {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support suspend"));
+  }
+
+  virtual void resume() {
+    TORCH_CHECK(
+        false,
+        c10::str("Backend ", getBackendName(), " does not support resume"));
+  }
+
+  virtual std::unordered_map<std::string, uint64_t> getMemoryStats() {
+    TORCH_CHECK(
+        false,
+        c10::str(
+            "Backend ", getBackendName(), " does not support getMemoryStats"));
+  }
+
+ protected:
+  // Implementations of this interface need to call this to setup
+  // appropriate logging etc.
+  void init();
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const int rank_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const int size_;
+  // Debug level setting. It is parsed once when ProcessGroup is constructed and
+  // remains the same across use of this process group.
+  DebugLevel dist_debug_level_;
+  std::string pg_uid_;
+  std::string pg_desc_;
+
+  std::function<void(std::shared_ptr<WorkInfo>)> onCompletionHook_;
+
+  std::optional<at::Device> bound_device_id_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Backoff.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Backoff.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..329dc1f97857e93886096a508bcde1b20c75146d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Backoff.hpp
@@ -0,0 +1,57 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <chrono>
+#include <random>
+#include <thread>
+
+#include <c10/macros/Macros.h>
+
+namespace c10d {
+
+class TORCH_API Backoff {
+ public:
+  virtual ~Backoff() = default;
+
+  virtual std::chrono::milliseconds nextBackoff() = 0;
+  virtual void reset() = 0;
+
+  void sleepBackoff() {
+    std::this_thread::sleep_for(nextBackoff());
+  }
+};
+
+class TORCH_API ExponentialBackoffWithJitter : public Backoff {
+ public:
+  ExponentialBackoffWithJitter();
+
+  std::chrono::milliseconds nextBackoff() override;
+  void reset() override;
+
+ public:
+  std::chrono::milliseconds initialInterval{500};
+  double randomizationFactor{0.5};
+  double multiplier{1.5};
+  std::chrono::milliseconds maxInterval{60000};
+
+ private:
+  std::mt19937 gen_;
+  std::chrono::milliseconds currentInterval_{0};
+};
+
+class TORCH_API FixedBackoff : public Backoff {
+ public:
+  FixedBackoff(std::chrono::milliseconds interval);
+
+  std::chrono::milliseconds nextBackoff() override;
+  void reset() override;
+
+ private:
+  std::chrono::milliseconds interval_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FakeProcessGroup.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FakeProcessGroup.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..453d2f4e64f83cf9466ce951e8be8a4f666a1393
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FakeProcessGroup.hpp
@@ -0,0 +1,258 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/utils.h>
+
+namespace c10d {
+
+class FakeWork : public Work {
+ public:
+  int seq_id = -1;
+  bool wait(std::chrono::milliseconds timeout = kNoTimeout) override {
+    return true;
+  }
+
+  c10::intrusive_ptr<c10::ivalue::Future> getFuture() override {
+    auto fut = c10::make_intrusive<c10::ivalue::Future>(c10::NoneType::get());
+    fut->markCompleted();
+    return fut;
+  }
+};
+
+class FakeProcessGroup : public Backend {
+ public:
+  struct Options : Backend::Options {
+    explicit Options() : Backend::Options("fake") {}
+
+    int fake_option = 0;
+    bool error_on_collective = false;
+  };
+
+  // Static factory method for official APIs
+  static c10::intrusive_ptr<FakeProcessGroup> _create_internal(
+      int rank,
+      int size,
+      c10::intrusive_ptr<Options> options = c10::make_intrusive<Options>()) {
+    return c10::make_intrusive<FakeProcessGroup>(
+        rank, size, std::move(options));
+  }
+
+  const std::string getBackendName() const override {
+    return "fake";
+  }
+
+  c10::intrusive_ptr<Backend::Options> getBackendOptions() override {
+    return c10::static_intrusive_pointer_cast<Backend::Options>(options_);
+  }
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& /* tensors */,
+      const BroadcastOptions& /* opts */ = BroadcastOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& /* tensors */,
+      const AllreduceOptions& /* opts */ = AllreduceOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> allreduce_sparse(
+      std::vector<at::Tensor>& /* tensors */,
+      const AllreduceOptions& /* opts */ = AllreduceOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& /* tensors */,
+      const AllreduceCoalescedOptions& /* opts */ =
+          AllreduceCoalescedOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& /* tensors */,
+      const ReduceOptions& /* opts */ = ReduceOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  // NOTE [allgather on FakeProcessGroup]
+  // Assume each rank have the same input tensor so we just copy to the results
+  // since it's not a real allgather, we simply make this copying logic to let
+  // some simple validation works (i.e. calling allgather to see if each rank
+  // have the same tensor or not).
+  //
+  // NOTE: in general it's not good form to try to make FakeProcessGroup work
+  // with real data, but the reasoning here is that we want FakeProcessGroup to
+  // work with DeviceMesh's init code that have the data validation, which
+  // makes it worth the tradeoff.
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) override {
+    checkCollectiveError();
+    for (auto& tensor : outputTensors[0]) {
+      tensor.copy_(inputTensors[0]);
+    }
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) override {
+    checkCollectiveError();
+    auto chunks = outputBuffer.chunk(size_);
+    for (auto& tensor : chunks) {
+      tensor.copy_(inputBuffer);
+    }
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> allgather_coalesced(
+      std::vector<std::vector<at::Tensor>>& /* outputTensorLists */,
+      std::vector<at::Tensor>& /* inputTensors */,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const AllgatherOptions& /* opts */ = AllgatherOptions()) override {
+    checkCollectiveError();
+    for (size_t i = 0; i < outputs.size(); ++i) {
+      auto chunks = outputs[i].chunk(size_);
+      for (auto& chunk : chunks) {
+        chunk.copy_(inputs[i]);
+      }
+    }
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& /* outputTensors */,
+      std::vector<at::Tensor>& /* inputTensors */,
+      const GatherOptions& /* opts */ = GatherOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& /* outputTensors */,
+      std::vector<std::vector<at::Tensor>>& /* inputTensors */,
+      const ScatterOptions& /* opts */ = ScatterOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& /* outputTensors */,
+      std::vector<std::vector<at::Tensor>>& /* inputTensors */,
+      const ReduceScatterOptions& /* opts */ =
+          ReduceScatterOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& /* outputBuffer */,
+      at::Tensor& /* inputBuffer */,
+      const ReduceScatterOptions& /* opts */ =
+          ReduceScatterOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& /* outputs */,
+      std::vector<at::Tensor>& /* inputs */,
+      const ReduceScatterOptions& /* opts */ =
+          ReduceScatterOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& /* outputBuffer */,
+      at::Tensor& /* inputBuffer */,
+      std::vector<int64_t>& /* outputSplitSizes */,
+      std::vector<int64_t>& /* inputSplitSizes */,
+      const AllToAllOptions& /* opts */ = AllToAllOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& /* outputTensors */,
+      std::vector<at::Tensor>& /* inputTensors */,
+      const AllToAllOptions& opts = AllToAllOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& /* tensors */,
+      int /* dstRank */,
+      int /* tag */) override {
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& /* tensors */,
+      int /* srcRank */,
+      int /* tag */) override {
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> recvAnysource(
+      std::vector<at::Tensor>& /* tensors */,
+      int /* tag */) override {
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  void startCoalescing() override {
+    // No-op
+  }
+
+  c10::intrusive_ptr<Work> endCoalescing(OpType /* optype */) {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> endCoalescing() override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& /* opts */ = BarrierOptions()) override {
+    checkCollectiveError();
+    return c10::make_intrusive<FakeWork>();
+  }
+
+  // Private constructor used by official APIs
+  FakeProcessGroup(int rank, int size, c10::intrusive_ptr<Options> options)
+      : Backend(rank, size), options_(std::move(options)) {}
+  c10::intrusive_ptr<Options> options_;
+
+ private:
+  void checkCollectiveError() {
+    TORCH_CHECK(
+        !options_ || !options_->error_on_collective,
+        "FakeProcessGroup collective operation error (error_on_collective=true)");
+  }
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FileStore.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FileStore.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..7725d9856572a906d688992dd6eaa284d019887d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FileStore.hpp
@@ -0,0 +1,72 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <sys/types.h>
+
+#include <mutex>
+#include <unordered_map>
+
+#include <torch/csrc/distributed/c10d/Store.hpp>
+
+namespace c10d {
+
+class TORCH_API FileStore : public Store {
+ public:
+  explicit FileStore(std::string path, int numWorkers);
+
+  c10::intrusive_ptr<Store> clone() override;
+
+  ~FileStore() override;
+
+  void set(const std::string& key, const std::vector<uint8_t>& value) override;
+
+  std::vector<uint8_t> compareSet(
+      const std::string& key,
+      const std::vector<uint8_t>& expectedValue,
+      const std::vector<uint8_t>& desiredValue) override;
+
+  std::vector<uint8_t> get(const std::string& key) override;
+
+  int64_t add(const std::string& key, int64_t value) override;
+
+  int64_t getNumKeys() override;
+
+  bool deleteKey(const std::string& key) override;
+
+  bool check(const std::vector<std::string>& keys) override;
+
+  void wait(const std::vector<std::string>& keys) override;
+
+  void wait(
+      const std::vector<std::string>& keys,
+      const std::chrono::milliseconds& timeout) override;
+
+  // Returns the path used by the FileStore.
+  const std::string& getPath() const noexcept {
+    return path_;
+  }
+
+  std::vector<std::string> listKeys() override;
+
+ protected:
+  int64_t addHelper(const std::string& key, int64_t i);
+
+  std::string path_;
+  off_t pos_{0};
+
+  int numWorkers_;
+  const std::string cleanupKey_;
+  const std::string refCountKey_;
+  const std::string regularPrefix_;
+  const std::string deletePrefix_;
+
+  std::unordered_map<std::string, std::vector<uint8_t>> cache_;
+
+  std::mutex activeFileOpLock_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FlightRecorder.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FlightRecorder.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..60985a09d1e7153c4949737ee82ffb1cf66c104a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FlightRecorder.hpp
@@ -0,0 +1,336 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <cstdio>
+#include <cstdlib>
+
+#include <memory>
+#include <mutex>
+
+#include <ATen/ATen.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/distributed/c10d/TraceUtils.h>
+#include <torch/csrc/distributed/c10d/logger.hpp>
+#include <optional>
+
+namespace c10d {
+
+#define DEFINE_CONSTANT(name, value) \
+  static c10::IValue name = value;   \
+  static std::string name##_str = value;
+// Update whenever changing contents or formatting of the dump
+// (minor when adding fields, major when changing existing fields)
+// Also update both JSON and Pickle dumps to make use of the newly defined
+// field(s).
+DEFINE_CONSTANT(version_val, "2.10")
+DEFINE_CONSTANT(entries_key, "entries")
+DEFINE_CONSTANT(nccl_comm_key, "nccl_comm_state")
+DEFINE_CONSTANT(comm_lib_version_key, "comm_lib_version")
+DEFINE_CONSTANT(version_key, "version")
+DEFINE_CONSTANT(pg_config_key, "pg_config")
+DEFINE_CONSTANT(pg_status_key, "pg_status")
+DEFINE_CONSTANT(record_id_key, "record_id")
+DEFINE_CONSTANT(pg_id_key, "pg_id")
+DEFINE_CONSTANT(pg_name_key, "process_group")
+DEFINE_CONSTANT(collective_seq_id_key, "collective_seq_id")
+DEFINE_CONSTANT(p2p_seq_id_key, "p2p_seq_id")
+DEFINE_CONSTANT(is_p2p_key, "is_p2p")
+DEFINE_CONSTANT(op_id_key, "op_id")
+DEFINE_CONSTANT(profiling_name_key, "profiling_name")
+DEFINE_CONSTANT(input_sizes_key, "input_sizes")
+DEFINE_CONSTANT(input_dtypes_key, "input_dtypes")
+DEFINE_CONSTANT(output_sizes_key, "output_sizes")
+DEFINE_CONSTANT(output_dtypes_key, "output_dtypes")
+DEFINE_CONSTANT(time_created_key, "time_created_ns")
+DEFINE_CONSTANT(duration_key, "duration_ms")
+DEFINE_CONSTANT(timeout_key, "timeout_ms")
+DEFINE_CONSTANT(frames_key, "frames")
+DEFINE_CONSTANT(state_key, "state")
+DEFINE_CONSTANT(line_key, "line")
+DEFINE_CONSTANT(name_key, "name")
+DEFINE_CONSTANT(filename_key, "filename")
+DEFINE_CONSTANT(retired_key, "retired")
+DEFINE_CONSTANT(time_discovered_started_key, "time_discovered_started_ns")
+DEFINE_CONSTANT(time_discovered_completed_key, "time_discovered_completed_ns")
+DEFINE_CONSTANT(completed_state, "completed")
+DEFINE_CONSTANT(scheduled_state, "scheduled")
+DEFINE_CONSTANT(started_state, "started")
+DEFINE_CONSTANT(thread_id_key, "thread_id")
+DEFINE_CONSTANT(thread_name_key, "thread_name")
+#undef DEFINE_CONSTANT
+
+// Write NCCL debug info to local disk or any storage users define.
+// There are some constrains we set for the debug info writer:
+// 1. The writer should only be registered once.
+// 2. Once registered, users cannot change it including un-register.
+// 3. It is recommended to register the customized writer in the trainer setup,
+//    If users don't register before calling launchAsyncDebugDump, then users
+//    lose the chance to register (and the default writer will be
+//    auto-registered).
+class TORCH_API DebugInfoWriter {
+ public:
+  virtual ~DebugInfoWriter() = default;
+  virtual void write(const std::string& trace);
+  static DebugInfoWriter& getWriter(int rank);
+  static void registerWriter(std::unique_ptr<DebugInfoWriter> writer);
+  virtual std::string getWriterTarget() {
+    return filename_;
+  }
+
+ protected:
+  DebugInfoWriter(
+      const std::string& namePrefix,
+      int rank,
+      bool enableDynamicFilename = false) {
+    filename_ = c10::str(namePrefix, rank);
+    enable_dynamic_filename_ = enableDynamicFilename;
+    rank_ = rank;
+  }
+  std::string filename_;
+  int rank_;
+  bool enable_dynamic_filename_;
+
+ private:
+  static std::unique_ptr<DebugInfoWriter> writer_;
+  static std::atomic<bool> hasWriterRegistered_;
+};
+
+template <typename EventType>
+struct FlightRecorder {
+  static FlightRecorder<EventType>* get() {
+    // intentionally leak on exit
+    // because this will hold python state that may get destructed
+    static FlightRecorder<EventType>* instance =
+        new FlightRecorder<EventType>();
+    return instance;
+  }
+  FlightRecorder() {
+    // NOTE: This default value (2000) is duplicated in ProcessGroupNCCL.cpp
+    // and ProcessGroupNCCL.hpp because they cannot directly query max_entries_
+    // (no public accessor). Keep these values in sync.
+    max_entries_ = getCvarInt(
+        {"TORCH_FR_BUFFER_SIZE", "TORCH_NCCL_TRACE_BUFFER_SIZE"}, 2000);
+    capture_cpp_stack_ = getCvarBool(
+        {"TORCH_FR_CPP_STACK", "TORCH_NCCL_TRACE_CPP_STACK"}, false);
+    enabled_ = max_entries_ > 0;
+    reset_epoch_start_idx_[0] = 0;
+  }
+  struct Entry {
+    size_t id_; // incremented id in the trace buffer
+                // used to figure out where in the circular entries
+                // buffer this entry will be located to
+                // update state information
+    size_t reset_epoch_; // epoch when this entry was created
+    size_t pg_id_;
+    std::tuple<std::string, std::string> pg_name_; // <group_name, group_desc>
+
+    // collective_seq_id and p2p_seq_id refer to actual kernel launches (e.g. 1
+    // per coalesced group).
+    // collective_seq_id only increments for true collective operations (over
+    // all ranks in the group). p2p_seq_id only increments over non-collective
+    // operations in the group. op_id refers to logical operations (e.g. one per
+    // op inside coalesced group)
+    size_t collective_seq_id_;
+    size_t p2p_seq_id_;
+    size_t op_id_;
+    std::string profiling_name_;
+
+    std::shared_ptr<torch::CapturedTraceback> traceback_;
+    // we borrow pointers to start_ and end_ so we can query the state
+    // on reporting. However, once the event is completed, the call
+    // to `complete` will clear these.
+    EventType *start_, *end_;
+
+    // timestamp when the entry was created, likely close to the time the work
+    // was 'enqueued'- not necessarily started
+    c10::time_t time_created_;
+
+    // configured timeout for this entry
+    c10::time_t timeout_ms_;
+
+    // Is this a P2P event?
+    bool isP2P_;
+
+    std::optional<float> duration_;
+
+    // timestamp when our CPU threads discovered that the kernel started.
+    // will always be _after_ it actually started, and can be very late
+    // if the watchdog thread got stuck on CUDA APIs.
+    std::optional<c10::time_t> time_discovered_started_;
+
+    // timestamp when our CPU threads discovered that the kernel completed.
+    // will always be _after_ it actually completed, and can be the same time
+    // as the discovery of the start if the watchdog thread is stuck on CUDA
+    // APIs
+    std::optional<c10::time_t> time_discovered_completed_;
+
+    // size information for input/output tensors
+    c10::SmallVector<int64_t, 4> input_dims_;
+    std::vector<c10::ScalarType> input_dtypes_;
+    c10::SmallVector<int64_t, 4> output_dims_;
+    std::vector<c10::ScalarType> output_dtypes_;
+    c10::SmallVector<int64_t, 8> sizes_; // flattened from inputs, outputs
+    std::thread::id thread_id_;
+    std::string thread_name_;
+    bool retired_ = false; // is this work entry no longer in the workMetaList_?
+                           // a retired but not completed event has timed out
+
+    // Returns the traceback of current entry, in string form.
+    // Note: `getTraceback` invokes `torch::symbolize`, which may need to
+    // acquire the GIL. If you don't want to block the current thread or take
+    // the risk of a GIL deadlock, you can use an asynchronous calling mechanism
+    // like std::async.
+    TORCH_API std::string getTraceback();
+  };
+
+  bool enabled_ = false;
+  bool capture_cpp_stack_ = false;
+  std::mutex mutex_;
+  std::vector<Entry> entries_;
+  size_t max_entries_ = 0;
+  size_t next_ = 0;
+  size_t id_ = 0;
+  size_t reset_epoch_ = 0;
+  std::unordered_map<size_t, size_t>
+      reset_epoch_start_idx_; // maps reset_epoch to the idx where it starts
+  std::map<size_t, std::shared_ptr<ProcessGroupStatus>> all_pg_status_;
+  std::map<std::tuple<std::string, std::string>, std::vector<uint64_t>>
+      pg_name_to_ranks_;
+  std::string comm_lib_version_;
+
+  struct TraceIdentifier {
+    std::optional<size_t> id;
+    std::optional<size_t> reset_epoch;
+  };
+
+  TraceIdentifier recordWithResetEnabled(
+      size_t pg_id,
+      const std::tuple<std::string, std::string>& pg_name,
+      size_t collective_seq_id,
+      size_t p2p_seq_id,
+      size_t op_id,
+      std::string profiling_name,
+      const std::vector<at::Tensor>& inputs,
+      const std::vector<at::Tensor>& outputs,
+      EventType* start,
+      EventType* end,
+      std::chrono::milliseconds timeout_ms,
+      std::shared_ptr<ProcessGroupStatus> pg_status,
+      bool isP2P);
+
+  std::optional<size_t> record(
+      size_t pg_id,
+      const std::tuple<std::string, std::string>& pg_name,
+      size_t collective_seq_id,
+      size_t p2p_seq_id,
+      size_t op_id,
+      std::string profiling_name,
+      const std::vector<at::Tensor>& inputs,
+      const std::vector<at::Tensor>& outputs,
+      EventType* start,
+      EventType* end,
+      std::chrono::milliseconds timeout_ms,
+      std::shared_ptr<ProcessGroupStatus> pg_status,
+      bool isP2P);
+
+  TORCH_API void record_pg_ranks(
+      const std::tuple<std::string, std::string>& pg_name,
+      std::vector<uint64_t> ranks);
+
+  void record_accelerator_version(const std::string comm_lib_version);
+
+  void update_state(Entry& r);
+
+  std::vector<Entry> dump_entries();
+
+  // Returns the index in entries_ for the given id and reset_epoch.
+  // Caller must hold mutex_lock before calling this method.
+  size_t getIdxFromId(size_t id, size_t reset_epoch) const;
+
+  // Returns the entry with the given id and reset_epoch, if it exists.
+  // Otherwise, returns std::nullopt.
+  TORCH_API std::optional<Entry> getEntry(
+      std::optional<size_t> id,
+      std::optional<size_t> reset_epoch);
+
+  TORCH_API std::optional<Entry> getEntry(std::optional<size_t> id);
+
+  /*
+  Mark an Event as completed and free its events.
+  This is called by the watchdog thread, and is asynchronous from the
+  perspective of the main thread.
+  compute_duration defaults to true since retire_id is only called in the
+  watchdog thread, which is currently a place we call cuda APIs which may hang,
+  but care should be taken to avoid computing duration in any function that must
+  never hang. (timing must also be enabled for compute_duration - see
+  TORCH_NCCL_ENABLE_TIMING).
+  */
+  TORCH_API void retire_id(
+      std::optional<size_t> id,
+      std::optional<size_t> reset_epoch,
+      bool compute_duration = true);
+
+  TORCH_API void retire_id(
+      std::optional<size_t> id,
+      bool compute_duration = true);
+
+  TORCH_API void reset_all();
+
+  const c10::List<c10::IValue> getCollectiveTrace(
+      bool includeStacktraces,
+      bool onlyActive);
+
+  // dump pg_entries
+  const c10::Dict<c10::IValue, c10::IValue> getPgConfig();
+
+  const std::map<std::string, std::map<std::string, std::string>>
+  getPgConfigJson();
+
+  // dump pg_status
+  const c10::Dict<c10::IValue, c10::IValue> getPgStatus();
+
+  const std::map<std::string, std::map<std::string, std::string>>
+  getPgStatusJson();
+
+  std::string dump_json(
+      const std::optional<std::unordered_map<
+          std::string,
+          std::unordered_map<std::string, std::string>>>& extraDumpMap,
+      bool includeCollectives,
+      bool onlyActive);
+
+  std::string dump(
+      const std::optional<std::unordered_map<
+          std::string,
+          std::unordered_map<std::string, std::string>>>& extraDumpMap,
+      bool includeCollectives,
+      bool includeStackTraces,
+      bool onlyActive);
+};
+
+// Whether to include stack trace in the Flight Recorder trace (default true)
+static std::vector<std::string> TORCH_INCLUDE_STACK_TRACE = {
+    "TORCH_INCLUDE_STACK_TRACE"};
+
+// Whether to include only active collectives in the Flight Recorder trace
+// (default false)
+static std::vector<std::string> TORCH_INCLUDE_ONLY_ACTIVE = {
+    "TORCH_INCLUDE_ONLY_ACTIVE"};
+
+// Dumps the fr traces and additional information about the Process
+// Group.
+TORCH_API std::string dump_fr_trace(
+    bool includeCollectives,
+    bool includeStackTraces,
+    bool onlyActive);
+
+// Dumps the fr traces and additional information about the Process
+// Group in JSON formatted string.
+// We don't include stack traces in JSON format as it is far too much data.
+TORCH_API std::string dump_fr_trace_json(
+    bool includeCollectives,
+    bool onlyActive);
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..89b9987f011d005e95dd208dbd53a4c6b6954b81
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp
@@ -0,0 +1,641 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <nlohmann/json.hpp>
+
+#include <c10/util/WaitCounter.h>
+#include <c10/util/thread_name.h>
+
+#include <torch/csrc/distributed/c10d/FlightRecorder.hpp>
+
+namespace c10d {
+
+template <typename EventType>
+float getDurationFromEvent(EventType& start, EventType& end);
+
+// Returns the traceback of current entry, in string form.
+// Note: `getTraceback` invokes `torch::symbolize`, which may need to acquire
+// the GIL. If you don't want to block the current thread or take the risk of a
+// GIL deadlock, you can use an asynchronous calling mechanism like std::async.
+template <typename EventType>
+std::string FlightRecorder<EventType>::Entry::getTraceback() {
+  torch::CapturedTraceback* traceback = traceback_.get();
+  torch::SymbolizedTracebacks s_tbs = torch::symbolize({traceback});
+  // We use 0 because we only have one traceback here.
+  const auto& s_tb = s_tbs.tracebacks.at(0);
+  std::stringstream oss;
+  for (auto idx : c10::irange(s_tb.size())) {
+    auto frame_id = s_tb[idx];
+    const auto& frame = s_tbs.all_frames.at(frame_id);
+    oss << '#' << idx << ' ' << frame.funcname << " from " << frame.filename
+        << ':' << frame.lineno << '\n';
+  }
+  /* Resulted format is like:
+    #0 all_reduce from pytorch/torch/distributed/distributed_c10d.py:2696
+    #1 wrapper from pytorch/torch/distributed/c10d_logger.py:83
+    #2 bar from /home/user/repro.py:15
+    #3 foo from /home/user/repro.py:24
+    #4 main from /home/user/repro.py:34
+    #5 <module> from /home/user/repro.py:40
+  */
+  return oss.str();
+}
+
+template <typename EventType>
+std::optional<size_t> FlightRecorder<EventType>::record(
+    size_t pg_id,
+    const std::tuple<std::string, std::string>& pg_name,
+    size_t collective_seq_id,
+    size_t p2p_seq_id,
+    size_t op_id,
+    std::string profiling_name,
+    const std::vector<at::Tensor>& inputs,
+    const std::vector<at::Tensor>& outputs,
+    EventType* start,
+    EventType* end,
+    std::chrono::milliseconds timeout_ms,
+    std::shared_ptr<ProcessGroupStatus> pg_status,
+    bool isP2P) {
+  auto result = recordWithResetEnabled(
+      pg_id,
+      pg_name,
+      collective_seq_id,
+      p2p_seq_id,
+      op_id,
+      std::move(profiling_name),
+      inputs,
+      outputs,
+      start,
+      end,
+      timeout_ms,
+      std::move(pg_status),
+      isP2P);
+  return result.id;
+}
+
+template <typename EventType>
+typename FlightRecorder<EventType>::TraceIdentifier FlightRecorder<EventType>::
+    recordWithResetEnabled(
+        size_t pg_id,
+        const std::tuple<std::string, std::string>& pg_name,
+        size_t collective_seq_id,
+        size_t p2p_seq_id,
+        size_t op_id,
+        std::string profiling_name,
+        const std::vector<at::Tensor>& inputs,
+        const std::vector<at::Tensor>& outputs,
+        EventType* start,
+        EventType* end,
+        std::chrono::milliseconds timeout_ms,
+        std::shared_ptr<ProcessGroupStatus> pg_status,
+        bool isP2P) {
+  if (!enabled_) {
+    return TraceIdentifier{std::nullopt, std::nullopt};
+  }
+  auto traceback =
+      torch::CapturedTraceback::gather(true, true, capture_cpp_stack_);
+  std::lock_guard<std::mutex> guard(mutex_);
+  if (all_pg_status_.find(pg_id) == all_pg_status_.end()) {
+    // Current pg_status is not in FR.
+    all_pg_status_[pg_id] = std::move(pg_status);
+  }
+
+  TORCH_CHECK(
+      reset_epoch_start_idx_.find(reset_epoch_) !=
+      reset_epoch_start_idx_.end());
+
+  auto te = Entry{
+      id_,
+      reset_epoch_,
+      pg_id,
+      pg_name,
+      collective_seq_id,
+      p2p_seq_id,
+      op_id,
+      std::move(profiling_name),
+      std::move(traceback),
+      start,
+      end,
+      c10::getTime(),
+      timeout_ms.count(),
+      isP2P,
+      std::nullopt,
+      std::nullopt,
+      std::nullopt,
+      {},
+      {},
+      {},
+      {},
+      {},
+      std::this_thread::get_id(),
+      c10::getThreadName(),
+      false};
+
+  for (const auto& input : inputs) {
+    c10::IntArrayRef sizes = input.sizes();
+    te.input_dtypes_.push_back(input.dtype().toScalarType());
+    te.input_dims_.push_back(static_cast<int64_t>(sizes.size()));
+    te.sizes_.insert(te.sizes_.end(), sizes.begin(), sizes.end());
+  }
+
+  for (const auto& output : outputs) {
+    c10::IntArrayRef sizes = output.sizes();
+    te.output_dtypes_.push_back(output.dtype().toScalarType());
+    te.output_dims_.push_back(static_cast<int64_t>(sizes.size()));
+    te.sizes_.insert(te.sizes_.end(), sizes.begin(), sizes.end());
+  }
+
+  const auto next = next_++;
+
+  if (entries_.size() < max_entries_) {
+    entries_.emplace_back(std::move(te));
+  } else {
+    entries_[next] = std::move(te);
+  }
+
+  if (next_ == max_entries_) {
+    next_ = 0;
+  }
+
+  const auto id = id_++;
+  return TraceIdentifier{id, reset_epoch_};
+}
+
+template <typename EventType>
+void FlightRecorder<EventType>::record_pg_ranks(
+    const std::tuple<std::string, std::string>& pg_name,
+    std::vector<uint64_t> ranks) {
+  if (!enabled_) {
+    return;
+  }
+  std::lock_guard<std::mutex> guard(mutex_);
+  pg_name_to_ranks_[pg_name] = std::move(ranks);
+}
+
+template <typename EventType>
+void FlightRecorder<EventType>::record_accelerator_version(
+    const std::string comm_lib_version) {
+  if (!enabled_) {
+    return;
+  }
+  std::lock_guard<std::mutex> guard(mutex_);
+  comm_lib_version_ = std::move(comm_lib_version);
+}
+
+template <typename EventType>
+void FlightRecorder<EventType>::update_state(Entry& r) {
+  try {
+    if (r.start_ != nullptr) {
+      bool started = r.start_->query();
+      if (started && !r.time_discovered_started_) {
+        r.time_discovered_started_ = c10::getTime();
+      }
+    }
+    if (r.end_ != nullptr) {
+      bool completed = r.end_->query();
+      if (completed && !r.time_discovered_completed_) {
+        r.time_discovered_completed_ = c10::getTime();
+      }
+    }
+  } catch (std::exception& e) {
+    LOG(ERROR) << "Failed to update state for entry " << r.id_ << ": "
+               << r.profiling_name_ << " with error: " << e.what();
+  }
+}
+
+template <typename EventType>
+std::vector<typename FlightRecorder<EventType>::Entry> FlightRecorder<
+    EventType>::dump_entries() {
+  std::vector<Entry> result;
+  {
+    std::lock_guard<std::mutex> guard(mutex_);
+    // Filter entries during insertion - only keep entries from current epoch
+    auto filter = [this](const Entry& e) {
+      return e.reset_epoch_ == reset_epoch_;
+    };
+    std::copy_if(
+        entries_.begin() + static_cast<std::ptrdiff_t>(next_),
+        entries_.end(),
+        std::back_inserter(result),
+        filter);
+    std::copy_if(
+        entries_.begin(),
+        entries_.begin() + static_cast<std::ptrdiff_t>(next_),
+        std::back_inserter(result),
+        filter);
+  }
+  // query any remaining events
+  for (auto& r : result) {
+    update_state(r);
+    r.start_ = r.end_ = nullptr;
+  }
+  return result;
+}
+
+template <typename EventType>
+// Returns the index in entries_ for the given id and reset_epoch.
+// Caller must hold mutex_lock before calling this method.
+size_t FlightRecorder<EventType>::getIdxFromId(size_t id, size_t reset_epoch)
+    const {
+  // Look up the starting idx for the given reset epoch
+  auto it = reset_epoch_start_idx_.find(reset_epoch);
+  TORCH_CHECK(it != reset_epoch_start_idx_.end());
+  // Calculate idx based on where the epoch started
+  return (it->second + id) % max_entries_;
+}
+
+template <typename EventType>
+// Returns the entry with the given id and reset_epoch, if it exists. Otherwise,
+// returns std::nullopt.
+std::optional<typename FlightRecorder<EventType>::Entry> FlightRecorder<
+    EventType>::
+    getEntry(std::optional<size_t> id, std::optional<size_t> reset_epoch) {
+  if (!enabled_ || !id || !reset_epoch) {
+    return std::nullopt;
+  }
+
+  std::unique_lock<std::mutex> guard(mutex_);
+  Entry entry = entries_.at(getIdxFromId(*id, *reset_epoch));
+  if (entry.id_ == *id && entry.reset_epoch_ == *reset_epoch) {
+    return entry;
+  }
+  return std::nullopt;
+}
+
+template <typename EventType>
+std::optional<typename FlightRecorder<EventType>::Entry> FlightRecorder<
+    EventType>::getEntry(std::optional<size_t> id) {
+  return getEntry(id, 0);
+}
+
+template <typename EventType>
+void FlightRecorder<EventType>::retire_id(
+    std::optional<size_t> id,
+    std::optional<size_t> reset_epoch,
+    bool compute_duration) {
+  if (!enabled_ || !id || !reset_epoch) {
+    return;
+  }
+
+  bool can_compute_duration = false;
+  EventType* startEvent = nullptr;
+  EventType* endEvent = nullptr;
+  std::optional<float> duration = std::nullopt;
+
+  std::unique_lock<std::mutex> guard(mutex_);
+
+  Entry* entry = &entries_.at(getIdxFromId(*id, *reset_epoch));
+  if (entry->id_ == *id && entry->reset_epoch_ == *reset_epoch) {
+    update_state(*entry);
+
+    if (compute_duration) {
+      can_compute_duration = entry->time_discovered_completed_.has_value() &&
+          entry->start_ && entry->end_;
+      startEvent = entry->start_;
+      endEvent = entry->end_;
+    }
+    entry->retired_ = true;
+    entry->start_ = entry->end_ = nullptr;
+  }
+
+  if (can_compute_duration) {
+    // Compute duration without without holding the lock, because
+    // cudaEventDuration() can hang, and we need to acquire the lock before we
+    // can dump(), which we never want to block.
+    guard.unlock();
+    duration = getDurationFromEvent<EventType>(*startEvent, *endEvent);
+    guard.lock();
+
+    // Refresh the entry pointer, see if the entry has been overwritten
+    entry = &entries_.at(getIdxFromId(*id, *reset_epoch));
+    if (!(entry->id_ == *id && entry->reset_epoch_ == *reset_epoch)) {
+      LOG(INFO) << "retire_id abandoned for id " << *id
+                << ", event was overwritten while waiting to compute duration.";
+      return;
+    }
+    if (duration.has_value()) {
+      entry->duration_ = duration;
+    }
+  }
+}
+
+template <typename EventType>
+void FlightRecorder<EventType>::retire_id(
+    std::optional<size_t> id,
+    bool compute_duration) {
+  retire_id(id, 0, compute_duration);
+}
+
+template <typename EventType>
+void FlightRecorder<EventType>::reset_all() {
+  std::lock_guard<std::mutex> guard(mutex_);
+  if (!entries_.empty()) {
+    // Soft delete: increment epoch to mark all existing entries as old
+    // Store where the new epoch starts in the circular buffer
+    reset_epoch_++;
+    reset_epoch_start_idx_[reset_epoch_] = next_;
+    id_ = 0;
+  }
+}
+
+template <typename EventType>
+const c10::List<c10::IValue> FlightRecorder<EventType>::getCollectiveTrace(
+    bool includeStacktraces,
+    bool onlyActive) {
+  auto entries = new_list();
+  // Entries are returned in the order they were recorded
+  auto result = dump_entries();
+  std::vector<torch::CapturedTraceback*> tracebacks;
+  torch::SymbolizedTracebacks stracebacks;
+  std::vector<c10::IValue> all_frames;
+  if (includeStacktraces) {
+    for (auto& e : result) {
+      tracebacks.push_back(e.traceback_.get());
+    }
+    stracebacks = torch::symbolize(tracebacks);
+    for (const auto& f : stracebacks.all_frames) {
+      auto d = new_dict();
+      d.insert(name_key, f.funcname);
+      d.insert(filename_key, f.filename);
+      d.insert(line_key, int64_t(f.lineno));
+      all_frames.emplace_back(std::move(d));
+    }
+  }
+  for (auto i : c10::irange(result.size())) {
+    auto dict = new_dict();
+    auto& e = result.at(i);
+    // Skip completed events
+    if (onlyActive && e.time_discovered_completed_.has_value()) {
+      continue;
+    }
+    if (includeStacktraces) {
+      auto& tb = stracebacks.tracebacks.at(i);
+      auto frames = new_list();
+      for (auto frame : tb) {
+        frames.push_back(all_frames.at(frame));
+      }
+      dict.insert(frames_key, frames);
+    }
+
+    dict.insert(record_id_key, int64_t(e.id_));
+    dict.insert(pg_id_key, int64_t(e.pg_id_));
+    dict.insert(pg_name_key, e.pg_name_);
+    dict.insert(thread_name_key, e.thread_name_);
+    dict.insert(thread_id_key, c10::str(e.thread_id_));
+    dict.insert(collective_seq_id_key, int64_t(e.collective_seq_id_));
+    dict.insert(p2p_seq_id_key, int64_t(e.p2p_seq_id_));
+    dict.insert(op_id_key, int64_t(e.op_id_));
+    dict.insert(profiling_name_key, e.profiling_name_);
+    dict.insert(time_created_key, int64_t(e.time_created_));
+    if (e.duration_) {
+      dict.insert(duration_key, *e.duration_);
+    }
+
+    auto it = e.sizes_.begin();
+    auto read_sizes = [&](const c10::SmallVector<int64_t, 4>& dims) {
+      auto sizes = new_list();
+      for (auto dim : dims) {
+        auto arg_sizes = new_list();
+        for ([[maybe_unused]] auto i : c10::irange(dim)) {
+          arg_sizes.push_back(*it++);
+        }
+        sizes.push_back(arg_sizes);
+      }
+      return sizes;
+    };
+
+    dict.insert(input_sizes_key, read_sizes(e.input_dims_));
+    std::vector<std::string> input_dtypes_strs;
+    input_dtypes_strs.reserve(e.input_dtypes_.size());
+    for (const auto& input_dtype : e.input_dtypes_) {
+      input_dtypes_strs.emplace_back(c10::toString(input_dtype));
+    }
+    dict.insert(input_dtypes_key, input_dtypes_strs);
+    dict.insert(output_sizes_key, read_sizes(e.output_dims_));
+    std::vector<std::string> output_dtypes_strs;
+    output_dtypes_strs.reserve(e.output_dtypes_.size());
+    for (const auto& output_dtype : e.output_dtypes_) {
+      output_dtypes_strs.emplace_back(c10::toString(output_dtype));
+    }
+    dict.insert(output_dtypes_key, output_dtypes_strs);
+    if (e.time_discovered_completed_.has_value()) {
+      dict.insert(state_key, completed_state);
+    } else if (e.time_discovered_started_.has_value()) {
+      dict.insert(state_key, started_state);
+    } else {
+      dict.insert(state_key, scheduled_state);
+    }
+
+    dict.insert(
+        time_discovered_started_key,
+        e.time_discovered_started_.has_value()
+            ? int64_t(*e.time_discovered_started_)
+            : c10::IValue());
+    dict.insert(
+        time_discovered_completed_key,
+        e.time_discovered_completed_.has_value()
+            ? int64_t(*e.time_discovered_completed_)
+            : c10::IValue());
+    dict.insert(retired_key, e.retired_);
+    dict.insert(timeout_key, e.timeout_ms_);
+    dict.insert(is_p2p_key, e.isP2P_);
+
+    entries.push_back(dict);
+  }
+  return entries;
+}
+
+template <typename EventType>
+const c10::Dict<c10::IValue, c10::IValue> FlightRecorder<
+    EventType>::getPgConfig() {
+  auto pg_config = new_dict();
+  for (const auto& [pg_name, ranks] : pg_name_to_ranks_) {
+    auto pg_info = new_dict();
+    pg_info.insert("name", std::get<0>(pg_name));
+    pg_info.insert("desc", std::get<1>(pg_name));
+    pg_info.insert("ranks", ranks_str(ranks));
+    pg_config.insert(std::get<0>(pg_name), pg_info);
+  }
+  return pg_config;
+}
+
+template <typename EventType>
+const std::map<std::string, std::map<std::string, std::string>> FlightRecorder<
+    EventType>::getPgConfigJson() {
+  std::map<std::string, std::map<std::string, std::string>> result;
+  for (const auto& [pg_name, ranks] : pg_name_to_ranks_) {
+    auto pg_info = std::map<std::string, std::string>();
+    pg_info["name"] = std::get<0>(pg_name);
+    pg_info["desc"] = std::get<1>(pg_name);
+    pg_info["ranks"] = ranks_str(ranks);
+    result.emplace(std::get<0>(pg_name), pg_info);
+  }
+  return result;
+}
+
+template <typename EventType>
+const c10::Dict<c10::IValue, c10::IValue> FlightRecorder<
+    EventType>::getPgStatus() {
+  auto all_pg_status = new_dict();
+  for (const auto& [pg_id, status] : all_pg_status_) {
+    auto pg_status = new_dict();
+    pg_status.insert("last_enqueued_collective", status->lastEnqueuedSeq);
+    pg_status.insert("last_started_collective", status->lastStartedSeq);
+    pg_status.insert("last_completed_collective", status->lastCompletedSeq);
+    all_pg_status.insert(std::to_string(pg_id), pg_status);
+  }
+  return all_pg_status;
+}
+
+template <typename EventType>
+const std::map<std::string, std::map<std::string, std::string>> FlightRecorder<
+    EventType>::getPgStatusJson() {
+  std::map<std::string, std::map<std::string, std::string>> result;
+  for (const auto& [pg_id, status] : all_pg_status_) {
+    auto pg_status = std::map<std::string, std::string>();
+    pg_status["last_enqueued_collective"] =
+        std::to_string(status->lastEnqueuedSeq);
+    pg_status["last_started_collective"] =
+        std::to_string(status->lastStartedSeq);
+    pg_status["last_completed_collective"] =
+        std::to_string(status->lastCompletedSeq);
+    result[std::to_string(pg_id)] = pg_status;
+  }
+  return result;
+}
+
+using json = nlohmann::json;
+template <typename EventType>
+std::string FlightRecorder<EventType>::dump_json(
+    const std::optional<std::unordered_map<
+        std::string,
+        std::unordered_map<std::string, std::string>>>& extraDumpMap,
+    bool includeCollectives,
+    bool onlyActive) {
+  json result;
+  result[version_key_str] = version_val_str;
+  result[comm_lib_version_key_str] = comm_lib_version_;
+  result[pg_config_key_str] = getPgConfigJson();
+  result[pg_status_key_str] = getPgStatusJson();
+
+  // collective trace
+  if (includeCollectives) {
+    std::list<json> entries;
+    for (auto& e : dump_entries()) {
+      json j;
+      if (onlyActive && e.time_discovered_completed_.has_value()) {
+        continue;
+      }
+      j[record_id_key_str] = int64_t(e.id_);
+      j[pg_id_key_str] = int64_t(e.pg_id_);
+      j[pg_name_key_str] = e.pg_name_;
+      j[thread_name_key_str] = e.thread_name_;
+      j[thread_id_key_str] = c10::str(e.thread_id_);
+      j[collective_seq_id_key_str] = int64_t(e.collective_seq_id_);
+      j[p2p_seq_id_key_str] = int64_t(e.p2p_seq_id_);
+      j[op_id_key_str] = int64_t(e.op_id_);
+      j[profiling_name_key_str] = e.profiling_name_;
+      j[time_created_key_str] = int64_t(e.time_created_);
+      if (e.duration_) {
+        j[duration_key_str] = *e.duration_;
+      }
+      auto it = e.sizes_.begin();
+      auto read_sizes = [&](const c10::SmallVector<int64_t, 4>& dims) {
+        auto sizes = std::list<std::list<int64_t>>();
+        for (auto dim : dims) {
+          auto arg_sizes = std::list<int64_t>();
+          for (auto i : c10::irange(dim)) {
+            (void)i;
+            arg_sizes.push_back(*it++);
+          }
+          sizes.push_back(arg_sizes);
+        }
+        return sizes;
+      };
+      j[input_sizes_key_str] = read_sizes(e.input_dims_);
+      std::vector<std::string> input_dtypes_strs;
+      input_dtypes_strs.reserve(e.input_dtypes_.size());
+      for (const auto& input_dtype : e.input_dtypes_) {
+        input_dtypes_strs.emplace_back(c10::toString(input_dtype));
+      }
+      j[input_dtypes_key_str] = input_dtypes_strs;
+      j[output_sizes_key_str] = read_sizes(e.output_dims_);
+      std::vector<std::string> output_dtypes_strs;
+      output_dtypes_strs.reserve(e.output_dtypes_.size());
+      for (const auto& output_dtype : e.output_dtypes_) {
+        output_dtypes_strs.emplace_back(c10::toString(output_dtype));
+      }
+      j[output_dtypes_key_str] = output_dtypes_strs;
+      if (e.time_discovered_completed_.has_value()) {
+        j[state_key_str] = completed_state_str;
+      } else if (e.time_discovered_started_.has_value()) {
+        j[state_key_str] = started_state_str;
+      } else {
+        j[state_key_str] = scheduled_state_str;
+      }
+      j[time_discovered_started_key_str] =
+          e.time_discovered_started_.has_value()
+          ? int64_t(*e.time_discovered_started_)
+          : 0;
+      j[time_discovered_completed_key_str] =
+          e.time_discovered_completed_.has_value()
+          ? int64_t(*e.time_discovered_completed_)
+          : 0;
+      j[retired_key_str] = e.retired_;
+      j[timeout_key_str] = e.timeout_ms_;
+      j[is_p2p_key_str] = e.isP2P_;
+      entries.emplace_back(j);
+    }
+
+    if (!entries.empty()) {
+      result[entries_key_str] = entries;
+    }
+  }
+
+  if (extraDumpMap.has_value()) {
+    result[nccl_comm_key_str] = extraDumpMap.value();
+  }
+  return result.dump();
+}
+
+template <typename EventType>
+std::string FlightRecorder<EventType>::dump(
+    const std::optional<std::unordered_map<
+        std::string,
+        std::unordered_map<std::string, std::string>>>& extraDumpMap,
+    bool includeCollectives,
+    bool includeStackTraces,
+    bool onlyActive) {
+  STATIC_SCOPED_WAIT_COUNTER(pytorch.wait_counter.FlightRecorder__dump);
+  auto result = new_dict();
+  // common values
+  result.insert(version_key, version_val);
+  result.insert(pg_config_key, getPgConfig());
+  result.insert(comm_lib_version_key_str, comm_lib_version_);
+  result.insert(pg_status_key, getPgStatus());
+
+  // collective trace
+  if (includeCollectives) {
+    result.insert(
+        entries_key, getCollectiveTrace(includeStackTraces, onlyActive));
+  }
+
+  // convert extraDumpMap into a dictionary
+  auto per_comm_dict = new_dict();
+  if (extraDumpMap.has_value()) {
+    for (const auto& [ncclId, ncclDump] : extraDumpMap.value()) {
+      auto inner_dict = new_dict();
+      for (const auto& [key, value] : ncclDump) {
+        inner_dict.insert(key, value);
+      }
+      per_comm_dict.insert(ncclId, inner_dict);
+    }
+  }
+  if (!per_comm_dict.empty()) {
+    result.insert(nccl_comm_key, per_comm_dict);
+  }
+  return pickle_str(result);
+}
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Functional.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Functional.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..1970c4abb9fa4ef0192a7814e2eebc8637c11410
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Functional.hpp
@@ -0,0 +1,182 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+
+namespace c10d {
+
+C10_EXPORT at::Tensor& all_reduce_(
+    at::Tensor& input,
+    std::string reduce_op,
+    std::string group_name);
+
+C10_EXPORT at::Tensor& all_reduce_(
+    at::Tensor& input,
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor all_reduce(
+    const at::Tensor& input,
+    std::string reduce_op,
+    std::string group_name);
+
+C10_EXPORT at::Tensor all_reduce(
+    const at::Tensor& input,
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced_(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced_(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string group_name);
+
+C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced(
+    std::vector<at::Tensor> inputs,
+    std::string reduce_op,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT std::vector<at::Tensor> all_reduce_coalesced(
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::vector<at::Tensor> inputs,
+    std::string reduce_op,
+    std::string group_name);
+
+C10_EXPORT std::vector<at::Tensor> all_gather_into_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    int64_t group_size,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string group_name);
+
+C10_EXPORT std::vector<at::Tensor> all_gather_into_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor all_gather_into_tensor(
+    const at::Tensor& input,
+    int64_t group_size,
+    std::string group_name);
+
+C10_EXPORT at::Tensor all_gather_into_tensor(
+    const at::Tensor& input,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor& all_gather_into_tensor_out(
+    at::Tensor& input,
+    int64_t group_size,
+    const std::string& group_name,
+    at::Tensor& output);
+
+C10_EXPORT at::Tensor& all_gather_into_tensor_out(
+    at::Tensor& input,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group,
+    at::Tensor& output);
+
+C10_EXPORT std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT std::vector<at::Tensor> reduce_scatter_tensor_coalesced(
+    std::vector<at::Tensor> inputs,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string reduce_op,
+    int64_t group_size,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string group_name);
+
+C10_EXPORT at::Tensor reduce_scatter_tensor(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor reduce_scatter_tensor(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    std::string group_name);
+
+C10_EXPORT at::Tensor reduce_scatter_tensor_out(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    c10::intrusive_ptr<ProcessGroup> group,
+    at::Tensor& output);
+
+C10_EXPORT at::Tensor reduce_scatter_tensor_out(
+    const at::Tensor& input,
+    std::string reduce_op,
+    int64_t group_size,
+    std::string group_name,
+    at::Tensor& output);
+
+C10_EXPORT at::Tensor all_to_all_single(
+    const at::Tensor& input,
+    at::SymIntArrayRef output_split_sizes,
+    at::SymIntArrayRef input_split_sizes,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    std::string group_name);
+
+C10_EXPORT at::Tensor all_to_all_single(
+    const at::Tensor& input,
+    at::SymIntArrayRef output_split_sizes,
+    at::SymIntArrayRef input_split_sizes,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor& broadcast_(
+    at::Tensor& input,
+    int64_t src,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor& broadcast_(
+    at::Tensor& input,
+    int64_t src,
+    std::string group_name);
+
+C10_EXPORT at::Tensor broadcast(
+    const at::Tensor& input,
+    int64_t src,
+    c10::intrusive_ptr<ProcessGroup> group);
+
+C10_EXPORT at::Tensor broadcast(
+    const at::Tensor& input,
+    int64_t src,
+    std::string group_name);
+
+C10_EXPORT at::Tensor isend(
+    at::Tensor& send_buf,
+    int64_t dst,
+    int64_t tag,
+    std::string group_name);
+
+C10_EXPORT at::Tensor irecv(
+    at::Tensor& recv_buf,
+    int64_t src,
+    int64_t tag,
+    std::string group_name);
+
+C10_EXPORT std::vector<at::Tensor> batch_p2p_ops(
+    std::vector<std::string> op_list,
+    std::vector<int64_t> peer_list,
+    std::vector<int64_t> tag_list,
+    std::vector<at::Tensor> tensors_for_op,
+    std::string group_name);
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/GlooDeviceFactory.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/GlooDeviceFactory.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..9770bbe9f3822bc68fa8c3344e92c17b95004dff
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/GlooDeviceFactory.hpp
@@ -0,0 +1,40 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_GLOO
+
+#include <string>
+
+#include <c10/util/Registry.h>
+#include <gloo/config.h>
+#include <gloo/transport/device.h>
+
+namespace c10d {
+
+class TORCH_API GlooDeviceFactory {
+ public:
+  // Create new device instance for specific interface.
+  static std::shared_ptr<::gloo::transport::Device> makeDeviceForInterface(
+      const std::string& interface,
+      bool lazyInit);
+
+  // Create new device instance for specific hostname or address.
+  static std::shared_ptr<::gloo::transport::Device> makeDeviceForHostname(
+      const std::string& hostname,
+      bool lazyInit);
+};
+
+TORCH_DECLARE_SHARED_REGISTRY(
+    GlooDeviceRegistry,
+    ::gloo::transport::Device,
+    const std::string&, /* interface */
+    const std::string&, /* hostname */
+    bool /* lazyInit */);
+
+} // namespace c10d
+
+#endif // USE_C10D_GLOO
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/GroupRegistry.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/GroupRegistry.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..4df50c8b39e5fbfc27e74477651d865a5d3dc06d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/GroupRegistry.hpp
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+
+namespace c10d {
+
+C10_EXPORT void set_thread_isolation_mode(bool enable);
+
+bool get_thread_isolation_mode();
+
+C10_EXPORT void register_process_group(
+    const std::string& group_name,
+    const c10::intrusive_ptr<c10d::ProcessGroup>& group);
+
+C10_EXPORT c10::intrusive_ptr<c10d::ProcessGroup> resolve_process_group(
+    const std::string& group_name);
+
+C10_EXPORT void unregister_process_group(const std::string& group_name);
+
+C10_EXPORT void unregister_all_process_groups();
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/HashStore.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/HashStore.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..e9b6186ef3bbf413ade67ddc83b764c43e655fa6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/HashStore.hpp
@@ -0,0 +1,86 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <condition_variable>
+#include <mutex>
+#include <unordered_map>
+
+#include <torch/csrc/distributed/c10d/Store.hpp>
+
+namespace c10d {
+
+class TORCH_API HashStore : public Store {
+ public:
+  c10::intrusive_ptr<Store> clone() override;
+
+  ~HashStore() override = default;
+
+  void set(const std::string& key, const std::vector<uint8_t>& data) override;
+
+  std::vector<uint8_t> compareSet(
+      const std::string& key,
+      const std::vector<uint8_t>& expectedValue,
+      const std::vector<uint8_t>& desiredValue) override;
+
+  std::vector<uint8_t> get(const std::string& key) override;
+
+  void wait(const std::vector<std::string>& keys) override {
+    wait(keys, timeout_);
+  }
+
+  void wait(
+      const std::vector<std::string>& keys,
+      const std::chrono::milliseconds& timeout) override;
+
+  int64_t add(const std::string& key, int64_t value) override;
+
+  int64_t getNumKeys() override;
+
+  bool check(const std::vector<std::string>& keys) override;
+
+  bool deleteKey(const std::string& key) override;
+
+  void append(const std::string& key, const std::vector<uint8_t>& value)
+      override;
+
+  std::vector<std::vector<uint8_t>> multiGet(
+      const std::vector<std::string>& keys) override;
+
+  void multiSet(
+      const std::vector<std::string>& keys,
+      const std::vector<std::vector<uint8_t>>& values) override;
+
+  // Returns true if this store support append, multiGet and multiSet
+  bool hasExtendedApi() const override;
+
+  void queuePush(const std::string& key, const std::vector<uint8_t>& value)
+      override;
+
+  std::vector<uint8_t> queuePop(const std::string& key, bool block) override;
+
+  int64_t queueLen(const std::string& key) override;
+
+  std::vector<std::string> listKeys() override;
+
+ protected:
+  bool checkLocked(
+      const std::unique_lock<std::mutex>& lock,
+      const std::vector<std::string>& keys);
+
+  void waitLocked(
+      std::unique_lock<std::mutex>& lock,
+      const std::vector<std::string>& keys,
+      const std::chrono::milliseconds& timeout);
+
+ protected:
+  std::unordered_map<std::string, std::vector<uint8_t>> map_;
+  std::unordered_map<std::string, std::deque<std::vector<uint8_t>>> queues_;
+  std::mutex m_;
+  std::condition_variable cv_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/NCCLUtils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/NCCLUtils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..c8dd7f59bb97fbd0b1a9ccbce42bf0cf50ade5d1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/NCCLUtils.hpp
@@ -0,0 +1,458 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_NCCL
+
+#include <sched.h>
+#include <cstdio>
+#include <cstdlib>
+
+#include <memory>
+#include <mutex>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAEvent.h>
+#include <c10/util/Exception.h>
+#include <nccl.h>
+#include <torch/csrc/cuda/nccl.h>
+#include <optional>
+
+constexpr int64_t kCommInitBusyWaitMillis = 2;
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 14, 0)
+#define NCCL_HAS_COMM_NONBLOCKING
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 18, 0)
+#define NCCL_HAS_COMM_SPLIT
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 23, 0)
+#define NCCL_HAS_INIT_RANK_SCALABLE
+#endif
+
+// ncclGetLastError() is enabled only for NCCL versions 2.13+
+// ncclRemoteError only exists in NCCL versions 2.13+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 13, 0)
+#define ENABLE_NCCL_GET_LAST_ERROR
+#define NCCL_REMOTE_ERROR
+#endif
+
+static_assert(
+    NCCL_VERSION_CODE >= NCCL_VERSION(2, 7, 0),
+    "NCCL version must be 2.7 or later");
+// The following macros represent features supported prior to NCCL 2.7,
+// therefore we can define them unconditionally, given the static_assert above.
+// TODO: remove these macros from code.
+#define ENABLE_NCCL_ERROR_CHECKING
+#define ENABLE_NCCL_P2P_SUPPORT
+// End of macros for NCCL 2.7 and below.
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 11, 0)
+#define ENABLE_NCCL_PREMUL_SUM_SUPPORT
+#endif
+
+// Note: the first version that supports ncclConfig_t is 2.14. Here we
+// fast-forward the version requirement to 2.17 where ncclConfig_t has CTA and
+// CGA fields because they have already been pybinded out.
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 17, 0)
+#define NCCL_HAS_CONFIG
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0)
+#define NCCL_HAS_COMM_REGISTER
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 27, 0)
+#define NCCL_HAS_COMM_WINDOW_REGISTER
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0)
+#define NCCL_HAS_MEM_ALLOC
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 26, 0)
+#define NCCL_HAS_QOS
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 24, 0)
+#define NCCL_SUPPORTS_FP8
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 27, 0)
+#define NCCL_HAS_COLLNET
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 27, 0)
+#define NCCL_HAS_CTA_POLICY
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 27, 0)
+#define NCCL_HAS_NVLS_CTAS
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 27, 0)
+#define NCCL_HAS_COMM_SHRINK
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 7)
+#define NCCL_HAS_COMM_OFFLOAD
+#endif
+
+// Macro to throw on a non-successful NCCL return value.
+#define C10D_NCCL_CHECK(cmd, failureReason)                                   \
+  do {                                                                        \
+    ncclResult_t result = cmd;                                                \
+    if (result != ncclSuccess) {                                              \
+      std::string err = "NCCL error in: " + std::string(__FILE__) + ":" +     \
+          std::to_string(__LINE__) + ", " + ncclGetErrorWithVersion(result) + \
+          "\n" + getNcclErrorDetailStr(result, failureReason);                \
+      TORCH_CHECK_WITH(DistBackendError, false, err);                         \
+    }                                                                         \
+  } while (0)
+
+// Macro to throw on a non-successful NCCL return value for NONBLOCKING calls.
+#define C10D_NCCL_CHECK_NONBLOCKING(cmd, failureReason)                       \
+  do {                                                                        \
+    ncclResult_t result = cmd;                                                \
+    if (result != ncclSuccess && result != ncclInProgress) {                  \
+      std::string err = "NCCL error in: " + std::string(__FILE__) + ":" +     \
+          std::to_string(__LINE__) + ", " + ncclGetErrorWithVersion(result) + \
+          "\n" + getNcclErrorDetailStr(result, failureReason);                \
+      TORCH_CHECK_WITH(DistBackendError, false, err);                         \
+    }                                                                         \
+  } while (0)
+
+// Error out if (current time - startTime) is greater than timeout (sec).
+#define C10D_CHECK_TIMEOUT(startTime, timeout)                              \
+  do {                                                                      \
+    auto currentTime = std::chrono::steady_clock::now();                    \
+    auto timeElapsed = std::chrono::duration_cast<std::chrono::seconds>(    \
+                           currentTime - startTime)                         \
+                           .count();                                        \
+    if (timeElapsed > timeout) {                                            \
+      std::string err = "NCCL timeout in: " + std::string(__FILE__) + ":" + \
+          std::to_string(__LINE__);                                         \
+      TORCH_CHECK_WITH(DistBackendError, false, err);                       \
+    }                                                                       \
+  } while (0)
+
+// Macro to throw on a non-successful NCCL return value, non-blocking.
+// Thread-safe: uses NCCLComm wrapper's getAsyncError() which acquires mutex
+// before calling ncclCommGetAsyncError to prevent race conditions between
+// watchdog and main threads.
+#define C10D_NCCL_CHECK_TIMEOUT_BASE(                                         \
+    cmd, commWrapper, failureReason, yield_fn)                                \
+  do {                                                                        \
+    ncclResult_t result = cmd;                                                \
+    auto startTimepoint = std::chrono::steady_clock::now();                   \
+    auto timeout = nccl_nonblocking_timeout();                                \
+    while (result == ncclInProgress) {                                        \
+      C10D_CHECK_TIMEOUT(startTimepoint, timeout);                            \
+      yield_fn;                                                               \
+      commWrapper->getAsyncError(&result);                                    \
+    }                                                                         \
+    if (result != ncclSuccess) {                                              \
+      std::string err = "NCCL error in: " + std::string(__FILE__) + ":" +     \
+          std::to_string(__LINE__) + ", " + ncclGetErrorWithVersion(result) + \
+          "\n" + getNcclErrorDetailStr(result, failureReason);                \
+      TORCH_CHECK_WITH(DistBackendError, false, err);                         \
+    }                                                                         \
+  } while (0)
+
+// Sleep for kCommInitBusyWaitMillis milliseconds.
+#define C10D_SCHED_SLEEP()     \
+  std::this_thread::sleep_for( \
+      std::chrono::milliseconds(kCommInitBusyWaitMillis))
+
+// Macro to throw exception on a non-successful NCCL return value or timeout.
+// This macro uses sched_yield() to yield the CPU.
+// Thus suitable for NCCL calls that would quickly turn ncclSuccess, e.g.
+// collectives.
+#define C10D_NCCL_CHECK_TIMEOUT(cmd, commWrapper, failureReason) \
+  C10D_NCCL_CHECK_TIMEOUT_BASE(cmd, commWrapper, failureReason, sched_yield())
+
+// Macro to throw exception on a non-successful NCCL return value or timeout.
+// This macro uses sleep to yield the CPU.
+// Thus suitable for NCCL calls that would take longer to turn ncclSuccess, e.g.
+// ncclCommInitRankConfig, ncclCommFinalize, etc.
+#define C10D_NCCL_CHECK_TIMEOUT_SLEEP(cmd, commWrapper, failureReason) \
+  C10D_NCCL_CHECK_TIMEOUT_BASE(                                        \
+      cmd, commWrapper, failureReason, C10D_SCHED_SLEEP())
+
+#define C10D_NCCL_CHECK_TIMEOUT_GROUPEND(cmd, comm, failureReason)           \
+  do {                                                                       \
+    ncclResult_t state = cmd;                                                \
+    auto startTimepoint = std::chrono::steady_clock::now();                  \
+    auto timeout = nccl_nonblocking_timeout();                               \
+    if (state == ncclInProgress) {                                           \
+      do {                                                                   \
+        C10D_CHECK_TIMEOUT(startTimepoint, timeout);                         \
+        sched_yield();                                                       \
+        comm->getAsyncError(&state);                                         \
+      } while (state == ncclInProgress);                                     \
+    }                                                                        \
+    if (state != ncclSuccess) {                                              \
+      std::string err = "NCCL error in: " + std::string(__FILE__) + ":" +    \
+          std::to_string(__LINE__) + ", " + ncclGetErrorWithVersion(state) + \
+          "\n" + getNcclErrorDetailStr(state, failureReason);                \
+      TORCH_CHECK_WITH(DistBackendError, false, err);                        \
+    }                                                                        \
+  } while (0)
+
+// Macro to print and abort on a non-successful NCCL return value.
+#define C10D_NCCL_ASSERT(cmd)                            \
+  do {                                                   \
+    ncclResult_t result = cmd;                           \
+    if (result != ncclSuccess) {                         \
+      std::string err = ncclGetErrorWithVersion(result); \
+      fprintf(                                           \
+          stderr,                                        \
+          "NCCL error in: %s:%d, %s\n",                  \
+          __FILE__,                                      \
+          __LINE__,                                      \
+          err.c_str());                                  \
+      abort();                                           \
+    }                                                    \
+  } while (0)
+
+namespace c10d {
+
+// NCCL type typing
+static std::map<at::ScalarType, ncclDataType_t> ncclDataType = {
+    {at::kChar, ncclInt8},
+    {at::kByte, ncclUint8},
+    {at::kFloat, ncclFloat},
+    {at::kDouble, ncclDouble},
+    {at::kInt, ncclInt32},
+    {at::kLong, ncclInt64},
+    {at::kHalf, ncclHalf},
+    {at::kBool, ncclUint8},
+#ifdef NCCL_SUPPORTS_FP8
+    {at::kFloat8_e5m2, ncclFloat8e5m2},
+    {at::kFloat8_e4m3fn, ncclFloat8e4m3},
+#else
+    {at::kFloat8_e5m2, ncclUint8},
+    {at::kFloat8_e4m3fn, ncclUint8},
+#endif
+    // NVIDIA GPUs does not support the UZ version standing for "no negative
+    // zero".  See https://onnx.ai/onnx/technical/float8.html
+    {at::kFloat8_e4m3fnuz, ncclUint8},
+    {at::kFloat8_e5m2fnuz, ncclUint8},
+#if HAS_NCCL_BF16_DATATYPE
+    {at::kBFloat16, ncclBfloat16},
+#endif // HAS_NCCL_BF16_DATATYPE
+};
+
+TORCH_API size_t hashTensors(const std::vector<at::Tensor>& tensors);
+TORCH_API std::string getNcclVersion();
+TORCH_API std::tuple<int, int, int> getNcclVersionTuple();
+TORCH_API int getNcclVersionNumber();
+TORCH_API std::string ncclGetErrorWithVersion(ncclResult_t error);
+int nccl_nonblocking_timeout();
+
+// Provides additional detail into NCCL error codes based on when these are
+// thrown in the NCCL codebase.
+TORCH_API std::string getNcclErrorDetailStr(
+    ncclResult_t error,
+    std::optional<std::string> processGroupFailureReason = std::nullopt);
+
+// Helper function that gets the data type and issues error if not supported
+ncclDataType_t getNcclDataType(at::ScalarType type);
+
+// RAII wrapper for NCCL communicator
+class NCCLComm {
+  using MutexType = std::recursive_mutex;
+  using LockType = std::unique_lock<MutexType>;
+
+ public:
+  explicit NCCLComm(ncclComm_t ncclComm);
+
+  NCCLComm() = default;
+
+  ~NCCLComm() noexcept;
+
+  void setUniqueHash(ncclUniqueId ncclId);
+  void setUniqueHash(std::string hash);
+  std::string getUniqueHash();
+
+  static std::shared_ptr<NCCLComm> create(
+      int numRanks,
+      int rank,
+      ncclUniqueId commId,
+      at::DeviceIndex deviceIndex);
+
+#ifdef NCCL_HAS_CONFIG
+  static std::shared_ptr<NCCLComm> create(
+      int numRanks,
+      int rank,
+      ncclUniqueId commId,
+      at::DeviceIndex deviceIndex,
+      ncclConfig_t& config);
+#ifdef NCCL_HAS_INIT_RANK_SCALABLE
+  static std::shared_ptr<NCCLComm> create_scalable(
+      int numRanks,
+      int rank,
+      std::vector<ncclUniqueId>& commIds,
+      at::DeviceIndex deviceIndex,
+      ncclConfig_t& config);
+#endif // NCCL_HAS_INIT_RANK_SCALABLE
+#endif // NCCL_HAS_CONFIG
+
+#ifdef NCCL_HAS_COMM_SPLIT
+  static std::shared_ptr<NCCLComm> split(
+      NCCLComm* source,
+      int color_id,
+      int rank,
+      ncclConfig_t& config);
+#endif // NCCL_HAS_COMM_SPLIT
+
+#ifdef NCCL_HAS_COMM_SHRINK
+  static std::shared_ptr<NCCLComm> shrink(
+      NCCLComm* source,
+      std::vector<int>& ranks_to_exclude,
+      ncclConfig_t* config,
+      int shrinkFlags = 0);
+#endif // NCCL_HAS_COMM_SHRINK
+
+#if (defined(IS_NCCLX) || defined(USE_ROCM)) && defined(NCCL_COMM_DUMP)
+  std::unordered_map<std::string, std::string> ncclCommDump();
+#endif
+
+  at::DeviceIndex getDeviceIndex();
+
+  // Must not be copyable
+  NCCLComm(const NCCLComm&) = delete;
+  NCCLComm& operator=(const NCCLComm&) = delete;
+
+  // Do not support move assignment as there is no valid use case
+  NCCLComm& operator=(NCCLComm&& other) = delete;
+
+  // Move constructable
+  // NOLINTNEXTLINE(*-noexcept-move-*)
+  NCCLComm(NCCLComm&& other);
+
+  ncclComm_t getNcclComm();
+
+  // Wait for the communicator to be ready. This is a blocking function.
+  // Useful in nonblocking mode: NCCL requires the communicator to be ready
+  // before issuing a second command.
+  // Arguments:
+  //   longInterval: if true, wait with sleep of an interval; otherwise, wait
+  //   with `sched_yield` which is faster (but acquires CPU more frequently).
+  //   Use `longInterval=true` when waiting for initialization or finalize to
+  //   complete. Use `longInterval=false` when waiting collective call to return
+  //   ncclSuccess.
+  void waitReady(bool longInterval);
+
+  std::optional<std::string> getNcclCommFailureReason() const;
+
+  void abort(std::optional<std::string> commFailureReason = std::nullopt);
+
+  // Finalize a communicator -- asking it to flush its operations. When the
+  // communicator is marked as nonblocking, this is a nonblocking function;
+  // otherwise, it will block till all operations complete.
+  void finalize();
+
+  // Destroy a communicator. This is a blocking function.
+  void destroy();
+
+  bool isInitialized() const;
+
+  bool isAborted() const;
+
+  uint64_t getCommSplitCounter() const;
+
+  ncclResult_t checkForNcclError();
+
+  // Thread-safe wrapper for ncclCommGetAsyncError that acquires the mutex
+  // before calling the NCCL API. This is needed because NCCL does not provide
+  // thread-safety guarantees for ncclCommGetAsyncError, and both the main
+  // thread and watchdog thread may call it concurrently.
+  ncclResult_t getAsyncError(ncclResult_t* asyncError);
+
+  ncclResult_t registerSegment(
+      void* ptr,
+      size_t size,
+      bool errorOnRereg = true,
+      bool window = false);
+
+  ncclResult_t deregisterSegment(void* ptr, bool window = false);
+
+  std::string repr() const;
+
+  // APIs related to memory offload (require NCCL 2.29.7+ at runtime)
+  void suspend();
+
+  void resume();
+
+  std::unordered_map<std::string, uint64_t> getMemoryStats();
+
+  friend class ProcessGroupNCCL;
+
+ protected:
+  // Unique hash for this communicator.
+  std::string uniqueHash_;
+  bool aborted_{false};
+  uint64_t ncclCommSplitCounter_{0};
+  ncclResult_t ncclAsyncErr_{ncclSuccess};
+  mutable MutexType mutex_;
+  // Rank that this communicator corresponds to.
+  int rank_{};
+  // Optional reason for communicator failure, provided by ProcessGroupNCCL for
+  // better error messaging.
+  std::optional<std::string> commFailureReason_;
+  bool initialized_{false};
+  // Whether this communicator is using nonblocking mode. Recorded during comm
+  // creation or split. For safety, we give a default value of true (more
+  // protection).
+  bool nonBlocking_{true};
+  // Device index for which the NCCL comm is created
+  at::DeviceIndex deviceIndex_{-1};
+#ifdef NCCL_HAS_COMM_REGISTER
+  // Stores handlers for tensors registered by NCCL
+  std::unordered_map<void*, void*> registeredSegmentHandles_;
+#endif // NCCL_HAS_COMM_REGISTER
+
+ private:
+  ncclComm_t ncclComm_{nullptr};
+};
+
+// Helper that automatically cleans up premul sums.
+struct ncclRedOpRAII {
+  ncclRedOpRAII() = default;
+  ncclRedOpRAII(ncclRedOp_t op) : op_(op) {}
+  ncclRedOpRAII(ncclRedOp_t op, ncclComm_t comm)
+      : op_(op), comm_(comm), premul_sum_(true) {}
+  ncclRedOpRAII(const ncclRedOpRAII&) = delete;
+  ncclRedOpRAII& operator=(const ncclRedOpRAII&) = delete;
+  ncclRedOpRAII(ncclRedOpRAII&& tmp) noexcept : ncclRedOpRAII() {
+    std::swap(tmp.op_, this->op_);
+    std::swap(tmp.comm_, this->comm_);
+    std::swap(tmp.premul_sum_, this->premul_sum_);
+  }
+#if defined(ENABLE_NCCL_PREMUL_SUM_SUPPORT)
+  ~ncclRedOpRAII() {
+    if (premul_sum_) {
+      ncclRedOpDestroy(op_, comm_);
+    }
+  }
+#endif // ENABLE_NCCL_PREMUL_SUM_SUPPORT
+  operator ncclRedOp_t() const {
+    return op_;
+  }
+  ncclRedOp_t op_{};
+  ncclComm_t comm_{};
+  bool premul_sum_ = false;
+};
+
+void printNcclCommProxyTrace(
+    const std::string& dumpReason,
+    const std::unordered_map<std::string, std::string>& dumpMap);
+} // namespace c10d
+
+#endif // USE_C10D_NCCL
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/NanCheck.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/NanCheck.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..a912bfeed4ee795f41b1e26cd540caea5d40e36a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/NanCheck.hpp
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <torch/csrc/Export.h>
+
+namespace c10d {
+
+// Check for NaNs in a tensor. If any are found, throw an error.
+// Dispatches to device-specific implementations via the c10d::check_for_nan op.
+TORCH_API void checkForNan(const at::Tensor& tensor);
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ParamCommsUtils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ParamCommsUtils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..9cec8cbdb9ad81a4ccf6b35e5879e13acef4722b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ParamCommsUtils.hpp
@@ -0,0 +1,260 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <ATen/record_function.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/ThreadLocalDebugInfo.h>
+#include <string>
+#include <tuple>
+#include <vector>
+
+namespace torch {
+
+class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase {
+ public:
+  ParamCommsDebugInfo() = default;
+  ParamCommsDebugInfo(
+      std::tuple<std::string, std::string> pgName,
+      int rank,
+      std::string&& collName,
+      int64_t inNelems,
+      int64_t outNelems,
+      at::ScalarType dType,
+      std::vector<int64_t> inSplitSizes,
+      std::vector<int64_t> outSplitSizes,
+      int globalRankStart,
+      int globalRankStride,
+      int worldSize,
+      bool isAsynchronizedOp = true);
+
+  ~ParamCommsDebugInfo() override = default;
+
+  const std::string getProcessGroupName() const {
+    return std::get<0>(pgName_);
+  }
+
+  const std::string getProcessGroupDesc() const {
+    return std::get<1>(pgName_);
+  }
+
+  int getRank() const {
+    return rank_;
+  }
+
+  int getWorldSize() const {
+    return worldSize_;
+  }
+
+  int getGlobalRankStart() const {
+    return globalRankStart_;
+  }
+
+  int getGlobalRankStride() const {
+    return globalRankStride_;
+  }
+
+  const std::string getCollectiveName() const {
+    return collectiveName_;
+  }
+
+  int64_t getInMessageNelems() const {
+    return inMessageNelems_;
+  }
+
+  int64_t getOutMessageNelems() const {
+    return outMessageNelems_;
+  }
+
+  at::ScalarType getDType() const {
+    return dType_;
+  }
+
+  const std::vector<int64_t>& getInputSplitSizes() const {
+    return inputSplitSizes_;
+  }
+
+  const std::vector<int64_t>& getOutputSplitSizes() const {
+    return outputSplitSizes_;
+  }
+
+  const std::vector<int64_t>& getGroupRanks() const {
+    return groupRanks_;
+  }
+
+  bool isAsynchronizedOp() const {
+    return isAsynchronizedOp_;
+  }
+
+  int64_t getSequenceNumber() const {
+    return sequenceNumber_;
+  }
+
+  bool getIsP2P() const {
+    return isP2P_;
+  }
+
+  void setSequenceInfo(int64_t seqNum, bool isP2P) {
+    sequenceNumber_ = seqNum;
+    isP2P_ = isP2P;
+  }
+
+ private:
+  std::tuple<std::string, std::string> pgName_; // <group_name, group_desc>
+  int rank_{};
+  int worldSize_{};
+  std::string collectiveName_;
+  int64_t inMessageNelems_{};
+  int64_t outMessageNelems_{};
+  at::ScalarType dType_ = at::kByte;
+  std::vector<int64_t> inputSplitSizes_;
+  std::vector<int64_t> outputSplitSizes_;
+  int globalRankStart_{};
+  int globalRankStride_{};
+  std::vector<int64_t> groupRanks_;
+  bool isAsynchronizedOp_{};
+  int64_t sequenceNumber_{-1};
+  bool isP2P_{false};
+};
+
+// Helper to set sequence info from tuple-typed seq arguments (NCCL backend).
+// No-op fallback for backends that pass non-tuple seq types (e.g., XPU/XCCL).
+template <typename A, typename B>
+inline void maybeSetSequenceInfo(
+    const std::shared_ptr<ParamCommsDebugInfo>& info,
+    const std::tuple<A, B>& seq) {
+  info->setSequenceInfo(std::get<0>(seq), std::get<1>(seq));
+}
+
+template <typename T>
+inline void maybeSetSequenceInfo(
+    const std::shared_ptr<ParamCommsDebugInfo>&,
+    const T&) {}
+
+#define RECORD_PARAM_COMMS(                                                    \
+    seq,                                                                       \
+    pgName,                                                                    \
+    rank,                                                                      \
+    collName,                                                                  \
+    inNelems,                                                                  \
+    outNelems,                                                                 \
+    dType,                                                                     \
+    inSplitSizes,                                                              \
+    outSplitSizes,                                                             \
+    globalRankStart,                                                           \
+    globalRankStride,                                                          \
+    worldSize)                                                                 \
+  auto paramCommsInfo = std::make_shared<torch::ParamCommsDebugInfo>(          \
+      pgName,                                                                  \
+      rank,                                                                    \
+      collName,                                                                \
+      inNelems,                                                                \
+      outNelems,                                                               \
+      dType,                                                                   \
+      inSplitSizes,                                                            \
+      outSplitSizes,                                                           \
+      globalRankStart,                                                         \
+      globalRankStride,                                                        \
+      worldSize,                                                               \
+      false);                                                                  \
+  torch::maybeSetSequenceInfo(paramCommsInfo, seq);                            \
+  c10::DebugInfoGuard g(c10::DebugInfoKind::PARAM_COMMS_INFO, paramCommsInfo); \
+  std::initializer_list<const c10::IValue> paramList = {                       \
+      seq,                                                                     \
+      pgName,                                                                  \
+      rank,                                                                    \
+      collName,                                                                \
+      inSplitSizes,                                                            \
+      outSplitSizes,                                                           \
+      globalRankStart,                                                         \
+      globalRankStride,                                                        \
+      worldSize,                                                               \
+      false};                                                                  \
+  c10::ArrayRef<const c10::IValue> paramInputs(paramList);                     \
+  RECORD_FUNCTION(at::kParamCommsCallName, paramInputs);
+
+#define RECORD_PARAM_COMMS_DATA(         \
+    seq,                                 \
+    pgName,                              \
+    InputTensors,                        \
+    OutputTensors,                       \
+    rank,                                \
+    collName,                            \
+    inNelems,                            \
+    outNelems,                           \
+    dType,                               \
+    inSplitSizes,                        \
+    outSplitSizes,                       \
+    globalRankStart,                     \
+    globalRankStride,                    \
+    worldSize)                           \
+  RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP( \
+      seq,                               \
+      pgName,                            \
+      InputTensors,                      \
+      OutputTensors,                     \
+      rank,                              \
+      collName,                          \
+      inNelems,                          \
+      outNelems,                         \
+      dType,                             \
+      inSplitSizes,                      \
+      outSplitSizes,                     \
+      globalRankStart,                   \
+      globalRankStride,                  \
+      worldSize,                         \
+      true);
+
+#define RECORD_PARAM_COMMS_DATA_WITH_ASYNC_OP(                                 \
+    seq,                                                                       \
+    pgName,                                                                    \
+    InputTensors,                                                              \
+    OutputTensors,                                                             \
+    rank,                                                                      \
+    collName,                                                                  \
+    inNelems,                                                                  \
+    outNelems,                                                                 \
+    dType,                                                                     \
+    inSplitSizes,                                                              \
+    outSplitSizes,                                                             \
+    globalRankStart,                                                           \
+    globalRankStride,                                                          \
+    worldSize,                                                                 \
+    isAsyncOp)                                                                 \
+  auto paramCommsInfo = std::make_shared<torch::ParamCommsDebugInfo>(          \
+      pgName,                                                                  \
+      rank,                                                                    \
+      collName,                                                                \
+      inNelems,                                                                \
+      outNelems,                                                               \
+      dType,                                                                   \
+      inSplitSizes,                                                            \
+      outSplitSizes,                                                           \
+      globalRankStart,                                                         \
+      globalRankStride,                                                        \
+      worldSize,                                                               \
+      isAsyncOp);                                                              \
+  torch::maybeSetSequenceInfo(paramCommsInfo, seq);                            \
+  c10::DebugInfoGuard g(c10::DebugInfoKind::PARAM_COMMS_INFO, paramCommsInfo); \
+  std::initializer_list<const c10::IValue> paramList = {                       \
+      c10::IValue(InputTensors),                                               \
+      seq,                                                                     \
+      pgName,                                                                  \
+      rank,                                                                    \
+      collName,                                                                \
+      inSplitSizes,                                                            \
+      outSplitSizes,                                                           \
+      globalRankStart,                                                         \
+      globalRankStride,                                                        \
+      worldSize,                                                               \
+      isAsyncOp};                                                              \
+  c10::ArrayRef<const c10::IValue> paramInputs(paramList);                     \
+  RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(                                         \
+      at::kParamCommsCallName,                                                 \
+      paramInputs,                                                             \
+      std::vector<c10::IValue>(1, c10::IValue(OutputTensors)));
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/PrefixStore.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/PrefixStore.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..48411dd8182e462c0e2592f30e0714a70769cc63
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/PrefixStore.hpp
@@ -0,0 +1,82 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Store.hpp>
+
+namespace c10d {
+
+class TORCH_API PrefixStore : public Store {
+ public:
+  explicit PrefixStore(std::string prefix, c10::intrusive_ptr<Store> store);
+
+  c10::intrusive_ptr<Store> clone() override;
+
+  using Store::set;
+  void set(const std::string& key, const std::vector<uint8_t>& value) override;
+
+  using Store::compareSet;
+  std::vector<uint8_t> compareSet(
+      const std::string& key,
+      const std::vector<uint8_t>& expectedValue,
+      const std::vector<uint8_t>& desiredValue) override;
+
+  std::vector<uint8_t> get(const std::string& key) override;
+
+  int64_t add(const std::string& key, int64_t value) override;
+
+  bool deleteKey(const std::string& key) override;
+
+  int64_t getNumKeys() override;
+
+  bool check(const std::vector<std::string>& keys) override;
+
+  void wait(const std::vector<std::string>& keys) override;
+
+  void wait(
+      const std::vector<std::string>& keys,
+      const std::chrono::milliseconds& timeout) override;
+
+  const std::chrono::milliseconds& getTimeout() const noexcept override;
+
+  void setTimeout(const std::chrono::milliseconds& timeout) override;
+
+  void append(const std::string& key, const std::vector<uint8_t>& value)
+      override;
+
+  std::vector<std::vector<uint8_t>> multiGet(
+      const std::vector<std::string>& keys) override;
+
+  void multiSet(
+      const std::vector<std::string>& keys,
+      const std::vector<std::vector<uint8_t>>& values) override;
+
+  // Returns true if this store support append, multiGet and multiSet
+  bool hasExtendedApi() const override;
+
+  void queuePush(const std::string& key, const std::vector<uint8_t>& value)
+      override;
+
+  std::vector<uint8_t> queuePop(const std::string& key, bool block) override;
+
+  int64_t queueLen(const std::string& key) override;
+
+  c10::intrusive_ptr<Store> getUnderlyingStore();
+
+  // Recursively to fetch the store before layers of wrapping with PrefixStore.
+  c10::intrusive_ptr<Store> getUnderlyingNonPrefixStore();
+
+  std::vector<std::string> listKeys() override;
+
+ protected:
+  std::string prefix_;
+  c10::intrusive_ptr<Store> store_;
+
+  std::string joinKey(const std::string& key);
+  std::vector<std::string> joinKeys(const std::vector<std::string>& keys);
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroup.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroup.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..ee183245b7d345bf1d26ce85b49c839608583d27
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroup.hpp
@@ -0,0 +1,1037 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/Work.hpp>
+#include <memory>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <c10/macros/Macros.h>
+
+// *************************************************************************
+// PROCESS GROUP collective communication API IS BEING CHANGED BETWEEN
+// versions 1.7 and 1.8.
+// PLEASE DO NOT ADD ANY DEPENDENCIES.
+// SEE RFC: https://github.com/pytorch/pytorch/issues/39662
+// *************************************************************************
+
+constexpr auto kProcessGroupDefaultTimeout =
+    std::chrono::milliseconds(30 * 60 * 1000);
+
+namespace c10d {
+
+// We only call `register_work()` in two cases:
+// 1. If the work object is created from a functional collective call.
+// 2. If the work object is created from a non-functional collective call within
+//    the `with allow_inflight_collective_as_graph_input_ctx()` context manager.
+C10_EXPORT void register_work(
+    const at::Tensor& tensor,
+    const c10::intrusive_ptr<c10d::Work>& work);
+
+C10_EXPORT at::Tensor wait_tensor(const at::Tensor& tensor);
+
+// We only call `unregister_work()` in one case:
+// 1. If the work object is created from a non-functional collective call within
+//    the `with allow_inflight_collective_as_graph_input_ctx()` context manager.
+//
+// Q: What about the functional collective case?
+// A: The unregistration of work object for functional collective is done in
+//    the required user-side explicit call to `wait_tensor()`.
+C10_EXPORT void unregister_work(const c10::intrusive_ptr<c10d::Work>& work);
+
+C10_EXPORT size_t get_work_registry_size();
+
+C10_EXPORT void set_allow_inflight_collective_as_graph_input(bool value);
+
+C10_EXPORT bool allow_inflight_collective_as_graph_input();
+
+// ProcessGroup is a base class that captures collective and point to
+// point communication in a fixed set of processes.
+//
+// The functions specified in the class below describe the API alone;
+// implementations are provided in subclasses.
+//
+// Every function that performs I/O is executed asynchronously by a
+// thread pool owned by the ProcessGroup (by default). They return an
+// object that can be used to wait for completion or error.
+//
+// The ProcessGroup can instantiate subgroups with fewer or an equal
+// number of members. Implementations must take care that multiple
+// process groups can be used in parallel and synchronize accordingly.
+//
+// The ProcessGroup assumes a fixed set of processes. If the set
+// changes, existing instances must be destructed and instantiation
+// and initialization must start from scratch. For members of the
+// process group to find each other (referred to as rendezvous from
+// hereon)
+//
+class TORCH_API ProcessGroup : public torch::CustomClassHolder {
+ public:
+  struct TORCH_API MergeOptions : torch::CustomClassHolder {
+    explicit MergeOptions(
+        const std::chrono::milliseconds timeout = kProcessGroupDefaultTimeout,
+        const std::optional<std::string> group_name = std::nullopt,
+        const std::optional<std::string> group_desc = std::nullopt)
+        : timeout(timeout), group_name(group_name), group_desc(group_desc) {}
+    ~MergeOptions() override = default;
+    MergeOptions(const MergeOptions&) = delete;
+    MergeOptions& operator=(const MergeOptions&) = delete;
+
+    std::chrono::milliseconds timeout;
+    std::optional<std::string> group_name;
+    std::optional<std::string> group_desc;
+  };
+
+  enum BackendType : uint8_t {
+    UNDEFINED = 0,
+    GLOO = 1,
+    NCCL = 2,
+    UCC = 3,
+    MPI = 4,
+    XCCL = 5,
+    CUSTOM = 6,
+  };
+
+  static std::string backendTypeToString(const BackendType& type) {
+    switch (type) {
+      case BackendType::GLOO:
+        return "gloo";
+      case BackendType::NCCL:
+        return "nccl";
+      case BackendType::XCCL:
+        return "xccl";
+      case BackendType::UCC:
+        return "ucc";
+      case BackendType::MPI:
+        return "mpi";
+      case BackendType::UNDEFINED:
+        return "undefined";
+      case BackendType::CUSTOM:
+        return "custom";
+      default:
+        TORCH_CHECK(false, "THis should never happen!");
+    }
+  }
+
+  static BackendType strToBackendType(const std::string& backend) {
+    if (backend == "undefined") {
+      return BackendType::UNDEFINED;
+    } else if (backend == "gloo") {
+      return BackendType::GLOO;
+    } else if (backend == "nccl") {
+      return BackendType::NCCL;
+    } else if (backend == "xccl") {
+      return BackendType::XCCL;
+    } else if (backend == "ucc") {
+      return BackendType::UCC;
+    } else if (backend == "mpi") {
+      return BackendType::MPI;
+    } else {
+      return BackendType::CUSTOM;
+    }
+  }
+
+  // Not used, set for backwards compatibility and only used for TypeDef in
+  // Ops.cpp
+  explicit ProcessGroup(int rank, int size);
+
+  explicit ProcessGroup(
+      c10::intrusive_ptr<::c10d::Store> store,
+      int rank,
+      int size);
+  ~ProcessGroup() override;
+
+  virtual int getRank() const {
+    return rank_;
+  }
+
+  virtual int getSize() const {
+    return size_;
+  }
+
+  // Returns an unique opaque ID of this process group object.
+  int64_t getID() const {
+    return reinterpret_cast<std::intptr_t>(this);
+  }
+
+  // Returns an unique opaque ID of a backend for the specific backend type
+  // that can correlate with this process group's collectives.
+  int64_t getBackendID(BackendType backend_type) const {
+    return reinterpret_cast<std::intptr_t>(getBackend(backend_type).get());
+  }
+
+  virtual const std::string getBackendName() const {
+    return backendTypeToString(backendType_);
+  }
+
+  BackendType getBackendType() const {
+    return backendType_;
+  }
+
+  inline bool backendSupportsSequenceNumbers(BackendType backendType) {
+    if (backendType == BackendType::GLOO || backendType == BackendType::NCCL ||
+        backendType == BackendType::XCCL || backendType == BackendType::UCC)
+      return true;
+    return false;
+  }
+
+  virtual void setTimeout(std::chrono::milliseconds timeout) {
+    for (auto& backend : backendTypeToBackend_) {
+      backend.second->setTimeout(timeout);
+    }
+  }
+
+  int64_t incrementSplitCount() {
+    return splitCounter_++;
+  }
+
+  virtual void startCoalescing(c10::DeviceType deviceType) {
+    // only nccl has implemented startCoalescing so only execute for nccl
+    // backends
+    auto backend = getBackend(deviceType);
+    backend->startCoalescing();
+  }
+
+  virtual c10::intrusive_ptr<Work> endCoalescing(c10::DeviceType deviceType) {
+    // only nccl has implemented endCoalescing so only execute for nccl
+    // backends
+    auto backend = getBackend(deviceType);
+    auto work = backend->endCoalescing();
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& tensors,
+      const BroadcastOptions& opts = BroadcastOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::broadcast_", "")
+            .typed<
+                std::tuple<std::vector<at::Tensor>, c10::intrusive_ptr<Work>>(
+                    at::TensorList,
+                    const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                    int64_t,
+                    int64_t,
+                    bool,
+                    int64_t)>();
+    // It's awakward to unbox the opts here and box them again in the custom C++
+    // op. But it's also complicated to make opts as a CustomClassHolder. Leave
+    // it as it is now.
+    auto work = std::get<1>(op.call(
+        tensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.rootRank,
+        opts.rootTensor,
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : tensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::allreduce_", "")
+            .typed<
+                std::tuple<std::vector<at::Tensor>, c10::intrusive_ptr<Work>>(
+                    at::TensorList,
+                    const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                    const c10::intrusive_ptr<::c10d::ReduceOp>&,
+                    const std::optional<at::Tensor>& sparse_indices,
+                    bool,
+                    int64_t)>();
+
+    auto work = std::get<1>(op.call(
+        tensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        c10::make_intrusive<ReduceOp>(opts.reduceOp),
+        opts.sparseIndices,
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : tensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceCoalescedOptions& opts = AllreduceCoalescedOptions()) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::allreduce_coalesced_", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             at::TensorList,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             const c10::intrusive_ptr<::c10d::ReduceOp>&,
+                             bool,
+                             int64_t)>();
+
+    auto work = op.call(
+        tensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        c10::make_intrusive<ReduceOp>(opts.reduceOp),
+        opts.asyncOp,
+        opts.timeout.count());
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : tensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& tensors,
+      const ReduceOptions& opts = ReduceOptions()) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::reduce_", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             at::TensorList,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             const c10::intrusive_ptr<::c10d::ReduceOp>&,
+                             int64_t,
+                             int64_t,
+                             bool,
+                             int64_t)>();
+    auto work = op.call(
+        tensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        c10::make_intrusive<ReduceOp>(opts.reduceOp),
+        opts.rootRank,
+        opts.rootTensor,
+        opts.asyncOp,
+        opts.timeout.count());
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : tensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::allgather_", "")
+                         .typed<std::tuple<
+                             std::vector<std::vector<at::Tensor>>,
+                             c10::intrusive_ptr<Work>>(
+                             const std::vector<std::vector<at::Tensor>>&,
+                             at::TensorList,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             bool,
+                             int64_t)>();
+
+    auto work = std::get<1>(op.call(
+        outputTensors,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor_list : outputTensors) {
+        for (const auto& tensor : tensor_list) {
+          c10d::register_work(tensor, work);
+        }
+      }
+    }
+    return work;
+  }
+
+  // Gathers a single tensor inputBuffer into a single buffer outputBuffer that
+  // is interpreted as a contiguous collection of size inputBuffer * WORLD_SIZE.
+  // For implementers of ProcessGroup API and advanced users only.
+  // Note: this function will be deprecated in near future.
+  virtual c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      const AllgatherOptions& opts = AllgatherOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::_allgather_base_", "")
+            .typed<std::tuple<at::Tensor, c10::intrusive_ptr<Work>>(
+                at::Tensor&,
+                at::Tensor&,
+                const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                bool,
+                int64_t)>();
+
+    auto work = std::get<1>(op.call(
+        outputBuffer,
+        inputBuffer,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      c10d::register_work(outputBuffer, work);
+    }
+    return work;
+  }
+
+  // This function is deprecated and will be moved out of ProcessGroup to comms:
+  // * do not add dependencies on this function,
+  // * do not implement it in your ProcessGroup, implement _allgather_base
+  //   instead.
+  virtual c10::intrusive_ptr<Work> allgather_coalesced(
+      std::vector<std::vector<at::Tensor>>& outputTensorLists,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::allgather_coalesced_", "")
+                         .typed<c10::intrusive_ptr<Work>(
+                             const std::vector<std::vector<at::Tensor>>&,
+                             const at::TensorList&,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             bool)>();
+
+    auto work = op.call(
+        outputTensorLists,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.asyncOp);
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor_list : outputTensorLists) {
+        for (const auto& tensor : tensor_list) {
+          c10d::register_work(tensor, work);
+        }
+      }
+    }
+    return work;
+  }
+
+  // This function is a coalesced version of `allgather_into_tensor` (currently
+  // still named as `_allgather_base`). Each tensor in the vector corresponds to
+  // an input/output of one `allgather_into_tensor` operation.
+  virtual c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::allgather_into_tensor_coalesced_", "")
+            .typed<c10::intrusive_ptr<Work>(
+                const at::TensorList,
+                const at::TensorList,
+                const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                bool)>();
+
+    auto work = op.call(
+        outputTensors,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.asyncOp);
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : outputTensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const GatherOptions& opts = GatherOptions()) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::gather_", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             const std::vector<std::vector<at::Tensor>>&,
+                             const at::TensorList&,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             int64_t,
+                             bool,
+                             int64_t)>();
+    auto work = op.call(
+        outputTensors,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.rootRank,
+        opts.asyncOp,
+        opts.timeout.count());
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor_list : outputTensors) {
+        for (const auto& tensor : tensor_list) {
+          c10d::register_work(tensor, work);
+        }
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ScatterOptions& opts = ScatterOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::scatter_", "")
+            .typed<
+                std::tuple<std::vector<at::Tensor>, c10::intrusive_ptr<Work>>(
+                    const at::TensorList&,
+                    const std::vector<std::vector<at::Tensor>>&,
+                    const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                    int64_t,
+                    bool,
+                    int64_t)>();
+    auto work = std::get<1>(op.call(
+        outputTensors,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.rootRank,
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : outputTensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::reduce_scatter_", "")
+            .typed<
+                std::tuple<std::vector<at::Tensor>, c10::intrusive_ptr<Work>>(
+                    const at::TensorList&,
+                    const std::vector<std::vector<at::Tensor>>&,
+                    const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                    const c10::intrusive_ptr<::c10d::ReduceOp>&,
+                    bool,
+                    int64_t)>();
+    auto work = std::get<1>(op.call(
+        outputTensors,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        c10::make_intrusive<::c10d::ReduceOp>(opts.reduceOp),
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : outputTensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::_reduce_scatter_base_", "")
+            .typed<std::tuple<at::Tensor, c10::intrusive_ptr<Work>>(
+                at::Tensor&,
+                at::Tensor&,
+                const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                const c10::intrusive_ptr<::c10d::ReduceOp>&,
+                bool,
+                int64_t)>();
+    auto work = std::get<1>(op.call(
+        outputBuffer,
+        inputBuffer,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        c10::make_intrusive<::c10d::ReduceOp>(opts.reduceOp),
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      c10d::register_work(outputBuffer, work);
+    }
+    return work;
+  }
+
+  // This function is a coalesced version of `reduce_scatter_tensor` (currently
+  // still named as `_reduce_scatter_base`). Each tensor in the vector
+  // corresponds to an input/output of one `reduce_scatter_tensor` operation.
+  virtual c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::reduce_scatter_tensor_coalesced_", "")
+            .typed<c10::intrusive_ptr<Work>(
+                const at::TensorList,
+                const at::TensorList,
+                const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                const c10::intrusive_ptr<::c10d::ReduceOp>&,
+                bool,
+                int64_t)>();
+
+    auto work = op.call(
+        outputTensors,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        c10::make_intrusive<::c10d::ReduceOp>(opts.reduceOp),
+        opts.asyncOp,
+        opts.timeout.count());
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : outputTensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      std::vector<int64_t>& outputSplitSizes,
+      std::vector<int64_t>& inputSplitSizes,
+      const AllToAllOptions& opts = AllToAllOptions()) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::alltoall_base_", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             at::Tensor&,
+                             at::Tensor&,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             std::vector<int64_t>,
+                             std::vector<int64_t>,
+                             bool,
+                             int64_t)>();
+    auto work = op.call(
+        outputBuffer,
+        inputBuffer,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        outputSplitSizes,
+        inputSplitSizes,
+        opts.asyncOp,
+        opts.timeout.count());
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      c10d::register_work(outputBuffer, work);
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllToAllOptions& opts = AllToAllOptions()) {
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow("c10d::alltoall_", "")
+            .typed<
+                std::tuple<std::vector<at::Tensor>, c10::intrusive_ptr<Work>>(
+                    const at::TensorList&,
+                    const at::TensorList&,
+                    const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                    bool,
+                    int64_t)>();
+    auto work = std::get<1>(op.call(
+        outputTensors,
+        inputTensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.asyncOp,
+        opts.timeout.count()));
+
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : outputTensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual void monitoredBarrier(
+      const BarrierOptions& opts,
+      bool wait_all_ranks = false) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::monitored_barrier_", "")
+                         .typed<void(
+                             at::Tensor,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             const std::vector<int64_t>&,
+                             int64_t,
+                             bool)>();
+    // Default to using cpu implementation, monitored barrier is only for GLOO
+    at::Tensor tensor = at::empty({0}, at::TensorOptions().device(at::kCPU));
+    op.call(
+        tensor,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.device_ids,
+        opts.timeout.count(),
+        wait_all_ranks);
+  }
+
+  // Agrees on an initial sequence number for the whole group by having rank 0
+  // create it and broadcast it to other ranks using the store. Only implemented
+  // for GLOO and NCCL backends currently.
+  virtual void setSequenceNumberForGroup() {
+    auto backendType = getBackendType();
+    // TODO: HACK for backend name to get sequence number for that backend.
+    if (backendSupportsSequenceNumbers(backendType)) {
+      getDefaultBackend()->setSequenceNumberForGroup();
+    } else {
+      TORCH_CHECK(
+          false,
+          c10::str(
+              "ProcessGroup ",
+              getBackendName(),
+              " does not yet support sequence numbers."));
+    }
+  }
+
+  // Retrieves the current sequence number for the whole group, which should be
+  // in sync. If the returned number is not consistent across the group, it
+  // may indicate that there is some sort of collective desynchronization.
+  virtual uint64_t getSequenceNumberForGroup() {
+    auto backendType = getBackendType();
+
+    // TODO: HACK for backend name to get sequence number for that backend.
+    if (backendSupportsSequenceNumbers(backendType)) {
+      return getDefaultBackend()->getSequenceNumberForGroup();
+    } else {
+      TORCH_CHECK(
+          false,
+          c10::str(
+              "ProcessGroup ",
+              getBackendName(),
+              " does not yet support sequence numbers."));
+    }
+  }
+
+  virtual c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::send", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             at::TensorList,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             int64_t,
+                             int64_t)>();
+    auto work = op.call(
+        tensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        dstRank,
+        tag);
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : tensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::recv_", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             at::TensorList,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             int64_t,
+                             int64_t)>();
+    auto work = op.call(
+        tensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        srcRank,
+        tag);
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : tensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> recvAnysource(
+      std::vector<at::Tensor>& tensors,
+      int tag) {
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::recv_any_source_", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             at::TensorList,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             int64_t)>();
+    auto work = op.call(
+        tensors,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        tag);
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      for (const auto& tensor : tensors) {
+        c10d::register_work(tensor, work);
+      }
+    }
+    return work;
+  }
+
+  virtual c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) {
+    at::Tensor tensor;
+    // TODO: if nccl was specified then use it
+    auto device = opts.device;
+    if (device.has_value()) {
+      // set device tensor from argument
+      tensor = at::empty(
+          {1}, at::TensorOptions().device(device.value()).dtype(at::kByte));
+    } else if (backendType_ == c10d::ProcessGroup::BackendType::NCCL) {
+      // set cuda tensor
+      tensor = at::empty(
+          {1},
+          at::TensorOptions().device(at::DeviceType::CUDA).dtype(at::kByte));
+    } else if (backendType_ == c10d::ProcessGroup::BackendType::XCCL) {
+      // set xpu tensor for override cpu dispatch
+      tensor = at::empty(
+          {1},
+          at::TensorOptions().device(at::DeviceType::XPU).dtype(at::kByte));
+    } else {
+      // Default to using cpu implementation
+      tensor = at::empty(
+          {1},
+          at::TensorOptions().device(at::DeviceType::CPU).dtype(at::kByte));
+    }
+
+    static auto op = c10::Dispatcher::singleton()
+                         .findSchemaOrThrow("c10d::barrier", "")
+                         .typed<c10::intrusive_ptr<::c10d::Work>(
+                             at::Tensor,
+                             const c10::intrusive_ptr<::c10d::ProcessGroup>&,
+                             const std::vector<int64_t>&,
+                             bool,
+                             int64_t)>();
+
+    auto work = op.call(
+        tensor,
+        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
+        opts.device_ids,
+        opts.asyncOp,
+        opts.timeout.count());
+    if (c10d::allow_inflight_collective_as_graph_input()) {
+      c10d::register_work(tensor, work);
+    }
+    return work;
+  }
+
+  bool hasBackends() {
+    return !deviceTypeToBackendType_.empty();
+  }
+
+  void setBackend(
+      c10::DeviceType deviceType,
+      BackendType backendType,
+      const std::optional<c10::intrusive_ptr<Backend>>& backend) {
+    // TODO: should we add these entries after the backend setting succeeds?
+    deviceTypeToBackendType_[deviceType] = backendType;
+    deviceTypes_.insert(deviceType);
+    // if the backendType is already set then reuse it for this device
+    if (backendTypeToBackend_.find(backendType) !=
+        backendTypeToBackend_.end()) {
+      auto existingBackend = backendTypeToBackend_.at(backendType);
+      deviceTypeToBackend_[deviceType] = existingBackend;
+      TORCH_CHECK(
+          existingBackend->getBoundDeviceId() ==
+          (*backend)->getBoundDeviceId());
+    } else {
+      // check if backend has value
+      if (backend.has_value()) {
+        deviceTypeToBackend_[deviceType] = backend.value();
+        backendTypeToBackend_[backendType] = backend.value();
+        (*backend)->setBoundDeviceId(bound_device_id_);
+      }
+    }
+  }
+
+  c10::intrusive_ptr<Backend> getDefaultBackend() const {
+    auto backend_iter = backendTypeToBackend_.find(backendType_);
+    TORCH_CHECK(
+        backend_iter != backendTypeToBackend_.end(),
+        "Could not find the default backend type ",
+        uint16_t(backendType_),
+        " for Process Group with name ",
+        getBackendName(),
+        ".");
+    return backend_iter->second;
+  }
+
+  void setDefaultBackend(const BackendType& backendType) {
+    backendType_ = backendType;
+  }
+
+  void setDefaultBackend(const std::string& backend) {
+    backendType_ = strToBackendType(backend);
+  }
+
+  c10::intrusive_ptr<Backend> getBackend(c10::DeviceType deviceType);
+
+  c10::intrusive_ptr<Backend> getBackend(BackendType backendType) const {
+    TORCH_CHECK(
+        backendTypeToBackend_.find(backendType) != backendTypeToBackend_.end(),
+        "Could not find backend type ",
+        uint16_t(backendType),
+        " for Process Group with name ",
+        backendTypeToString(backendType),
+        ".");
+    return backendTypeToBackend_.at(backendType);
+  }
+
+  // Return device types supported by this ProcessGroup.
+  // Note: the return type is `Device` rather than `DeviceType` for the purpose
+  // of easy comparison at Python level. The `Device` will have default index
+  // (-1).
+  std::vector<c10::Device> getDeviceTypes() const {
+    std::vector<c10::Device> devices;
+    devices.reserve(deviceTypes_.size());
+    for (auto& dt : deviceTypes_) {
+      devices.emplace_back(dt);
+    }
+    return devices;
+  }
+
+  void registerOnCompletionHook(
+      std::function<void(std::shared_ptr<WorkInfo>)>&& hook) {
+    getDefaultBackend()->registerOnCompletionHook(std::move(hook));
+  }
+
+  void waitForPendingWorks() {
+    getDefaultBackend()->waitForPendingWorks();
+  }
+
+  virtual void shutdown() {
+    for (auto& backend : backendTypeToBackend_) {
+      backend.second->shutdown();
+    }
+  }
+
+  virtual void abort() {
+    for (auto& backend : backendTypeToBackend_) {
+      backend.second->abort();
+    }
+  }
+
+  bool hasHooks() const {
+    auto backend_iter = backendTypeToBackend_.find(backendType_);
+    if (backend_iter == backendTypeToBackend_.end()) {
+      TORCH_WARN(
+          "No backend of type ",
+          uint16_t(backendType_),
+          " found for Process Group with name ",
+          getBackendName(),
+          ". Assuming no hooks are registered.");
+      return false;
+    }
+
+    return backend_iter->second->hasHooks();
+  }
+
+  virtual const std::string& getGroupName() const;
+  virtual void setGroupName(const std::string& name);
+  virtual const std::string& getGroupDesc() const;
+  virtual void setGroupDesc(const std::string& name);
+  void enableCollectivesTiming();
+
+  void release_resources() override;
+
+  // ProcessGroups optionally can be "bound" to a specific device.
+  // Currently this is only for nccl and allows for some opt-in
+  // optimizations such as automatic use of ncclCommSplit.  The device
+  // is specified in `init_process_group` and eventually makes it
+  // here and then down into the actual backend instances.
+  std::optional<at::Device> getBoundDeviceId() const {
+    return bound_device_id_;
+  }
+
+  c10::intrusive_ptr<c10d::Store> getStore() const {
+    return store_;
+  }
+
+  void setBoundDeviceId(std::optional<at::Device> device) {
+    if (device) {
+      TORCH_CHECK(device->has_index(), "setBoundDeviceId must have an index");
+    }
+    bound_device_id_ = device;
+  }
+
+  // This creates a new subgroup using the specified ranks.
+  // The current rank must be included in the list of new_ranks.
+  virtual c10::intrusive_ptr<ProcessGroup> splitGroup(
+      const std::vector<int>& ranks,
+      const std::optional<std::chrono::milliseconds>& timeout,
+      const std::optional<c10::intrusive_ptr<Backend::Options>>& opts,
+      const std::optional<std::string>& name,
+      const std::optional<std::string>& groupDesc);
+
+  // This creates a new subgroup using the specified ranks.
+  // The current rank must be included in the list of new_ranks.
+  virtual c10::intrusive_ptr<ProcessGroup> mergeRemoteGroup(
+      const c10::intrusive_ptr<Store>& store,
+      const MergeOptions& opts,
+      const int& size);
+
+ protected:
+  // Implementations of this interface need to call this to setup
+  // appropriate logging etc.
+  void init();
+
+  c10::intrusive_ptr<c10d::Store> store_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const int rank_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const int size_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  BackendType backendType_;
+  std::string pg_desc_;
+  int64_t splitCounter_;
+
+  // Debug level setting. It is parsed once when ProcessGroup is constructed and
+  // remains the same across use of this process group.
+  DebugLevel dist_debug_level_{DebugLevel::Off};
+
+  // Backend classes for this ProcessGroup
+  std::unordered_set<c10::DeviceType> deviceTypes_;
+  // This mapping is ordered, as splitGroup must call split on the underlying
+  // backends in a consistent order.
+  std::map<c10::DeviceType, BackendType> deviceTypeToBackendType_;
+  std::unordered_map<c10::DeviceType, c10::intrusive_ptr<Backend>>
+      deviceTypeToBackend_;
+  std::unordered_map<BackendType, c10::intrusive_ptr<Backend>>
+      backendTypeToBackend_;
+
+  std::optional<at::Device> bound_device_id_;
+};
+
+// Thread local functions for managing the currently active process group.
+TORCH_API c10::intrusive_ptr<ProcessGroup>& currentProcessGroup();
+TORCH_API void setProcessGroup(c10::intrusive_ptr<ProcessGroup> processGroup);
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..b90e183651105c4dd2518056dbd16c6bb4c82923
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp
@@ -0,0 +1,508 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_GLOO
+
+#include <condition_variable>
+#include <deque>
+#include <mutex>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include <gloo/algorithm.h>
+#include <gloo/common/error.h>
+#include <gloo/context.h>
+#include <gloo/rendezvous/store.h>
+#include <gloo/transport/device.h>
+
+#include <c10/util/hash.h>
+
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+#include <torch/csrc/distributed/c10d/logger.hpp>
+
+#include <ATen/ThreadLocalState.h>
+
+namespace c10d {
+
+constexpr const char* GLOO_BACKEND_NAME = "gloo";
+
+// Control whether or not connections are established in a full mesh or lazily
+// as needed.
+static std::vector<std::string> TORCH_GLOO_LAZY_INIT = {"TORCH_GLOO_LAZY_INIT"};
+
+// Returns default value for lazyInit.
+bool TORCH_API getDefaultGlooLazyInit();
+
+// ProcessGroupGloo implements Gloo bindings for c10d.
+//
+// All functions on this class are expected to be called in the same
+// order across processes in the group. This is the only way that we
+// can guarantee to match up the same calls across processes. For
+// multi-threaded usage of process groups, you can consider using
+// multiple process group instances.
+//
+class TORCH_API ProcessGroupGloo : public Backend {
+ public:
+  // AsyncWork is the Gloo specific superclass for asynchronous work items.
+  // We can split asynchronous work into 3 phases:
+  // 1) Sanity checks and prepare input (e.g. memcpy)
+  // 2) Run operation on background thread
+  // 3) Synchronize with completion on foreground thread
+  //
+  // There is state to be shared between these 3 phases and all of this state
+  // is captured in the AsyncWork class and its derivatives.
+  //
+  // Note: while we are porting operations to use new style collectives, there
+  // is a split between operations using the existing caching approach and
+  // operations using the new AsyncWork base class. Over time we will port
+  // all operations and perform needed cleanup.
+  //
+  // FIXME: This probably should be called WorkGloo since the work is executed
+  // in sync mode by a background thread.
+  class TORCH_API AsyncWork : public Work {
+   public:
+    explicit AsyncWork(
+        std::shared_ptr<gloo::Context> context,
+        std::vector<std::vector<at::Tensor>> outputTensors,
+        OpType opType,
+        uint64_t seq,
+        std::chrono::milliseconds timeout,
+        const char* profilingTitle = nullptr,
+        const std::optional<std::vector<at::Tensor>>& inputTensors =
+            std::nullopt);
+
+    ~AsyncWork() override = default;
+
+    static void execute(const c10::intrusive_ptr<AsyncWork>& work);
+
+    virtual void run() = 0;
+
+    std::vector<at::Tensor> result() override;
+
+    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
+    uint64_t getSequencenumber() const override;
+    std::chrono::milliseconds getTimeout() const;
+    virtual const std::vector<at::Tensor> getInputTensors() = 0;
+    virtual const std::vector<at::Tensor> getOutputTensors() = 0;
+    inline std::string getProfilerTitle() const {
+      return profilingTitle_;
+    }
+    inline at::ThreadLocalState getTLS() const {
+      return tls_;
+    }
+
+   protected:
+    friend class ProcessGroupGloo;
+    // unique id used to tell the trace buffer that this
+    // work has completed
+    std::optional<uint64_t> trace_id_;
+    std::optional<uint64_t> trace_reset_epoch_;
+    std::shared_ptr<gloo::Context> context_;
+    const std::chrono::milliseconds timeout_;
+
+   private:
+    void finishWorkGloo();
+    void finishWorkGlooError(const std::exception_ptr& eptr);
+    inline void recordAsyncWorkProfilingInfo(
+        const char* profilingTitle,
+        const std::optional<std::vector<at::Tensor>>& inputTensors);
+
+    const std::vector<std::vector<at::Tensor>> outputTensors_;
+    c10::intrusive_ptr<at::ivalue::Future> future_;
+    std::function<void()> recordFunctionBeforeCallback_;
+    const uint64_t seq_;
+    std::string profilingTitle_;
+    at::ThreadLocalState tls_;
+  };
+
+  // Wrap c10d store as Gloo store
+  class TORCH_API GlooStore : public ::gloo::rendezvous::Store {
+   public:
+    GlooStore(c10::intrusive_ptr<::c10d::Store> store)
+        : store_(std::move(store)) {}
+
+    void setUint(const std::string& key, const std::vector<uint8_t>& value) {
+      store_->set(key, value);
+    }
+
+    void set(const std::string& key, const std::vector<char>& value) override {
+      std::vector<uint8_t> tmp(value.begin(), value.end());
+      store_->set(key, tmp);
+    }
+
+    std::vector<uint8_t> getUint(const std::string& key) {
+      auto value = store_->get(key);
+      return value;
+    }
+
+    std::vector<char> get(const std::string& key) override {
+      auto value = store_->get(key);
+      return std::vector<char>(value.begin(), value.end());
+    }
+
+    void wait(const std::vector<std::string>& keys) override {
+      store_->wait(keys, ::c10d::Store::kDefaultTimeout);
+    }
+
+    void wait(
+        const std::vector<std::string>& keys,
+        const std::chrono::milliseconds& timeout) override {
+      store_->wait(keys, timeout);
+    }
+
+#ifdef GLOO_STORE_HAS_STORE_V2
+    bool has_v2_support() override {
+      return store_->hasExtendedApi();
+    }
+
+    std::vector<std::vector<char>> multi_get(
+        const std::vector<std::string>& keys) override {
+      std::vector<std::vector<char>> res;
+      for (auto& value : store_->multiGet(keys)) {
+        res.emplace_back(value.begin(), value.end());
+      }
+      return res;
+    }
+
+    void multi_set(
+        const std::vector<std::string>& keys,
+        const std::vector<std::vector<char>>& values) override {
+      std::vector<std::vector<uint8_t>> u_values;
+      u_values.reserve(values.size());
+      for (auto& value : values) {
+        u_values.emplace_back(value.begin(), value.end());
+      }
+      store_->multiSet(keys, u_values);
+    }
+
+    void append(const std::string& key, const std::vector<char>& value)
+        override {
+      std::vector<uint8_t> tmp(value.begin(), value.end());
+      return store_->append(key, tmp);
+    }
+
+    int64_t add(const std::string& key, int64_t value) override {
+      return store_->add(key, value);
+    }
+#endif
+
+    const c10::intrusive_ptr<::c10d::Store>& _getStore() const {
+      return store_;
+    }
+
+   protected:
+    c10::intrusive_ptr<::c10d::Store> store_;
+  };
+
+  // For send and recv operations there is no need to pass them to the
+  // thread pool as they are entirely completed by the device thread.
+  // This work object is used to synchronize completion of the send or
+  // recv operation. It keeps a reference to the tensor it is
+  // operating on to prevent it from being deallocated while the
+  // operation is still in flight.
+  class TORCH_API SendWork : public Work {
+   public:
+    explicit SendWork(
+        at::Tensor& tensor,
+        std::unique_ptr<::gloo::transport::UnboundBuffer> buffer,
+        uint64_t seq);
+
+    bool wait(std::chrono::milliseconds timeout = kNoTimeout) override;
+
+    void abort() override;
+
+    uint64_t getSequencenumber() const override;
+
+   protected:
+    at::Tensor tensor_;
+    std::unique_ptr<::gloo::transport::UnboundBuffer> buffer_;
+    const uint64_t seq_;
+  };
+
+  class TORCH_API RecvWork : public Work {
+   public:
+    explicit RecvWork(
+        at::Tensor& tensor,
+        std::unique_ptr<::gloo::transport::UnboundBuffer> buffer,
+        OpType opType,
+        uint64_t seq,
+        const char* profilingTitle = nullptr);
+
+    int sourceRank() const override;
+
+    bool wait(std::chrono::milliseconds timeout = kNoTimeout) override;
+
+    void abort() override;
+
+    uint64_t getSequencenumber() const override;
+
+   protected:
+    at::Tensor tensor_;
+    std::unique_ptr<::gloo::transport::UnboundBuffer> buffer_;
+    int srcRank_{-1};
+    const uint64_t seq_;
+  };
+
+  struct TORCH_API Options : public Backend::Options {
+    explicit Options(
+        std::chrono::milliseconds timeout = kBackendDefaultTimeout);
+
+    // return intrusive_ptr of the object
+    static c10::intrusive_ptr<Options> create(
+        std::chrono::milliseconds timeout = kBackendDefaultTimeout) {
+      return c10::make_intrusive<Options>(timeout);
+    }
+
+    static c10::intrusive_ptr<Options> create_default(
+        std::chrono::milliseconds timeout = kBackendDefaultTimeout);
+
+    std::vector<std::shared_ptr<::gloo::transport::Device>> devices;
+    int threads{2};
+  };
+
+  const std::string getBackendName() const override {
+    return std::string(GLOO_BACKEND_NAME);
+  }
+
+  bool supportsSplitting() const override {
+    return true;
+  }
+
+  // Helper functions to create a new device object.
+  // They are static functions on this class to keep them logically
+  // separate from the rest of the code base (e.g. torch/csrc/distributed).
+
+  // Create new device instance for specific interface.
+  static std::shared_ptr<::gloo::transport::Device> createDeviceForInterface(
+      const std::string& interface,
+      bool lazyInit = false);
+
+  // Create new device instance for specific hostname or address.
+  static std::shared_ptr<::gloo::transport::Device> createDeviceForHostname(
+      const std::string& hostname,
+      bool lazyInit = false);
+
+  // Create new device instance.
+  // It tries to resolve this machine's hostname and bind to that address.
+  // If that fails (i.e. the hostname doesn't resolve to an address), it
+  // falls back to binding to the loopback address.
+  static std::shared_ptr<::gloo::transport::Device> createDefaultDevice(
+      bool lazyInit = false);
+
+  explicit ProcessGroupGloo(
+      const c10::intrusive_ptr<Store>& store,
+      int rank,
+      int size,
+      c10::intrusive_ptr<Options> options = Options::create());
+
+  ~ProcessGroupGloo() override;
+
+  c10::intrusive_ptr<Options> getOptions() {
+    return options_;
+  }
+
+  void setTimeout(std::chrono::milliseconds timeout) override {
+    options_->timeout = timeout;
+    for (auto& context : contexts_) {
+      context->setTimeout(timeout);
+    }
+  }
+
+  c10::intrusive_ptr<Backend::Options> getBackendOptions() override {
+    return c10::static_intrusive_pointer_cast<Backend::Options>(options_);
+  }
+
+  c10::intrusive_ptr<Backend> split(
+      const c10::intrusive_ptr<Store>& store,
+      const std::vector<int>& ranks,
+      const c10::intrusive_ptr<Backend::Options>& opts) override;
+
+  c10::intrusive_ptr<Backend> merge(
+      const c10::intrusive_ptr<Store>& store,
+      const c10::intrusive_ptr<Backend::Options>& opts,
+      const int& rank,
+      const int& size) override;
+
+  const std::vector<uint64_t>& groupRanks() const;
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& tensors,
+      const BroadcastOptions& opts = BroadcastOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_sparse(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceCoalescedOptions& opts =
+          AllreduceCoalescedOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& tensors,
+      const ReduceOptions& opts = ReduceOptions()) override;
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& output_tensor,
+      at::Tensor& input_tensor,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_coalesced(
+      std::vector<std::vector<at::Tensor>>& output_lists,
+      std::vector<at::Tensor>& input_list,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const GatherOptions& opts = GatherOptions()) override;
+
+  c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& outputs,
+      std::vector<std::vector<at::Tensor>>& inputs,
+      const ScatterOptions& opts = ScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputs,
+      std::vector<std::vector<at::Tensor>>& inputs,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      std::vector<int64_t>& outputCounts,
+      std::vector<int64_t>& inputCounts,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recvAnysource(
+      std::vector<at::Tensor>& tensors,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) override;
+
+  void enableCollectivesTiming() override;
+
+  const std::shared_ptr<::gloo::rendezvous::Store>& _getStore() const {
+    return store_;
+  }
+
+  // Similar to barrier(), but blocks rank 0 until all other ranks have
+  // acknowledged that they are alive (through send/recv from rank 0). Rank 0
+  // is able to report all failed ranks if waitAllRanks = true, otherwise
+  // reports the first rank it detected as failed.
+  void monitoredBarrier(
+      const BarrierOptions& opts = BarrierOptions(),
+      bool waitAllRanks = false) override;
+
+  // Agrees on an initial sequence number for the whole group by having rank 0
+  // create it and broadcast it to other ranks using the store.
+  void setSequenceNumberForGroup() override;
+
+  // Retrieves the current sequence number for the whole group, which should be
+  // in sync. If the returned number is not consistent across the group, it
+  // may indicate that there is some sort of collective desynchronization.
+  uint64_t getSequenceNumberForGroup() override;
+
+  int getNumThreads() {
+    return options_->threads;
+  }
+
+ protected:
+  std::shared_ptr<::gloo::rendezvous::Store> store_;
+  const c10::intrusive_ptr<Options> options_;
+
+  // Every Gloo context represents a set of connections to its peers.
+  // In order to use more than one device (or allow for parallelism on
+  // a single device), you need multiple contexts.
+  std::vector<std::shared_ptr<::gloo::Context>> contexts_;
+  std::vector<std::thread> threads_;
+  bool stop_{false};
+
+  // Incremented for every collective we kick off.
+  // The value is used as tag for collective operations. Collectives are kicked
+  // off in identical order across processes. Therefore the tag can be used
+  // to match up operations during concurrent execution.
+  uint32_t collectiveCounter_{0};
+
+  // Returns next collective tag to use (uses collectiveCounter_).
+  uint32_t nextTag();
+
+  // Returns the context to use for the specified tag.
+  // With `nextTag` returning an increasing number, this should lead
+  // to contexts being used in a round-robin fashion.
+  std::shared_ptr<::gloo::Context> getContext(uint32_t tag);
+
+  // Entrypoint for worker threads.
+  void runLoop(int workerIndex);
+
+  // Queue work to run on worker thread.
+  void enqueue(c10::intrusive_ptr<AsyncWork> work);
+
+  // Keep both a queue of pending work, and a vector with in progress work.
+  // Both of these can only be mutated when holding the queue lock.
+  // We keep both around instead of just the queue, so we can grab a weak_ptr
+  // to all in progress and pending work when executing a barrier.
+  // When executing a barrier, we need to ensure that all prior work
+  // has completed before completing itself.
+  std::deque<c10::intrusive_ptr<AsyncWork>> workQueue_;
+  std::vector<c10::intrusive_ptr<AsyncWork>> workInProgress_;
+  std::mutex workMutex_;
+  std::condition_variable workProduceCV_;
+  std::condition_variable workConsumeCV_;
+  uint64_t seq_{0};
+  size_t local_id_;
+  std::shared_ptr<ProcessGroupStatus> pgStatus_ =
+      std::make_shared<ProcessGroupStatus>();
+};
+
+} // namespace c10d
+
+#endif // USE_C10D_GLOO
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupGlooDetail.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupGlooDetail.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..98edd8124f9e6735f5025b16e00b20481b41373e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupGlooDetail.hpp
@@ -0,0 +1,679 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_GLOO
+
+#include <c10/util/Registry.h>
+#include <torch/csrc/distributed/c10d/ProcessGroupGloo.hpp>
+
+#include <gloo/allgather.h>
+#include <gloo/allgatherv.h>
+#include <gloo/allreduce.h>
+#include <gloo/alltoall.h>
+#include <gloo/alltoallv.h>
+#include <gloo/barrier.h>
+#include <gloo/broadcast.h>
+#include <gloo/gather.h>
+#include <gloo/reduce.h>
+#include <gloo/scatter.h>
+
+#ifdef _WIN32
+#define GENERATE_ALL_TYPES(type, func, ...)      \
+  switch (type) {                                \
+    case ::at::ScalarType::Float:                \
+      func<float>(__VA_ARGS__);                  \
+      break;                                     \
+    case ::at::ScalarType::Double:               \
+      func<double>(__VA_ARGS__);                 \
+      break;                                     \
+    case ::at::ScalarType::Half:                 \
+      func<c10::Half>(__VA_ARGS__);              \
+      break;                                     \
+    case ::at::ScalarType::BFloat16:             \
+      func<c10::BFloat16>(__VA_ARGS__);          \
+      break;                                     \
+    case ::at::ScalarType::Char:                 \
+      func<int8_t>(__VA_ARGS__);                 \
+      break;                                     \
+    case ::at::ScalarType::Byte:                 \
+    case ::at::ScalarType::Bool:                 \
+      func<uint8_t>(__VA_ARGS__);                \
+      break;                                     \
+    case ::at::ScalarType::Int:                  \
+      func<int32_t>(__VA_ARGS__);                \
+      break;                                     \
+    case ::at::ScalarType::Long:                 \
+      func<int64_t>(__VA_ARGS__);                \
+      break;                                     \
+    default:                                     \
+      TORCH_CHECK(false, "Invalid scalar type"); \
+  }
+
+#define HOST_NAME_MAX 256
+#else
+#define GENERATE_ALL_TYPES(type, func, args...)  \
+  switch (type) {                                \
+    case ::at::ScalarType::Float:                \
+      func<float>(args);                         \
+      break;                                     \
+    case ::at::ScalarType::Double:               \
+      func<double>(args);                        \
+      break;                                     \
+    case ::at::ScalarType::Half:                 \
+      func<c10::Half>(args);                     \
+      break;                                     \
+    case ::at::ScalarType::BFloat16:             \
+      func<c10::BFloat16>(args);                 \
+      break;                                     \
+    case ::at::ScalarType::Char:                 \
+      func<int8_t>(args);                        \
+      break;                                     \
+    case ::at::ScalarType::Byte:                 \
+    case ::at::ScalarType::Bool:                 \
+      func<uint8_t>(args);                       \
+      break;                                     \
+    case ::at::ScalarType::Int:                  \
+      func<int32_t>(args);                       \
+      break;                                     \
+    case ::at::ScalarType::Long:                 \
+      func<int64_t>(args);                       \
+      break;                                     \
+    default:                                     \
+      TORCH_CHECK(false, "Invalid scalar type"); \
+  }
+#endif
+
+namespace c10d {
+
+TORCH_DECLARE_TYPED_REGISTRY(
+    GlooAllreduceRegistry,
+    c10::DeviceType,
+    ProcessGroupGloo::AsyncWork,
+    c10::intrusive_ptr,
+    std::shared_ptr<gloo::Context>,
+    std::vector<at::Tensor>&,
+    ReduceOp,
+    uint32_t,
+    uint64_t,
+    std::chrono::milliseconds);
+
+// This function initializes a vector of CUDA streams, one for every
+// tensor in the input tensor vector, and ensures that these streams are
+// synchronized with the current default streams. This is needed so
+// that new work on the new streams is serialized w.r.t. all operations
+// on the tensors.
+TORCH_API void initializeStreamsEvents(
+    const std::vector<at::Tensor>& tensors,
+    std::vector<c10::Stream>& streams,
+    std::vector<c10::Event>& events);
+
+// This function initializes a vector of CUDA streams, one per device,
+// and ensures that these streams are synchronized with the current default
+// streams. It is assumed that the tensors in the nested tensor vectors are
+// on the same device.
+TORCH_API void initializeStreamsEvents(
+    std::vector<std::vector<at::Tensor>>& tensors,
+    std::vector<c10::Stream>& streams,
+    std::vector<c10::Event>& events);
+
+typedef void (*ReduceFunc)(void*, const void*, const void*, size_t);
+
+template <typename T, std::enable_if_t<!std::is_integral_v<T>, int> = 0>
+ReduceFunc toFunction(const ReduceOp& r) {
+  switch (r) {
+    case ReduceOp::SUM:
+    case ReduceOp::AVG:
+      return ReduceFunc(&::gloo::sum<T>);
+    case ReduceOp::PRODUCT:
+      return ReduceFunc(&::gloo::product<T>);
+    case ReduceOp::MIN:
+      return ReduceFunc(&::gloo::min<T>);
+    case ReduceOp::MAX:
+      return ReduceFunc(&::gloo::max<T>);
+    case ReduceOp::BAND:
+      TORCH_CHECK(false, "Cannot use ReduceOp.BAND with non-integral dtype");
+      break;
+    case ReduceOp::BOR:
+      TORCH_CHECK(false, "Cannot use ReduceOp.BOR with non-integral dtype");
+      break;
+    case ReduceOp::BXOR:
+      TORCH_CHECK(false, "Cannot use ReduceOp.BXOR with non-integral dtype");
+      break;
+    case ReduceOp::PREMUL_SUM:
+      TORCH_CHECK(false, "Cannot use ReduceOp.PREMUL_SUM with Gloo");
+      break;
+    case ReduceOp::UNUSED:
+    default:
+      break;
+  }
+
+  TORCH_CHECK(false, "Unhandled ReduceOp");
+}
+
+// Bitwise AND with SFINAE guard for integral types.
+template <typename T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
+void band(void* c, const void* a, const void* b, size_t n) {
+  auto tc = static_cast<T*>(c);
+  auto ta = static_cast<const T*>(a);
+  auto tb = static_cast<const T*>(b);
+  for (const auto i : c10::irange(n)) {
+    tc[i] = ta[i] & tb[i];
+  }
+}
+
+// Bitwise OR with SFINAE guard for integral types.
+template <typename T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
+void bor(void* c, const void* a, const void* b, size_t n) {
+  auto tc = static_cast<T*>(c);
+  auto ta = static_cast<const T*>(a);
+  auto tb = static_cast<const T*>(b);
+  for (const auto i : c10::irange(n)) {
+    tc[i] = ta[i] | tb[i];
+  }
+}
+
+// Bitwise XOR with SFINAE guard for integral types.
+template <typename T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
+void bxor(void* c, const void* a, const void* b, size_t n) {
+  auto tc = static_cast<T*>(c);
+  auto ta = static_cast<const T*>(a);
+  auto tb = static_cast<const T*>(b);
+  for (const auto i : c10::irange(n)) {
+    tc[i] = ta[i] ^ tb[i];
+  }
+}
+
+template <typename T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
+ReduceFunc toFunction(const ReduceOp& r) {
+  switch (r) {
+    case ReduceOp::SUM:
+    case ReduceOp::AVG:
+      return ReduceFunc(&::gloo::sum<T>);
+    case ReduceOp::PRODUCT:
+      return ReduceFunc(&::gloo::product<T>);
+    case ReduceOp::MIN:
+      return ReduceFunc(&::gloo::min<T>);
+    case ReduceOp::MAX:
+      return ReduceFunc(&::gloo::max<T>);
+    case ReduceOp::BAND:
+      return ReduceFunc(&band<T>);
+    case ReduceOp::BOR:
+      return ReduceFunc(&bor<T>);
+    case ReduceOp::BXOR:
+      return ReduceFunc(&bxor<T>);
+    case ReduceOp::PREMUL_SUM:
+      TORCH_CHECK(false, "Cannot use ReduceOp.PREMUL_SUM with Gloo");
+      break;
+    case ReduceOp::UNUSED:
+    default:
+      break;
+  }
+
+  TORCH_CHECK(false, "Unhandled ReduceOp");
+}
+
+template <typename T, typename O>
+void setInputs(O& opts, std::vector<at::Tensor>& tensors) {
+  opts.setInputs(getDataPointers<T>(tensors), tensors[0].numel());
+}
+
+template <typename T, typename O>
+void setInput(O& opts, at::Tensor& tensor) {
+  opts.setInput(getDataPointer<T>(tensor), tensor.numel());
+}
+
+template <typename T, typename O>
+void setInput(O& opts, at::Tensor& tensor, std::vector<size_t>& counts) {
+  opts.setInput(getDataPointer<T>(tensor), counts);
+}
+
+template <typename T, typename O>
+void setInput(O& opts, at::Tensor& tensor, std::vector<int64_t>& counts) {
+  opts.setInput(getDataPointer<T>(tensor), counts);
+}
+
+template <typename T, typename O>
+void setOutputs(O& opts, std::vector<at::Tensor>& tensors, int64_t count) {
+  opts.setOutputs(getDataPointers<T>(tensors), count);
+}
+
+template <typename T, typename O>
+void setOutput(O& opts, at::Tensor& tensor) {
+  opts.setOutput(getDataPointer<T>(tensor), tensor.numel());
+}
+
+template <typename T, typename O>
+void setOutput(O& opts, at::Tensor& tensor, std::vector<size_t>& counts) {
+  opts.setOutput(getDataPointer<T>(tensor), counts);
+}
+
+template <typename T, typename O>
+void setOutput(O& opts, at::Tensor& tensor, std::vector<int64_t>& counts) {
+  opts.setOutput(getDataPointer<T>(tensor), counts);
+}
+
+static at::Tensor pinnedLike(at::Tensor& tensor) {
+  auto* allocator = at::detail::getCUDAHooks().getPinnedMemoryAllocator();
+  auto storage = c10::Storage(
+      c10::Storage::use_byte_size_t(),
+      static_cast<int64_t>(at::detail::computeStorageNbytes(
+          tensor.sizes(), tensor.strides(), tensor.dtype().itemsize())),
+      allocator,
+      /*resizable=*/false);
+  return at::empty({0}, tensor.options().device(at::kCPU))
+      .set_(storage, 0, tensor.sizes(), tensor.strides());
+}
+
+class AsyncAllreduceWork : public ProcessGroupGloo::AsyncWork {
+ public:
+  AsyncAllreduceWork(
+      std::shared_ptr<gloo::Context> context,
+      std::vector<at::Tensor>& inputs,
+      ReduceOp reduceOp,
+      uint32_t tag,
+      uint64_t seq,
+      std::chrono::milliseconds timeout)
+      : ProcessGroupGloo::AsyncWork(
+            std::move(context),
+            {inputs},
+            OpType::ALLREDUCE,
+            seq,
+            timeout,
+            "gloo:all_reduce",
+            inputs),
+        inputs(inputs),
+        reduceOp(std::move(reduceOp)),
+        tag(tag) {}
+
+  std::vector<at::Tensor> inputs;
+  const ReduceOp reduceOp;
+  const uint32_t tag;
+
+  void allreduce(std::vector<at::Tensor>& tensors) {
+    auto tensor = tensors[0];
+    if (tensor.is_complex()) {
+      TORCH_CHECK(
+          c10d::isComplexViewAsRealAllowed(reduceOp),
+          "all_reduce does not support",
+          reduceOp,
+          "on complex tensors");
+      tensor = at::view_as_real(tensor);
+    }
+    gloo::AllreduceOptions opts(context_);
+    const auto& scalarType = tensor.scalar_type();
+    opts.setReduceFunction(getFunction(scalarType, reduceOp));
+    opts.setTag(tag);
+    opts.setTimeout(getTimeout());
+    // Use tensor.numel() instead of tensors[0].numel() to
+    // get the right number of elements when tensors[0] is complex
+    GENERATE_ALL_TYPES(scalarType, setOutputs, opts, tensors, tensor.numel());
+    gloo::allreduce(opts);
+
+    // Gloo doesn't support AVG so we use SUM + division.
+    if (reduceOp == ReduceOp::AVG) {
+      tensors[0] /= context_->size;
+    }
+  }
+
+  const std::vector<at::Tensor> getInputTensors() override {
+    return inputs;
+  }
+
+  const std::vector<at::Tensor> getOutputTensors() override {
+    return inputs;
+  }
+
+  void run() override {
+    allreduce(inputs);
+  }
+
+  template <typename T>
+  void getFunction(gloo::AllreduceOptions::Func& fn, const ReduceOp op) {
+    fn = toFunction<T>(op);
+  }
+
+  gloo::AllreduceOptions::Func getFunction(
+      const at::ScalarType& dtype,
+      const ReduceOp& op) {
+    gloo::AllreduceOptions::Func fn;
+    GENERATE_ALL_TYPES(dtype, getFunction, fn, op);
+    return fn;
+  }
+};
+
+class AsyncAllreduceCoalescedWork : public AsyncAllreduceWork {
+ public:
+  AsyncAllreduceCoalescedWork(
+      const std::shared_ptr<gloo::Context>& context,
+      std::vector<at::Tensor>& inputs,
+      ReduceOp reduceOp,
+      uint32_t tag,
+      uint64_t seq,
+      std::chrono::milliseconds timeout)
+      : AsyncAllreduceWork(
+            context,
+            inputs,
+            std::move(reduceOp),
+            tag,
+            seq,
+            timeout) {}
+
+  void run() override {
+    allreduceCoalesced(inputs);
+  }
+
+ private:
+  void allreduceCoalesced(std::vector<at::Tensor>& tensors) {
+    // reduce coalesced, flattened tensors.
+    at::Tensor coalescedTensor = flattenDenseTensors(tensors);
+    std::vector<at::Tensor> allreduceInput = {coalescedTensor};
+    allreduce(allreduceInput);
+
+    // separate and reshape tensors.
+    size_t offset = 0;
+    for (at::Tensor& tensor : tensors) {
+      const int64_t tensorNumel = tensor.numel();
+      const c10::IntArrayRef tensorShape = tensor.sizes();
+      tensor.copy_(coalescedTensor.slice(0, offset, offset + tensorNumel)
+                       .view(tensorShape));
+      offset += tensorNumel;
+    }
+  }
+};
+
+class AsyncSparseAllreduceWork : public ProcessGroupGloo::AsyncWork {
+ public:
+  AsyncSparseAllreduceWork(
+      std::shared_ptr<gloo::Context> context,
+      std::vector<at::Tensor>& inputs,
+      uint32_t tag,
+      uint64_t seq,
+      std::chrono::milliseconds timeout)
+      : ProcessGroupGloo::AsyncWork(
+            std::move(context),
+            {inputs},
+            OpType::_ALLREDUCE_SPARSE,
+            seq,
+            timeout,
+            "gloo:sparse_all_reduce",
+            inputs),
+        inputs(inputs),
+        tag(tag) {}
+
+  std::vector<at::Tensor> inputs;
+  const uint32_t tag;
+
+  // We share dimensionality about the sparse tensors before collecting
+  // their contents. We assume here that the maximum number of sparse
+  // and dense dimensions is 4. This is stored in a contiguous piece of
+  // memory so that we can easily run allgather on it.
+  //
+  // The layout of this memory is as follows:
+  //
+  //   - [0:4]: sparse dims
+  //   - [4:8]: dense dims
+  //   -   [8]: nnz
+  //
+  class SparseTensorMetadata {
+   public:
+    static constexpr auto dim = 9;
+
+    // Construct from an existing metadata tensor to facilitate structured
+    // access to metadata from peers, after gathering it.
+    explicit SparseTensorMetadata(at::Tensor metadata)
+        : metadata_(std::move(metadata)),
+          data_(metadata_.mutable_data_ptr<int64_t>()) {
+      AT_ASSERT(metadata_.scalar_type() == at::kLong);
+      AT_ASSERT(metadata_.dim() == 1);
+      AT_ASSERT(metadata_.size(0) == dim);
+    }
+
+    // Populate the metadata.
+    void populate_from_sparse_tensor(const at::Tensor& tensor) {
+      const auto sparse_dim = tensor.sparse_dim();
+      AT_ASSERT(sparse_dim <= 4);
+      for (const auto i : c10::irange(4)) {
+        if (i < sparse_dim) {
+          data_[i] = tensor.size(i);
+        }
+      }
+      const auto dense_dim = tensor.dense_dim();
+      AT_ASSERT(dense_dim <= 4);
+      for (const auto i : c10::irange(4)) {
+        if (i < dense_dim) {
+          data_[i + 4] = tensor.size(sparse_dim + i);
+        }
+      }
+      data_[8] = tensor._nnz();
+    }
+
+    std::vector<int64_t> sizes() const {
+      std::vector<int64_t> sizes;
+      // Sparse sizes
+      for (const auto i : c10::irange(4)) {
+        if (data_[i] <= 0) {
+          break;
+        }
+        sizes.push_back(data_[i]);
+      }
+      // Dense sizes
+      for (const auto i : c10::irange(4, 8)) {
+        if (data_[i] <= 0) {
+          break;
+        }
+        sizes.push_back(data_[i]);
+      }
+      return sizes;
+    }
+
+    int64_t nnz() const {
+      return data_[8];
+    }
+
+   protected:
+    at::Tensor metadata_;
+    int64_t* data_;
+  };
+
+  // Sparse allreduce is implemented with allgather on indices and values.
+  // Every process then sums the resulting sparse tensors locally.
+  // The nnz for sparse tensors may be different across processes, so first
+  // we run allgather on the nnz, and then allgather with max(nnz).
+  at::Tensor allreduce(std::vector<at::Tensor>& tensors) {
+    // TODO: This is a massive hack!  There is some confusion about
+    // Variable/Tensor inside the body of this function.  Turning off
+    // grad smooths over the confusion for now.  This fixes
+    // test/test_c10d_gloo.py ProcessGroupGlooTest.test_sparse_allreduce_basics
+    //
+    // The correct fix is to stop allocating tensors that are not variables,
+    // but to conveniently do this c10d must depend on torch not ATen
+    at::AutoDispatchBelowAutograd guard;
+    auto input = tensors[0];
+
+    // Perform local reduction if we have multiple inputs.
+    for (const auto i : c10::irange(1, tensors.size())) {
+      input += tensors[i];
+    }
+
+    // Need to coalesce before we can access indices and values.
+    input = input.coalesce();
+
+    // Gather metadata information from all ranks.
+    auto metadata = allgather_metadata(input);
+
+    // Sanity check dimensionality across ranks.
+    {
+      const auto expected = metadata[context_->rank].sizes();
+      for (const auto i : c10::irange(context_->size)) {
+        if (i == context_->rank) {
+          continue;
+        }
+        const auto actual = metadata[i].sizes();
+        TORCH_CHECK(actual == expected, "Sparse dimensions do not match");
+      }
+    }
+
+    // Gather all indices and all values.
+    auto indices = allgather_indices(input, metadata);
+    auto values = allgather_values(input, metadata);
+
+    // Perform global reduction.
+    AT_ASSERT(static_cast<int>(indices.size()) == context_->size);
+    AT_ASSERT(static_cast<int>(values.size()) == context_->size);
+    auto output = at::sparse_coo_tensor(
+        indices[0], values[0], input.sizes(), input.options());
+    for (const auto i : c10::irange(1, context_->size)) {
+      output += at::sparse_coo_tensor(
+          indices[i], values[i], input.sizes(), input.options());
+    }
+
+    // Coalesce for good measure.
+    return output.coalesce();
+  }
+
+  void run() override {
+    auto output = allreduce(inputs);
+
+    // This copy is needed when we run a multi-gpu version of reduce (multiple
+    // inputs per rank).
+    for (const auto i : c10::irange(inputs.size())) {
+      inputs[i].copy_(output);
+    }
+  }
+
+  const std::vector<at::Tensor> getInputTensors() override {
+    return inputs;
+  }
+
+  const std::vector<at::Tensor> getOutputTensors() override {
+    return inputs;
+  }
+
+ private:
+  std::vector<SparseTensorMetadata> allgather_metadata(
+      const at::Tensor& tensor) {
+    auto buffer =
+        at::zeros({context_->size, SparseTensorMetadata::dim}, at::kLong);
+
+    // Prepare metadata vector (1 entry per rank)
+    std::vector<SparseTensorMetadata> metadata;
+    metadata.reserve(context_->size);
+    for (const auto i : c10::irange(context_->size)) {
+      metadata.emplace_back(buffer.select(0, i));
+    }
+
+    // Populate data for this rank
+    metadata[context_->rank].populate_from_sparse_tensor(tensor);
+
+    // Allgather metadata
+    gloo::AllgatherOptions opts(context_);
+    opts.setOutput(buffer.mutable_data_ptr<int64_t>(), buffer.numel());
+    opts.setTag(tag);
+    opts.setTimeout(getTimeout());
+    gloo::allgather(opts);
+
+    return metadata;
+  }
+
+  std::vector<at::Tensor> allgather_indices(
+      const at::Tensor& tensor,
+      const std::vector<SparseTensorMetadata>& metadata) {
+    const auto sparseDim = tensor.sparse_dim();
+
+    std::vector<size_t> counts(context_->size);
+    size_t totalSize = 0;
+    for (const auto i : c10::irange(metadata.size())) {
+      counts[i] = metadata[i].nnz() * sparseDim;
+      totalSize += counts[i];
+    }
+
+    auto output = at::empty({static_cast<int64_t>(totalSize)}, at::kLong);
+
+    // tensors copied from cuda may not be contiguous, get a contiguous
+    // tensor before use its data_ptr
+    auto input = tensor.indices().contiguous();
+
+    // Allgatherv indices.
+    gloo::AllgathervOptions opts(context_);
+    opts.setInput(
+        // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
+        const_cast<int64_t*>(input.const_data_ptr<int64_t>()),
+        input.numel());
+    opts.setOutput(output.mutable_data_ptr<int64_t>(), counts);
+    opts.setTag(tag);
+    opts.setTimeout(getTimeout());
+    gloo::allgatherv(opts);
+
+    // Compile indices tensor per rank.
+    std::vector<at::Tensor> indices;
+    indices.reserve(metadata.size());
+    int64_t offset = 0;
+    for (const auto& i : metadata) {
+      const auto nnz = i.nnz();
+      const auto numel = sparseDim * nnz;
+      indices.push_back(
+          output.narrow(0, offset, numel).reshape({sparseDim, nnz}));
+      offset += numel;
+    }
+
+    return indices;
+  }
+
+  std::vector<at::Tensor> allgather_values(
+      const at::Tensor& tensor,
+      const std::vector<SparseTensorMetadata>& metadata) {
+    // There are nnz #dense_dim()-dimensional tensors per rank.
+    const auto valueShape = tensor.sizes().slice(tensor.sparse_dim());
+    int64_t denseNumel = 1;
+    for (auto dim : valueShape) {
+      denseNumel *= dim;
+    }
+
+    std::vector<size_t> counts(context_->size);
+    int64_t totalSize = 0;
+    for (const auto i : c10::irange(metadata.size())) {
+      counts[i] = metadata[i].nnz() * denseNumel;
+      totalSize += static_cast<int64_t>(counts[i]);
+    }
+
+    auto output = at::empty({totalSize}, tensor.scalar_type());
+
+    // Allgatherv indices.
+    gloo::AllgathervOptions opts(context_);
+    // tensors copied from cuda may not be contiguous, get a contiguous
+    // tensor before use its data_ptr
+    at::Tensor valueTensor = tensor.values().contiguous();
+    GENERATE_ALL_TYPES(valueTensor.scalar_type(), setInput, opts, valueTensor);
+    GENERATE_ALL_TYPES(
+        valueTensor.scalar_type(), setOutput, opts, output, counts);
+    opts.setTag(tag);
+    opts.setTimeout(getTimeout());
+    gloo::allgatherv(opts);
+
+    // Compile values tensor per rank.
+    std::vector<at::Tensor> values;
+    values.reserve(metadata.size());
+    int64_t offset = 0;
+    for (const auto& i : metadata) {
+      const auto nnz = i.nnz();
+      const auto numel = denseNumel * nnz;
+      auto tensorShape = std::vector<int64_t>({(int64_t)nnz});
+      std::copy(
+          valueShape.begin(),
+          valueShape.end(),
+          std::back_inserter(tensorShape));
+      values.push_back(output.narrow(0, offset, numel).reshape(tensorShape));
+      offset += numel;
+    }
+
+    return values;
+  }
+};
+
+} // namespace c10d
+
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupMPI.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupMPI.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..467ec0c08c5339947f514e6d1ae243fe2520ce29
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupMPI.hpp
@@ -0,0 +1,278 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_MPI
+
+#include <condition_variable>
+#include <deque>
+#include <exception>
+#include <memory>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+#include <ATen/core/ivalue.h>
+#include <ATen/core/ivalue_inl.h>
+
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+
+#include <mpi.h>
+
+namespace c10d {
+
+constexpr const char* MPI_BACKEND_NAME = "mpi";
+
+// WorkEntry is the state associated with a single MPI run instance.
+// It include the source Tensor list and destination Tensor list, as well as
+// The actual run function that will operate either on src or dst or both.
+struct WorkEntry {
+  explicit WorkEntry(
+      std::vector<at::Tensor>* srcPtr,
+      std::vector<at::Tensor>* dstPtr,
+      std::function<void(std::unique_ptr<WorkEntry>&)> run)
+      : dst(dstPtr ? *dstPtr : std::vector<at::Tensor>()), run(std::move(run)) {
+    if (srcPtr) {
+      src = *srcPtr;
+    }
+  }
+
+  // Not copyable
+  WorkEntry(const WorkEntry&) = delete;
+  // Not copy assignable
+  WorkEntry& operator=(const WorkEntry&) = delete;
+
+  // For input and output tensors (in-place), we will always use src
+  std::vector<at::Tensor> src;
+
+  // Copy of user provided outputs.
+  const std::vector<at::Tensor> dst;
+
+  // src rank returned, for recv only
+  int* srcRank = nullptr;
+  std::function<void(std::unique_ptr<WorkEntry>&)> run;
+};
+
+// ProcessGroupMPI implements MPI bindings for c10d.
+//
+// All functions on this class are expected to be called in the same
+// order across processes in the group. This is the only way that we
+// can guarantee to match up the same calls across processes.
+//
+// All MPI functions provided by this class is asynchronously scheduled on a
+// Worker thread. Therefore, ProcessGroupMPI requires the MPI implementation
+// that is used to have a minimum thread support value of MPI_THREAD_SERIALIZED.
+// That is, The process may be multi-threaded, and multiple threads may make
+// MPI calls, but only one at a time: MPI calls are not made concurrently from
+// two distinct threads (all MPI calls are serialized). However, with
+// MPI_THREAD_SERIALIZED, ProcessGroupMPI will only support a single process
+// group. In other words, no more than 1 process group can be created globally.
+//
+// If you would like to use multiple ProcessGroupMPI, it requires your MPI
+// implementation to have a thread support value of MPI_THREAD_MULTIPLE, that
+// is, multiple threads may call MPI, with no restriction.
+//
+// Also note that ProcessGroupMPI only supports a single Tensor operation. In
+// other words, the size of the input Tensor vector should always be 1.
+//
+// CUDA tensor can be supported if the MPI used is CUDA-aware MPI, and
+// ProcessGroupMPI will automatically detect this support.
+class TORCH_API ProcessGroupMPI : public Backend {
+ public:
+  class WorkMPI : public Work {
+   public:
+    explicit WorkMPI(
+        std::vector<at::Tensor> outputTensors,
+        const char* profilingTitle = nullptr,
+        const std::optional<std::vector<at::Tensor>>& inputTensors =
+            std::nullopt)
+        : Work(-1, OpType::UNKNOWN, profilingTitle, inputTensors),
+          outputTensors_(std::move(outputTensors)),
+          future_(c10::make_intrusive<at::ivalue::Future>(
+              c10::ListType::create(c10::TensorType::get()))) {}
+
+    std::vector<at::Tensor> result() override;
+
+    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
+
+   protected:
+    friend class ProcessGroupMPI;
+
+   private:
+    void finishWorkMPI();
+    void finishWorkMPIError(const std::exception_ptr& eptr);
+
+    std::vector<at::Tensor> outputTensors_;
+    c10::intrusive_ptr<at::ivalue::Future> future_;
+  };
+
+  class AsyncWork : public Work {
+   public:
+    AsyncWork(
+        MPI_Request request,
+        std::vector<at::Tensor> outputTensors,
+        const char* profilingTitle = nullptr,
+        const std::optional<std::vector<at::Tensor>>& inputTensors =
+            std::nullopt);
+
+    ~AsyncWork() override;
+
+    bool isCompleted() override;
+
+    bool isSuccess() const override;
+
+    int sourceRank() const override;
+
+    bool wait(std::chrono::milliseconds timeout = kUnsetTimeout) override;
+
+    void abort() override;
+
+    std::vector<at::Tensor> result() override;
+
+   protected:
+    void populateException();
+
+   private:
+    const std::vector<at::Tensor> outputTensors_;
+    MPI_Request request_;
+    MPI_Status status_{};
+  };
+
+  // Constructor will spawn up the worker thread loop
+  explicit ProcessGroupMPI(int rank, int size, MPI_Comm pgComm);
+
+  ~ProcessGroupMPI() override;
+
+  // Abort the MPI program, needs to be called when exception is detected
+  void abort() override;
+
+  const std::string getBackendName() const override {
+    return std::string(MPI_BACKEND_NAME);
+  }
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& data,
+      const BroadcastOptions& opts = BroadcastOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceCoalescedOptions& opts =
+          AllreduceCoalescedOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& tensors,
+      const ReduceOptions& opts = ReduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& outputbuffer,
+      at::Tensor& inputbuffer,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_coalesced(
+      std::vector<std::vector<at::Tensor>>& outputTensorLists,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const GatherOptions& opts = GatherOptions()) override;
+
+  c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ScatterOptions& opts = ScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      std::vector<int64_t>& outputSplitSizes,
+      std::vector<int64_t>& inputSplitSizes,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recvAnysource(
+      std::vector<at::Tensor>& tensor,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) override;
+
+  // Creating a new ProcessGroupMPI, will initialize MPI if not initialized
+  static c10::intrusive_ptr<ProcessGroupMPI> createProcessGroupMPI(
+      std::vector<int> ranks = {});
+
+ protected:
+  using WorkType =
+      std::tuple<std::unique_ptr<WorkEntry>, c10::intrusive_ptr<WorkMPI>>;
+  // Worker thread loop
+  void runLoop();
+  // Helper function that is called by the destructor
+  void destroy();
+
+  c10::intrusive_ptr<Work> enqueue(
+      std::unique_ptr<WorkEntry> entry,
+      const char* profilingTitle = nullptr,
+      const std::optional<std::vector<at::Tensor>>& inputTensors =
+          std::nullopt);
+
+  bool stop_{false};
+
+  std::mutex pgMutex_;
+  std::thread workerThread_;
+
+  std::deque<WorkType> queue_;
+  std::condition_variable queueProduceCV_;
+  std::condition_variable queueConsumeCV_;
+
+  // Global states
+  static void initMPIOnce();
+  static void mpiExit();
+
+  static std::mutex pgGlobalMutex_;
+  static int mpiThreadSupport_;
+
+  MPI_Comm pgComm_;
+};
+
+} // namespace c10d
+
+#endif // USE_C10D_MPI
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..ba4265b0613193c88987a766416ca749b54182fa
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp
@@ -0,0 +1,1547 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_NCCL
+
+#if defined(__linux__)
+#include <fcntl.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+#endif
+
+#include <atomic>
+#include <chrono>
+#include <deque>
+#include <future>
+#include <iostream>
+#include <list>
+#include <mutex>
+#include <thread>
+#include <unordered_map>
+
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
+#include <torch/csrc/distributed/c10d/PrefixStore.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/cuda/CUDAEventCache.hpp>
+#include <torch/csrc/distributed/c10d/logger.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp>
+
+#include <ATen/DynamicLibrary.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDAEvent.h>
+#include <ATen/cuda/MemPool.h>
+#include <c10/core/Stream.h>
+#include <c10/core/StreamGuard.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#include <torch/custom_class.h>
+
+namespace c10d {
+
+// Control broadcasting of NCCL uniqueId
+static std::vector<std::string> TORCH_NCCL_BCAST_UNIQUEID = {
+    "TORCH_NCCL_BCAST_UNIQUEID"};
+
+// Control EagerInit P2P serialization warning
+static std::vector<std::string>
+    TORCH_NCCL_SHOW_EAGER_INIT_P2P_SERIALIZATION_WARNING = {
+        "TORCH_NCCL_SHOW_EAGER_INIT_P2P_SERIALIZATION_WARNING"};
+
+// Control whether to always use high priority streams
+static std::vector<std::string> TORCH_NCCL_HIGH_PRIORITY = {
+    "TORCH_NCCL_HIGH_PRIORITY"};
+
+// Control whether or not wait() is blocking or non-blocking.
+static std::vector<std::string> TORCH_NCCL_BLOCKING_WAIT = {
+    "TORCH_NCCL_BLOCKING_WAIT",
+    "NCCL_BLOCKING_WAIT"};
+
+// TODO: We want to eventually remove this variable and make users to use
+// the default value (3 - SkipCleanUp).
+// Control whether or not we perform Async Error Handling with NCCL.
+static std::vector<std::string> TORCH_NCCL_ASYNC_ERROR_HANDLING = {
+    "TORCH_NCCL_ASYNC_ERROR_HANDLING",
+    "NCCL_ASYNC_ERROR_HANDLING"};
+
+// Control whether dumping debug info on watchdog
+// timeout is enabled. This variable must be set together with
+// TORCH_NCCL_ENABLE_MONITORING=1 and TORCH_NCCL_TRACE_BUFFER_SIZE > 0.
+static std::vector<std::string> TORCH_NCCL_DUMP_ON_TIMEOUT = {
+    "TORCH_NCCL_DUMP_ON_TIMEOUT"};
+
+// Control whether to propagate NCCL errors to all ranks through TCPStore.
+static std::vector<std::string> TORCH_NCCL_PROPAGATE_ERROR = {
+    "TORCH_NCCL_PROPAGATE_ERROR"};
+
+// Control whether Desync Debug is enabled. This variable must be set
+// together with TORCH_NCCL_ASYNC_ERROR_HANDLING.
+static std::vector<std::string> TORCH_NCCL_DESYNC_DEBUG = {
+    "TORCH_NCCL_DESYNC_DEBUG",
+    "NCCL_DESYNC_DEBUG"};
+
+// Enable recording start-events for all ProcessGroupNCCL collectives, and
+// compute accurate collective timing per-collective. (Note: end-events are
+// recorded by default. Turn on this flag can increase chances of a watchdog
+// hang due to performing a CUDA event query which eventually calls
+// cudaEventElapsedTime() API.
+static std::vector<std::string> TORCH_NCCL_ENABLE_TIMING = {
+    "TORCH_NCCL_ENABLE_TIMING",
+    "NCCL_ENABLE_TIMING"};
+
+// Enable monitoring thread which aborts the process when the ProcessGroupNCCL
+// Watchdog thread gets stuck and no heartbeat is detected after
+// TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC. This can happen due to calling CUDA/NCCL
+// APIs that may hang. It is Useful to prevent jobs being stuck for a prolonged
+// time than necessary tying up cluster resources.
+static std::vector<std::string> TORCH_NCCL_ENABLE_MONITORING = {
+    "TORCH_NCCL_ENABLE_MONITORING"};
+
+// Control the watchdog heartbeat timeout period after which the monitoring
+// thread will abort the process.
+static std::vector<std::string> TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC = {
+    "TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC"};
+
+// Whether to rethrow CUDA Errors in the watchdog (default true)
+static std::vector<std::string> TORCH_NCCL_RETHROW_CUDA_ERRORS = {
+    "TORCH_NCCL_RETHROW_CUDA_ERRORS"};
+
+// The maximum number of events we store in the flight recorder's ring buffer.
+// (One event could be the start or end of a collective, for example).
+static std::vector<std::string> TORCH_NCCL_TRACE_BUFFER_SIZE = {
+    "TORCH_NCCL_TRACE_BUFFER_SIZE"};
+
+// Control how much extra time we will wait for dumping the debugging info
+// before we exit and throws timeout exception.
+static std::vector<std::string> TORCH_NCCL_WAIT_TIMEOUT_DUMP_MILSEC = {
+    "TORCH_NCCL_WAIT_TIMEOUT_DUMP_MILSEC"};
+
+// Control the interval inside the monitoring thread to check the coordinated
+// signal from other ranks, e.g. to dump the debugging information.
+static std::vector<std::string> TORCH_NCCL_COORD_CHECK_MILSEC = {
+    "TORCH_NCCL_COORD_CHECK_MILSEC"};
+
+// Whether to log C++ stack traces on unclean shutdown (default true)
+static std::vector<std::string> TORCH_NCCL_LOG_CPP_STACK_ON_UNCLEAN_SHUTDOWN = {
+    "TORCH_NCCL_LOG_CPP_STACK_ON_UNCLEAN_SHUTDOWN"};
+
+// Whether to include only active collectives in the Flight Recorder trace
+// (default false)
+static std::vector<std::string> TORCH_NCCL_EXTRA_DUMP_ON_EXEC = {
+    "TORCH_NCCL_EXTRA_DUMP_ON_EXEC"};
+
+// Control whether to use CudaEventCache for the collective in watchdog thread.
+// We noticed in the past when cuda global lock is held, destroying CudaEvent
+// can cause a hang.
+static std::vector<std::string> TORCH_NCCL_CUDA_EVENT_CACHE = {
+    "TORCH_NCCL_CUDA_EVENT_CACHE"};
+
+// Control the number of ranks each root can cover during NCCL comm init.
+static std::vector<std::string> TORCH_NCCL_RANKS_PER_ROOT = {
+    "TORCH_NCCL_RANKS_PER_ROOT"};
+
+static std::vector<std::string> TORCH_NCCL_NAN_CHECK = {"TORCH_NCCL_NAN_CHECK"};
+
+constexpr const char* NCCL_BACKEND_NAME = "nccl";
+
+constexpr const char* kStoreDumpKey = "exception_dump";
+
+constexpr const char* kStoreErrorSignalKey = "remote_error";
+
+constexpr const int kWorkStatusUpdatePeriodMs = 30 * 1000; // 30 seconds
+
+constexpr auto kProcessGroupNCCLDefaultTimeout =
+    std::chrono::milliseconds(10 * 60 * 1000);
+
+// NoHandling: do not handle asynchronous NCCL errors
+// TearDown: tear down process upon error, see `WorkNCCL::handleException`
+// CleanUpOnly: just clean up collectives and abort communicators without
+// tearing down process SkipCleanUp: (this is a temporary option and can be
+// removed in future) tear down process without cleaning up NCCL communicators.
+// This should be used as a last resort in case `ncclCommAbort` itself is
+// hanging
+enum ErrorHandlingMode {
+  NoHandling = 0,
+  TearDown = 1,
+  CleanUpOnly = 2,
+  SkipCleanUp = 3
+};
+
+#define SHOULD_CLEAN_UP(a) (a != NoHandling && a != SkipCleanUp)
+
+#define SHOULD_TEAR_DOWN(a) (a != NoHandling && a != CleanUpOnly)
+
+#define PRINT_COLLECTIVE_HASH_SIGNATURE(phase, opType, numel, hashValue)      \
+  LOG(WARNING) << logPrefix() << "Hash of " << phase << " to NCCL " << opType \
+               << " with size " << numel << " is " << hashValue;
+
+// If set, ProcessGroupNCCL doesn't use recordStream calls to ensure
+// caching allocator safety for tensors used on both user-facing and
+// internal comm streams.
+// Instead, it stashes live references to those tensors until after
+// user-facing streams are synced with comm streams.
+// See stashed_for_allocator_safety_ below.
+static std::vector<std::string> TORCH_NCCL_AVOID_RECORD_STREAMS = {
+    "TORCH_NCCL_AVOID_RECORD_STREAMS"};
+
+// If set, ProcessGroupNCCL registers postAlloc and preFree hooks to cuda cache
+// allocator so that whenever a tensor is allocated or freed, ProcessGroupNCCL
+// can register/deregister the tensor on all available NCCL communicators.
+static std::vector<std::string> TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK =
+    {"TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK",
+     "NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK"};
+
+#if defined(__linux__)
+struct DumpPipe {
+  DumpPipe(int rank, const std::string& fileStem, int traceBufferSize) {
+    if (fileStem.empty() || traceBufferSize <= 0) {
+      return;
+    }
+    std::string filename = c10::str(fileStem, rank, ".pipe");
+    TORCH_CHECK(
+        unlink(filename.c_str()) != -1 || errno == ENOENT,
+        "Error removing existing named pipe ",
+        filename,
+        ", Error: ",
+        std::strerror(errno));
+    TORCH_CHECK(
+        mkfifo(filename.c_str(), 0666) != -1,
+        "Error creating named pipe ",
+        filename,
+        ", Error: ",
+        std::strerror(errno));
+    fd_ = open(filename.c_str(), O_RDONLY | O_NONBLOCK);
+    LOG(INFO) << "Pipe file " << filename
+              << " has been opened, write to it to trigger NCCL Debug Dump.";
+    TORCH_CHECK(fd_ != -1, "Error opening named pipe ", filename);
+  }
+  bool shouldDump() {
+    if (fd_ == -1) {
+      return false;
+    }
+    // NOLINTNEXTLINE(*array*)
+    char buf[128]{};
+    // non-blocking from O_NONBLOCK above.
+    // Ignore EINTR because we already will poll this
+    // again later.
+    ssize_t bytesRead = read(fd_, &buf, 128);
+    return bytesRead > 0;
+  }
+  ~DumpPipe() {
+    if (fd_ != -1) {
+      close(fd_);
+    }
+  }
+
+ private:
+  int fd_ = -1;
+};
+#else
+struct DumpPipe {
+  DumpPipe(int rank) {}
+  bool shouldDump() {
+    return false;
+  }
+};
+#endif
+
+// A shelf for stashing tensors between op call and `work.wait()`.
+// Used in case of async ops.
+class TensorShelf {
+ public:
+  // Stash tensors so that CachingAllocator cannot recycle them prematurely.
+  void stash(std::vector<at::Tensor>& tensors);
+  // Stash tensors from another shelf.
+  void stash(TensorShelf& other);
+  // Unstage the stashed tensors so that CachingAllocator can recycle them.
+  // Same as `clear()`.
+  void unstash();
+  // Whether shelf is empty.
+  bool empty();
+  // Clear the shelf.
+  void clear();
+
+ protected:
+  // Get the inner tensor vector. Use with caution as it is not protected by
+  // mutex.
+  std::vector<at::Tensor>& get();
+
+ private:
+  std::vector<at::Tensor> tVector_;
+  // Need a mutex to protect `tVector_` because it can be potentially accessed
+  // from both main thread and watchdog thread.
+  std::mutex mutex_;
+};
+
+// ProcessGroupNCCL implements NCCL bindings for c10d.
+//
+// All functions of the class are expected to be called in the same order
+// across all processes in the process group.  This is the only way that we
+// can guarantee to match up the same calls among all processes.
+//
+// All NCCL functions provided by this class are asynchronous functions. More
+// specifically, each NCCL call is scheduled on a separate CUDA stream that is
+// different from the current CUDA stream. This is for the purpose of
+// achieving potentially concurrency and better performance. As a result,
+// it is the callers' responsibility to make sure that the CUDA stream their
+// code works on needs to wait for the NCCL operation from
+// this class.
+//
+// This can be done by calling:
+//
+// either WorkNCCL::wait() or WorkNCCL::synchronize(), both achieves the same
+// functionality and are synonyms.
+//
+// Also note that WorkNCCL::finishedGPUExecution() is a helper function only
+// provided by ProcessGroupNCCL to check if the NCCL operation of WorkNCCL has
+// finished execution on the GPU (not just scheduled).
+//
+// Example on using the NCCL process group
+//
+//   ProcessGroupNCCL pg(store, rank, size);
+//   std::shared_ptr<WorkNCCL> work = pg.allreduce(tensors);
+//
+//   // At this point, NCCL kernel has already by queued successfully
+//   // Now, let current stream wait for the NCCL to finish, this function is
+//   // async operation as well
+//
+//   work->wait()
+//
+//   // Now continue on other work in the current stream.
+class TORCH_API ProcessGroupNCCL : public Backend {
+ public:
+  class WorkNCCL : public Work, public std::enable_shared_from_this<WorkNCCL> {
+   public:
+    friend struct WorkInfo;
+
+    // Constructor takes a list of CUDA devices
+    WorkNCCL(
+        std::string pgUID,
+        std::string pgDesc,
+        at::Device& device,
+        int rank,
+        OpType opType,
+        uint64_t seq,
+        bool isP2P = false,
+        const char* profilingTitle = nullptr,
+        const std::optional<std::vector<at::Tensor>>& inputs = std::nullopt,
+        bool enableTiming = false,
+        bool cudaEventCacheEnabled = false,
+        DebugLevel distDebugLevel = DebugLevel::Off);
+    // Copy constructor doing partial copy without outputs_. Cleanup thread
+    // monitors and removes finished works. However it will deadlock when
+    // destructs outputs_ tensors who are view tensors in autograd graph.
+    WorkNCCL(const WorkNCCL& w);
+
+    ~WorkNCCL() override = default;
+
+    // Checks if the NCCL kernel has started to execute.
+    bool isStarted();
+
+    // Checks if request has completed. In this specific case of NCCL, it checks
+    // if the NCCL operation has completed on the GPU in its own NCCL stream.
+    // Non-blocking operation.
+    bool isCompleted() override;
+
+    bool isSuccess() const override;
+
+    // Same as calling synchronize() for NCCL work if timeout is not set.
+    // Otherwise, it will block the CPU thread until the NCCL work is completed
+    // or timed out. If timeout, exception will be thrown.
+    bool wait(std::chrono::milliseconds timeout = kNoTimeout) override;
+
+    void blockCurrentStream() override {
+      synchronize();
+    }
+
+    void abort() override;
+
+    // Let current stream wait on the completion of the NCCL work
+    // Throws on exceptions.
+    void synchronize() override;
+
+    // Synchronize streams by blocking each on the NCCL stream
+    void synchronizeStream();
+
+    // Helper function to handle exception (throw if needed).
+    void handleException(ErrorHandlingMode asyncErrorHandling);
+
+    // Helper function that checks if the NCCL kernels have finished
+    // execution on the GPUs
+    bool finishedGPUExecution();
+
+    // Get a Future object that will be marked as completed internally.
+    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
+
+    // Get a Future result of each work (e.g. success, different error types).
+    // instead of the tensor output.
+    c10::intrusive_ptr<c10::ivalue::Future> getFutureResult() override;
+
+    float getDuration() const override;
+
+    uint64_t getSequencenumber() const override;
+
+    const std::string& logPrefix() const;
+
+    // Helper function that sets an exception_ptr on the WorkNCCL object.
+    void setException(std::exception_ptr exception_ptr);
+
+    // Helper function that returns True if the WorkNCCL object has timed out
+    // and False otherwise.
+    // In case of timeout, set exception on the WorkNCCL object.
+    bool checkTimeout(
+        std::optional<std::chrono::milliseconds> timeout = std::nullopt);
+
+    // Print the traceback of the collective at call time
+    void printTraceback() const;
+
+    std::string getTraceback() const;
+
+    std::vector<at::Tensor> result() override;
+
+   protected:
+    // The process group unique id
+    std::string pgUID_;
+
+    // The process group description
+    std::string pgDesc_;
+
+    // The cached list of CUDA devices to operate on
+    at::Device device_;
+
+    // The start CUDA event of NCCL operator tracking this work item. These
+    // start CUDA events are needed by desync debugging if enabled.
+    std::shared_ptr<at::cuda::CUDAEvent> ncclStartEvent_;
+
+    // The end CUDA event of NCCL operator tracking this work item.
+    std::shared_ptr<at::cuda::CUDAEvent> ncclEndEvent_;
+
+    // The NCCL communicator used for this work item.
+    std::shared_ptr<NCCLComm> ncclComm_;
+
+    // whether this work is a barrier op
+    bool isBarrierOp_{false};
+
+    // Clone of blockingWait_ from ProcessGroupNCCL.
+    bool blockingWait_{false};
+
+    // Clone of opTimeout_ from ProcessGroupNCCL.
+    std::chrono::milliseconds opTimeout_{};
+
+    // Ephemeral timeouts are owned by exactly one work,
+    // and reset after that work completes.
+    // There may be more than one ephemeral timeout active at the same time,
+    // and this variable is used to track the ownership of ephemeral timeout.
+    std::chrono::milliseconds ownedEphermeralTimeout_ =
+        std::chrono::milliseconds(0);
+
+    // Time point representing when the work started.
+    std::chrono::time_point<std::chrono::steady_clock> workStartTime_;
+
+    // Record the sequential number of collective or p2p.
+    uint64_t seq_;
+    bool isP2P_;
+
+    // Indicates if the nccl start event has been updated to the store trace.
+    // This will be used by desync debug.
+    bool startTraceUpdated_{false};
+
+    // Record collective sizes for debug. We only record the size on the first
+    // device as multi-device per process is deprecated
+    size_t numelIn_ = 0;
+    size_t numelOut_ = 0;
+
+    // Wrapper method for the static checkForNCCLErrors which can be overridden
+    // for tests.
+    virtual std::exception_ptr checkForNCCLErrors();
+
+    friend std::ostream& operator<<(
+        std::ostream& output,
+        const WorkNCCL& workNCCL);
+
+    // Checks for NCCL errors and sets an appropriate exception_ptr.
+    void checkAndSetException();
+
+    // Just checks whether GPU execution has started, without modifying
+    // exception_ptr.
+    bool startedGPUExecutionInternal() const;
+
+    // Just checks whether GPU execution has completed, without modifying
+    // exception_ptr.
+    bool finishedGPUExecutionInternal() const;
+
+    // Reference to the store so that we can write aborted communicators
+    // to the store.
+    c10::intrusive_ptr<Store> store_;
+
+    // Store a reference to NCCL collective's outputs, used by result and to
+    // give a more descriptive message when representing the Work as a string.
+    std::shared_ptr<std::vector<at::Tensor>> outputs_;
+
+    // TORCH_NCCL_AVOID_RECORD_STREAMS implementation helper.
+    // Stores references to participating non-output tensors (ie inputs,
+    // flattened intermediates).
+    // We'll clear this list in synchronizeStream, just after user-facing
+    // stream(s) are synced with the nccl work stream(s).
+    // By keeping these refs (as well as outputs_) alive until after the
+    // collective's work rejoins the user-facing streams, we achieve
+    // caching allocator safety without any recordStream calls.
+    // For in-place collectives, some refs stashed here may alias outputs_,
+    // but that doesn't do any harm.
+    std::shared_ptr<TensorShelf> stashed_for_allocator_safety_;
+
+    // The future returned by getFuture.
+    c10::intrusive_ptr<at::ivalue::Future> future_;
+
+    // the future result (e.g., success or failure) of the work
+    c10::intrusive_ptr<at::ivalue::Future> futureWorkResult_;
+
+    bool timingEnabled_;
+    // unique id used to tell the trace buffer that this
+    // work has completed
+    std::optional<uint64_t> trace_id_;
+    std::optional<uint64_t> trace_reset_epoch_;
+    DebugLevel distDebugLevel_;
+    friend class ProcessGroupNCCL;
+  };
+
+  struct Options : Backend::Options {
+    // NOTE: timeout in ProcessGroupNCCL::Options denote the timeout for
+    // operations. This is only used when blockingWait_ is enabled.
+    explicit Options(bool is_high_priority_stream = false);
+    Options(const Options&) = default;
+    Options(Options&&) noexcept = default;
+    Options& operator=(const Options&) = delete;
+    Options& operator=(Options&&) noexcept = delete;
+    ~Options() override = default;
+
+    // return intrusive_ptr of the object
+    static c10::intrusive_ptr<Options> create(
+        bool is_high_priority_stream = false) {
+      return c10::make_intrusive<Options>(is_high_priority_stream);
+    }
+
+    // Schedule NCCL operations on high priority CUDA streams
+    bool is_high_priority_stream;
+
+#ifdef NCCL_HAS_CONFIG
+    // Configure ranks
+    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+#endif
+
+    // Optional "parent" backend and color to create communicators from
+    // via `ncclCommSplit`
+    c10::intrusive_ptr<ProcessGroupNCCL> split_from;
+    // Color to use for `ncclCommSplit`, values:
+    // * Non-negative value: in group;
+    // * NCCL_SPLIT_NOCOLOR (-1): not in group;
+    // * NCCL_SPLIT_NOCOLOR - 1: uninitialized.
+    // [Note 1]: the type must be `int` instead of `int64_t` because NCCL API
+    // accepts int. Otherwise, an implicit conversion may happen at the API call
+    // and the value may become negative.
+    // [Note 2]: this member is pybinded to Python, the value passed from Python
+    // must be within the numerical range of C++ int. Otherwise, Python will
+    // raise a RuntimeError saying type is incompatible. See also
+    // `_process_group_color` in `distributed_c10d.py`.
+#ifdef NCCL_HAS_COMM_SPLIT
+    int split_color{NCCL_SPLIT_NOCOLOR - 1};
+#else
+    // [Note 3]: for older NCCL versions, NCCL_SPLIT_NOCOLOR is not defined. But
+    // `split_color` is pybinded to Python, so we need to define it. So we use
+    // the int value of `NCCL_SPLIT_NOCOLOR` (-1) instead.
+    int split_color{-2};
+#endif
+  };
+
+  // Helper class related to TORCH_NCCL_DESYNC_DEBUG
+  class DesyncDebugger {
+   public:
+    // Initialize and enable DesyncDebugger
+    void init(
+        int rank,
+        int size,
+        int globalRank,
+        int pgId,
+        c10::intrusive_ptr<Store> store);
+
+    // Run desync debug. This function is called by watchdog at time of timeout.
+    void run();
+
+    // Log work start to store.
+    void logWorkStart(WorkNCCL& work);
+
+    // Log work end to store.
+    void logWorkEnd(WorkNCCL& work);
+
+   private:
+    // Whether desync debug is enabled.
+    // If false, all functions are no-op.
+    bool enabled_{false};
+
+    // From ProcessGroupNCCL
+    int rank_;
+    int size_;
+    int globalRank_;
+    int pgId_;
+
+    // Reference to the store so that we can log start/end event.
+    c10::intrusive_ptr<Store> store_;
+
+    // The store keys to trace the last NCCL collective kernel CUDA events -
+    // start event and end event respectively. These are used to do desync root
+    // cause analysis.
+    std::string traceKeyStart_;
+    std::string traceKeyEnd_;
+  };
+
+  // Class that runs as a separate thread aside from watchdog
+  // thread because we need to check the heartbeat from watchdog thread
+  // so that when we get stuck in some NCCL/CUDA calls,
+  // we can dump the debugging information and abort the process.
+  class HeartbeatMonitor {
+   public:
+    HeartbeatMonitor(ProcessGroupNCCL* pg);
+    virtual ~HeartbeatMonitor() = default;
+
+    // Start the heartbeat monitor thread.
+    void start();
+
+    // Join the heartbeat monitor thread.
+    void join();
+
+    // Run the actual loop to check watchdog heartbeat.
+    virtual void runLoop();
+
+    // Set the terminal flag and notify the heartbeat monitor thread to stop.
+    void stop();
+
+    // Set the last update time of watchdog thread.
+    void setLastWorkListUpdateTime(
+        std::chrono::time_point<std::chrono::steady_clock> time);
+
+    int getDumpTimeout() const;
+
+    // Util function to get the timeout error message
+    std::string getNCCLWatchdogTimeoutErrorMsg(const std::string& extraMsg);
+
+    // Util function to get the timeout exit message
+    std::string getNCCLWatchdogTimeoutExitMsg(const std::string& exitReason);
+
+   protected:
+    // We need to keep a reference to the PG instance so that we can access
+    // the member functions of the PG instance. We store a raw pointer on
+    // purpose because the heartbeat monitor thread now still lives within the
+    // lifetime of the PG instance.
+    ProcessGroupNCCL* pg_;
+
+   private:
+    // Whether or not to print C++ stack traces to logs on unclean shutdown.
+    bool logCppStackOnUncleanShutdown_;
+
+    // The time interval used for deciding whether there is no watchdog
+    // heartbeat.
+    int heartbeatTimeoutInSec_;
+
+    // timeout for the dump to finish.
+    int waitTimeoutDumpInMilSec_;
+
+    // Interval of check coordinated signals in ProcessGroupNCCL from other
+    // ranks e.g., trigger the dump of the debugging info for timeout when
+    // notified.
+    int coordCheckIntervalMilSec_;
+
+    // We gate the heartbeat monitor thread so that we can roll it out
+    // gradually.
+    bool watchdogHeartbeatMonitorEnabled_;
+
+    // Monitor thread which checks the heartbeat of Watchdog thread.
+    // If the monitor thread finds there is no heartbeat, it will dump debug
+    // info and then kill the watchdog thread to avoid hang.
+    std::thread ncclHeartbeatMonitorThread_;
+
+    // Whether or not we should terminate the heartbeat monitoring threads.
+    std::atomic<bool> terminateHeartbeatMonitorThread_{false};
+
+    // Condition Variable for monitor thread to wake up early
+    std::condition_variable monitorWakeUpCV_;
+
+    // Whether or not to dump debug info on exception including both watchdog
+    // timeout and nccl errors.
+    bool dumpOnTimeoutOrEx_;
+
+    // Mutex to Guard monitorWakeUpCV_
+    std::mutex monitorMutex_;
+
+    // The last update time of WorkList inside watchdog thread.
+    std::chrono::time_point<std::chrono::steady_clock> lastWorkListUpdateTime_;
+  };
+
+  // Class that runs as a side thread to check whether the NCCL collective
+  // is timed out or errors on the cached NCCL communicators.
+  class Watchdog {
+   public:
+    Watchdog(ProcessGroupNCCL* pg);
+    virtual ~Watchdog() = default;
+
+    // Start the watchdog thread.
+    void start();
+
+    // Join the watchdog thread.
+    void join();
+
+    // Function that runs as part of a separate thread and checks for errors on
+    // NCCL communicators. We need a separate thread to check for NCCL errors
+    // since we can't rely on the user calling certain methods like wait(),
+    // isCompleted() etc. to detect and remediate errors. In addition to this,
+    // we need a mechanism to safely abort and remove NCCL communicators from
+    // our cache. This can be done cleanly by having a thread for the
+    // ProcessGroupNCCL class. Attempting to modify the communicator cache from
+    // the WorkNCCL class might run into issues with object lifetime since the
+    // ProcessGroupNCCL object might get destroyed before the WorkNCCL object.
+    void run();
+
+    // Watchdog's inside loop.
+    // Takes care of cleaning up completed work, and aborting upon failure or
+    // timeout.
+    void runLoop();
+
+    // Notify the loop inside watchdog.
+    void notify();
+
+    void checkAndSetRemoteError();
+
+    // A helper function to get the src rank of a signal from the Store. This is
+    // nonblocking function returning -1 if the signal is not available yet.
+    int getSignalSrcRank(
+        c10::intrusive_ptr<Store>& store,
+        const std::string& signal);
+
+    uint64_t getHeartbt() const;
+
+    void setDesyncDebug(bool desyncDebug);
+
+   private:
+    std::thread ncclCommWatchdogThread_;
+
+    // We need to keep a reference to the PG instance so that we can access
+    // the member functions of the PG instance. We store a raw pointer on
+    // purpose because the watchdog thread now still lives within the
+    // lifetime of the PG instance.
+    ProcessGroupNCCL* pg_;
+
+    // Whether the NCCL watchdog should rethrow CUDA errors.
+    bool rethrowCUDAErrors_ = false;
+
+    std::exception_ptr watchDogException_ = nullptr;
+
+    // Condition Variable for watchdog thread sleep
+    std::condition_variable workMetaListCV_;
+
+    // Heartbeat of watchdog thread.
+    std::atomic_uint64_t heartbeat_;
+
+    // Whether or not to propagate detected errors to all ranks in the same PG
+    // through TCPStore.
+    bool propagatePgError_;
+
+    // Whether or not to enable timeout root cause analysis.
+    bool desyncDebug_;
+
+    DesyncDebugger desyncDebugger_;
+  };
+
+  // If you wish to create multiple process groups, each with a potentially
+  // different rank and size, you can do so by passing a new store instance
+  // to each one. If you have only a single store object, you can
+  // use the `c10d::PrefixStore` to derive scoped instances.
+  // This is also what the Python API in torch.distributed does.
+  //
+  // The process group instance keeps a reference to the store because
+  // it may be used long after the constructor runs. In fact, the constructor
+  // doesn't create any NCCL communicators. A single NCCL communicator can
+  // only be used on a specific set of devices, and are therefore created
+  // on-demand when a collective runs. If another collective is executed later,
+  // against a different set of devices, the process group creates another NCCL
+  // communicator. These NCCL communicators are cached and reused if possible.
+  //
+  ProcessGroupNCCL(
+      c10::intrusive_ptr<Store> store,
+      int rank,
+      int size,
+      c10::intrusive_ptr<Options> options = Options::create());
+
+  // This constructor includes the deprecated `groupName` argument.
+  // If you have existing code that uses the `groupName`, you can replace
+  // it by specifying a `c10d::PrefixStore(groupName, store)` for store.
+  C10_DEPRECATED ProcessGroupNCCL(
+      const c10::intrusive_ptr<Store>& store,
+      int rank,
+      int size,
+      const std::string& groupName,
+      c10::intrusive_ptr<Options> options = Options::create())
+      : ProcessGroupNCCL(store, rank, size, std::move(options)) {}
+
+  ~ProcessGroupNCCL() override;
+
+  // This function returns a local uid for ProcessGroupNCCL.
+  uint64_t getUid() {
+    return static_cast<uint64_t>(local_id_);
+  }
+
+  c10::intrusive_ptr<Options> getOptions() {
+    return options_;
+  }
+
+  c10::intrusive_ptr<Backend::Options> getBackendOptions() override {
+    return c10::static_intrusive_pointer_cast<Backend::Options>(options_);
+  }
+
+  const std::string getBackendName() const override {
+    return std::string(NCCL_BACKEND_NAME);
+  }
+
+  bool supportsSplitting() const override {
+    return true;
+  }
+
+  bool supportsCoalescing() const override {
+    return true;
+  }
+
+  bool supportsTimeEstimation() const override {
+#ifdef NCCL_SIM_INFO_INITIALIZER
+    return true;
+#else
+    return false;
+#endif
+  }
+
+  void setTimeout(std::chrono::milliseconds timeout) override {
+    options_->timeout = timeout;
+  }
+
+  void startCoalescing() override;
+
+  c10::intrusive_ptr<Work> endCoalescing() override;
+
+  void startTimeEstimate();
+
+  float endTimeEstimate();
+
+  // For specifying a composite optype, such as ALLGATHER and REDUCE_SCATTER
+  c10::intrusive_ptr<Work> endCoalescing(OpType optype);
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& tensors,
+      const BroadcastOptions& opts = BroadcastOptions()) override;
+
+  c10::intrusive_ptr<Work> _broadcast_oop(
+      at::Tensor& outputTensors,
+      at::Tensor& inputTensors,
+      const BroadcastOptions& opts = BroadcastOptions());
+
+  c10::intrusive_ptr<Work> allreduce_sparse(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceCoalescedOptions& opts =
+          AllreduceCoalescedOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& tensors,
+      const ReduceOptions& opts = ReduceOptions()) override;
+
+  c10::intrusive_ptr<Work> _reduce_oop(
+      at::Tensor& outputTensors,
+      at::Tensor& inputTensors,
+      const ReduceOptions& opts = ReduceOptions());
+
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& outputbuffer,
+      at::Tensor& inputbuffer,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_coalesced(
+      std::vector<std::vector<at::Tensor>>& outputTensorLists,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      std::vector<int64_t>& outputSplitSizes,
+      std::vector<int64_t>& inputSplitSizes,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) override;
+
+  int64_t getCommPtr();
+
+  void groupStart();
+
+  void groupEnd();
+
+  void groupEndNonblocking(const std::shared_ptr<NCCLComm>& comm);
+
+  c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const GatherOptions& opts = GatherOptions()) override;
+
+  c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ScatterOptions& opts = ScatterOptions()) override;
+
+  // Unsupported Ops
+  c10::intrusive_ptr<Work> recvAnysource(
+      std::vector<at::Tensor>& tensors,
+      int tag) override;
+
+  // Agrees on an initial sequence number for the whole group by having rank 0
+  // create it and broadcast it to other ranks using the store.
+  void setSequenceNumberForGroup() override;
+
+  // Retrieves the current sequence number for the whole group, which should be
+  // in sync. If the returned number is not consistent across the group, it
+  // may indicate that there is some sort of collective desynchronization.
+  uint64_t getSequenceNumberForGroup() override;
+
+  // Return the total number of splits the communicators held by this process
+  // group have performed.  Counts ncclCommCreateFromRanks() for ncclx v2.21.5+
+  uint64_t getCommSplitCounter() const;
+
+  void registerOnCompletionHook(
+      std::function<void(std::shared_ptr<WorkInfo>)>&& hook) override;
+  void waitForPendingWorks() override;
+
+  void enableCollectivesTiming() override;
+
+  c10::intrusive_ptr<Backend> split(
+      const c10::intrusive_ptr<Store>& store,
+      const std::vector<int>& ranks,
+      const c10::intrusive_ptr<Backend::Options>& opts) override;
+
+  c10::intrusive_ptr<Backend> merge(
+      const c10::intrusive_ptr<Store>& store,
+      const c10::intrusive_ptr<Backend::Options>& opts,
+      const int& rank,
+      const int& size) override;
+
+  // Helper function for iteratively aborting communicators in the provided map
+  void abortCommsFromMap(
+      std::unordered_map<std::string, std::shared_ptr<NCCLComm>>& ncclCommsMap,
+      const std::optional<std::string>& abortReason);
+
+  c10::intrusive_ptr<intra_node_comm::IntraNodeComm> initIntraNodeComm();
+
+  // Destroy (shutdown) this backend -- normal exit.
+  void shutdown() override;
+
+  // Provides an API to abort the ProcessGroup (similar to ncclCommAbort)
+  // instead of relying on ProcessGroupNCCL destructor.
+  void abort() override;
+
+  void eagerConnectSingleDevice(at::Device device) override;
+
+  void performNocolorSplit(at::Device device);
+
+  // If all comms on this PG are fully initialized, return true.
+  bool isInitialized();
+
+  ErrorType getError() override;
+
+  bool supportsShrinking() const override {
+#ifdef NCCL_HAS_COMM_SHRINK
+    return true;
+#else
+    return false;
+#endif
+  }
+
+  // Backend-style shrink override that returns a Backend instance.
+  c10::intrusive_ptr<Backend> shrink(
+      const std::vector<int64_t>& ranks_to_exclude,
+      int shrink_flags = 0,
+      const c10::intrusive_ptr<Backend::Options>& opts_override =
+          nullptr) override;
+
+  std::shared_ptr<c10::Allocator> getMemAllocator() override;
+
+  // Allocate tensor from communication-optimized memory pool
+  at::Tensor allocateTensor(long size, at::TensorOptions options = {}) override;
+
+  // Whether tensor allocation from NCCL memory pool is supported
+  bool supportsTensorAlloc(c10::DeviceIndex deviceIdx) override;
+
+  // Performs NCCL user buffer registration for all buffers in
+  // the given MemPool
+  void registerMemPool(at::cuda::MemPool* pool, bool symm = false);
+
+  // Performs NCCL user buffer de-registration for all buffers in
+  // the given MemPool
+  void deregisterMemPool(at::cuda::MemPool* pool);
+
+  // This method adds a temporary extension for the timeout period,
+  // applying to all collectives between the calling of this API and
+  // the completion of the first collective on the GPU. While this feature
+  // provides flexibility in specific scenarios, it introduces statefulness
+  // to timeout setting. Therefore, it is advisable to use this API sparingly
+  // and consider alternative approaches, such as directly setting the timeout
+  // or utilizing a barrier collective (one can set any timeout to the barrier),
+  // whenever feasible.
+  void addEphemeralTimeout(const std::chrono::milliseconds& timeout);
+
+  // This function is only intended for testing purposes because we don't
+  // want to expose the `WorkNCCL` via pybind. It verifies whether the
+  // `opTimeout_` of the provided WorkNCCL instance is the same as the specified
+  // timeout.
+  bool verifyWorkTimeoutForTest(
+      const c10::intrusive_ptr<Work>& work,
+      const std::chrono::milliseconds& timeout);
+
+  void setEnableNanCheck(bool enableNanCheck);
+
+  // APIs related to memory offload (require NCCL 2.29.7+ at runtime)
+  void suspend() override;
+
+  void resume() override;
+
+  std::unordered_map<std::string, uint64_t> getMemoryStats() override;
+
+ protected:
+  uint64_t getWatchdogHeartbt() const;
+
+  // Instance of the heartbeat monitor thread.
+  std::unique_ptr<HeartbeatMonitor> heartbeatMonitor_;
+
+  // Instance of the watchdog thread.
+  std::unique_ptr<Watchdog> watchdog_;
+
+  // Helper that broadcasts nccl unique ID to all ranks through the store
+  void broadcastUniqueNCCLID(
+      ncclUniqueId* ncclID,
+      bool isSingleP2POp,
+      const std::string& devicesKey,
+      int p2pRank);
+
+  // Helper that allgathers nccl unique IDs to all ranks through the store
+  void allgatherUniqueNCCLIDs(
+      int rootIdx,
+      ncclUniqueId* ncclID,
+      std::vector<ncclUniqueId>& ncclIDs);
+
+  // Helper that looks up the cached NCCL communicators only
+  std::shared_ptr<NCCLComm> getNCCLComm(const std::string& deviceKey);
+
+  std::shared_ptr<NCCLComm> initNCCLComm(
+      const std::string& deviceKey,
+      at::Device& device,
+      OpType opType,
+      int p2pRank = 0,
+      bool isSendRecvSelf = false);
+
+  // Initialize device-specific state (comm, stream, event, bookkeeping) for a
+  // given communicator on this process group instance.
+  void initializeDeviceStateForComm(
+      const at::Device& device,
+      std::shared_ptr<NCCLComm> comm);
+
+  // Wrapper method which can be overridden for tests.
+  virtual std::exception_ptr checkForNCCLErrors(
+      std::shared_ptr<NCCLComm>& ncclComm);
+
+  // Ensure thaht if record is True, the work obj will be enqueued via
+  // workEnqueue
+  virtual c10::intrusive_ptr<ProcessGroupNCCL::WorkNCCL> initWork(
+      at::Device& device,
+      int rank,
+      OpType opType,
+      bool isP2P,
+      const char* profilingTitle = nullptr,
+      const std::vector<at::Tensor>& inputs = {},
+      const std::vector<at::Tensor>& outputs = {},
+      bool record = false);
+
+  // In the timeout case and we will dump debug info such as the NCCL flight
+  // recorder to storage. Down the road, if we have more complicated or blocking
+  // operations, we might need to use a side thread to do it.
+  bool dumpDebuggingInfo(
+      bool includeStackTrace = true,
+      bool onlyActive = false);
+
+  void dumpExtraDebuggingInfo();
+
+  // Abort all communicators on this rank.
+  bool abortComms(const std::optional<std::string>& abortReason = std::nullopt);
+
+  // A helper function to check if nonblocking API mode should be used.
+  // Use this helper instead of directly checking `useNonblocking_` variable.
+  bool useNonblocking();
+
+ protected:
+  int globalRankStart_{};
+  int globalRankStride_{};
+
+ private:
+  bool eagerInit_{false};
+  bool showSerializationWarning_{true};
+
+  // Helper that encapsulates work shared across all collective communication
+  // primitives.  The callbacks have the following signatures:
+  //
+  //    ncclResult_t fn(at::Tensor& input, at::Tensor& output,
+  //                    ncclComm_t, at::cuda::CUDAStream&);
+  //    void {pre,post}(std::vector<at::cuda::CUDAStream&>);
+  template <typename Fn>
+  c10::intrusive_ptr<Work> collective(
+      at::Tensor& input,
+      at::Tensor& output,
+      Fn fn,
+      OpType opType,
+      bool asyncOp,
+      const char* profilingTitle = nullptr,
+      bool nanCheck = true);
+
+  template <typename Fn, typename PreProcess, typename PostProcess>
+  c10::intrusive_ptr<Work> collective(
+      at::Tensor& input,
+      at::Tensor& output,
+      Fn fn,
+      PreProcess pre,
+      PostProcess post,
+      OpType opType,
+      bool asyncOp,
+      const char* profilingTitle = nullptr,
+      bool nanCheck = true);
+
+  template <typename Fn, typename PreProcess, typename PostProcess>
+  c10::intrusive_ptr<Work> collective(
+      std::vector<at::Tensor>& inputs,
+      std::vector<at::Tensor>& outputs,
+      Fn fn,
+      PreProcess pre,
+      PostProcess post,
+      OpType opType,
+      bool asyncOp,
+      const char* profilingTitle = nullptr,
+      bool nanCheck = true);
+
+  template <typename Fn>
+  c10::intrusive_ptr<Work> collectiveCoalesced(
+      std::vector<at::Tensor>& input,
+      std::vector<at::Tensor>& output,
+      Fn fn,
+      OpType opType,
+      bool asyncOp,
+      const char* profilingTitle = nullptr);
+
+  // Helper that encapsulates work shared across point-to-point communication
+  // primitives. It is the same structure as the helper used for collective
+  // communication primitives.
+  template <typename Fn>
+  c10::intrusive_ptr<Work> pointToPoint(
+      at::Tensor& tensor,
+      Fn fn,
+      int peer,
+      OpType opType,
+      const char* profilingTitle = nullptr);
+
+  template <typename Fn, typename PreProcess, typename PostProcess>
+  c10::intrusive_ptr<Work> pointToPoint(
+      at::Tensor& tensor,
+      Fn fn,
+      int peer,
+      OpType opType,
+      PreProcess pre,
+      PostProcess post,
+      const char* profilingTitle);
+
+  c10::intrusive_ptr<Work> allreduce_impl(
+      at::Tensor& tensor,
+      const char* profilingTitle = "nccl:all_reduce",
+      const AllreduceOptions& opts = AllreduceOptions());
+
+  // Checks for NCCL errors on each of the communicators and returns an
+  // appropriate exception_ptr (nullptr if no errors).
+  static std::exception_ptr checkForNCCLErrorsInternal(
+      std::shared_ptr<NCCLComm>& ncclComm);
+
+  // Return the CUDA device most likely associated with this backend.
+  // If we aren't bound to a specific device, there is no strict
+  // guarantee that this heuristic is the correct assignment of ranks
+  // to GPUs that Python layers use, but in practice it tends to be.
+  // Fortunately we don't rely on this for correctness of any tensor
+  // operations, just for ancillary uses like barriers.
+  at::Device guessDeviceForRank() const;
+
+  // Destroys initialized NCCL communicators in devNCCLComMap_ given by input
+  // key. Throws if there are no communicators to destroy. Also removes
+  // communicators from the cache and clears used device indices.
+  void destroyNCCLComms(const std::string& devNCCLCommMapKey);
+
+  void runHookLoop();
+
+  // Generates a prefix that is unique to this process group and rank, for
+  // disambiguating logs
+  std::string createLogPrefix() const;
+
+  // Returns the unique prefix created in createLogPrefix
+  const std::string& logPrefix() const;
+
+  // Returns the global rank of the device. This function assumes that users
+  // always create a default global process group(PG) which includes all
+  // devices. It is called in the constructor of ProcessGroupNCCL, so it always
+  // return the rank_ of the very first PG created, aka, default global PG.
+  const int& globalRank() const;
+
+  const c10::intrusive_ptr<Store>& globalStore() const;
+
+  // Returns the global ranks of a PG.
+  const std::vector<uint64_t>& groupRanks() const;
+
+  // Util function to assign timeout to each work.
+  void assignTimeoutToWork(
+      const c10::intrusive_ptr<ProcessGroupNCCL::WorkNCCL>& work,
+      const c10::intrusive_ptr<Options>& option);
+
+  // Broadcast flight-recorder dump signal
+  void broadcastDumpSignal();
+
+  // A helper function to broadcast a signal (key) from a src rank to all other
+  // ranks using the specified store.
+  void broadcastSignal(
+      c10::intrusive_ptr<Store>& store,
+      const std::string& signal,
+      int srcRank);
+
+ protected:
+  // Function that directly trigger std::abort so that the whole process
+  // gets terminated.
+  virtual void terminateProcess(const std::string& errMsg);
+
+  // A helper function to wait for a future to complete or timeout.
+  // Returns true if the future completes before timeout, false otherwise.
+  bool waitForFutureOrTimeout(
+      std::future<bool>& fut,
+      const std::chrono::milliseconds& timeOutMilSec,
+      const std::string& futDescription,
+      ::c10d::C10dLoggingData& debugLog,
+      bool throwException = false);
+
+  // A helper function to guess the device id of the current rank, based on
+  // bounded device or used device. Do not use this function if you already know
+  // the device id to operate on.
+  c10::DeviceIndex guessDeviceId() const;
+
+  static const int64_t kWatchdogThreadSleepMillis;
+
+  // The store is used to broadcast the NCCL unique ID of rank 0. This store
+  // comes with prefix and it is different across ProcessGroup NCCL instances
+  // (aka, different ProcessGroups).
+  c10::intrusive_ptr<Store> store_;
+
+  // Reference to the store without prefix so that keys are same across all
+  // ProcessGroup NCCL instances and (key, value) pairs written to the store are
+  // global.
+  c10::intrusive_ptr<Store> globalStore_;
+
+  // The lock which protects the write/read of
+  // ephemeralTimeoutActive_/ephemeralTimeoutInflight_.
+  // TODO(fduwjj): We need to have an audit on all mutexes we are adding here.
+  // And consolidate them if possible.
+  std::mutex mtxTimeoutExtension_;
+
+  // The ephemeral timeout added on top of existing timeout for works issued
+  // before first work finishes.
+  std::chrono::milliseconds ephemeralTimeoutActive_ =
+      std::chrono::milliseconds(0);
+
+  // The ephemeral timeout addition which has been already applied to work.
+  std::chrono::milliseconds ephemeralTimeoutInflight_ =
+      std::chrono::milliseconds(0);
+
+  const c10::intrusive_ptr<Options> options_;
+
+  // The number of NCCL communicators that have been created during
+  // the lifetime of this process group. This sequence number is
+  // used to scope keys used in the store.
+  uint64_t ncclCommCounter_{0};
+
+  // The NCCL communicator that the process group has cached.
+  //
+  // For collective operations:
+  // The key is a list of GPU devices that an operation is operating on
+  // The GPU devices are stored in a device sequence and the cache NCCL
+  // communicator is associated with this GPU device sequence
+  //
+  // e.g. If the process group op only uses device 0, then the value of
+  // the used device string stored (value of the hashmap) would be "0".
+  //
+  //      If the process group op uses device 0 - 7 and the each tensor of the
+  //      input tensor list is on device, 0, 1, 2, 3, 4, 5, 6, 7 separately,
+  //      then the value of the used device string (key) stored would be
+  //      "0,1,2,3,4,5,6,7"
+  //
+  //      If the process group op uses device 0 - 7 and the each tensor of the
+  //      input tensor list is on device, 0, 4, 5, 6, 7, 1, 2, 3 separately,
+  //      then the value of the used device string stored would be
+  //      "0,4,5,6,7,1,2,3"
+  //
+  //      Note that the order of the device for the tensor list matters.
+  //
+  // For point-to-point operations:
+  // The key is a string of my current rank and the peer process rank.
+  // e.g. If process 1 and process 2 are involved in a point-to-point
+  // communication, the key will be "1:2" on both processes. Note: this is for
+  // the scenario where there is only 1 GPU per process. When it comes to
+  // multiple GPUs per process, this part may need to redesigned.
+  // TODO: we probably need a separate map for P2P comms
+  std::unordered_map<std::string, std::shared_ptr<NCCLComm>> devNCCLCommMap_;
+
+  // The NCCL communicators currently in process of being initialized.
+  std::unordered_map<std::string, std::shared_ptr<NCCLComm>>
+      inInitializationCommMap_;
+
+  // Mutex to guard maps like devNCCLCommMap_.
+  std::mutex mutex_;
+
+  // Size of ring buffer where we store NCCL Traces for debugging.
+  int traceBufferSize_;
+
+  // Stores TORCH_NCCL_DEBUG_INFO_PIPE_FILE
+  std::string debugInfoPipeFile_;
+
+  // We gate the cudaEventCache so that we can roll it out gradually.
+  std::atomic<bool> cudaEventCacheEnabled_;
+
+  std::thread onCompletionHookThread_;
+
+  // Whether or not we should terminate the watchdog and workCleanup threads.
+  std::atomic<bool> terminateProcessGroup_;
+
+  // Whether there are hooks pending to be fired
+  std::atomic<bool> hasPendingHooks_;
+
+  // This is the signal from watchdog threads to indicate whether the monitor
+  // thread should dump. Making it static so that it is accessible from all the
+  // PGs. With this flag, monitor thread would dump debug info under any one of
+  // the three conditions:
+  //
+  // 1: watchdog thread of any PG detects a collective timeout.
+  // 2: timeout signal is received from other ranks through tcpstore.
+  // 3: current PG's watchdog heartbeat timeout occurs.
+  //
+  // Note that only the monitor thread from PG0 will dump the debug info for
+  // case one and two so that the debug info is only dumped once.
+  static std::atomic<bool> shouldDump_;
+
+  // Mutex to Guard workMetaList_
+  std::mutex workMetaListMutex_;
+
+  bool writeDebugInfo_ = false;
+
+  // Vector to store WorkNCCL pointers
+  std::list<ProcessGroupNCCL::WorkNCCL> workMetaList_;
+
+  // Mutex to Guard workMetaList_
+  std::mutex completedWorkListMutex_;
+
+  // Condition Variable for watchdog thread sleep
+  std::condition_variable completedWorkListCV_;
+
+  std::list<ProcessGroupNCCL::WorkNCCL> completedWorkList_;
+
+  // Add Work Pointer to workVector
+  void workEnqueue(
+      const c10::intrusive_ptr<ProcessGroupNCCL::WorkNCCL>& /*work*/);
+
+  // The CUDA streams used by NCCL kernels
+  std::unordered_map<std::string, at::cuda::CUDAStream> ncclStreams_;
+
+  // The CUDA events used to sync NCCL streams
+  std::unordered_map<std::string, at::cuda::CUDAEvent> ncclEvents_;
+
+  // Device Indexes used for all collectives in this group
+  std::set<c10::DeviceIndex> usedDeviceIdxs_;
+
+  // Flag to denote if a coalescing groupStart/groupEnd block is active
+  int coalescing_state_ = 0;
+
+  // Stores device indexes for all collectives run inside a coalescing block
+  at::Device coalescedDevice_ = at::Device("cuda");
+
+  // Stores communicators for all collectives run inside a coalescing block
+  std::shared_ptr<NCCLComm> coalescedComm_ = nullptr;
+
+  // Whether the coalesced calls are sync or async.
+  bool coalescedAsync_{};
+
+  // keeps track of input and output tensors when coalescing is in flight.  Will
+  // hand over these tensors to WorkNCCL's stash when coalescing is ended.
+  TensorShelf coalescedTensors_;
+
+  // Some ops may have completed, but user still hasn't called `work.wait()`.
+  // When watchdog detects this, it transfers the TensorShelf from `work` to
+  // this `shelves` structure. Next time we execute ProcessGroupNCCL's methods
+  // on main thread, we clear the `shelves` in one shot. This is mainly because
+  // watchdog (a side thread) unstashing the shelf directly seems to cause some
+  // problem.
+  std::vector<std::shared_ptr<TensorShelf>> shelvesToUnstash_;
+  std::mutex shelvesMutex_;
+
+  // Whether or not wait() and synchronize() are blocking operations that wait
+  // for the operation to complete.
+  bool blockingWait_ = false;
+
+  // Whether or not the workCleanupThread is used to perform async error
+  // handling.
+  ErrorHandlingMode asyncErrorHandling_ = NoHandling;
+
+  ErrorType error_ = ErrorType::SUCCESS;
+
+  std::mutex errorMutex_;
+
+  // Whether or not to sleep after an exception is thrown in the watchdog.
+  bool sleepAfterException_{};
+
+  // Whether or not to enable nan check for input tensors to collectives.
+  bool enableNanCheck_;
+
+  // Whether or not to create start CUDAEvent and enable timing for start
+  // and end events. Note that enableTiming_ is always true if desyncDebug_
+  // is set to true.
+  std::atomic<bool> enableTiming_;
+
+  // Flag to enable the print of hash value of input/output of collectives for
+  // verification.
+  std::atomic<bool> enableCollectiveHashDebug_;
+
+  // Whether or not TORCH_NCCL_AVOID_RECORD_STREAMS was set
+  bool avoidRecordStreams_ = false;
+
+  // The number of active ncclGroupStart() calls. This counter will be increased
+  // by 1 when ncclGroupStart() is called and decreased by 1 when ncclGroupEnd()
+  // is called.
+  static thread_local uint64_t ncclActiveGroupCounter_;
+
+  // Counting for the sequential number of NCCL collective call.
+  // (specifically, how many actual kernels we launched, which differs from
+  // op_id_ when coalescing is enabled)
+  uint64_t seqCollective_{0};
+
+  // Counting for the sequential number of NCCL P2P calls.
+  uint64_t seqP2P_{0};
+
+  // Incrementing counter for logical operations (collective or p2p) issued on
+  // the ProcessGroup
+  uint64_t op_id_{0};
+
+  // The number of ProcessGroupNCCL created on the current rank.
+  size_t local_id_;
+
+  std::string logPrefix_;
+
+  c10::intrusive_ptr<intra_node_comm::IntraNodeComm> intraNodeComm_;
+
+  // Number of devices on this node.
+  int localDeviceCount_{0};
+
+  std::shared_ptr<ProcessGroupStatus> pgStatus_ =
+      std::make_shared<ProcessGroupStatus>();
+
+  // Internal cached value: use NCCL non-blocking API mode or not.
+  // Use `useNonblocking()` method instead of accessing this variable directly.
+  std::optional<bool> useNonblocking_{std::nullopt};
+
+  // Communication-optimized memory pool associated with this PG
+  std::unique_ptr<at::cuda::MemPool> memPool_ = nullptr;
+};
+
+// Reset the flighrecorder recordings for the current rank.
+TORCH_API void reset_nccl_trace();
+
+// Dumps the NCCL comm traces and additional information about the Process
+// Group.
+TORCH_API std::string dump_nccl_trace(
+    bool includeCollectives,
+    bool includeStackTraces,
+    bool onlyActive);
+
+// Dumps the NCCL comm traces and additional information about the Process
+// Group in JSON formatted string.
+// We don't include stack traces in JSON format as it is far too much data.
+TORCH_API std::string dump_nccl_trace_json(
+    bool includeCollectives,
+    bool onlyActive);
+
+// Gets a mutable reference to a global optional function.Heartbeat Monitor
+// will use this function to dump traces, if available. Inside fbcode, we
+// store a function here that uses an internal tool for process tracing
+TORCH_API std::optional<
+    std::function<void(std::function<void(const std::string&)>)>>&
+get_cpp_trace_dumper();
+
+// Similar to get_cpp_trace_dumper, this stores a function defined in
+// torch-python layer that lets us check whether the GIL can be acquired,
+// helpful for instrumenting in cases where a hang was observed.
+typedef bool (*gil_checker_t)();
+
+TORCH_API gil_checker_t& get_gil_checker();
+} // namespace c10d
+
+#endif // USE_C10D_NCCL
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupUCC.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupUCC.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..12e737b61df23a0fc5503b60fc7a6136d311aaeb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupUCC.hpp
@@ -0,0 +1,363 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_UCC
+
+#include <torch/csrc/distributed/c10d/UCCUtils.hpp>
+
+#include <exception>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <thread>
+#include <vector>
+
+#include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+#ifdef USE_CUDA
+#include <ATen/cuda/CUDAEvent.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
+
+namespace c10d {
+
+#define TORCH_UCC_DEVICE_NOT_SET -2
+
+#ifdef USE_CUDA
+#define SAVE_TENSORS(_TENSORS, _DATA)                       \
+  do {                                                      \
+    if ((_TENSORS)[0].device().is_cuda()) {                 \
+      for (const auto i : c10::irange((_TENSORS).size())) { \
+        c10::cuda::CUDACachingAllocator::recordStream(      \
+            (_TENSORS)[i].storage().data_ptr(), (*stream)); \
+      }                                                     \
+    } else {                                                \
+      (_DATA) = (_TENSORS);                                 \
+    }                                                       \
+  } while (0)
+
+#else
+#define SAVE_TENSORS(_TENSORS, _DATA) (_DATA) = (_TENSORS);
+#endif
+
+constexpr const char* UCC_BACKEND_NAME = "ucc";
+
+struct event_pool_t {
+#ifdef USE_CUDA
+  std::queue<std::unique_ptr<at::cuda::CUDAEvent>> event_pool;
+#endif
+  std::mutex event_pool_mutex;
+};
+
+class Comm;
+
+// UCC does not support multiple CUDA devices per process.
+class TORCH_API ProcessGroupUCC : public Backend {
+ private:
+  void set_timeout(ucc_coll_args_t& args);
+
+ public:
+  class WorkData {
+   public:
+    std::vector<at::Tensor> src;
+    std::vector<at::Tensor> dst;
+    std::vector<at::Tensor> flat;
+    WorkData() {}
+    virtual ~WorkData() = default;
+  };
+  class AlltoallWorkData : public WorkData {
+   public:
+    AlltoallWorkData(int size)
+        : send_lengths(size),
+          send_offsets(size),
+          recv_lengths(size),
+          recv_offsets(size) {}
+    std::vector<uint64_t> send_lengths;
+    std::vector<uint64_t> send_offsets;
+    std::vector<uint64_t> recv_lengths;
+    std::vector<uint64_t> recv_offsets;
+  };
+
+  class AllgathervWorkData : public WorkData {
+   public:
+    AllgathervWorkData(int size) : recv_lengths(size), recv_offsets(size) {}
+    std::vector<uint64_t> recv_lengths;
+    std::vector<uint64_t> recv_offsets;
+  };
+
+  class ScattervWorkData : public WorkData {
+   public:
+    ScattervWorkData(int size) : send_lengths(size), send_offsets(size) {}
+    std::vector<uint64_t> send_lengths;
+    std::vector<uint64_t> send_offsets;
+  };
+
+  class ProgressEntry {
+    friend class ProcessGroupUCC;
+    friend class Comm;
+
+   public:
+    ProgressEntry(CommBase* comm, ucc_coll_req_h request)
+        : status_(UCC_INPROGRESS), comm_(comm), request_(request) {}
+    // Finalizes UCC status or exception of collective request.
+    void finalize(std::exception_ptr eptr = nullptr);
+    ucc_status_t status_;
+    CommBase* comm_;
+    ucc_coll_req_h request_;
+    std::unique_ptr<WorkData> data;
+    c10::intrusive_ptr<c10::ivalue::Future> future_;
+    std::exception_ptr eptr_;
+  };
+
+  class WorkUCC : public Work {
+    friend class ProcessGroupUCC;
+    friend class Comm;
+
+   public:
+    WorkUCC(
+        OpType opType,
+        uint64_t seq,
+        const char* prof_title,
+        const std::optional<std::vector<at::Tensor>>& inputs,
+        const c10::intrusive_ptr<ProcessGroupUCCLogger>& logger)
+        : Work(-1, opType, prof_title, inputs), logger_(logger), seq_(seq) {}
+    ~WorkUCC();
+    void setException();
+    void setAndThrowException();
+    bool isCompleted() override;
+    bool isSuccess() const override;
+    bool wait(std::chrono::milliseconds timeout = kUnsetTimeout) override;
+    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
+    std::vector<at::Tensor> result() override;
+    int sourceRank() const override;
+#ifdef USE_CUDA
+    std::unique_ptr<at::cuda::CUDAEvent> fence = nullptr;
+    event_pool_t* ep = nullptr;
+#endif
+    int sourceRank_;
+
+   protected:
+    std::shared_ptr<ProgressEntry> entry_;
+    c10::intrusive_ptr<ProcessGroupUCCLogger> logger_;
+    uint64_t seq_;
+
+   private:
+    // The future returned by getFuture.
+    c10::intrusive_ptr<at::ivalue::Future> future_;
+    // Store a reference to collective's outputs, used by result
+    std::shared_ptr<std::vector<at::Tensor>> outputs_;
+  };
+
+  explicit ProcessGroupUCC(
+      const c10::intrusive_ptr<Store>& store,
+      int rank = -1,
+      int size = -1,
+      std::chrono::duration<float> timeout = kBackendDefaultTimeout);
+
+  void initComm(c10::Device dev);
+
+  ~ProcessGroupUCC() override;
+
+  const std::string getBackendName() const override {
+    return std::string(UCC_BACKEND_NAME);
+  }
+
+#ifdef USE_CUDA
+  std::unique_ptr<at::cuda::CUDAEvent> getPooledEvent();
+#endif
+
+  // Performs a health check by initializing dummy UCC & UCX communicators and
+  // then destroying them. This will help indicate and signal any
+  // UCC/UCX-related issues prior to the first collective. The actual
+  // initialization and subsequent destruction is ran on a separate thread and
+  // the main thread is signalled about timeouts/errors to report to the
+  // application.
+  void runHealthCheck();
+
+  template <typename PreProcess, typename PostProcess>
+  c10::intrusive_ptr<Work> collective_post(
+      OpType opType,
+      PreProcess preproc,
+      PostProcess postproc,
+      ucc_coll_args_t& coll,
+      std::unique_ptr<ProcessGroupUCC::WorkData> data,
+      c10::Device dev,
+      std::vector<at::Tensor>& inputTensors,
+      std::vector<at::Tensor>& outputTensors,
+      const char* prof_title);
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& data,
+      const BroadcastOptions& opts = BroadcastOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceCoalescedOptions& opts =
+          AllreduceCoalescedOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& tensors,
+      const ReduceOptions& opts = ReduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) override;
+
+  c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const GatherOptions& opts = GatherOptions()) override;
+
+  c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ScatterOptions& opts = ScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      std::vector<int64_t>& outputSplitSizes,
+      std::vector<int64_t>& inputSplitSizes,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) override;
+
+  // Counting for the sequential number of UCC collective_post call.
+  uint64_t seq_{0};
+
+  // Agrees on an initial sequence number for the whole group by having rank 0
+  // create it and broadcast it to other ranks using the store.
+  void setSequenceNumberForGroup() override;
+
+  // Retrieves the current sequence number for the whole group, which should be
+  // in sync. If the returned number is not consistent across the group, it
+  // may indicate that there is some sort of collective desynchronization.
+  uint64_t getSequenceNumberForGroup() override;
+
+  static c10::intrusive_ptr<Backend> createProcessGroupUCC(
+      const c10::intrusive_ptr<::c10d::Store>& store,
+      int rank,
+      int size,
+      const std::chrono::duration<float>& timeout);
+
+ protected:
+  const std::chrono::duration<float> timeout_;
+  std::shared_ptr<torch_ucc_oob_coll_info_t> oob;
+  std::shared_ptr<Comm> comm = {nullptr};
+  uint32_t comm_id;
+  ucc_team_h team{nullptr};
+  ucc_ee_h cuda_ee{nullptr};
+  ucc_ee_h cuda_ee_p2p[2]{nullptr, nullptr};
+
+#ifdef USE_CUDA
+  std::unique_ptr<at::cuda::CUDAStream> stream = nullptr;
+  std::unique_ptr<at::cuda::CUDAStream> stream_p2p[2] = {nullptr, nullptr};
+  event_pool_t ep;
+#endif
+  c10::intrusive_ptr<ProcessGroupUCCLogger> logger;
+};
+
+class Comm {
+  c10::intrusive_ptr<ProcessGroupUCCLogger> logger;
+  std::shared_ptr<torch_ucc_oob_coll_info_t> oob;
+  CommUCC ucc_comm;
+  std::mutex mutex;
+  std::thread progress_thread;
+  std::condition_variable queue_produce_cv;
+  std::condition_variable queue_consume_cv;
+  std::deque<std::shared_ptr<ProcessGroupUCC::ProgressEntry>> progress_queue;
+  bool stop_progress_loop;
+  bool collective_inprogress;
+  torch_ucc_phase_t finalize_phase;
+
+ public:
+  c10::DeviceIndex cuda_device_index;
+  Comm(
+      const c10::intrusive_ptr<ProcessGroupUCCLogger>& logger,
+      std::shared_ptr<torch_ucc_oob_coll_info_t> oob,
+      c10::Device dev,
+      bool is_health_check);
+
+  ~Comm();
+
+  void ucc_create_team(
+      ucc_team_h& team,
+      std::shared_ptr<torch_ucc_oob_coll_info_t> oob);
+
+  void ucc_destroy_team(ucc_team_h& team);
+
+  c10::intrusive_ptr<Work> enqueue_p2p(
+      OpType opType,
+      ucc_coll_req_h request,
+      const char* prof_title);
+
+#ifdef USE_CUDA
+  void enqueue_cuda_collective(
+      std::unique_ptr<ProcessGroupUCC::WorkData> data,
+      c10::intrusive_ptr<ProcessGroupUCC::WorkUCC> work,
+      ucc_coll_args_t& coll,
+      ucc_team_h team,
+      ucc_ee_h ee);
+#endif
+
+  void enqueue_collective(
+      std::unique_ptr<ProcessGroupUCC::WorkData> data,
+      c10::intrusive_ptr<ProcessGroupUCC::WorkUCC> work,
+      ucc_coll_args_t& coll,
+      ucc_team_h team);
+
+  static std::shared_ptr<Comm> get_comm(
+      uint32_t& id,
+      c10::Device dev,
+      std::shared_ptr<torch_ucc_oob_coll_info_t> oob,
+      const c10::intrusive_ptr<ProcessGroupUCCLogger>& logger,
+      bool is_health_check = false);
+
+  void progress_loop();
+};
+
+} // namespace c10d
+
+#endif // USE_C10D_UCC
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..71f2ffed254a8441143bd0ecf3f921984f9dedf6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp
@@ -0,0 +1,204 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_GLOO
+
+#include <torch/csrc/distributed/c10d/ProcessGroupGloo.hpp>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+
+namespace c10d {
+
+// ProcessGroupWrapper wraps a Backend for debugging purposes. It intercepts
+// collective operations to verify consistency across ranks before dispatching
+// to the wrapped backend.
+//
+// IMPORTANT: This wrapper must forward all Backend virtual methods to backend_.
+// When adding new virtual methods to Backend that are overridden by backends
+// like ProcessGroupNCCL, you must also add forwarding methods here. Otherwise,
+// those methods will fail when TORCH_DISTRIBUTED_DEBUG=DETAIL is set.
+// See https://github.com/pytorch/pytorch/issues/173538 for an example.
+class TORCH_API ProcessGroupWrapper : public Backend {
+ public:
+  explicit ProcessGroupWrapper(
+      const c10::intrusive_ptr<Backend>& backend,
+      c10::intrusive_ptr<Backend> glooBackend);
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& data,
+      const BroadcastOptions& opts = BroadcastOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& data,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_sparse(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceCoalescedOptions& opts =
+          AllreduceCoalescedOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce(
+      std::vector<at::Tensor>& tensors,
+      const ReduceOptions& opts = ReduceOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> _allgather_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  // This function is deprecated and will be moved out of ProcessGroup to comms:
+  // * do not add dependencies on this function,
+  // * do not implement it in your ProcessGroup, implement _allgather_base
+  //   instead.
+  c10::intrusive_ptr<Work> allgather_coalesced(
+      std::vector<std::vector<at::Tensor>>& outputTensorLists,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const AllgatherOptions& opts = AllgatherOptions()) override;
+
+  c10::intrusive_ptr<Work> gather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const GatherOptions& opts = GatherOptions()) override;
+
+  c10::intrusive_ptr<Work> scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ScatterOptions& opts = ScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> _reduce_scatter_base(
+      at::Tensor& inputBuffer,
+      at::Tensor& outputBuffer,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& outputs,
+      std::vector<at::Tensor>& inputs,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& outputTensor,
+      at::Tensor& inputTensor,
+      std::vector<int64_t>& outputSplitSizes,
+      std::vector<int64_t>& inputSplitSizes,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  c10::intrusive_ptr<Work> alltoall(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllToAllOptions& opts = AllToAllOptions()) override;
+
+  void monitoredBarrier(const BarrierOptions& opts, bool waitAllRanks = false)
+      override;
+
+  // Agrees on an initial sequence number for the whole group by having rank 0
+  // create it and broadcast it to other ranks using the store. Only implemented
+  // for GLOO and NCCL backends currently.
+  // dont implement this
+  void setSequenceNumberForGroup() override;
+
+  // Retrieves the current sequence number for the whole group, which should be
+  // in sync. If the returned number is not consistent across the group, it
+  // may indicate that there is some sort of collective desynchronization.
+  uint64_t getSequenceNumberForGroup() override; // just call underlying
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> recvAnysource(
+      std::vector<at::Tensor>& tensors,
+      int tag) override;
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) override;
+  void registerOnCompletionHook(
+      std::function<void(std::shared_ptr<WorkInfo>)>&& hook) override;
+
+  void waitForPendingWorks() override;
+  void enableCollectivesTiming() override;
+
+  c10::intrusive_ptr<Backend> split(
+      const c10::intrusive_ptr<Store>& store,
+      const std::vector<int>& ranks,
+      const c10::intrusive_ptr<Options>& opts) override;
+
+  c10::intrusive_ptr<Backend> merge(
+      const c10::intrusive_ptr<Store>& store,
+      const c10::intrusive_ptr<Options>& opts,
+      const int& rank,
+      const int& size) override;
+
+  // Forward methods to wrapped backend
+  bool supportsSplitting() const override;
+  bool supportsCoalescing() const override;
+  bool supportsTimeEstimation() const override;
+  bool supportsShrinking() const override;
+  c10::intrusive_ptr<Backend> shrink(
+      const std::vector<int64_t>& ranks_to_exclude,
+      int shrink_flags = 0,
+      const c10::intrusive_ptr<Options>& opts_override = nullptr) override;
+  void setTimeout(std::chrono::milliseconds timeout) override;
+  void startCoalescing() override;
+  c10::intrusive_ptr<Work> endCoalescing() override;
+  const std::string getBackendName() const override;
+  c10::intrusive_ptr<Options> getBackendOptions() override;
+  std::shared_ptr<c10::Allocator> getMemAllocator() override;
+  at::Tensor allocateTensor(long size, at::TensorOptions options = {}) override;
+  bool supportsTensorAlloc(c10::DeviceIndex deviceIdx) override;
+  void abort() override;
+  void shutdown() override;
+  void suspend() override;
+  void resume() override;
+  std::unordered_map<std::string, uint64_t> getMemoryStats() override;
+
+  ErrorType getError() override;
+  void eagerConnectSingleDevice(at::Device device) override;
+
+  c10::intrusive_ptr<Backend> getWrappedPg() const;
+
+ private:
+  // Underlying process group that actual application collectives will be
+  // dispatched to
+  c10::intrusive_ptr<Backend> backend_;
+  // Gloo process group responsible for internal coordination such as monitored
+  // barrier, sequence number checking, collective fingerprint collecting.
+  c10::intrusive_ptr<Backend> glooBackend_;
+  // Conducts several checks to ensure that the underlying collective is well
+  // formed with the goal of notifying the user about incorrect collective use
+  // in the application.
+  void runCollectiveChecks(
+      OpType op_type,
+      const std::vector<at::Tensor>& tensors);
+};
+} // namespace c10d
+
+#endif // USE_C10D_GLOO
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/PyProcessGroup.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/PyProcessGroup.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..dd974d19037a5e7ae362dd5c4218c4bfcf4ea34f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/PyProcessGroup.hpp
@@ -0,0 +1,361 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace c10d {
+
+// PyProcessGroup is a pybind11 trampoline class to allow a Python
+// class to inherit from torch.distributed.ProcessGroup
+class PyProcessGroup : public ProcessGroup {
+ public:
+  // PyWork is a pybind11 trampoline class to allow a Python
+  // class to inherit from torch.distributed.Work
+  class TORCH_PYTHON_API PyWork : public Work {
+   public:
+    PyWork() = default;
+
+    bool wait(std::chrono::milliseconds timeout = kNoTimeout) override {
+      PYBIND11_OVERRIDE(
+          bool, /* Return type */
+          Work, /* Parent class */
+          wait, /* Name of function in C++ */
+          timeout);
+    }
+
+    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override {
+      // We cannot use PYBIND11_OVERRIDE because:
+      // 1. We have to >MANUALLY< unwrap the PyFutureWrapper and
+      // 2. The python name is get_future
+      pybind11::gil_scoped_acquire gil;
+      auto override =
+          pybind11::get_override(static_cast<const Work*>(this), "get_future");
+
+      if (override) {
+        py::object o = override();
+        auto futWrapper =
+            o.cast<std::shared_ptr<torch::jit::PythonFutureWrapper>>();
+        return futWrapper->fut;
+      }
+
+      return Work::getFuture();
+    }
+  };
+
+#define WORK_OVERRIDE(cname, name, ...)                                 \
+  do {                                                                  \
+    pybind11::gil_scoped_acquire gil;                                   \
+    pybind11::function override =                                       \
+        pybind11::get_override(static_cast<const cname*>(this), #name); \
+    if (override) {                                                     \
+      auto o = override(__VA_ARGS__);                                   \
+      return c10::make_intrusive<PyWorkHolder>(o);                      \
+    }                                                                   \
+    return cname::name(__VA_ARGS__);                                    \
+  } while (false)
+
+  // This class is used to wrap a PyWork trampoline with it's corresponding
+  // Python object to prevent the Python object from being garbage collected.
+  class PyWorkHolder : public Work {
+   public:
+    PyWorkHolder(const c10::intrusive_ptr<Work>& work, py::object pyWork)
+        : work_(work), pyWork_(std::move(pyWork)) {}
+
+    PyWorkHolder(py::object pyWork)
+        : work_(pyWork.cast<c10::intrusive_ptr<Work>>()),
+          pyWork_(std::move(pyWork)) {}
+
+    ~PyWorkHolder() override {
+      // GIL must be held when freeing python objects.
+      py::gil_scoped_acquire gil;
+      pyWork_ = py::object();
+    }
+
+    bool wait(std::chrono::milliseconds timeout = kNoTimeout) override {
+      return work_->wait(timeout);
+    }
+
+    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override {
+      return work_->getFuture();
+    }
+
+   private:
+    c10::intrusive_ptr<Work> work_;
+    py::object pyWork_;
+  };
+
+  using ProcessGroup::ProcessGroup;
+
+  const std::string getBackendName() const override {
+    PYBIND11_OVERRIDE(
+        std::string, /* Return type */
+        ProcessGroup, /* Parent class */
+        getBackendName, /* Name of function in C++ */
+    );
+  }
+
+  int getRank() const override {
+    PYBIND11_OVERRIDE(
+        int, /* Return type */
+        ProcessGroup, /* Parent class */
+        getRank, /* Name of function in C++ */
+    );
+  }
+
+  int getSize() const override {
+    PYBIND11_OVERRIDE(
+        int, /* Return type */
+        ProcessGroup, /* Parent class */
+        getSize, /* Name of function in C++ */
+    );
+  }
+
+  void abort() override {
+    PYBIND11_OVERRIDE(
+        void, /* Return type */
+        ProcessGroup, /* Parent class */
+        abort, /* Name of function in C++ */
+    );
+  }
+
+  const std::string& getGroupName() const override {
+    PYBIND11_OVERRIDE(
+        const std::string&, /* Return type */
+        ProcessGroup, /* Parent class */
+        getGroupName, /* Name of function in C++ */
+    );
+  }
+
+  void setGroupName(const std::string& group_name) override {
+    PYBIND11_OVERRIDE(
+        void, /* Return type */
+        ProcessGroup, /* Parent class */
+        setGroupName, /* Name of function in C++ */
+        group_name);
+  }
+
+  const std::string& getGroupDesc() const override {
+    PYBIND11_OVERRIDE(
+        const std::string&, /* Return type */
+        ProcessGroup, /* Parent class */
+        getGroupDesc, /* Name of function in C++ */
+    );
+  }
+
+  void setGroupDesc(const std::string& group_desc) override {
+    PYBIND11_OVERRIDE(
+        void, /* Return type */
+        ProcessGroup, /* Parent class */
+        setGroupDesc, /* Name of function in C++ */
+        group_desc);
+  }
+
+  c10::intrusive_ptr<ProcessGroup> splitGroup(
+      const std::vector<int>& ranks,
+      const std::optional<std::chrono::milliseconds>& timeout,
+      const std::optional<c10::intrusive_ptr<Backend::Options>>& opts,
+      const std::optional<std::string>& group_name,
+      const std::optional<std::string>& group_desc) override {
+    PYBIND11_OVERRIDE(
+        c10::intrusive_ptr<ProcessGroup>, /* Return type */
+        ProcessGroup, /* Parent class */
+        splitGroup, /* Name of function in C++ */
+        ranks,
+        timeout,
+        opts,
+        group_name,
+        group_desc);
+  }
+
+  c10::intrusive_ptr<ProcessGroup> mergeRemoteGroup(
+      const c10::intrusive_ptr<c10d::Store>& store,
+      const MergeOptions& opts,
+      const int& size) override {
+    PYBIND11_OVERRIDE(
+        c10::intrusive_ptr<ProcessGroup>, /* Return type */
+        ProcessGroup, /* Parent class */
+        mergeRemoteGroup, /* Name of function in C++ */
+        store,
+        opts,
+        size);
+  }
+
+  c10::intrusive_ptr<Work> allgather(
+      std::vector<std::vector<at::Tensor>>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        allgather, /* Name of function in C++ */
+        outputTensors,
+        inputTensors,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> allgather_into_tensor_coalesced(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const AllgatherOptions& opts = AllgatherOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        allgather_into_tensor_coalesced, /* Name of function in C++ */
+        outputTensors,
+        inputTensors,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> allreduce(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceOptions& opts = AllreduceOptions()) override {
+    WORK_OVERRIDE(
+        // py::object, /* Return type */
+        ProcessGroup, /* Parent class */
+        allreduce, /* Name of function in C++ */
+        tensors,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> allreduce_coalesced(
+      std::vector<at::Tensor>& tensors,
+      const AllreduceCoalescedOptions& opts =
+          AllreduceCoalescedOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        allreduce_coalesced, /* Name of function in C++ */
+        tensors,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> alltoall_base(
+      at::Tensor& outputBuffer,
+      at::Tensor& inputBuffer,
+      std::vector<int64_t>& outputSplitSizes,
+      std::vector<int64_t>& inputSplitSizes,
+      const AllToAllOptions& opts = AllToAllOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        alltoall_base, /* Name of function in C++ */
+        outputBuffer,
+        inputBuffer,
+        outputSplitSizes,
+        inputSplitSizes,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> barrier(
+      const BarrierOptions& opts = BarrierOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        barrier, /* Name of function in C++ */
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> broadcast(
+      std::vector<at::Tensor>& tensors,
+      const BroadcastOptions& opts = BroadcastOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        broadcast, /* Name of function in C++ */
+        tensors,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> reduce_scatter(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<std::vector<at::Tensor>>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        reduce_scatter, /* Name of function in C++ */
+        outputTensors,
+        inputTensors,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> reduce_scatter_tensor_coalesced(
+      std::vector<at::Tensor>& outputTensors,
+      std::vector<at::Tensor>& inputTensors,
+      const ReduceScatterOptions& opts = ReduceScatterOptions()) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        reduce_scatter_tensor_coalesced, /* Name of function in C++ */
+        outputTensors,
+        inputTensors,
+        opts);
+  }
+
+  c10::intrusive_ptr<Work> send(
+      std::vector<at::Tensor>& tensors,
+      int dstRank,
+      int tag) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        send, /* Name of function in C++ */
+        tensors,
+        dstRank,
+        tag);
+  }
+
+  c10::intrusive_ptr<Work> recv(
+      std::vector<at::Tensor>& tensors,
+      int srcRank,
+      int tag) override {
+    WORK_OVERRIDE(
+        ProcessGroup, /* Parent class */
+        recv, /* Name of function in C++ */
+        tensors,
+        srcRank,
+        tag);
+  }
+};
+
+class TORCH_PYTHON_API PythonOnCompletionHook {
+ public:
+  // Wraps a py::object hook and acquires Python GIL in dtor before
+  // destructing the hook object.
+  PythonOnCompletionHook(py::object hook) : hook_(std::move(hook)) {}
+  PythonOnCompletionHook(const PythonOnCompletionHook&) = default;
+
+  // NOLINTNEXTLINE(bugprone-exception-escape)
+  ~PythonOnCompletionHook() {
+    py::gil_scoped_acquire ag;
+    hook_.dec_ref();
+    // Explicitly set hook_ to nullptr to prevent py::object's dtor
+    // to decref on the PyObject again.
+    // See Note [Destructing py::object] in python_ivalue.h
+    hook_.ptr() = nullptr;
+  }
+
+  void operator()(const std::shared_ptr<WorkInfo>& workInfo) const {
+    std::exception_ptr eptr;
+    {
+      py::gil_scoped_acquire acquire;
+      try {
+        hook_(workInfo);
+      } catch (py::error_already_set& e) {
+        // py::error_already_set requires GIL to destruct, take
+        // special care.
+        eptr = std::make_exception_ptr(std::runtime_error(e.what()));
+        e.restore();
+        PyErr_Clear();
+      } catch (std::exception&) {
+        eptr = std::current_exception();
+      }
+    }
+    // No more Python-related stuff at this point, i.e., this
+    // exception can be captured and handled by PG backend.
+    if (eptr)
+      std::rethrow_exception(eptr);
+  }
+
+ private:
+  py::object hook_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/RankLocal.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/RankLocal.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..fbbde7e8ee787a1287dfe173f14128cf53e9c622
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/RankLocal.hpp
@@ -0,0 +1,94 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+
+#pragma once
+
+#include <shared_mutex>
+
+#include <torch/csrc/autograd/function.h>
+
+namespace c10d {
+
+// `RankLocal` maintains a unique instance of T for each non-autograd thread.
+// For non-autograd threads, `RankLocal<T>::get()` functions similar to
+// thread_local. For autograd threads, `RankLocal<T>::get()` returns the
+// instance of T corresponding to the enqueuing non-autograd thread. The
+// mechanism allows for rank-specific context shared between forward and
+// backward. It works for both the one-rank-per-process and one-rank-per-thread
+// scenarios.
+//
+// NOTE: RankLocal doesn't make the underlying objects thread-safe.
+template <typename T>
+class RankLocal {
+ public:
+  RankLocal(const RankLocal&) = delete;
+  RankLocal& operator=(const RankLocal&) = delete;
+
+  static T& get() {
+    // Fast path: non-autograd threads can simply return
+    // the object reference cached in TLS.
+    if (cached_ != nullptr) {
+      return *cached_;
+    }
+    const auto node = torch::autograd::get_current_node();
+    auto fwd_thread_id = node == nullptr ? at::RecordFunction::currentThreadId()
+                                         : node->thread_id();
+    // Optimistically acquire the read lock first, since most likely we are in
+    // an autograd thread and the object has already been constructed.
+    {
+      std::shared_lock read_lock(lock_);
+      auto it = thread_id_to_rank_local_.find(fwd_thread_id);
+      if (it != thread_id_to_rank_local_.end()) {
+        // Cache for non-autograd threads
+        if (node == nullptr) {
+          cached_ = &it->second;
+        }
+        return it->second;
+      }
+    }
+
+    std::unique_lock write_lock(lock_);
+    auto [it, _] = thread_id_to_rank_local_.try_emplace(fwd_thread_id);
+    // Cache for non-autograd threads
+    if (node == nullptr) {
+      cached_ = &it->second;
+    }
+    return it->second;
+  }
+
+  // Apply a function to all thread-local instances and return the first
+  // non-empty result. This is useful for cross-thread lookups when we need
+  // to find data that may have been registered on a different thread.
+  // The function should have signature: std::optional<R>(T&)
+  template <typename F>
+  static auto find_across_all(F&& func) -> decltype(func(std::declval<T&>())) {
+    std::shared_lock read_lock(lock_);
+    for (auto& [thread_id, instance] : thread_id_to_rank_local_) {
+      auto result = func(instance);
+      if (result) {
+        return result;
+      }
+    }
+    return decltype(func(std::declval<T&>()))();
+  }
+
+ private:
+  RankLocal() = default;
+  thread_local static T* cached_;
+  static std::unordered_map<uint64_t, T> thread_id_to_rank_local_;
+  static std::shared_mutex lock_;
+};
+
+template <typename T>
+thread_local T* RankLocal<T>::cached_ = nullptr;
+
+template <typename T>
+std::unordered_map<uint64_t, T> RankLocal<T>::thread_id_to_rank_local_;
+
+template <typename T>
+std::shared_mutex RankLocal<T>::lock_;
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Store.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Store.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..d3a30355e12486bb302a9172caaf8d5f7cf8f283
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Store.hpp
@@ -0,0 +1,172 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <chrono>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <c10/macros/Macros.h>
+#include <torch/custom_class.h>
+
+namespace c10d {
+
+// callback function will be given arguments (std::optional<string> oldValue,
+// std::optional<string> newValue)
+using WatchKeyCallback =
+    std::function<void(std::optional<std::string>, std::optional<std::string>)>;
+
+class TORCH_API Store : public torch::CustomClassHolder {
+ public:
+  static constexpr std::chrono::milliseconds kDefaultTimeout =
+      std::chrono::seconds(300);
+  static constexpr std::chrono::milliseconds kNoTimeout =
+      std::chrono::milliseconds::zero();
+
+  Store() : timeout_(kDefaultTimeout) {}
+
+  explicit Store(const std::chrono::milliseconds& timeout)
+      : timeout_(timeout) {}
+
+  Store(const Store&) = default;
+  Store(Store&&) noexcept = default;
+
+  ~Store() override = default;
+
+  // Clone a thread safe copy of this store object that points to the same
+  // underlying store.
+  virtual c10::intrusive_ptr<Store> clone() = 0;
+
+  void set(const std::string& key, const std::string& value);
+
+  virtual void set(
+      const std::string& key,
+      const std::vector<uint8_t>& value) = 0;
+
+  std::string compareSet(
+      const std::string& key,
+      const std::string& currentValue,
+      const std::string& newValue);
+
+  virtual std::vector<uint8_t> compareSet(
+      const std::string& key,
+      const std::vector<uint8_t>& currentValue,
+      const std::vector<uint8_t>& newValue) {
+    C10_THROW_ERROR(NotImplementedError, "Not implemented.");
+  }
+
+  std::string get_to_str(const std::string& key);
+
+  virtual std::vector<uint8_t> get(const std::string& key) = 0;
+
+  virtual int64_t add(const std::string& key, int64_t value) = 0;
+
+  virtual bool deleteKey(const std::string& key) = 0;
+
+  virtual bool check(const std::vector<std::string>& keys) = 0;
+
+  virtual int64_t getNumKeys() = 0;
+
+  virtual void wait(const std::vector<std::string>& keys) = 0;
+
+  virtual void wait(
+      const std::vector<std::string>& keys,
+      const std::chrono::milliseconds& timeout) = 0;
+
+  virtual const std::chrono::milliseconds& getTimeout() const noexcept;
+
+  virtual void setTimeout(const std::chrono::milliseconds& timeout);
+
+  // watchKey() is deprecated and no longer supported.
+  virtual void watchKey(
+      const std::string& /* unused */,
+      // NOLINTNEXTLINE(performance-unnecessary-value-param)
+      WatchKeyCallback /* unused */) {
+    C10_THROW_ERROR(
+        NotImplementedError,
+        "watchKey is deprecated, no implementation support it.");
+  }
+
+  virtual void append(
+      const std::string& key,
+      const std::vector<uint8_t>& value);
+
+  virtual std::vector<std::vector<uint8_t>> multiGet(
+      const std::vector<std::string>& keys);
+
+  virtual void multiSet(
+      const std::vector<std::string>& keys,
+      const std::vector<std::vector<uint8_t>>& values);
+
+  // Returns true if this store support append, multiGet and multiSet
+  virtual bool hasExtendedApi() const;
+
+  virtual void queuePush(
+      const std::string& key,
+      const std::vector<uint8_t>& value) {
+    C10_THROW_ERROR(NotImplementedError, "queue support is not implemented.");
+  }
+
+  virtual std::vector<uint8_t> queuePop(const std::string& key, bool block) {
+    C10_THROW_ERROR(NotImplementedError, "queue support is not implemented.");
+  }
+
+  virtual int64_t queueLen(const std::string& key) {
+    C10_THROW_ERROR(NotImplementedError, "queue support is not implemented.");
+  }
+
+  virtual std::vector<std::string> listKeys() {
+    C10_THROW_ERROR(
+        NotImplementedError, "listKeys support is not implemented.");
+  }
+
+  // Barrier operation that blocks until world_size workers have reached it.
+  // This is an optimized operation that combines increment and wait into a
+  // single operation, reducing network round trips compared to using
+  // separate add() and wait() calls.
+  virtual void barrier(
+      const std::string& key,
+      int64_t world_size,
+      const std::chrono::milliseconds& timeout);
+
+  void barrier(const std::string& key, int64_t world_size) {
+    barrier(key, world_size, timeout_);
+  }
+
+ protected:
+  std::chrono::milliseconds timeout_;
+};
+
+/*
+StoreTimeoutGuard is a RAII guard that will set the store timeout and restore it
+when it returns.
+*/
+class StoreTimeoutGuard {
+ public:
+  explicit StoreTimeoutGuard(
+      Store& store,
+      const std::chrono::milliseconds& timeout)
+      : store_(store), oldTimeout_(store.getTimeout()) {
+    store.setTimeout(timeout);
+  }
+
+  ~StoreTimeoutGuard() {
+    store_.setTimeout(oldTimeout_);
+  }
+
+  /* Disabling copy and move semantics */
+  StoreTimeoutGuard(const StoreTimeoutGuard&) = delete;
+  StoreTimeoutGuard& operator=(const StoreTimeoutGuard&) = delete;
+  StoreTimeoutGuard(StoreTimeoutGuard&&) = delete;
+  StoreTimeoutGuard& operator=(StoreTimeoutGuard&&) = delete;
+
+ private:
+  Store& store_;
+  std::chrono::milliseconds oldTimeout_{};
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TCPStore.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TCPStore.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..2b3bec0f2fbb045cd227911a919e74d81ea5e67a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TCPStore.hpp
@@ -0,0 +1,181 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include <torch/csrc/distributed/c10d/Store.hpp>
+
+namespace c10d {
+namespace detail {
+
+// TCPStore is a key-value store used by PyTorch mainly for distributed
+// rendezvous, but for other purposes as well. (e.g., a centralized storage for
+// synchronization among different processes.)
+//
+// It is run via a classic client-server architecture, where the server runs
+// a separate background thread (alternatively we call it daemon thread). The
+// client and server communicate via TCP sockets.
+//
+// Currently we have two types of server backends:
+// 1. TCPStoreBackend: a single thread to handle all incoming request
+// synchronously.
+// 2. LibUVTCPStoreBackend: an event-driven asynchronous stream processing that
+// leverages libuv library (https://github.com/libuv/libuv) for better
+// performance. And this backend now is recommended to users. (We set the
+// default value of `useLibUV` inside `TCPStoreOptions` to true now, so users
+// should get it by default).
+//
+// Code structure:
+// ├── TCPStore client side API and server setup code:
+// │   TCPStore.hpp/TCPStore.cpp
+// ├── TCPStoreBackend server side API implementation code:
+// │   TCPStoreBackend.hpp/TCPStoreBackend.cpp
+// |   (actual class:`TCPStoreMasterDaemon`)
+// ├── LibUVTCPStoreBackend
+// │   TCPStoreLibUvBackend.cpp
+// |   (actual class: `LibUVStoreDaemon`)
+
+class TCPServer;
+
+class TCPClient;
+
+struct SocketAddress {
+  std::string host;
+  std::uint16_t port{};
+};
+
+} // namespace detail
+
+struct TCPStoreOptions {
+  static constexpr std::uint16_t kDefaultPort = 29500;
+
+  std::uint16_t port = kDefaultPort;
+  bool isServer = false;
+  std::optional<std::size_t> numWorkers = std::nullopt;
+  bool waitWorkers = true;
+  std::chrono::milliseconds timeout = Store::kDefaultTimeout;
+
+  // A boolean value indicating whether multiple store instances can be
+  // initialized with the same host:port pair.
+  bool multiTenant = false;
+
+  // If specified, and if isServer is true, the underlying TCPServer will take
+  // over the bound socket associated to this fd. This option is useful to avoid
+  // port assignment races in certain scenarios.
+  std::optional<int> masterListenFd = std::nullopt;
+
+  // A boolean value indicating whether to use the experimental libUV backend.
+  bool useLibUV = true;
+};
+
+class TORCH_API TCPStore : public Store {
+ public:
+  static constexpr std::chrono::milliseconds kConnectRetryDelay{1000};
+
+  explicit TCPStore(std::string host, const TCPStoreOptions& opts = {});
+
+  ~TCPStore() override;
+
+  c10::intrusive_ptr<Store> clone() override;
+
+  void set(const std::string& key, const std::vector<uint8_t>& value) override;
+
+  std::vector<uint8_t> compareSet(
+      const std::string& key,
+      const std::vector<uint8_t>& expectedValue,
+      const std::vector<uint8_t>& desiredValue) override;
+
+  std::vector<uint8_t> get(const std::string& key) override;
+
+  int64_t add(const std::string& key, int64_t value) override;
+
+  bool deleteKey(const std::string& key) override;
+
+  bool check(const std::vector<std::string>& keys) override;
+
+  int64_t getNumKeys() override;
+
+  void wait(const std::vector<std::string>& keys) override;
+
+  void wait(
+      const std::vector<std::string>& keys,
+      const std::chrono::milliseconds& timeout) override;
+
+  void append(const std::string& key, const std::vector<uint8_t>& value)
+      override;
+
+  std::vector<std::vector<uint8_t>> multiGet(
+      const std::vector<std::string>& keys) override;
+
+  void multiSet(
+      const std::vector<std::string>& keys,
+      const std::vector<std::vector<uint8_t>>& values) override;
+
+  bool hasExtendedApi() const override;
+
+  void queuePush(const std::string& key, const std::vector<uint8_t>& value)
+      override;
+
+  std::vector<uint8_t> queuePop(const std::string& key, bool block) override;
+
+  int64_t queueLen(const std::string& key) override;
+
+  std::vector<std::string> listKeys() override;
+
+  void barrier(
+      const std::string& key,
+      int64_t world_size,
+      const std::chrono::milliseconds& timeout) override;
+
+  // Waits for all workers to join.
+  void waitForWorkers();
+
+  // Returns the hostname used by the TCPStore.
+  const std::string& getHost() const noexcept {
+    return addr_.host;
+  }
+
+  // Returns the port used by the TCPStore.
+  std::uint16_t getPort() const noexcept {
+    return addr_.port;
+  }
+
+  bool isLibUvBackend() const noexcept {
+    return usingLibUv_;
+  }
+
+  // note(xilunwu): this function is only for internal testing
+  void _splitSet(const std::string& key, const std::vector<uint8_t>& data);
+
+  std::string repr() const;
+
+ private:
+  int64_t incrementValueBy(const std::string& key, int64_t delta);
+
+  void ping();
+  void validate();
+
+  std::vector<uint8_t> doGet(const std::string& key);
+
+  void doWait(
+      c10::ArrayRef<std::string> keys,
+      std::chrono::milliseconds timeout);
+
+  detail::SocketAddress addr_;
+  std::shared_ptr<detail::TCPServer> server_;
+  std::unique_ptr<detail::TCPClient> client_;
+  std::optional<std::size_t> numWorkers_;
+
+  const std::string initKey_ = "init/";
+  const std::string keyPrefix_ = "/";
+  std::mutex activeOpLock_;
+  bool usingLibUv_ = true;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TCPStoreBackend.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TCPStoreBackend.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..531f4bdd2e3dd93751a835148225453bbe218ee1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TCPStoreBackend.hpp
@@ -0,0 +1,84 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <thread>
+
+#include <torch/csrc/distributed/c10d/TCPStore.hpp>
+#include <torch/csrc/distributed/c10d/socket.h>
+
+#ifdef _WIN32
+#include <io.h>
+#include <winsock2.h>
+#else
+#include <poll.h>
+#include <unistd.h>
+#endif
+
+namespace c10d::detail {
+
+// Magic number for client validation.
+static const uint32_t validationMagicNumber = 0x3C85F7CE;
+
+enum class QueryType : uint8_t {
+  VALIDATE,
+  SET,
+  COMPARE_SET,
+  GET,
+  ADD,
+  CHECK,
+  WAIT,
+  GETNUMKEYS,
+  DELETE_KEY,
+  APPEND,
+  MULTI_GET,
+  MULTI_SET,
+  CANCEL_WAIT,
+  PING,
+  QUEUE_PUSH,
+  QUEUE_POP,
+  QUEUE_LEN,
+  LIST_KEYS,
+  BARRIER,
+};
+
+enum class CheckResponseType : uint8_t { READY, NOT_READY };
+
+enum class WaitResponseType : uint8_t { STOP_WAITING, WAIT_CANCELED };
+
+// Abstract base class to handle thread state for TCPStoreMasterDaemon.
+// Contains the windows/unix implementations to signal a
+// shutdown sequence for the thread
+class BackgroundThread {
+ public:
+  explicit BackgroundThread();
+
+  virtual ~BackgroundThread() = 0;
+  virtual std::uint16_t port() const = 0;
+
+  void start();
+  bool stop_requested();
+
+ protected:
+  void dispose();
+  virtual void run() = 0;
+  virtual void stop() = 0;
+  bool is_running() {
+    return is_running_.load();
+  }
+
+ private:
+  std::atomic<bool> is_running_{false};
+  std::thread daemonThread_;
+};
+
+std::unique_ptr<BackgroundThread> create_tcpstore_backend(
+    const TCPStoreOptions& opts);
+std::unique_ptr<BackgroundThread> create_libuv_tcpstore_backend(
+    const TCPStoreOptions& opts);
+bool is_libuv_tcpstore_backend_available();
+
+} // namespace c10d::detail
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TraceUtils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TraceUtils.h
new file mode 100644
index 0000000000000000000000000000000000000000..8e584ac92c5c2a8aaf092cd25fc273f1fe1740f2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/TraceUtils.h
@@ -0,0 +1,324 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/ScalarType.h>
+#include <c10/util/ApproximateClock.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+#include <torch/csrc/jit/serialization/pickler.h>
+#include <torch/csrc/profiler/combined_traceback.h>
+
+#include <fmt/compile.h>
+#include <fmt/core.h>
+#include <fmt/ostream.h> // optional, for ostream fallback
+#include <fmt/ranges.h> // for fmt::join
+
+#include <sys/types.h>
+#include <cstdlib>
+#include <cstring>
+#include <iterator>
+#include <string>
+#include <vector>
+
+namespace c10d {
+
+inline std::string getTraceStartKey(const std::string& pgName, int rank) {
+  return fmt::format(FMT_COMPILE("{}_{}_trace_start"), pgName, rank);
+}
+
+inline std::string getTraceEndKey(const std::string& pgName, int rank) {
+  return fmt::format(FMT_COMPILE("{}_{}_trace_end"), pgName, rank);
+}
+
+inline bool traceUpdate(
+    c10::intrusive_ptr<Store>& store,
+    const std::string& key,
+    uint64_t seq,
+    const std::string& col) {
+  std::vector<uint8_t> value(col.size() + sizeof(seq) + 1);
+  std::memcpy(value.data(), &seq, sizeof(seq));
+  std::memcpy(value.data() + sizeof(seq), col.data(), col.size());
+  try {
+    store->set(key, value);
+    return true;
+  } catch (...) {
+    LOG(ERROR) << "Store is down while updating #" << seq << " with key "
+               << key;
+    return false;
+  }
+  return true;
+}
+
+enum TraceDebugEvent {
+  kEventStart,
+  kEventEnd,
+};
+// <seq, <rank, <col, start/end>>>
+using TraceMap =
+    std::map<uint64_t, std::map<int, std::pair<std::string, TraceDebugEvent>>>;
+
+inline std::string ranksToString(const std::vector<int>& ranks) {
+  return fmt::to_string(fmt::join(ranks, ", "));
+}
+
+inline std::string ranksFromTrace(
+    const std::vector<std::pair<int, std::string>>& items) {
+  fmt::memory_buffer buf;
+  bool first = true;
+  for (const auto& [rank, _] : items) {
+    if (!first) {
+      fmt::format_to(std::back_inserter(buf), ", ");
+    }
+    fmt::format_to(std::back_inserter(buf), "{}", rank);
+    first = false;
+  }
+  return fmt::to_string(buf);
+}
+
+inline std::string analyzeMissingRanks(const std::vector<int>& missingRanks) {
+  return c10::str(
+      "\n\t - To our best knowledge, ranks [",
+      ranksToString(missingRanks),
+      "] are the lagging ranks that caused this timeout. "
+      "They never joined any collectives");
+}
+
+inline std::string analyzeLaggingRanks(const TraceMap& traceMap) {
+  uint64_t lagSeq = traceMap.begin()->first;
+  std::vector<int> startRanks;
+  std::vector<int> endRanks;
+  for (auto& p : traceMap.begin()->second) {
+    if (p.second.second == kEventStart) {
+      startRanks.push_back(p.first);
+    } else {
+      endRanks.push_back(p.first);
+    }
+  }
+  std::string report =
+      "\n\t - To our best knowledge, the lagging/dead/mismatched ranks "
+      "that caused the desync are:";
+  if (!startRanks.empty()) {
+    report += c10::str(
+        "\n\t   - [",
+        ranksToString(startRanks),
+        "] joined but didn't finish collective #",
+        lagSeq,
+        " (count from 1)");
+  }
+  if (!endRanks.empty()) {
+    report += c10::str(
+        "\n\t     [",
+        ranksToString(endRanks),
+        "] finished collective #",
+        lagSeq,
+        ", but didn't join collective #",
+        lagSeq + 1,
+        " (count from 1)");
+  }
+  return report;
+}
+
+inline std::string dumpSnapshot(TraceMap& traceMap) {
+  std::string report = "\n\t - Snapshot of ranks' latest states:";
+  for (auto& tracePair : traceMap) {
+    uint64_t seq = tracePair.first;
+    std::map<int, std::pair<std::string, TraceDebugEvent>>& subMap =
+        tracePair.second;
+
+    std::unordered_map<std::string, std::vector<int>> collectivesStart;
+    std::unordered_map<std::string, std::vector<int>> collectivesEnd;
+    for (const auto& p : subMap) {
+      int rank = p.first;
+      const std::string& col = p.second.first;
+      if (p.second.second == kEventStart) {
+        collectivesStart[col].push_back(rank);
+      } else {
+        collectivesEnd[col].push_back(rank);
+      }
+    }
+
+    if (!collectivesStart.empty()) {
+      report += c10::str("\n\t   #", seq, " started ranks:");
+      for (auto& mapPair : collectivesStart) {
+        report += c10::str(
+            "\n\t     [",
+            ranksToString(mapPair.second),
+            "] started ",
+            mapPair.first);
+      }
+    }
+    if (!collectivesEnd.empty()) {
+      report += c10::str("\n\t   #", seq, " finished ranks:");
+      for (auto& mapPair : collectivesEnd) {
+        report += c10::str(
+            "\n\t     [",
+            ranksToString(mapPair.second),
+            "] finished ",
+            mapPair.first);
+      }
+    }
+  }
+  return report;
+}
+
+inline bool parseTraceValue(
+    c10::intrusive_ptr<Store>& store,
+    const std::string& key,
+    uint64_t& seq,
+    std::string& col) {
+  try {
+    std::vector<uint8_t> traceValue = store->get(key);
+    std::memcpy(&seq, traceValue.data(), sizeof(seq));
+    std::string colName((char*)traceValue.data() + sizeof(seq));
+    col = colName;
+    return true;
+  } catch (...) {
+    LOG(ERROR) << "Store is down while getting key " << key;
+    return false;
+  }
+  return true;
+}
+
+inline std::string retrieveDesyncReport(
+    c10::intrusive_ptr<Store>& store,
+    const std::string& pgName,
+    int myRank,
+    int worldSize) {
+  std::string report;
+
+  uint64_t thisSeq = 0;
+  std::string thisCol;
+
+  std::vector<int> missingRanks;
+  TraceMap traceMap;
+
+  for (const auto rank : c10::irange(worldSize)) {
+    // Build traceMapStart.
+    uint64_t seqStart = 0;
+    {
+      std::string traceKeyStart = getTraceStartKey(pgName, rank);
+      if (!store->check({traceKeyStart})) {
+        missingRanks.push_back(rank);
+        continue;
+      }
+      std::string col;
+      if (!parseTraceValue(store, traceKeyStart, seqStart, col)) {
+        return report;
+      }
+      traceMap[seqStart].emplace(rank, std::make_pair(col, kEventStart));
+      if (rank == myRank) {
+        thisSeq = seqStart;
+        thisCol = std::move(col);
+      }
+    }
+
+    // Build traceMapEnd.
+    {
+      std::string traceKeyEnd = getTraceEndKey(pgName, rank);
+      if (!store->check({traceKeyEnd})) {
+        continue;
+      }
+      uint64_t seq = 0;
+      std::string col;
+      if (!parseTraceValue(store, traceKeyEnd, seq, col)) {
+        return report;
+      }
+      if (seq == seqStart) {
+        traceMap[seq][rank].second = kEventEnd;
+      }
+    }
+  }
+
+  TORCH_INTERNAL_ASSERT(
+      !missingRanks.empty() || !traceMap.empty(),
+      "Trace shouldn't be empty while enabled GLOO_ASYNC_TIMEOUT_DEBUG");
+  TORCH_INTERNAL_ASSERT(
+      !thisCol.empty(),
+      "Timeout rank [",
+      myRank,
+      "] must have collective tracking iteam in c10::Store trace");
+  TORCH_INTERNAL_ASSERT(
+      traceMap[thisSeq][myRank].second == kEventStart,
+      "Timeout rank [",
+      myRank,
+      "] last trace item must be kEventStart. thisSeq = ",
+      thisSeq,
+      ", col = ",
+      thisCol);
+
+  report += c10::str(
+      "\n\t - [", myRank, "] Timeout at collective: ", thisCol, ", #", thisSeq);
+
+  if (!missingRanks.empty()) {
+    report += analyzeMissingRanks(missingRanks);
+  } else {
+    report += analyzeLaggingRanks(traceMap);
+    report += dumpSnapshot(traceMap);
+  }
+
+  return report;
+}
+
+inline std::string pickle_str(const c10::IValue& v) {
+  std::vector<char> result;
+  {
+    auto writer = [&](const char* data, size_t size) {
+      result.insert(result.end(), data, data + size);
+    };
+    torch::jit::Pickler pickler(
+        writer, nullptr, nullptr, nullptr, nullptr, false);
+    pickler.protocol();
+    pickler.pushIValue(v);
+    pickler.stop();
+  }
+  return std::string(result.begin(), result.end());
+}
+
+inline std::string get_python_cpp_trace() {
+  // usage:
+  // LOG(INFO) << "stacktrace: "
+  //           << get_python_cpp_trace();
+  // warn: might be slow in getting cpp traces
+  // because of slow/broken addr2line
+  // in different system libs
+  std::shared_ptr<torch::CapturedTraceback> tb =
+      torch::CapturedTraceback::gather(
+          /*python=*/true, /*script=*/true, /*cpp=*/true);
+  torch::SymbolizedTracebacks s_tbs = torch::symbolize({tb.get()});
+  const auto& s_tb = s_tbs.tracebacks.at(0);
+  constexpr auto TB_FMT_CSTR = FMT_COMPILE("#{} {} from {}:{}\n");
+  fmt::memory_buffer buf;
+  auto buf_iter = std::back_inserter(buf);
+  for (auto idx : c10::irange(s_tb.size())) {
+    auto frame_id = s_tb[idx];
+    const auto& frame = s_tbs.all_frames.at(frame_id);
+    fmt::format_to(
+        buf_iter,
+        TB_FMT_CSTR,
+        idx,
+        frame.funcname,
+        frame.filename,
+        frame.lineno);
+  }
+  return fmt::to_string(buf);
+}
+
+inline c10::Dict<c10::IValue, c10::IValue> new_dict() {
+  return c10::Dict<c10::IValue, c10::IValue>(
+      c10::AnyType::get(), c10::AnyType::get());
+}
+
+inline c10::List<c10::IValue> new_list() {
+  return c10::List<c10::IValue>(c10::AnyType::get());
+}
+
+inline std::string ranks_str(const std::vector<uint64_t>& ranks) {
+  return fmt::format("[{}]", fmt::join(ranks, ", "));
+}
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Types.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Types.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..c02edd2ef0eb6f4a04b6abe2976d1c9cd20e2e90
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Types.hpp
@@ -0,0 +1,190 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Store.hpp>
+
+#include <chrono>
+#include <cstdint>
+
+#include <ATen/core/Tensor.h>
+#include <ATen/core/ivalue.h>
+
+#include <c10/macros/Macros.h>
+#include <c10/util/intrusive_ptr.h>
+
+namespace c10d {
+
+// Base class for supplementary data potentially needed by ReduceOps
+struct TORCH_API _SupplementBase : torch::CustomClassHolder {
+  ~_SupplementBase() override = default;
+};
+
+// Supplementary data specific to NCCL PREMUL_SUM
+// The point of use in ProcessGroupNCCL knows how to unpack it.
+struct NCCLPreMulSumSupplement : _SupplementBase {
+  double double_factor{0.0};
+  at::Tensor tensor_factor;
+  NCCLPreMulSumSupplement(double f) : double_factor{f} {}
+  NCCLPreMulSumSupplement(at::Tensor t) : tensor_factor{std::move(t)} {
+    TORCH_CHECK_EQ(tensor_factor.numel(), 1);
+  }
+};
+
+// Other ReduceOps that need different supplementary data can also
+// derive from _SupplementBase.
+struct TORCH_API ReduceOp : torch::CustomClassHolder {
+  // note(crcrpar): RedOpType could be defined outside of `ReduceOp`
+  enum RedOpType : uint8_t {
+    SUM = 0,
+    AVG = 1,
+    PRODUCT = 2,
+    MIN = 3,
+    MAX = 4,
+    BAND = 5, // Bitwise AND
+    BOR = 6, // Bitwise OR
+    BXOR = 7, // Bitwise XOR
+    PREMUL_SUM = 8, // Multiply by a user-supplied constant before summing.
+    UNUSED = 9
+  };
+
+  ReduceOp() = default;
+
+  ReduceOp(RedOpType op) : op_(op) {
+    TORCH_INTERNAL_ASSERT(
+        op_ != PREMUL_SUM,
+        "Use `torch.distributed._make_nccl_premul_sum` to create an instance of ReduceOp with PREMUL_SUM");
+  }
+
+  ReduceOp(
+      RedOpType op,
+      const c10::intrusive_ptr<_SupplementBase>& optional_supplement) {
+    if (optional_supplement) {
+      op_ = op;
+    } else {
+      supplement_ = optional_supplement;
+    }
+  }
+
+  // The heap resource supplement_, if it exists, is managed by a
+  // c10::intrusive_ptr, so constructors and operator= can be simple
+  ReduceOp(const ReduceOp& other) = default;
+  ReduceOp& operator=(const ReduceOp& other) = default;
+
+  ReduceOp(ReduceOp&& other) = default;
+  ReduceOp& operator=(ReduceOp&& other) = default;
+  ~ReduceOp() override = default;
+
+  operator RedOpType() const {
+    return op_;
+  }
+
+  bool operator==(const std::uint8_t other) {
+    TORCH_INTERNAL_ASSERT(other < 9, "Invalid other op value");
+    return other == op_;
+  }
+
+  bool operator==(const ReduceOp::RedOpType other) {
+    return *this == static_cast<std::uint8_t>(other);
+  }
+
+  // todo(crcrpar): Handle `RedOpType::PREMUL_SUM` with its scaling factor.
+  bool operator==(const ReduceOp& other) {
+    return *this == other.op_;
+  }
+
+  RedOpType op_ = SUM;
+  // supplement_ is "type-erased" storage for optional supplementary
+  // data the op might need.
+  // The point of use will know the derived type supplement_ really is,
+  // and downcast its pointer to extract the data as the needed type(s).
+  // Right now, only PREMUL_SUM needs supplementary data, but the same
+  // mechanism could extend to support other nontrivial reduce ops with
+  // different supplementary payloads.
+  c10::intrusive_ptr<_SupplementBase> supplement_;
+};
+
+template <typename T>
+ReduceOp makeNCCLPreMulSum(const T& factor) {
+  ReduceOp rop;
+  rop.op_ = ReduceOp::PREMUL_SUM;
+  rop.supplement_ = c10::make_intrusive<NCCLPreMulSumSupplement>(factor);
+  return rop;
+}
+
+TORCH_API bool isComplexViewAsRealAllowed(const ReduceOp& reduceOp);
+
+constexpr auto kUnsetTimeout = std::chrono::milliseconds(-1);
+
+struct BroadcastOptions {
+  int64_t rootRank = 0;
+  int64_t rootTensor = 0;
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+};
+
+struct AllreduceOptions {
+  ReduceOp reduceOp = ReduceOp::SUM;
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+  std::optional<at::Tensor> sparseIndices = std::nullopt;
+};
+
+struct AllreduceCoalescedOptions : AllreduceOptions {};
+
+struct ReduceOptions {
+  ReduceOp reduceOp = ReduceOp::SUM;
+  int64_t rootRank = 0;
+  int64_t rootTensor = 0;
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+};
+
+struct AllgatherOptions {
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+};
+
+struct GatherOptions {
+  int64_t rootRank = 0;
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+};
+
+struct ScatterOptions {
+  int64_t rootRank = 0;
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+};
+
+struct ReduceScatterOptions {
+  ReduceOp reduceOp = ReduceOp::SUM;
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+};
+
+struct AllToAllOptions {
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  bool asyncOp = true;
+};
+
+struct BarrierOptions {
+  std::vector<int64_t> device_ids;
+  std::chrono::milliseconds timeout = kUnsetTimeout;
+  std::optional<at::Device> device;
+  bool asyncOp = true;
+};
+
+struct DistributedBackendOptions {
+  c10::intrusive_ptr<::c10d::Store> store;
+  int group_rank;
+  int group_size;
+  std::chrono::duration<float> timeout;
+  std::string group_id;
+  std::vector<int64_t> global_ranks_in_group;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UCCTracing.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UCCTracing.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..6e53aa7a21fb3a67135d15b2add0fe344d05dab6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UCCTracing.hpp
@@ -0,0 +1,63 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_UCC
+
+#include <torch/csrc/distributed/c10d/UCCUtils.hpp>
+
+namespace c10d {
+
+#define RECORD_COMMS_TRACE(                                                    \
+    _comms_tracer, _work, _opType, _rank, _comm_size, _inTensors, _outTensors) \
+  do {                                                                         \
+    if (torch_ucc_config.enable_comms_logger) {                                \
+      _comms_tracer->recordComms(                                              \
+          opTypeToString(_opType),                                             \
+          (uintptr_t)_work.get(),                                              \
+          _rank,                                                               \
+          _comm_size,                                                          \
+          _inTensors,                                                          \
+          _outTensors);                                                        \
+    }                                                                          \
+  } while (0)
+
+// interfaces to collect communication traces
+class TORCH_API CommTraceLogger : public torch::CustomClassHolder {
+ private:
+  std::vector<std::string> comms_trace_;
+  std::vector<std::string> curBlocks_; /* unused */
+  std::vector<int64_t> curOutSplitSizes_;
+  std::vector<int64_t> curInSplitSizes_;
+  int curRoot_ = -1;
+  unsigned long seqnum = 0;
+
+ public:
+  void setCurBlock(const std::string& name); /* unused */
+  void popBlock(); /* unused */
+  // record root info if applicable, e.g., broadcast, gather, scatter
+  void recordOptionalInfo(int root = -1);
+  // record input/output splits of Alltoallv
+  void recordOptionalInfo(
+      const std::vector<int64_t>& outputSplitSizes = {},
+      const std::vector<int64_t>& inputSplitSizes = {});
+  // record essential comms information
+  void recordComms(
+      const std::string& collName,
+      const uintptr_t workReq = 0,
+      const int rank = -1,
+      const int world_size = -1,
+      const std::vector<at::Tensor>& inputTensors = {},
+      const std::vector<at::Tensor>& outputTensor = {});
+  // return collected comms traces
+  std::vector<std::string>& getCommsTrace() {
+    return comms_trace_;
+  }
+};
+
+} // namespace c10d
+
+#endif // USE_C10D_UCC
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UCCUtils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UCCUtils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..66357622a880f36d1cad25d5a1642478482bcf05
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UCCUtils.hpp
@@ -0,0 +1,192 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_C10D_UCC
+
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <ucc/api/ucc.h>
+
+namespace c10d {
+
+// Macro to generate the error message on a non-successful UCC return value.
+#define TORCH_UCC_GET_ERROR_MSG(_err, _error_msg, _result) \
+  do {                                                     \
+    _err = c10::str(                                       \
+        "[",                                               \
+        std::string(__FILE__),                             \
+        ":",                                               \
+        std::to_string(__LINE__),                          \
+        "] ",                                              \
+        logger->getLogPrefix(),                            \
+        _error_msg,                                        \
+        ", error code ",                                   \
+        _result,                                           \
+        ": ",                                              \
+        ucc_status_string(_result),                        \
+        ", system error code ",                            \
+        errno);                                            \
+  } while (0)
+
+// Macro to throw on a non-successful UCC return value.
+#define TORCH_UCC_CHECK(_cmd, _error_msg)               \
+  do {                                                  \
+    ucc_status_t result = _cmd;                         \
+    if (result != UCC_OK) {                             \
+      std::string err;                                  \
+      TORCH_UCC_GET_ERROR_MSG(err, _error_msg, result); \
+      TORCH_CHECK(false, err);                          \
+    }                                                   \
+  } while (0)
+
+// Macro and throw on a non-successful UCC return value and free its request.
+#define TORCH_UCC_CHECK_REQUEST(_request, _cmd, _error_msg) \
+  do {                                                      \
+    ucc_status_t result = _cmd;                             \
+    if (result != UCC_OK) {                                 \
+      std::string err;                                      \
+      TORCH_UCC_GET_ERROR_MSG(err, _error_msg, result);     \
+      if (_request != nullptr) {                            \
+        ucc_collective_finalize(_request);                  \
+      }                                                     \
+      TORCH_CHECK(false, err);                              \
+    }                                                       \
+  } while (0)
+
+// Macros to print logs with unified format
+#define TORCH_UCC_LOG_ERROR(_phase, _msg) \
+  LOG(ERROR) << logger->getLogPrefix(_phase) << "[ERROR] " << _msg;
+#define TORCH_UCC_LOG_INFO(_phase, _msg) \
+  LOG(INFO) << logger->getLogPrefix(_phase) << "[INFO] " << _msg;
+#define TORCH_UCC_LOG_DEBUG(_phase, _msg) \
+  VLOG(1) << logger->getLogPrefix(_phase) << "[DEBUG] " << _msg;
+
+enum torch_ucc_phase_t {
+  TORCH_UCC_UNKNOWN = -1,
+  TORCH_UCC_INIT,
+  TORCH_UCC_HEALTH_CHECK,
+  TORCH_UCC_READY,
+  TORCH_UCC_COLL_POST,
+  TORCH_UCC_COLL_PROGRESS,
+  TORCH_UCC_FINALIZE,
+};
+
+const std::map<torch_ucc_phase_t, std::string> ucc_phase_map = {
+    {TORCH_UCC_UNKNOWN, "UNKNOWN"},
+    {TORCH_UCC_INIT, "INIT"},
+    {TORCH_UCC_HEALTH_CHECK, "HEALTH_CHECK"},
+    {TORCH_UCC_READY, "READY"},
+    {TORCH_UCC_COLL_POST, "COLL_POST"},
+    {TORCH_UCC_COLL_PROGRESS, "COLL_PROGRESS"},
+    {TORCH_UCC_FINALIZE, "FINALIZE"},
+};
+
+class CommTraceLogger;
+
+class TORCH_API ProcessGroupUCCLogger : public torch::CustomClassHolder {
+ public:
+  ProcessGroupUCCLogger();
+  ProcessGroupUCCLogger(std::string log_prefix, torch_ucc_phase_t phase);
+
+  std::string getLogPrefix(torch_ucc_phase_t phase = TORCH_UCC_UNKNOWN);
+  void setLogPrefix(std::string log_prefix);
+  inline void setPhase(torch_ucc_phase_t phase) {
+    local_phase = phase;
+  }
+
+  void initCommsTracer();
+  void flushComms(int rank, int world_size);
+  std::shared_ptr<CommTraceLogger> trace_generator = nullptr;
+
+ protected:
+  std::string log_prefix;
+  torch_ucc_phase_t local_phase = TORCH_UCC_UNKNOWN;
+  bool initialized_CommTraceLogger = false;
+};
+
+struct torch_ucc_oob_coll_info_t {
+  c10::intrusive_ptr<Store> store;
+  uint32_t comm_id;
+  int rank;
+  int size;
+  void* rbuf;
+  size_t msglen;
+  std::string getKey(std::string key) {
+    return std::to_string(comm_id) + key;
+  }
+};
+
+class CommBase {
+ public:
+  CommBase(const c10::intrusive_ptr<ProcessGroupUCCLogger>& logger_)
+      : logger(logger_) {}
+  virtual void progress() = 0;
+  virtual void free_request(ucc_coll_req_h request) = 0;
+  virtual ~CommBase() {}
+  c10::intrusive_ptr<ProcessGroupUCCLogger> logger;
+};
+class CommUCC : public CommBase {
+ public:
+  ucc_lib_h lib{nullptr};
+  ucc_context_h context{nullptr};
+
+ public:
+  void progress() override;
+  CommUCC(
+      std::shared_ptr<torch_ucc_oob_coll_info_t> oob,
+      const c10::intrusive_ptr<ProcessGroupUCCLogger>& logger);
+  void free_request(ucc_coll_req_h request) override;
+  ~CommUCC();
+};
+
+ucc_status_t oob_allgather(
+    void* sbuf,
+    void* rbuf,
+    size_t msglen,
+    void* coll_info,
+    void** req);
+
+ucc_status_t oob_allgather_test(void* req);
+
+ucc_status_t oob_allgather_free(void* req);
+
+// trim: remove spaces before and after the string view
+// implementation borrowed from https://stackoverflow.com/a/17976541
+inline std::string_view trim(std::string_view s) {
+  auto wsfront = std::find_if_not(
+      s.begin(), s.end(), [](int c) { return std::isspace(c); });
+  auto wsback = std::find_if_not(s.rbegin(), s.rend(), [](int c) {
+                  return std::isspace(c);
+                }).base();
+  return (
+      wsback <= wsfront ? "" : s.substr(wsfront - s.begin(), wsback - wsfront));
+}
+
+inline std::string tolower(std::string_view s) {
+  std::string result;
+  result.reserve(s.size());
+  for (auto c : s) {
+    result.push_back(std::tolower(c));
+  }
+  return result;
+}
+
+inline std::vector<std::string> parse_list(std::string list) {
+  std::vector<std::string> result;
+  list = tolower(trim(list));
+  while (!list.empty()) {
+    const auto end_pos = list.find_first_of(',');
+    const auto token = trim(list.substr(0, end_pos));
+    result.push_back(std::string(token));
+    list = (end_pos != std::string_view::npos) ? list.substr(end_pos + 1) : "";
+  }
+  return result;
+}
+
+} // namespace c10d
+
+#endif // USE_C10D_UCC
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UnixSockUtils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UnixSockUtils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..4293911d1faea1f3595eea7af19f7094afcfac60
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/UnixSockUtils.hpp
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+
+namespace c10d::tcputil {
+
+#define CONNECT_SOCKET_OFFSET 2
+
+inline int poll(struct pollfd* fds, unsigned long nfds, int timeout) {
+  return ::poll(fds, nfds, timeout);
+}
+
+inline void addPollfd(
+    std::vector<struct pollfd>& fds,
+    int socket,
+    short events) {
+  fds.push_back({.fd = socket, .events = events});
+}
+
+inline struct ::pollfd getPollfd(int socket, short events) {
+  struct ::pollfd res = {.fd = socket, .events = events};
+  return res;
+}
+
+} // namespace c10d::tcputil
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Utils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Utils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..fedb64349bbbb088ea87d493c37c0e8efaef2e65
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Utils.hpp
@@ -0,0 +1,750 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/util/Exception.h>
+#include <c10/util/accumulate.h>
+#include <c10/util/env.h>
+#include <c10/util/error.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/distributed/c10d/Types.hpp>
+
+#ifdef _WIN32
+#include <winsock2.h>
+#include <ws2tcpip.h>
+typedef SSIZE_T ssize_t;
+#pragma comment(lib, "Ws2_32.lib")
+#else
+#include <fcntl.h>
+#include <netdb.h>
+#include <sys/poll.h>
+#include <sys/socket.h>
+#include <unistd.h>
+#endif
+
+#include <sys/types.h>
+
+#include <cstdint>
+#include <cstdlib>
+#include <functional>
+#include <string>
+#include <vector>
+
+namespace c10d {
+
+TORCH_API size_t getTensorsNumel(const std::vector<at::Tensor>& tensors);
+
+// Retrieve tensor shapes from a given tensor.
+TORCH_API std::vector<at::Tensor> getTensorShapes(
+    const std::vector<at::Tensor>& tensors);
+
+// Use -2 to represent unset state of env vars
+#define C10D_ENV_NOT_SET -2
+
+#define WARN_ENV_VAR_ONCE(deprecated_env, new_env)                        \
+  TORCH_WARN_ONCE(                                                        \
+      "Environment variable " + deprecated_env + " is deprecated; use " + \
+      new_env + " instead");
+
+// Turns at::IntArrayRef into "(1, 2, 3, 4)".
+inline std::string toString(at::IntArrayRef l) {
+  std::stringstream ss;
+  ss << '(';
+  for (const auto i : c10::irange(l.size())) {
+    if (i > 0) {
+      ss << ", ";
+    }
+    ss << l[i];
+  }
+  ss << ')';
+  return ss.str();
+}
+
+inline std::string toString(const c10::Layout& layout) {
+  std::stringstream ss;
+  ss << layout;
+  return ss.str();
+}
+
+inline void assertSameType(
+    const at::DeprecatedTypeProperties& type,
+    const std::vector<at::Tensor>& tensors) {
+  for (const auto i : c10::irange(tensors.size())) {
+    if (!tensors[i].options().type_equal(type.options())) {
+      const std::string expected = type.toString();
+      const std::string actual = tensors[i].toString();
+      throw std::invalid_argument(
+          // NOLINTNEXTLINE(performance-inefficient-string-concatenation)
+          "mixed types (" + expected + " and " + actual + ")");
+    }
+  }
+}
+
+inline std::vector<std::string> split(
+    char separator,
+    const std::string& string) {
+  std::vector<std::string> pieces;
+  std::stringstream ss(string);
+  std::string item;
+  while (std::getline(ss, item, separator)) {
+    pieces.push_back(std::move(item));
+  }
+  return pieces;
+}
+
+inline std::string getCvarString(
+    const std::vector<std::string>& env,
+    const char* def) {
+  std::string ret(def);
+
+  if (env.empty()) {
+    TORCH_CHECK(false, "No environment variables passed");
+    return ret;
+  }
+
+  /* parse environment variable in reverse order, so the early
+   * versions of a variable get higher priority than the latter
+   * versions of the same variable */
+  for (ssize_t i = static_cast<ssize_t>(env.size()) - 1; i >= 0; i--) {
+    auto val = c10::utils::get_env(env[i].c_str());
+    if (!val.has_value()) {
+      continue;
+    } else if (i) {
+      WARN_ENV_VAR_ONCE(env[i], env[0]);
+    }
+
+    ret = val.value();
+  }
+
+  return ret;
+}
+
+inline int getCvarInt(const std::vector<std::string>& env, int def) {
+  int ret = def;
+
+  if (env.empty()) {
+    TORCH_CHECK(false, "No environment variables passed");
+    return ret;
+  }
+
+  /* parse environment variable in reverse order, so the early
+   * versions of a variable get higher priority than the latter
+   * versions of the same variable */
+  for (ssize_t i = static_cast<ssize_t>(env.size()) - 1; i >= 0; i--) {
+    const auto val = c10::utils::get_env(env[i].c_str());
+    if (!val.has_value()) {
+      continue;
+    } else if (i) {
+      WARN_ENV_VAR_ONCE(env[i], env[0]);
+    }
+
+    try {
+      ret = std::stoi(val.value());
+    } catch (std::exception&) {
+      TORCH_CHECK(false, "Invalid value for environment variable: " + env[i]);
+    }
+  }
+
+  return ret;
+}
+
+inline bool getCvarBool(const std::vector<std::string>& env, bool def) {
+  bool ret = def;
+
+  if (env.empty()) {
+    TORCH_CHECK(false, "No environment variables passed");
+    return ret;
+  }
+
+  /* parse environment variable in reverse order, so the early
+   * versions of a variable get higher priority than the latter
+   * versions of the same variable */
+  for (ssize_t i = static_cast<ssize_t>(env.size()) - 1; i >= 0; i--) {
+    auto val = c10::utils::get_env(env[i].c_str());
+    if (!val.has_value()) {
+      continue;
+    } else if (i) {
+      WARN_ENV_VAR_ONCE(env[i], env[0]);
+    }
+
+    for (auto& x : val.value()) {
+      // NOLINTNEXTLINE(*-narrowing-conversions)
+      x = std::tolower(x);
+    }
+
+    if (val == "y" || val == "yes" || val == "1" || val == "t" ||
+        val == "true") {
+      ret = true;
+    } else if (
+        val == "n" || val == "no" || val == "0" || val == "f" ||
+        val == "false") {
+      ret = false;
+    } else {
+      TORCH_CHECK(false, "Invalid value for environment variable: " + env[i]);
+      return ret;
+    }
+  }
+
+  return ret;
+}
+
+inline void assertSameSizes(
+    const at::IntArrayRef& sizes,
+    const std::vector<at::Tensor>& tensors) {
+  for (const auto i : c10::irange(tensors.size())) {
+    if (!tensors[i].sizes().equals(sizes)) {
+      const auto expected = toString(sizes);
+      const auto actual = toString(tensors[i].sizes());
+      throw std::invalid_argument(
+          // NOLINTNEXTLINE(performance-inefficient-string-concatenation)
+          "mixed sizes (" + expected + " and " + actual + ")");
+    }
+  }
+}
+
+inline void assertSameSizeAndType(const std::vector<at::Tensor>& tensors) {
+  // Ensure we have at least one tensor
+  if (tensors.empty()) {
+    throw std::invalid_argument("argument is empty");
+  }
+
+  // Ensure all tensors have identical type and shape
+  auto options = tensors[0].options();
+  auto sizes = tensors[0].sizes();
+  for (const auto i : c10::irange(1, tensors.size())) {
+    if (!tensors[i].options().type_equal(options)) {
+      const auto expected = toString(options);
+      const auto actual = toString(tensors[i].options());
+      throw std::invalid_argument(
+          // NOLINTNEXTLINE(performance-inefficient-string-concatenation)
+          "argument contains mixed types (" + expected + " and " + actual +
+          ")");
+    }
+    if (!tensors[i].sizes().equals(sizes)) {
+      const auto expected = toString(sizes);
+      const auto actual = toString(tensors[i].sizes());
+      throw std::invalid_argument(
+          // NOLINTNEXTLINE(performance-inefficient-string-concatenation)
+          "argument contains mixed types (" + expected + " and " + actual +
+          ")");
+    }
+  }
+}
+
+inline void assertTypeMatch(
+    const std::function<void(const std::string&)>& fn,
+    const at::DeprecatedTypeProperties& type,
+    const at::ArrayRef<at::Tensor> tensors,
+    size_t index) {
+  if (!tensors[index].options().type_equal(type.options())) {
+    fn("invalid tensor type at index " + std::to_string(index) + " (expected " +
+       type.toString() + ", got " + tensors[index].toString() + ")");
+  }
+}
+
+inline void assertTypeMatch(
+    const std::function<void(const std::string&)>& fn,
+    const at::TensorOptions& options,
+    const at::ArrayRef<at::Tensor> tensors,
+    size_t index) {
+  if (!tensors[index].options().type_equal(options)) {
+    fn("invalid tensor type at index " + std::to_string(index) + " (expected " +
+       toString(options) + ", got " + toString(tensors[index].options()) + ")");
+  }
+}
+
+inline void assertSizesMatch(
+    const std::function<void(const std::string&)>& fn,
+    const at::IntArrayRef& sizes,
+    const at::ArrayRef<at::Tensor> tensors,
+    size_t index) {
+  if (tensors[index].sizes() != sizes) {
+    fn("invalid tensor size at index " + std::to_string(index) + " (expected " +
+       toString(sizes) + ", got " + toString(tensors[index].sizes()) + ")");
+  }
+}
+
+inline void assertLayoutMatch(
+    const std::function<void(const std::string&)>& fn,
+    const c10::Layout& expected,
+    const at::ArrayRef<at::Tensor> tensors,
+    size_t index) {
+  const auto& actual = tensors[index].layout();
+  if (actual != expected) {
+    fn("invalid tensor layout at index " + std::to_string(index) +
+       " (expected " + toString(expected) + ", got " + toString(actual) + ")");
+  }
+}
+
+inline void assertLayoutMatch(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  const auto& layout = tensors[0].layout();
+  for (const auto i : c10::irange(1, tensors.size())) {
+    assertLayoutMatch(fn, layout, tensors, i);
+  }
+}
+
+inline void assertNonEmpty(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  if (tensors.empty()) {
+    fn("requires non-empty tensor list");
+  }
+}
+
+inline void assertSingleElement(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  if (tensors.size() != 1) {
+    fn("requires a single-element tensor list");
+  }
+}
+
+inline void assertSingleElementInput(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  if (tensors.size() != 1) {
+    fn("requires a single-element input tensor list");
+  }
+}
+
+inline void assertSingleElementOutput(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  if (tensors.size() != 1) {
+    fn("requires a single-element output tensor list");
+  }
+}
+
+inline void assertRootRank(
+    const std::function<void(const std::string&)>& fn,
+    int64_t rank,
+    int64_t size) {
+  if (rank < 0 || rank >= size) {
+    fn("invalid root rank: " + std::to_string(rank));
+  }
+}
+
+inline void assertRootTensor(
+    const std::function<void(const std::string&)>& fn,
+    int64_t rank,
+    int64_t size) {
+  if (rank < 0 || rank >= size) {
+    fn("invalid root tensor: " + std::to_string(rank));
+  }
+}
+
+inline void assertDense(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  const auto& layout = tensors[0].layout();
+  if (layout != at::kStrided) {
+    fn("only supports dense tensors");
+  }
+}
+
+inline void assertCPU(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  const auto& device = tensors[0].device();
+  if (device.type() != at::kCPU) {
+    fn("only supports CPU tensors");
+  }
+}
+
+inline void assertSameDevice(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  if (tensors.size() < 2) {
+    return;
+  }
+  const auto& device = tensors[0].device();
+  for (const auto i : c10::irange(1, tensors.size())) {
+    if (tensors[i].device() != device) {
+      fn("tensors should be on the same device");
+    }
+  }
+}
+
+inline void assertTypeAndSizesMatch(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors,
+    const at::DeprecatedTypeProperties& type,
+    const at::IntArrayRef& sizes) {
+  for (const auto i : c10::irange(tensors.size())) {
+    assertTypeMatch(fn, type, tensors, i);
+    assertSizesMatch(fn, sizes, tensors, i);
+  }
+}
+
+inline void assertTypeAndSizesMatch(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors,
+    const at::TensorOptions& options,
+    const at::IntArrayRef& sizes) {
+  for (const auto i : c10::irange(tensors.size())) {
+    assertTypeMatch(fn, options, tensors, i);
+    assertSizesMatch(fn, sizes, tensors, i);
+  }
+}
+
+inline void assertTypeAndSizesMatch(
+    const std::function<void(const std::string&)>& fn,
+    const at::ArrayRef<at::Tensor> tensors) {
+  const auto& options = tensors[0].options();
+  const auto sizes = tensors[0].sizes();
+  assertTypeAndSizesMatch(fn, tensors.slice(1), options, sizes);
+}
+
+// Copied from ATen/core/functional.h.
+template <typename F, typename T>
+inline auto fmap(T& inputs, const F& fn)
+    -> std::vector<decltype(fn(*inputs.begin()))> {
+  std::vector<decltype(fn(*inputs.begin()))> r;
+  r.reserve(inputs.size());
+  for (auto& input : inputs) {
+    r.push_back(fn(input));
+  }
+  return r;
+}
+
+// Copied from torch/csrc/utils/tensor_flatten.h.
+inline at::Tensor flattenDenseTensors(at::TensorList tensors) {
+  static const auto flatten = [](const at::Tensor& t) {
+    return t.contiguous().view({-1});
+  };
+  if (tensors.size() == 1) {
+    return flatten(tensors[0]);
+  }
+  return at::cat(::c10d::fmap(tensors, flatten));
+}
+
+inline at::Tensor newLikeFlat(
+    std::vector<std::vector<at::Tensor>>& tensors,
+    size_t deviceIdx) {
+  if (tensors.empty() || tensors[0].empty()) {
+    TORCH_CHECK(false, "Received an empty list");
+  }
+  if (deviceIdx >= tensors.size()) {
+    TORCH_CHECK(false, "Invalid device index");
+  }
+  auto& t = tensors[deviceIdx][0];
+  auto device = t.device();
+  for (const auto i : c10::irange(1, tensors[deviceIdx].size())) {
+    if (tensors[deviceIdx][i].device() != device) {
+      TORCH_CHECK(false, "Expecting all tensors on the same device");
+    }
+  }
+  at::DeviceGuard gpuGuard(device);
+  std::vector<int64_t> sizes{static_cast<int64_t>(tensors[deviceIdx].size())};
+  std::vector<int64_t> strides{t.numel()};
+  sizes.insert(sizes.end(), t.sizes().begin(), t.sizes().end());
+  strides.insert(strides.end(), t.strides().begin(), t.strides().end());
+  return at::empty_strided(
+      sizes, strides, t.options().memory_format(std::nullopt));
+}
+
+inline at::Tensor newLikeFlat(std::vector<at::Tensor>& tensors) {
+  if (tensors.empty()) {
+    TORCH_CHECK(false, "Received an empty list");
+  }
+  auto& t = tensors[0];
+  at::DeviceGuard gpuGuard(t.device());
+  std::vector<int64_t> sizes{static_cast<int64_t>(tensors.size())};
+  sizes.insert(sizes.end(), t.sizes().begin(), t.sizes().end());
+  return at::empty(sizes, t.options());
+}
+
+inline std::vector<std::vector<int64_t>> getSizes(
+    const std::vector<at::Tensor>& tensors) {
+  std::vector<std::vector<int64_t>> sizes(tensors.size());
+  for (const auto i : c10::irange(tensors.size())) {
+    sizes[i] = tensors[i].sizes().vec();
+  }
+  return sizes;
+}
+
+inline std::vector<int> getDevices(const std::vector<at::Tensor>& tensors) {
+  std::vector<int> devices(tensors.size(), -1);
+  if (tensors[0].device().is_cuda()) {
+    for (const auto i : c10::irange(tensors.size())) {
+      // NOLINTNEXTLINE(bugprone-signed-char-misuse)
+      devices[i] = tensors[i].storage().device().index();
+    }
+  }
+  return devices;
+}
+
+template <typename T>
+inline T* getDataPointer(const at::Tensor& tensor) {
+  // This method is only used in ProcessGroupGloo for now. Call sites must make
+  // sure that the input tensor is contiguous. It is OK if the tensor does not
+  // start from the beginning of the storage. For example, it could come from
+  // chunk(..., dim=0)[1]. Hence, we need to use data_ptr() instead of
+  // tensor.storage().data()
+  // NB: not using tensor.data<T>() because tensor is not aware of gloo::TYPE
+  return static_cast<T*>(tensor.data_ptr());
+}
+
+template <typename T>
+std::vector<T*> getDataPointers(const std::vector<at::Tensor>& tensors) {
+  std::vector<T*> ptrs(tensors.size());
+  for (const auto i : c10::irange(tensors.size())) {
+    ptrs[i] = getDataPointer<T>(tensors[i]);
+  }
+  return ptrs;
+}
+
+// For alltoall split size sanity check
+inline void checkSplitSizes(
+    const std::vector<int64_t>& split_sizes,
+    const at::Tensor& tensor,
+    int group_size) {
+  if (split_sizes.empty()) {
+    TORCH_CHECK(
+        tensor.size(0) % group_size == 0,
+        "Tensor's dim 0 does not divide equally across group size");
+  } else {
+    TORCH_CHECK(
+        split_sizes.size() == static_cast<size_t>(group_size),
+        "Number of tensor splits not equal to group size");
+    const auto sum = c10::sum_integers(split_sizes);
+    TORCH_CHECK(
+        sum == tensor.size(0), "Split sizes doesn't match total dim 0 size");
+  }
+}
+
+// Compute alltoall lengths and offsets, handling multi-dimension tensors
+template <typename T>
+size_t computeLengthsAndOffsets(
+    const std::vector<int64_t>& split_sizes,
+    const at::Tensor& tensor,
+    std::vector<T>* lengths,
+    std::vector<T>* offsets) {
+  size_t group_size = lengths->size();
+  bool equal_splits = false;
+  size_t dim0_size = tensor.size(0);
+  size_t row_size = (dim0_size ? tensor.numel() / dim0_size : 1);
+  size_t split_size = 0;
+  size_t offset = 0;
+
+  if (split_sizes.empty()) {
+    equal_splits = true;
+    split_size = tensor.size(0) / group_size;
+  }
+  for (const auto i : c10::irange(group_size)) {
+    size_t length = row_size * (equal_splits ? split_size : split_sizes[i]);
+    (*lengths)[i] = length;
+    (*offsets)[i] = offset;
+    // TODO: see if we should add overflow protection for offset
+    offset += length;
+  }
+  return offset;
+}
+
+template <typename T>
+size_t computeLengthsAndOffsets(
+    const std::vector<at::Tensor>& tensors,
+    std::vector<T>* lengths,
+    std::vector<T>* offsets) {
+  size_t group_size = lengths->size();
+  size_t offset = 0;
+  for (const auto i : c10::irange(group_size)) {
+    size_t length = tensors[i].numel();
+    (*lengths)[i] = length;
+    (*offsets)[i] = offset;
+    offset += length;
+  }
+  return offset;
+}
+
+// Get the start and stride of the global rank from a list of global ranks
+// If the global ranks do not follow the consecutive rule, the stride will be -1
+void TORCH_API getGlobalRankStartAndStride(
+    const std::vector<uint64_t>& globalRanksInGroup,
+    int& globalRankStart,
+    int& globalRankStride);
+
+using RankType = uint32_t;
+using SizeType = uint64_t;
+
+// `errno` is only meaningful when it fails. E.g., a  successful `fork()` sets
+// `errno` to `EINVAL` in child process on some macos
+// (https://stackoverflow.com/a/20295079), and thus `errno` should really only
+// be inspected if an error occurred.
+//
+// `success_cond` is an expression used to check if an error has happened. So
+// for `fork()`, we can use `SYSCHECK(pid = fork(), pid != -1)`. The function
+// output is stored in variable `__output` and may be used in `success_cond`.
+#ifdef _WIN32
+#define SYSCHECK(expr, success_cond)                                           \
+  while (true) {                                                               \
+    auto __output = (expr);                                                    \
+    auto errno_local = WSAGetLastError();                                      \
+    (void)__output;                                                            \
+    if (!(success_cond)) {                                                     \
+      if (errno == EINTR) {                                                    \
+        continue;                                                              \
+      } else if (                                                              \
+          errno_local == WSAETIMEDOUT || errno_local == WSAEWOULDBLOCK) {      \
+        C10_THROW_ERROR(DistNetworkError, "Socket Timeout");                   \
+      } else {                                                                 \
+        C10_THROW_ERROR(DistNetworkError, c10::utils::str_error(errno_local)); \
+      }                                                                        \
+    } else {                                                                   \
+      break;                                                                   \
+    }                                                                          \
+  }
+#else
+#define SYSCHECK(expr, success_cond)                                     \
+  while (true) {                                                         \
+    auto __output = (expr);                                              \
+    (void)__output;                                                      \
+    if (!(success_cond)) {                                               \
+      if (errno == EINTR) {                                              \
+        continue;                                                        \
+      } else if (errno == EAGAIN || errno == EWOULDBLOCK) {              \
+        C10_THROW_ERROR(DistNetworkError, "Socket Timeout");             \
+      } else {                                                           \
+        C10_THROW_ERROR(DistNetworkError, c10::utils::str_error(errno)); \
+      }                                                                  \
+    } else {                                                             \
+      break;                                                             \
+    }                                                                    \
+  }
+#endif
+
+// Most functions indicate error by returning `-1`. This is a helper macro for
+// this common case with `SYSCHECK`.
+// Since SOCKET_ERROR = -1 in MSVC, so also leverage SYSCHECK_ERR_RETURN_NEG1
+#define SYSCHECK_ERR_RETURN_NEG1(expr) SYSCHECK(expr, __output != -1)
+
+namespace tcputil {
+
+// Send and receive
+template <typename T>
+void sendBytes(
+    int socket,
+    const T* buffer,
+    size_t length,
+    bool moreData = false) {
+  size_t bytesToSend = sizeof(T) * length;
+  if (bytesToSend == 0) {
+    return;
+  }
+
+  auto currentBytes = reinterpret_cast<const char*>(buffer);
+
+  int flags = 0;
+
+#ifdef MSG_MORE
+  if (moreData) { // there is more data to send
+    flags |= MSG_MORE;
+  }
+#endif
+
+// Ignore SIGPIPE as the send() return value is always checked for error
+#ifdef MSG_NOSIGNAL
+  flags |= MSG_NOSIGNAL;
+#endif
+
+  while (bytesToSend > 0) {
+    ssize_t bytesSent = 0;
+    SYSCHECK_ERR_RETURN_NEG1(
+        bytesSent = ::send(socket, currentBytes, bytesToSend, flags))
+    if (bytesSent == 0) {
+      C10_THROW_ERROR(
+          DistNetworkError,
+          "Failed to send, sent 0 bytes. "
+          "Connection was likely closed. "
+          "Did the remote server shutdown or crash?");
+    }
+
+    bytesToSend -= bytesSent;
+    currentBytes += bytesSent;
+  }
+}
+
+template <typename T>
+void recvBytes(int socket, T* buffer, size_t length) {
+  size_t bytesToReceive = sizeof(T) * length;
+  if (bytesToReceive == 0) {
+    return;
+  }
+
+  auto currentBytes = reinterpret_cast<char*>(buffer);
+
+  while (bytesToReceive > 0) {
+    ssize_t bytesReceived = 0;
+    SYSCHECK_ERR_RETURN_NEG1(
+        bytesReceived = recv(socket, currentBytes, bytesToReceive, 0))
+    if (bytesReceived == 0) {
+      C10_THROW_ERROR(
+          DistNetworkError,
+          "Failed to recv, got 0 bytes. "
+          "Connection was likely closed. "
+          "Did the remote server shutdown or crash?");
+    }
+
+    bytesToReceive -= bytesReceived;
+    currentBytes += bytesReceived;
+  }
+}
+
+// send a vector's length and data
+template <typename T>
+void sendVector(int socket, const std::vector<T>& vec, bool moreData = false) {
+  SizeType size = vec.size();
+  sendBytes<SizeType>(socket, &size, 1, true);
+  sendBytes<T>(socket, vec.data(), size, moreData);
+}
+
+// receive a vector as sent in sendVector
+template <typename T>
+std::vector<T> recvVector(int socket) {
+  SizeType valueSize = 0;
+  recvBytes<SizeType>(socket, &valueSize, 1);
+  std::vector<T> value(valueSize);
+  recvBytes<T>(socket, value.data(), value.size());
+  return value;
+}
+
+// this is only for convenience when sending rvalues
+template <typename T>
+void sendValue(int socket, const T& value, bool moreData = false) {
+  sendBytes<T>(socket, &value, 1, moreData);
+}
+
+template <typename T>
+T recvValue(int socket) {
+  T value;
+  recvBytes<T>(socket, &value, 1);
+  return value;
+}
+
+// send a string's length and data
+inline void sendString(
+    int socket,
+    const std::string& str,
+    bool moreData = false) {
+  SizeType size = str.size();
+  sendBytes<SizeType>(socket, &size, 1, true);
+  sendBytes<char>(socket, str.data(), size, moreData);
+}
+
+// receive a string as sent in sendString
+inline std::string recvString(int socket) {
+  SizeType valueSize = 0;
+  recvBytes<SizeType>(socket, &valueSize, 1);
+  std::vector<char> value(valueSize);
+  recvBytes<char>(socket, value.data(), value.size());
+  return std::string(value.data(), value.size());
+}
+
+} // namespace tcputil
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/WinSockUtils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/WinSockUtils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..4e94641b9b8ed2ce362b092398a6af0243aa4a81
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/WinSockUtils.hpp
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+
+namespace c10d::tcputil {
+
+#define CONNECT_SOCKET_OFFSET 1
+
+inline int poll(struct pollfd* fdArray, unsigned long fds, int timeout) {
+  return WSAPoll(fdArray, fds, timeout);
+}
+
+inline void addPollfd(
+    std::vector<struct pollfd>& fds,
+    int socket,
+    short events) {
+  fds.push_back({(SOCKET)socket, events});
+}
+
+inline struct ::pollfd getPollfd(int socket, short events) {
+  struct ::pollfd res = {(SOCKET)socket, events};
+  return res;
+}
+
+} // namespace c10d::tcputil
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Work.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Work.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..4f70b517c3f56898c8a33ca75a497c54ac956755
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/Work.hpp
@@ -0,0 +1,194 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <chrono>
+#include <mutex>
+#include <vector>
+
+constexpr auto kNoTimeout = std::chrono::milliseconds(0);
+
+namespace c10d {
+
+constexpr const char* const kSeqNumStoreKey = "SEQ_NUM_STORE_KEY";
+
+enum class OpType : std::uint8_t {
+  BROADCAST = 0,
+  ALLREDUCE = 1,
+  ALLREDUCE_COALESCED = 2,
+  REDUCE = 3,
+  ALLGATHER = 4,
+  _ALLGATHER_BASE = 5,
+  ALLGATHER_COALESCED = 6,
+  GATHER = 7,
+  SCATTER = 8,
+  REDUCE_SCATTER = 9,
+  ALLTOALL_BASE = 10,
+  ALLTOALL = 11,
+  SEND = 12,
+  RECV = 13,
+  RECVANYSOURCE = 14,
+  BARRIER = 15,
+  _REDUCE_SCATTER_BASE = 16,
+  COALESCED = 17,
+  _ALLREDUCE_SPARSE = 18,
+  REDUCE_SCATTER_TENSOR_COALESCED = 19,
+  UNKNOWN = 100,
+};
+
+// TODO: support different types of failures/errors
+enum class WorkResult : std::uint8_t {
+  SUCCESS = 0,
+  TIMEOUT = 1,
+  COMM_ERROR = 2,
+  UNKNOWN = 100,
+};
+
+// Converts OpType to human readable string.
+TORCH_API std::string opTypeToString(OpType opType);
+
+// Whether or not an OP is an p2p op (SEND, RECV, RECVANYSOURCE)
+TORCH_API bool isP2POp(OpType opType, bool batchP2P = false);
+
+// Please do not use Work API, it is going away, to be
+// replaced by ivalue::Future.
+// Python binding for this class might change, please do not assume
+// this will be bound using pybind.
+class TORCH_API Work : public torch::CustomClassHolder {
+ public:
+  Work(
+      int rank = -1,
+      OpType opType = OpType::UNKNOWN,
+      const char* profilingTitle = nullptr,
+      const std::optional<std::vector<at::Tensor>>& inputTensors =
+          std::nullopt);
+
+  ~Work() override;
+
+  // Checks if request has completed. Non-blocking operation.
+  virtual bool isCompleted();
+
+  // Returns if the work completed successfully.
+  // If false, the exception function can be called to get details.
+  virtual bool isSuccess() const;
+
+  // Returns exception if isSuccess() returned false.
+  virtual std::exception_ptr exception() const;
+
+  // Returns source rank if this objects represents a recv-from-any.
+  virtual int sourceRank() const;
+
+  // Returns result tensors, if applicable.
+  // If work is not supposed to have result, we return empty list.
+  virtual std::vector<at::Tensor> result();
+
+  // Ensures that operations on the output tensors that are invoked
+  // after this function returns are correctly sequenced after the
+  // asynchronous completion of this work.
+  //
+  // For CUDA tensors, it inserts stream synchronization such that
+  // the streams of the caller wait for completion of the
+  // asynchronous operations on the destination tensors.
+  //
+  // For CPU tensors, it is currently a nop.
+  //
+  // This function should only be used if the caller polls for
+  // completion through the `isCompleted` function, it has returned
+  // true, and the `isSuccess` function also has returned true.
+  //
+  virtual void synchronize();
+
+  // Waits until request completes. Blocking operation.
+  // Throws if the work completed with an exception.
+  // Returns false if the work is aborted.
+  // Otherwise, it always returns true, indicating the work is completed.
+  //
+  // Functionally equivalent to:
+  //
+  //   while (!isCompleted()) { /* nop */ }
+  //   auto success = isSuccess();
+  //   if (!success) { std::rethrow_exception(exception()); }
+  //   return success;
+  //
+  virtual bool wait(std::chrono::milliseconds timeout = kNoTimeout);
+
+  // Blocks the current stream until the work is completed.
+  // This is equivalent to synchronize for CUDA tensors but works for both CPU
+  // tensors and CUDA tensors by using a spinlock CUDA kernel.
+  // This will immediately return.
+  // If no stream is active it will throw an error.
+  virtual void blockCurrentStream();
+
+  virtual void abort();
+
+  // Returns a Future object that will be associated with the completion of
+  // work. Only NCCL backend is currently supported.
+  virtual c10::intrusive_ptr<c10::ivalue::Future> getFuture();
+
+  // Get a Future object that would be marked as either success or failure
+  // This API can be used by the user to track the completion of the work
+  // and handle the exception if any.
+  virtual c10::intrusive_ptr<c10::ivalue::Future> getFutureResult();
+
+  virtual float getDuration() const;
+
+  virtual uint64_t getSequencenumber() const;
+
+  OpType retrieveOpType() const;
+
+  static c10::intrusive_ptr<Work> create_from_future(
+      const c10::intrusive_ptr<c10::ivalue::Future>& /*future*/);
+
+ protected:
+  // Completes the work object and optionally sets the exception in a
+  // thread-safe manner. Notifies all waiting condition variables as well.
+  void finish(std::exception_ptr exception = nullptr);
+
+  // Similar to finish, but throws an exception if one is already set or
+  // provided by the user.
+  void finishAndThrow(std::exception_ptr exception);
+
+  mutable std::mutex mutex_;
+  std::condition_variable cv_;
+  bool completed_ = false;
+  std::exception_ptr exception_;
+
+  // Current rank of the node.
+  const int rank_;
+
+  // Operation type that this work object refers to.
+  OpType opType_;
+
+  // When profiling, the callback to record end of operation event. This
+  // callback needs to be called when collective operation is complete.
+  std::function<void()> recordFunctionEndCallback_;
+};
+
+struct TORCH_API WorkInfo {
+  WorkInfo(
+      const OpType& opType,
+      const uint64_t seq,
+      const std::chrono::time_point<std::chrono::steady_clock>& timeStarted,
+      const std::chrono::time_point<std::chrono::steady_clock>& timeFinished,
+      const std::chrono::duration<float>& activeDuration)
+      : opType(opType),
+        seq(seq),
+        timeStarted(timeStarted),
+        timeFinished(timeFinished),
+        activeDuration(activeDuration) {}
+
+  OpType opType;
+  uint64_t seq;
+  std::chrono::time_point<std::chrono::steady_clock> timeStarted;
+  std::chrono::time_point<std::chrono::steady_clock> timeFinished;
+  std::chrono::duration<float> activeDuration;
+};
+
+TORCH_API void set_comm_profiling_name(const std::string& name);
+TORCH_API const std::string& get_comm_profiling_name();
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/c10d.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/c10d.h
new file mode 100644
index 0000000000000000000000000000000000000000..45daf59b55ab849b2e72bdf275a6c09315413ef0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/c10d.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::distributed::c10d {
+
+PyMethodDef* python_functions();
+
+} // namespace torch::distributed::c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/comm.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/comm.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..e5099eb2980e5d3f38ab7aa723f5f92880b73d91
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/comm.hpp
@@ -0,0 +1,147 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <utility>
+
+namespace c10d {
+
+// Broadcast many tensors to all processes in the process group.
+TORCH_API void broadcast_coalesced(
+    const c10::intrusive_ptr<c10d::ProcessGroup>& process_group,
+    at::TensorList tensors,
+    size_t buffer_size,
+    int rank = 0);
+
+// This class passes bucket contents tensor to DDP communication hook.
+class TORCH_API GradBucket {
+ public:
+  explicit GradBucket(
+      size_t index,
+      size_t bucket_count,
+      at::Tensor tensor,
+      std::vector<size_t> offsets,
+      std::vector<size_t> lengths,
+      std::vector<c10::IntArrayRef> sizes_vec,
+      std::vector<at::Tensor> parameters,
+      std::optional<at::Tensor> sparse_grad_indices)
+      : index_(index),
+        bucket_count_(bucket_count),
+        buffer_(std::move(tensor)),
+        offsets_(std::move(offsets)),
+        lengths_(std::move(lengths)),
+        sizes_vec_(std::move(sizes_vec)),
+        parameters_(std::move(parameters)),
+        sparse_grad_indices_(std::move(sparse_grad_indices)) {}
+
+  // Returns the index of the bucket, which is unique across all the buckets.
+  size_t getIndex() const {
+    return index_;
+  }
+
+  const at::Tensor& getBuffer() const {
+    return buffer_;
+  }
+
+  // Returns a mutable buffer compared with the above method.
+  at::Tensor& getBufferRef() {
+    return buffer_;
+  }
+
+  // Overwrites the buffer at a specific index.
+  void setBuffer(at::Tensor& buffer) {
+    buffer_ = buffer;
+  }
+
+  // Each tensor in the list that getGradients corresponds to a
+  // parameter.
+  std::vector<at::Tensor> getGradients() const;
+
+  // Returns model parameters belonging to this bucket. They are returned in the
+  // same order as gradient tensors via getGradients(). For example,
+  // getParameters[i] will have its gradient stored in
+  // getGradients[i]
+  const std::vector<at::Tensor> getParameters() const {
+    return parameters_;
+  }
+
+  // Returns whether this bucket is the last bucket to allreduce in an
+  // iteration.
+  bool isLast() const {
+    return index_ == bucket_count_ - 1;
+  }
+
+  std::optional<at::Tensor>& getSparseGradIndices() {
+    return sparse_grad_indices_;
+  }
+
+ private:
+  size_t index_;
+  size_t bucket_count_;
+  at::Tensor buffer_;
+
+  // Per-variable info in buffer_.
+  std::vector<size_t> offsets_;
+  std::vector<size_t> lengths_;
+  std::vector<c10::IntArrayRef> sizes_vec_;
+
+  // Model parameters for this bucket.
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::vector<at::Tensor> parameters_;
+
+  // Predefined sparse indices for this bucket (only used for sparse tensors).
+  // The gradients will be updated to have indices with these tensor values
+  std::optional<at::Tensor> sparse_grad_indices_;
+};
+
+// Base class of both `PythonCommHook` and `CppCommHook`.
+// Requires implementing 1) `runHook` method that communicates gradients
+// asynchronously, and 2) `parseHookResult` method that converts the hook
+// result into a tensor.
+class TORCH_API CommHookInterface {
+ public:
+  virtual ~CommHookInterface() = default;
+
+  // Passes the input grad bucket to the registered communication hook.
+  // Once the tensor in the bucket are ready, kicks off the hook asynchronously
+  // and returns a future that holds the communication results.
+  virtual c10::intrusive_ptr<c10::ivalue::Future> runHook(
+      GradBucket& bucket) = 0;
+
+  // Returns the resulting tensor once the communication hook result is
+  // ready. The resulting tensor will then be copied to the grads of
+  // individual parameters.
+  virtual at::Tensor parseHookResult(const c10::IValue& result) = 0;
+};
+
+namespace detail {
+// This helper function is called both by CppCommHookInterface below and inside
+// reducer.
+TORCH_API at::Tensor parseCppCommHookResult(const c10::IValue& result);
+} // namespace detail
+
+// This CppCommHook interface only requires implementing runHook method that
+// potentially uses a state.
+template <typename T>
+class CppCommHookInterface : public CommHookInterface {
+ public:
+  explicit CppCommHookInterface(T state) : state_(std::move(state)) {}
+
+  ~CppCommHookInterface() override = default;
+
+  at::Tensor parseHookResult(const c10::IValue& result) override {
+    return detail::parseCppCommHookResult(result);
+  }
+
+ protected:
+  T state_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_collectives/ControlCollectives.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_collectives/ControlCollectives.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..5beaa289331f4c3b81d3a279dddb8a3be6b51b38
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_collectives/ControlCollectives.hpp
@@ -0,0 +1,64 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <chrono>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <c10/macros/Macros.h>
+#include <torch/custom_class.h>
+
+namespace c10d {
+
+using namespace std::chrono_literals;
+
+class TORCH_API ControlCollectives : public torch::CustomClassHolder {
+ public:
+  virtual void barrier(
+      const std::string& key,
+      std::chrono::milliseconds timeout = 5min,
+      bool block = true) = 0;
+
+  virtual void broadcastSend(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) = 0;
+  virtual std::vector<uint8_t> broadcastRecv(
+      const std::string& key,
+      std::chrono::milliseconds timeout = 5min) = 0;
+
+  virtual void gatherSend(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) = 0;
+  virtual std::vector<std::vector<uint8_t>> gatherRecv(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) = 0;
+
+  virtual std::vector<uint8_t> scatterSend(
+      const std::string& key,
+      const std::vector<std::vector<uint8_t>>& data,
+      std::chrono::milliseconds timeout = 5min) = 0;
+  virtual std::vector<uint8_t> scatterRecv(
+      const std::string& key,
+      std::chrono::milliseconds timeout = 5min) = 0;
+
+  virtual std::vector<std::vector<uint8_t>> allGather(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) = 0;
+
+  virtual int64_t allSum(
+      const std::string& key,
+      int64_t data,
+      std::chrono::milliseconds timeout = 5min) = 0;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_collectives/StoreCollectives.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_collectives/StoreCollectives.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..5782eb49f5755c1f2270428dd171260ad1d28a02
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_collectives/StoreCollectives.hpp
@@ -0,0 +1,73 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Macros.h>
+#include <c10/util/FbcodeMaps.h>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/control_collectives/ControlCollectives.hpp>
+
+namespace c10d {
+
+class TORCH_API StoreCollectives : public ControlCollectives {
+ public:
+  explicit StoreCollectives(
+      c10::intrusive_ptr<Store> store,
+      int rank,
+      int worldSize);
+
+  void barrier(
+      const std::string& key,
+      std::chrono::milliseconds timeout = 5min,
+      bool block = true) override;
+
+  void broadcastSend(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) override;
+  std::vector<uint8_t> broadcastRecv(
+      const std::string& key,
+      std::chrono::milliseconds timeout = 5min) override;
+
+  void gatherSend(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) override;
+  std::vector<std::vector<uint8_t>> gatherRecv(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) override;
+
+  std::vector<uint8_t> scatterSend(
+      const std::string& key,
+      const std::vector<std::vector<uint8_t>>& data,
+      std::chrono::milliseconds timeout = 5min) override;
+  std::vector<uint8_t> scatterRecv(
+      const std::string& key,
+      std::chrono::milliseconds timeout = 5min) override;
+
+  std::vector<std::vector<uint8_t>> allGather(
+      const std::string& key,
+      const std::vector<uint8_t>& data,
+      std::chrono::milliseconds timeout = 5min) override;
+
+  int64_t allSum(
+      const std::string& key,
+      int64_t data,
+      std::chrono::milliseconds timeout = 5min) override;
+
+ private:
+  void enforceUnique(const std::string& key);
+
+ private:
+  c10::intrusive_ptr<Store> store_;
+  int rank_;
+  int worldSize_;
+
+  c10::FastSet<std::string> seenKeys_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/Handlers.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/Handlers.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..bf79294f23d5515edbe5c98dc20b6b1ae16fdb5a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/Handlers.hpp
@@ -0,0 +1,82 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <functional>
+#include <map>
+#include <string>
+#include <utility>
+
+#include <c10/macros/Export.h>
+
+namespace c10d::control_plane {
+
+// Request represents a request to the handler. This conceptually maps to an
+// HTTP request but could be called via other transports.
+class TORCH_API Request {
+ public:
+  virtual ~Request() = default;
+
+  virtual const std::string& body() const = 0;
+
+  virtual const std::multimap<std::string, std::string>& params() const = 0;
+
+  std::string getParam(const std::string& key) const {
+    auto it = params().find(key);
+    if (it != params().end()) {
+      return it->second;
+    }
+    return "";
+  }
+};
+
+// Response represents a response to the handler. This conceptually maps to an
+// HTTP response but could be called via other transports.
+class TORCH_API Response {
+ public:
+  virtual ~Response() = default;
+
+  // Set the response body to the provided string.
+  // TODO: add support for chunked responses
+  virtual void setContent(
+      std::string&& content,
+      const std::string& content_type) = 0;
+
+  // Set the response status code.
+  // These should match standard HTTP status codes.
+  virtual void setStatus(int status) = 0;
+};
+
+using HandlerFunc = std::function<void(const Request&, Response&)>;
+
+// Registers a handler. The name needs to be unique and can be called by using
+// getHandler directly or via WorkerServer for remote requests.
+// These handlers are called from a background C++ thread concurrently with the
+// main thread. These handlers need to be thread safe and not cause issues
+// during Python training.
+TORCH_API void registerHandler(const std::string& name, HandlerFunc f);
+
+// Fetches a handler by name.
+TORCH_API HandlerFunc getHandler(const std::string& name);
+
+TORCH_API std::vector<std::string> getHandlerNames();
+
+// Registers a handler statically.
+// See registerHandler for more details.
+class TORCH_API RegisterHandler {
+ public:
+  RegisterHandler(const std::string& name, HandlerFunc f) {
+    registerHandler(name, std::move(f));
+  }
+
+  // disable move, copy
+  RegisterHandler(const RegisterHandler&) = delete;
+  RegisterHandler(RegisterHandler&&) = delete;
+  RegisterHandler& operator=(const RegisterHandler&) = delete;
+  RegisterHandler& operator=(RegisterHandler&&) = delete;
+};
+
+} // namespace c10d::control_plane
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/WaitCounterHandler.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/WaitCounterHandler.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..b6e73ec065fe5b1b094f4fabd97ad887a9b94b68
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/WaitCounterHandler.hpp
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+
+namespace c10d {
+namespace control_plane {
+
+// Returns all wait counter values as a JSON string
+std::string getWaitCounterValuesJson();
+
+// Ensures the wait counter backend is registered
+void ensureWaitCounterBackendRegistered();
+
+} // namespace control_plane
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/WorkerServer.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/WorkerServer.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..9ffaa77726734d694015d8900ae759b762deff6b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/control_plane/WorkerServer.hpp
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <thread>
+
+#include <c10/util/intrusive_ptr.h>
+#include <torch/csrc/distributed/c10d/control_plane/Handlers.hpp>
+
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wdeprecated-literal-operator")
+#include <httplib.h>
+C10_DIAGNOSTIC_POP()
+
+namespace c10d::control_plane {
+
+class TORCH_API WorkerServer : public c10::intrusive_ptr_target {
+ public:
+  WorkerServer(const std::string& hostOrFile, int port = -1);
+  ~WorkerServer() override;
+
+  void shutdown();
+
+  int port() {
+    return port_;
+  }
+
+ private:
+  httplib::Server server_;
+  std::thread serverThread_;
+  int port_;
+};
+
+} // namespace c10d::control_plane
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/CUDAEventCache.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/CUDAEventCache.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..e289e3d0aff7561ab197ce8c3e18ccf33d2c9ff9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/CUDAEventCache.hpp
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <deque>
+#include <memory>
+#include <mutex>
+
+#include <ATen/cuda/CUDAEvent.h>
+#include <c10/macros/Export.h>
+
+namespace c10d {
+
+class TORCH_API CUDAEventCache
+    : public std::enable_shared_from_this<CUDAEventCache> {
+ public:
+  CUDAEventCache();
+  std::shared_ptr<at::cuda::CUDAEvent> create(bool timing);
+  static std::shared_ptr<CUDAEventCache> get(at::DeviceIndex device);
+
+ private:
+  std::mutex cacheMutex_;
+  // NOTE: We intentionally store raw pointers so that
+  // we do not attempt to destroy the event objects on process exit,
+  // because cuda may be gone.
+  std::array<std::deque<at::cuda::CUDAEvent*>, 2>
+      eventsArray_; // 0 for timing=false, 1 for timing=true
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/StreamBlock.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/StreamBlock.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..8195a6faa33d4f86c771ff08b254022e43b3cacd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/StreamBlock.hpp
@@ -0,0 +1,43 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <chrono>
+#include <memory>
+
+#include <c10/util/Registry.h>
+
+namespace c10d::cuda {
+
+enum StreamBlockStatus : int32_t {
+  UNKNOWN = 0,
+  RUNNING = 1,
+  TIMED_OUT = 2,
+  ABORTED = 3,
+};
+
+/*
+StreamBlock implements a baton that will block a the active CUDA stream
+until aborted by the main process.
+*/
+class TORCH_API StreamBlock {
+ public:
+  virtual ~StreamBlock() = default;
+  virtual void abort() = 0;
+  virtual StreamBlockStatus status() = 0;
+};
+
+std::unique_ptr<StreamBlock> block_stream(std::chrono::milliseconds timeout);
+
+// Declare a registry so we can call the CUDA StreamBlock API from CPU only code
+// (i.e. ProcessGroup/Work objects in libtorch_cpu).
+// The implementation lives defined in StreamBlock.cu.
+TORCH_DECLARE_REGISTRY(
+    StreamBlockRegistry,
+    StreamBlock,
+    std::chrono::milliseconds);
+
+} // namespace c10d::cuda
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/utils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/utils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..d8a2520f4fd00d62e506f65e8645c1b035fee2c7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/cuda/utils.hpp
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// This file contains utility functions common for CUDA, which can be used by
+// ProcessGroupNCCL or SymmetricMemory.
+
+namespace c10d::cuda {
+
+bool deviceSupportsMulticast(int device_idx);
+
+} // namespace c10d::cuda
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/debug.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/debug.h
new file mode 100644
index 0000000000000000000000000000000000000000..cf343d0eef5a8b1a06ab96d0fbf0a19aaa61f87f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/debug.h
@@ -0,0 +1,28 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <c10/macros/Macros.h>
+
+namespace c10d {
+
+enum class DebugLevel { Off = 0, Info = 1, Detail = 2 };
+
+TORCH_API void setDebugLevel(DebugLevel level);
+
+// Sets the debug level based on the value of the `TORCH_DISTRIBUTED_DEBUG`
+// environment variable.
+TORCH_API void setDebugLevelFromEnvironment();
+
+TORCH_API DebugLevel debug_level() noexcept;
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/default_comm_hooks.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/default_comm_hooks.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..c56d6097a4d48dfe170ebab6b5e49d7f8504e83e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/default_comm_hooks.hpp
@@ -0,0 +1,57 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/distributed/c10d/comm.hpp>
+
+namespace c10d {
+
+enum class BuiltinCommHookType : uint8_t {
+  ALLREDUCE = 1,
+  FP16_COMPRESS = 2,
+};
+
+class AllReduceCommHook
+    : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {
+ public:
+  explicit AllReduceCommHook(const c10::intrusive_ptr<ProcessGroup>& state)
+      : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}
+
+  ~AllReduceCommHook() override = default;
+
+  c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;
+};
+
+class FP16CompressCommHook
+    : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {
+ public:
+  explicit FP16CompressCommHook(const c10::intrusive_ptr<ProcessGroup>& state)
+      : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}
+
+  ~FP16CompressCommHook() override = default;
+
+  c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;
+};
+
+// Almost same as AllReduceCommHook, but without division inside the hook.
+// This enables the optimization of fusing copy and division and saves one scan
+// over all the input parameters, when no communication hook is provided by the
+// user. Only used internally and not released as a public built-in
+// communication hook.
+class _AllReduceBySumCommHook
+    : public CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>> {
+ public:
+  explicit _AllReduceBySumCommHook(
+      const c10::intrusive_ptr<ProcessGroup>& state)
+      : CppCommHookInterface<c10::intrusive_ptr<ProcessGroup>>(state) {}
+
+  ~_AllReduceBySumCommHook() override = default;
+
+  c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/error.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/error.h
new file mode 100644
index 0000000000000000000000000000000000000000..e4e1a3e2d2be49904528720168ffc383b4332f9e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/error.h
@@ -0,0 +1,58 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <cstring>
+#include <system_error>
+
+#include <fmt/format.h>
+
+namespace fmt {
+
+template <>
+struct formatter<std::error_category> {
+  constexpr auto parse(format_parse_context& ctx) const {
+    return ctx.begin();
+  }
+
+  template <typename FormatContext>
+  auto format(const std::error_category& cat, FormatContext& ctx) const {
+    if (std::strcmp(cat.name(), "generic") == 0) {
+      return fmt::format_to(ctx.out(), "errno");
+    } else {
+      return fmt::format_to(ctx.out(), "{} error", cat.name());
+    }
+  }
+};
+
+template <>
+struct formatter<std::error_code> {
+  constexpr auto parse(format_parse_context& ctx) const {
+    return ctx.begin();
+  }
+
+  template <typename FormatContext>
+  auto format(const std::error_code& err, FormatContext& ctx) const {
+    return fmt::format_to(
+        ctx.out(), "({}: {} - {})", err.category(), err.value(), err.message());
+  }
+};
+
+} // namespace fmt
+
+namespace c10d::detail {
+
+inline std::error_code lastError() noexcept {
+  return std::error_code{errno, std::generic_category()};
+}
+
+} // namespace c10d::detail
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/exception.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/exception.h
new file mode 100644
index 0000000000000000000000000000000000000000..a79a93048797ceba3740f83fbef2b05ff5516121
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/exception.h
@@ -0,0 +1,45 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// @allow-raw-throw
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <c10/macros/Macros.h>
+#include <c10/util/Exception.h>
+
+// Utility macro similar to C10_THROW_ERROR, the major difference is that this
+// macro handles exception types defined in the c10d namespace, whereas
+// C10_THROW_ERROR requires an exception to be defined in the c10 namespace.
+#define C10D_THROW_ERROR(err_type, ...)                      \
+  throw ::c10d::err_type(                                    \
+      {__func__, __FILE__, static_cast<uint32_t>(__LINE__)}, \
+      c10::str(__VA_ARGS__))
+
+#define C10D_CHECK_WITH(error_t, cond, ...)                         \
+  if (C10_UNLIKELY_OR_CONST(!(cond))) {                             \
+    C10D_THROW_ERROR(                                               \
+        error_t, TORCH_CHECK_MSG(cond, "", c10::str(__VA_ARGS__))); \
+  }
+
+namespace c10d {
+
+using c10::DistNetworkError;
+using c10::DistStoreError;
+
+class TORCH_API SocketError : public DistNetworkError {
+  using DistNetworkError::DistNetworkError;
+};
+
+class TORCH_API TimeoutError : public DistNetworkError {
+  using DistNetworkError::DistNetworkError;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/logger.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/logger.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..9cb5cf16044d0e6fc6b6a8fdd8774a17c7cfcaf4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/logger.hpp
@@ -0,0 +1,176 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/Logging.h>
+#include <torch/csrc/distributed/c10d/reducer.hpp>
+
+#include <utility>
+
+namespace c10d {
+
+// A struct to hold the latest status of the process group.
+struct ProcessGroupStatus {
+  // the sequential number of the last collective enqueued into workMetaList_
+  // This is useful for identifying a rank that has not join a collective
+  // initialized to be -1 to indicate no collective has been enqueued
+  int64_t lastEnqueuedSeq{-1};
+  // the sequential number of the last collective started as the kernel
+  int64_t lastStartedSeq{-1};
+  // the sequential number of the last collective completed marked by
+  // the watchdog thread
+  // initialized to be -1 to indicate no collective has been completed
+  int64_t lastCompletedSeq{-1};
+
+  // the name of the last collective enqueued into workMetaList_
+  std::string lastEnqueuedWorkName;
+  // the name of the last collective started as the kernel
+  std::string lastStartedWorkName;
+  // the name of the last collective completed
+  std::string lastCompletedWorkName;
+
+  // the sizes of the last work enqueued
+  size_t lastEnqueuedNumelIn;
+  size_t lastEnqueuedNumelOut;
+  // the sizes of the last work completed
+  size_t lastCompletedNumelIn;
+  size_t lastCompletedNumelOut;
+  // the sizes of the last work started
+  size_t lastStartedNumelIn;
+  size_t lastStartedNumelOut;
+};
+
+class TORCH_API Logger {
+ public:
+  explicit Logger(std::shared_ptr<c10d::Reducer> reducer);
+  // Set logging data that can be got during DistributedDataParallel
+  // construction time.
+  void set_construction_data_and_log(
+      const std::string& module_name,
+      const std::vector<int>& device_ids,
+      int output_device,
+      bool broadcast_buffers,
+      bool has_sync_bn,
+      bool static_graph);
+
+  void set_static_graph();
+
+  // An interface for users to get DDPLoggingData and log them
+  // in the applications. Explanation of logging fields are in
+  // "struct DDPLoggingData" of "torch/c10/util/Logging.h".
+  at::DDPLoggingData get_ddp_logging_data();
+
+  // Stream insertion operator for logging data to stream under
+  // TORCH_DISTRIBUTED_DEBUG.
+  friend std::ostream& operator<<(std::ostream& output, const Logger& logger);
+
+  ~Logger() noexcept(false) {
+    // Log if DDP graph is static in Logger dtor instead of Reducer dtor since
+    // Logger is deleted before Reducer.
+    log_if_graph_static(reducer_->ddp_graph_static());
+  }
+
+  // Set environment variables.
+  void set_env_variables();
+  // Set parameters stats.
+  void set_parameter_stats();
+  // Get size of each bucket (Bytes).
+  std::vector<int64_t> get_bucket_sizes();
+  // Get variable indices for each bucket.
+  std::vector<std::vector<size_t>> get_per_bucket_variable_indices();
+  // Set comm. hook, if used
+  void set_comm_hook(const std::string& hook);
+  // Set running with uneven input detection (model.join() context manager)
+  void set_uneven_input_join();
+
+  // Reset performance stats at current iteration
+  void reset_performance_stats();
+
+  // Calculate avg stats using cpu timer and gpu timer
+  // that has been recorded in reducer.
+  void calculate_avg_time(
+      int64_t& avg_time,
+      int64_t& time_duration,
+      Timer& timer,
+      Timer::Event start_event,
+      Timer::Event end_event);
+
+  // Set the absolute time of the event that has been recorded in reducer.
+  void set_event_time(int64_t& event_time, Timer& timer, Timer::Event event);
+  // Set stats that can be collected only during
+  // training loop. It is called at the beginning of forward call
+  // to record the run time stats of sampled iterations that previously ran.
+  // GPU performance stats are collected only for single process
+  // single device program and single device module right now.
+  // TODO to support single process multiple devices and multi device modules,
+  // events need to be created and recorded on multiple devices.
+  void set_runtime_stats_and_log();
+
+  // Called when DDP/reducer is failing with an error. The
+  // logging data structure will have two fields filled: "has_error" indicating
+  // that this iteration encountered an error and other fields are not valid,
+  // and "error", a string which contains the error message that DDP failed
+  // with.
+  template <typename... Args>
+  void set_error_and_log(const std::string& ddp_error, const Args&... args) {
+    ddp_logging_data_->ints_map["has_error"] = 1;
+    auto err = c10::str(ddp_error, args...);
+    ddp_logging_data_->strs_map["error"] = err;
+    // Report the iteration we are erroring at so user knows how many examples
+    // successfully processed before this error was hit.
+    ddp_logging_data_->ints_map["iteration"] = reducer_->num_iterations_;
+    at::LogPyTorchDDPUsage(*ddp_logging_data_);
+  }
+
+  // When running without static graph, called when reducer is destroyed to log
+  // if graph was actually static and is a candidate for static graph
+  // optimization.
+  void log_if_graph_static(bool is_static);
+
+ private:
+  // ddp_logging_data_ is used to hold all the ddp related logging
+  // data fields.
+  std::unique_ptr<at::DDPLoggingData> ddp_logging_data_;
+  std::shared_ptr<c10d::Reducer> reducer_;
+  // track the number of iterations when runtime stats are collected so far.
+  long num_iterations_stats_recorded_ = 0;
+};
+
+// a generic logging data struct that holds different types of logging data.
+// starting with key value pairs of strings and integers,
+// It can be extended to more types as needed.
+struct C10dLoggingData {
+  // logging fields that are string types.
+  std::map<std::string, std::string> strings;
+  // logging fields that are int64_t types.
+  std::map<std::string, int64_t> integers;
+};
+
+class TORCH_API C10dLogger {
+ public:
+  C10dLogger(const C10dLogger&) = default;
+  C10dLogger(C10dLogger&&) = delete;
+  C10dLogger& operator=(const C10dLogger&) = default;
+  C10dLogger& operator=(C10dLogger&&) = delete;
+  virtual ~C10dLogger() = default;
+  virtual void log(const C10dLoggingData& data);
+  static C10dLogger* getLogger();
+  static void registerLogger(std::unique_ptr<C10dLogger> /*logger*/);
+
+ protected:
+  // singletion, hide constructor from the public
+  C10dLogger(std::string logDestination)
+      : logDestination_(std::move(logDestination)) {}
+
+  // the name of the destination this logger should log to
+  std::string logDestination_;
+
+ private:
+  static std::unique_ptr<C10dLogger> logger_;
+  static std::atomic<bool> registered_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/logging.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/logging.h
new file mode 100644
index 0000000000000000000000000000000000000000..596e4686212800c6ac419ac9fd77519d90257160
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/logging.h
@@ -0,0 +1,52 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <string>
+
+#include <c10/macros/Macros.h>
+#include <c10/util/Logging.h>
+#include <fmt/format.h>
+
+namespace c10d::detail {
+
+enum class LogLevel { Trace, Debug, Info, Warning, Error };
+
+TORCH_API bool isLogLevelEnabled(LogLevel level) noexcept;
+
+template <typename... T>
+// NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
+std::string formatLogMessage(fmt::string_view fmt, T&&... args) {
+  return fmt::vformat(fmt, fmt::make_format_args(args...));
+}
+
+} // namespace c10d::detail
+
+#define C10D_ERROR(...)                                               \
+  if (c10d::detail::isLogLevelEnabled(c10d::detail::LogLevel::Error)) \
+  LOG(ERROR) << "[c10d] " << c10d::detail::formatLogMessage(__VA_ARGS__)
+
+#define C10D_WARNING(...)                                               \
+  if (c10d::detail::isLogLevelEnabled(c10d::detail::LogLevel::Warning)) \
+  LOG(WARNING) << "[c10d] " << c10d::detail::formatLogMessage(__VA_ARGS__)
+
+#define C10D_INFO(...)                                               \
+  if (c10d::detail::isLogLevelEnabled(c10d::detail::LogLevel::Info)) \
+  LOG(INFO) << "[c10d] " << c10d::detail::formatLogMessage(__VA_ARGS__)
+
+#define C10D_DEBUG(...)                                               \
+  if (c10d::detail::isLogLevelEnabled(c10d::detail::LogLevel::Debug)) \
+  LOG(INFO) << "[c10d - debug] " << c10d::detail::formatLogMessage(__VA_ARGS__)
+
+#define C10D_TRACE(...)                                               \
+  if (c10d::detail::isLogLevelEnabled(c10d::detail::LogLevel::Trace)) \
+  LOG(INFO) << "[c10d - trace] " << c10d::detail::formatLogMessage(__VA_ARGS__)
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/python_callback_work.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/python_callback_work.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..27c25dc8235c65258a8fa05d88eca27ceb87416a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/python_callback_work.hpp
@@ -0,0 +1,33 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Work.hpp>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <torch/csrc/utils.h>
+
+namespace c10d {
+
+// PythonCallbackWork is a subclass of Work that wraps a Python callback
+// function that implements wait(). This allows asynchronous work to
+// be integrated with Python code, enabling custom completion logic or
+// post-processing in Python.
+class PythonCallbackWork : public Work {
+ public:
+  explicit PythonCallbackWork(py::function callback);
+
+  ~PythonCallbackWork() override;
+
+  bool wait(std::chrono::milliseconds timeout) override;
+
+  c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
+
+ private:
+  py::function callback_;
+  c10::intrusive_ptr<c10::ivalue::Future> future_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/python_comm_hook.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/python_comm_hook.h
new file mode 100644
index 0000000000000000000000000000000000000000..e6e985f87ced478c67147c8405303e682aac7403
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/python_comm_hook.h
@@ -0,0 +1,39 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/comm.hpp>
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/utils/pybind.h>
+
+namespace c10d {
+
+class TORCH_PYTHON_API PythonCommHook : public CommHookInterface {
+ public:
+  // Takes a state and a callable hook. The inputs are Python objects.
+  // The state is passed to the hook in runHook method, and it can be used to
+  // maintain and update any state information during the execution of the hook.
+  // The hook performs user-specified processing and returns a future indicating
+  // asynchronous communication of gradients.
+  PythonCommHook(py::object state, py::object hook)
+      : state_(std::move(state)), hook_(std::move(hook)) {}
+
+  ~PythonCommHook() override;
+
+  c10::intrusive_ptr<c10::ivalue::Future> runHook(GradBucket& bucket) override;
+
+  at::Tensor parseHookResult(const c10::IValue& result) override;
+
+ private:
+  // Only needed for stateful communication.
+  py::object state_;
+  py::object hook_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization.h
new file mode 100644
index 0000000000000000000000000000000000000000..08e8f2948af7ec89093fb8c7701bf7e2fe21a696
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <ATen/ATen.h>
+
+namespace torch::distributed::c10d::quantization {
+
+at::Tensor _float_to_bfloat16_cpu(const at::Tensor& input);
+at::Tensor _bfloat16_to_float_cpu(const at::Tensor& input);
+
+} // namespace torch::distributed::c10d::quantization
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization_gpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization_gpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..9c21a8f5a07cfab89ce47f08d201abb1f7e1bb32
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization_gpu.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <ATen/ATen.h>
+
+namespace torch::distributed::c10d::quantization {
+
+at::Tensor _float_to_bfloat16_cuda(const at::Tensor& input);
+at::Tensor _bfloat16_to_float_cuda(const at::Tensor& input);
+
+} // namespace torch::distributed::c10d::quantization
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..a2a6a171ba92a0741a7cc6f1990987917ff1b657
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/quantization/quantization_utils.h
@@ -0,0 +1,39 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <ATen/ATen.h>
+
+#include <typeinfo>
+
+inline std::string torch_tensor_device_name(const at::Tensor& ten) {
+  return c10::DeviceTypeName(ten.device().type());
+}
+
+#define TENSOR_NDIM_EQUALS(ten, dims)      \
+  TORCH_CHECK(                             \
+      (ten).ndimension() == (dims),        \
+      "Tensor '" #ten "' must have " #dims \
+      " dimension(s). "                    \
+      "Found ",                            \
+      (ten).ndimension())
+
+#define TENSOR_ON_CPU(x)                                      \
+  TORCH_CHECK(                                                \
+      !x.is_cuda(),                                           \
+      #x " must be a CPU tensor; it is currently on device ", \
+      torch_tensor_device_name(x))
+
+#define TENSOR_ON_CUDA_GPU(x)                                  \
+  TORCH_CHECK(                                                 \
+      x.is_cuda(),                                             \
+      #x " must be a CUDA tensor; it is currently on device ", \
+      torch_tensor_device_name(x))
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/reducer.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/reducer.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..ee1f83bea3da04a63ef43d4ecb8ad2442e9a9d64
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/reducer.hpp
@@ -0,0 +1,624 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/ScalarType.h>
+#include <atomic>
+#include <memory>
+#include <mutex>
+#include <tuple>
+#include <unordered_map>
+#include <vector>
+
+#include <ATen/core/ivalue_inl.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/intrusive_ptr.h>
+#include <torch/csrc/autograd/function.h>
+#include <torch/csrc/autograd/profiler.h>
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
+#include <torch/csrc/distributed/c10d/Utils.hpp>
+#include <torch/csrc/distributed/c10d/comm.hpp>
+#include <torch/csrc/distributed/c10d/debug.h>
+#include <torch/csrc/distributed/c10d/default_comm_hooks.hpp>
+#include <torch/csrc/distributed/c10d/reducer_timer.hpp>
+#ifndef _WIN32
+#include <torch/csrc/distributed/autograd/context/context.h>
+#endif
+
+namespace c10d {
+
+constexpr int kDefaultFirstBucketBytes = 1024 * 1024;
+constexpr int kDefaultBucketBytesCap = 25 * 1024 * 1024;
+// Collect runtime stats once for every kDDPRuntimeLoggingSampleRate iterations.
+constexpr int kDDPRuntimeLoggingSampleRate = 100;
+
+// Forward declaration
+class Logger;
+
+// Local accumulator type for a single bucket.
+struct BucketAccumulator {
+  std::vector<size_t> indices;
+  size_t size = 0;
+  size_t size_limit = 0;
+};
+
+class TORCH_API Reducer {
+ public:
+  // The constructor takes a list of variables (i.e. parameters) for this
+  // process's single model replica (as DDP assumes single-process
+  // single-device). The bucket assignment for this reducer, `bucket_indices`,
+  // is specified as a list of buckets, each of which is specified as a list of
+  // indices into the bucket's `variables` list.
+  explicit Reducer(
+      std::vector<at::Tensor> params,
+      std::vector<std::vector<size_t>> bucket_indices,
+      c10::intrusive_ptr<c10d::ProcessGroup> process_group,
+      std::vector<bool> expect_sparse_gradients,
+      int64_t bucket_bytes_cap,
+      bool find_unused_parameters,
+      bool gradient_as_bucket_view,
+      std::unordered_map<size_t, std::string> param_names,
+      int64_t first_bucket_bytes_cap,
+      bool skip_all_reduce_unused_params,
+      bool use_python_reducer,
+      std::vector<int64_t> bucket_bytes_cap_list,
+      bool batched_grad_copy = false);
+
+  ~Reducer() noexcept(false);
+
+  // To (re-)initialize bucket assignment, pass a list of buckets, each of
+  // which is specified by a list of indices in the bucket's `variables` list.
+  // This function performs validation that the variables within a bucket
+  // all live on the same device and have the same dimensionality.
+  void initialize_buckets(std::vector<std::vector<size_t>> bucket_indices);
+
+  void autograd_hook(size_t index);
+
+  // This function is called when the forward function has produced an output,
+  // and the user wishes to reduce gradients in the backwards pass.
+  // If they don't, and wish to accumulate gradients before reducing them,
+  // a call to this function can simply be omitted.
+  void prepare_for_backward(const std::vector<at::Tensor>& outputs);
+
+  // Called at the beginning of forward() inside DistributedDataParallel,
+  // right now it captures the starting time of forward in each iteration.
+  void prepare_for_forward();
+
+  // Returns the relative time in nanoseconds when gradients were ready,
+  // with respect to the time `prepare_for_backward` was called. The
+  // vector is for parameters for a single model replica.
+  std::vector<int64_t> get_backward_stats() const {
+    return backward_stats_;
+  }
+
+  // Registers a hook to the reducer. The hook is `CommHookInterface`
+  // type to allow both Python and CPP hooks. This function can only
+  // be called once before calling backward.
+  // Cannot combine with the call of `register_builtin_comm_hook`.
+  void register_comm_hook(std::unique_ptr<CommHookInterface> iface);
+
+  // Registers a built-in C++ comm hook to the reducer. This function can only
+  // be called once before calling backward.
+  // Cannot combine with the call of `register_comm_hook`.
+  void register_builtin_comm_hook(c10d::BuiltinCommHookType comm_hook_type);
+
+  // Informs reducer that optimizer is running in backward, so gradients
+  // don't need to be copied from buckets as the optimizer would've already
+  // been applied.
+  void set_optimizer_in_backward() {
+    optim_in_backward_ = true;
+  }
+
+  // Runs allreduce or installed communication hook given GradBucket instance.
+  c10::intrusive_ptr<c10::ivalue::Future> run_comm_hook(
+      GradBucket& grad_bucket);
+
+  // Runs default allreduce hook.
+  c10::intrusive_ptr<c10::ivalue::Future> run_allreduce_hook(
+      GradBucket& grad_bucket);
+
+  // Returns gradient buckets in sequential order of buckets_. This is the order
+  // in which buckets are reduced across processes. If return_zero_tensors=true,
+  // will return zero tensors of the same shape instead of the true tensors.
+  std::vector<c10d::GradBucket> get_grad_buckets(
+      bool return_zero_tensors = true) const;
+
+  // Rebuild buckets based on rebuilt_params_ and rebuilt_param_indices_
+  // according to when tensors received grads in the backward pass.
+  // TODO this function makes broadcast communication call and
+  // could be overlapped with next forward() call, thus
+  // it could be async. Will make it async when rebuilding buckets for
+  // find_unused_parameters = true case, as we could rebuild buckets more than
+  // once for find_unused_parameters = true case, where subgraphs are trained
+  // and parameter indices order may change more frequently.
+  // For find_unused_parameters = false case, buckets are only rebuilt once,
+  // the performance cost is negligible. Returns true if the buckets were
+  // rebuilt.
+  bool rebuild_buckets();
+
+  void setSparseMetadata(std::map<std::string, at::Tensor>& metadata);
+
+  // Install futures that should be awaited at end of backwards. Currently these
+  // are only used by user-defined custom buffer reduction hooks, but can be
+  // generalized to any user-originating futures that need to be awaited.
+  void install_futures(
+      const c10::List<c10::intrusive_ptr<c10::ivalue::Future>>& futs);
+
+  // Returns true if we should rebuild buckets, else false. We only rebuild
+  // buckets once after the first iteration and never rebuild them if
+  // find_unused_parameters_.
+  inline bool should_rebuild_buckets() const {
+    return (static_graph_ || !find_unused_parameters_) && !has_rebuilt_bucket_;
+  }
+
+  // Pushes all parameters to be rebuilt.
+  void push_rebuilt_params_for_all_indices();
+
+  // Creates and sets ForwardPassWorkHandle given a Work and the
+  // corresponding tensor being reduced.
+  void set_forward_pass_work_handle(
+      c10::intrusive_ptr<c10d::Work> forwardPassWorkHandle,
+      bool useStaticWorldSize);
+
+  // Retrieve on-device tensors used to track locally unused parameters. It is
+  // a tensor where index i = 1 if the Variable with that index has been used.
+  at::Tensor get_local_used_map_on_device() const;
+
+  // An function for users to set sample_rate of collecting
+  // runtime stats. The time stats will be recorded for the
+  // first 10 iterations, after 10 iterations time stats will be
+  // recorded once every "sample_rate" training iterations.
+  void set_ddp_runtime_logging_sample_rate(int sample_rate);
+
+  // Specify the training graph is static.
+  void set_static_graph();
+
+  // Delay all reduce to be after all gradients' calculation is complete.
+  void delay_all_reduce();
+
+  void set_mixed_precision_param_dtype(c10::ScalarType dtype);
+
+  // Weak reference to associated DDP logger. The reference is weak to avoid
+  // refcycle between reducer and logger.
+  void set_logger(std::weak_ptr<c10d::Logger> logger);
+
+  // When graph is not explicitly set by user as static and has unused
+  // parameters, this will return whether the graph has been static until the
+  // current iteration, which means unused params set has not changed.
+  bool ddp_graph_static();
+
+  // Removes autograd hooks registered by the Reducer on the model parameters.
+  void remove_autograd_hooks();
+
+  // Checks whether or not the reducer has finalized the current backward
+  // iteration.
+  void check_finalized();
+
+  // Updates the underlying process group used by DDP with the new process
+  // group.
+  void update_process_group(
+      c10::intrusive_ptr<c10d::ProcessGroup> new_process_group);
+
+  // Resets reducer state.
+  void reset_state();
+
+ protected:
+  // Forward declaration.
+  struct Bucket;
+
+  void push_rebuilt_params(const size_t& index);
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  mutable std::mutex mutex_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::vector<at::Tensor> params_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  c10::intrusive_ptr<::c10d::ProcessGroup> process_group_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::vector<bool> expect_sparse_gradients_;
+
+  std::vector<std::shared_ptr<torch::autograd::Node>>
+      grad_accumulators_; // NOLINT(cppcoreguidelines-non-private-member-variables-in-classes)
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::unordered_map<torch::autograd::Node*, size_t> gradAccToVariableMap_;
+  std::vector<std::pair<uintptr_t, std::shared_ptr<torch::autograd::Node>>>
+      hooks_; // NOLINT(cppcoreguidelines-non-private-member-variables-in-classes)
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  bool expect_autograd_hooks_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  bool require_finalize_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  size_t next_bucket_;
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  bool has_marked_unused_parameters_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const bool find_unused_parameters_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const bool gradient_as_bucket_view_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const bool batched_grad_copy_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::vector<size_t> unused_parameters_;
+  // Previous iteration's unused params, used for checking if unused parameters
+  // change between iterations. Only filled during the first backwards call.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::vector<size_t> prev_iteration_unused_parameters_;
+  // Whether graph is static or not. When user does not explicitly set static
+  // graph, the only possible dynamism is set of unused parameters changing
+  // between iterations which is tracked by this flag.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  bool ddp_graph_static_{true};
+  // Locally used parameter maps indicating if parameters are used locally
+  // during the current iteration or no_sync session if no_sync is on.
+  // Each map is a one-dim int32 tensor of number of parameters. These tensors
+  // are marked in autograd_hook to indicate the corresponding param has been
+  // used, and get allreduced in the end of backward step of current iteration
+  // or no_sync session for figuring out the globally unused parameters.
+  //
+  // local_used_map_:     CPU tensor for bookkeeping locally used params
+  // local_used_map_dev_: dev tensor for reducing globally unused params
+  at::Tensor local_used_map_;
+  at::Tensor local_used_map_dev_;
+  // Indicate that reduction is done and D2H copy is done as well.
+  bool local_used_map_reduced_;
+
+  // Weak pointer to associated DDP logger.
+  std::weak_ptr<c10d::Logger> logger_;
+  // List of futures installed by Reducer::install_futures that should be
+  // awaited at the end of backwards pass.
+  std::optional<c10::List<c10::intrusive_ptr<c10::ivalue::Future>>>
+      installed_futures_{std::nullopt};
+  // Mixed precision parameter dtype for bucket type checking.
+  std::optional<c10::ScalarType> mixed_precision_param_dtype_{std::nullopt};
+
+  // Work handle for allreduce on local_used_map_
+  c10::intrusive_ptr<c10d::Work> local_used_work_;
+
+  void mark_variable_ready_dense(size_t variable_index);
+
+  void mark_variable_ready_sparse(size_t variable_index);
+
+  void mark_variable_ready(size_t variable_index);
+
+  void mark_bucket_ready(size_t bucket_index);
+
+  void finalize_bucket_dense(Bucket& bucket);
+
+  void finalize_backward();
+
+  // Returns list of model parameters corresponding to the given bucket.
+  // bucket_index is a key to cache after buckets are rebuilt, after which this
+  // mapping never changes.
+  std::vector<at::Tensor> get_variables_for_bucket(
+      size_t bucket_index,
+      const Bucket& bucket) const;
+
+  // Asserts that the reduction for the previous iteration has finished before
+  // rebuilding buckets or kicking off the next one.
+  void ensure_prior_reduction_finished();
+
+  // Broadcast rebuilt buckets from rank 0 to other ranks before initializing
+  // the buckets
+  void sync_bucket_indices(std::vector<std::vector<size_t>>& bucket_indices);
+
+  // We'd like to use DistAutogradContext::GradCallback here but dist autograd
+  // doesn't exist under Windows. So we just directly use the concrete type but
+  // to preserve and enforce our original intent we do a static assert when dist
+  // autograd is available.
+  using GradCallback = std::function<bool(at::Tensor&)>;
+#ifndef _WIN32
+  static_assert(
+      std::is_same_v<
+          GradCallback,
+          torch::distributed::autograd::DistAutogradContext::GradCallback>);
+#endif
+  void runGradCallbackForVariable(at::Tensor& variable, const GradCallback& cb);
+
+  // Flushes deferred grad-to-bucket copies for a single bucket when
+  // batched_grad_copy_ is enabled. Called from mark_variable_ready (when
+  // bucket.pending == 0) and from delay_all_reduce (after all variables
+  // are marked ready).
+  void flush_deferred_copies(Bucket& bucket, size_t bucket_index);
+
+  // This function is called inside `initialize_buckets()`. It initializes both
+  // `bucket_views_in` and `bucket_views_out` with views for each variable's
+  // gradient into the bucket's flattened `gradients` tensor. Views serve as
+  // entry points to `copy_()` each grad's data in/out of the flattened
+  // `gradients` tensor.
+  void initialize_bucket_views(Bucket& bucket);
+
+  // This function is called inside `finalize_backward`, it happens only if
+  // DDP communication hook was registered to recreate just bucket_views_out
+  // with the result of `future_work`.
+  void populate_bucket_views_out(Bucket& bucket, at::Tensor& tensor);
+
+  // If gradient_as_bucket_view_ is false, after allreduce buckets,
+  // copy bucket results back to grads.
+  void copy_bucket_to_grad(
+      at::Tensor& variable,
+      Reducer::Bucket& bucket,
+      size_t intra_bucket_index,
+      bool global_unused);
+  // Check layout of grad and bucket_view before copying the grad to bucket.
+  void check_grad_layout(const at::Tensor& grad, const at::Tensor& bucket_view);
+
+  // A bucket contains [1..N] gradients to be reduced, where the gradients
+  // have the same dtype and device.
+  // Coalescing gradients together before reducing can result in lower overhead
+  // and/or faster time to completion. Coalescing requires the constituent
+  // gradients to have the same dtype and device, and the resulting flattened
+  // tensor uses that common dtype and device. The flattened tensor is filled
+  // as the corresponding gradients are computed (triggered by autograd hooks),
+  // and the buckets are reduced in a predetermined order consistent across
+  // processes.
+  struct Bucket {
+    // Gradients of the bucket flattened into a 1-dimensional tensor
+    at::Tensor gradients;
+
+    // Views into the `gradients` tensor for each individual gradient
+    // Each view is created with layout (size and stride) matching the
+    // gradient's expected layout (see the "Gradient Layout Contract" in
+    // torch/csrc/autograd/functions/accumulate_grad.h).
+    // `bucket_views_in[i].copy_(grad)` and `grad.copy_(bucket_views_out[i])`
+    // provide convenient ways to copy gradient data in/out of `gradients`,
+    // respectively.
+    // We keep both `bucket_views_in` and `bucket_views_out` because
+    // registering a DDP communication hook may re-initialize
+    // `bucket_views_out` with the value of the hook's `future_work` but we
+    // still need separate views into the bucket's original flattened gradient
+    // to copy in gradient data.
+    std::vector<at::Tensor> bucket_views_in;
+    std::vector<at::Tensor> bucket_views_out;
+
+    // Variables whose gradients are held in this bucket
+    // We use refcounted tensors here so that we can easily unflatten the
+    // bucket's flattened `gradients` tensor into the participating variables
+    // after reduction has completed.
+    std::vector<at::Tensor> variables;
+
+    // Per-variable offset/length into the flattened `gradients` tensor and
+    // the corresponding `GradBucket` instance for communication hooks
+    std::vector<size_t> offsets;
+    std::vector<size_t> lengths;
+
+    // Per-variable sizes slicing into the bucket's `gradients` tensor
+    std::vector<c10::IntArrayRef> sizes_vec;
+
+    // Number of gradients left to be computed before the bucket is ready to
+    // be reduced
+    size_t pending;
+
+    // Global indices of participating variables in the bucket
+    std::vector<size_t> variable_indices;
+
+    // Future work handle for DDP communication hook
+    // If no hook is registered, a temporary vanilla allreduce hook is used.
+    c10::intrusive_ptr<at::ivalue::Future> future_work;
+
+    // if this bucket contains complex parameters
+    bool is_complex_bucket = false;
+
+    // If this bucket should expect a single sparse gradient
+    // If `true`, then this implies that `bucket.variables.size() == 1`.
+    bool expect_sparse_gradient = false;
+
+    // Sparse indices tensor
+    std::optional<at::Tensor> sparse_tensor_indices = std::nullopt;
+
+    // TODO(@pietern)
+    // Memory copies from gradient tensors into the bucket are potentially
+    // done on different CUDA streams. We record an event for every copy
+    // so that we can synchronize with them prior to kicking off the reduction.
+    // std::vector<at::cuda::CUDAEvent> events;
+
+    // Intra-bucket indices of variables whose grad-to-bucket copies are
+    // deferred for batching. Flushed as _foreach_copy_ + flat div_ when
+    // pending == 0. Only used when batched_grad_copy is enabled.
+    std::vector<size_t> deferred_copy_indices;
+  };
+
+  std::vector<Bucket> buckets_;
+
+  // A variable locator locates a particular variable in the reducer's buckets
+  struct VariableLocator {
+    // Index of the bucket containing the variable in the `buckets_` vector
+    size_t bucket_index;
+    // Index of the variable in the bucket, which may be used consistently
+    // across `bucket_views_in`, `bucket_views_out`, `variables`, `offsets`,
+    // `lengths`, `sizes_vec`, and `variable_indices` in `Bucket`
+    size_t intra_bucket_index;
+
+    VariableLocator() = default;
+
+    VariableLocator(size_t bucket_index_, size_t intra_bucket_index_)
+        : bucket_index(bucket_index_),
+          intra_bucket_index(intra_bucket_index_) {}
+  };
+
+  // Map the index of a variable to its location in the bucket structure.
+  std::vector<VariableLocator> variable_locators_;
+
+  // track the number of iterations to synchronize grads in training so far.
+  long num_iterations_;
+  // track distinct iteration of backward call. This is distinct from
+  // num_iterations_, for example in the case of multiple forward before
+  // backward.
+  long num_bwd_calls_;
+  // whether the first autograd hook for a distinct backward pass has been
+  // called.
+  bool first_autograd_hook_called_;
+  // track the number of buckets that have been ready for
+  // communication calls like allReduce or communication hooks.
+  int num_buckets_ready_;
+  // track the number of buckets that have been reduced.
+  int num_buckets_reduced_;
+
+  // Timing information.
+  int64_t backward_compute_start_time_ = -1;
+  std::unique_ptr<Timer> timer_;
+
+  // We collect the relative timestamp of every gradient being ready
+  // when executing autograd. This can be used to derive a timeline of
+  // the point in time buckets were ready, or ideal bucket assignment/ordering.
+  std::vector<int64_t> backward_stats_;
+
+  bool should_collect_runtime_stats();
+  void record_forward_compute_start_time();
+  void record_backward_compute_start_time();
+  void record_backward_compute_end_time();
+  void record_backward_comm_start_time();
+  void record_backward_comm_end_time();
+
+  int get_ddp_runtime_logging_sample_rate();
+  int ddp_runtime_logging_sample_rate_ = kDDPRuntimeLoggingSampleRate;
+
+  bool is_multi_device_module_ = false;
+
+  // Following variables are to help build dynamic bucket order
+  bool has_rebuilt_bucket_;
+  std::vector<at::Tensor> rebuilt_params_;
+  std::vector<int64_t> rebuilt_param_indices_;
+  const int64_t bucket_bytes_cap_;
+
+#ifndef _WIN32
+  struct RpcContext {
+    using ContextPtr = torch::distributed::autograd::ContextPtr;
+    // The shared_ptr is to hold the context instance.
+    ContextPtr context_ptr_holder;
+    std::atomic<ContextPtr::element_type*> context_ptr{nullptr};
+
+    void set(ContextPtr&& new_context_ptr);
+  };
+  RpcContext rpc_context_;
+#endif
+
+  // A struct containing work handle and tensor for allreduce scheduled in
+  // forward pass, if applicable.
+  struct ForwardPassAllreduceWork {
+    c10::intrusive_ptr<c10d::Work> workHandle;
+    at::Tensor resultTensor;
+    // whether we should divide by the initial world_size or the no. of
+    // remaining DDP ranks.
+    bool useStaticWorldSize;
+  };
+
+  // Handle for the currently scheduled allreduce in the forward pass, if
+  // applicable.
+  ForwardPassAllreduceWork forwardPassWorkHandle_;
+
+  // Division factor for reduction of gradients.
+  // Equal to the process group size, with an exception of handling uneven
+  // input.
+  int div_factor_;
+
+  bool static_graph_;
+
+  bool skip_all_reduce_unused_params_;
+
+  // Key: size_t (index), Value: the number of times that a variable's
+  // autograd_hook() should be triggered before marking this variable's grad as
+  // ready for communication. Map will not change after 1st iteration.
+  std::unordered_map<size_t, int> numGradHooksTriggeredMap_;
+  // Key: size_t (index), Value: the number of times that a variable's
+  // autograd_hook() are left to be triggered before marking this variable's
+  // grad as ready for communication. Map will change after 1st iteration to
+  // track a grad is ready for communication or not.
+  std::unordered_map<size_t, int> numGradHooksTriggeredMapPerIteration_;
+
+ private:
+  // reset counting for buckets before backward starts
+  void reset_bucket_counting();
+  // search unused parameters beore backward starts
+  void search_unused_parameters(
+      const std::vector<torch::autograd::Variable>& outputs);
+  void set_divide_factor();
+  // kick off all reduce for the ready bucket
+  void all_reduce_bucket(Bucket& bucket);
+  // kick off all reduce to local used map, it can help find global unused
+  // parameters
+  void all_reduce_local_used_map();
+  // initialize locally used parameter maps
+  void initialize_local_used_map();
+  // get current cuda stream
+  const c10::Stream get_current_stream();
+  bool dynamic_graph_find_unused();
+  bool static_graph_first_iteration();
+  bool static_graph_after_first_iteration();
+
+  bool is_unused_bucket(Bucket& bucket);
+  bool should_skip_all_reduce_bucket(Bucket& bucket);
+
+  // comm_hook_ is used to access the DDP communication hook if registered.
+  std::unique_ptr<CommHookInterface> comm_hook_;
+
+  // Sparse metadata contains the indices that will be used
+  // when calling into sparse allreduce.
+  // This is only used in the sparse allreduce collective calls
+  std::unique_ptr<std::map<std::string, at::Tensor>> sparse_metadata_;
+
+  // Debug level setting. It is parsed once when Reducer is constructed, and
+  // remains the same across a single invocation of DDP training.
+  DebugLevel ddp_debug_level_;
+  // Mapping of variable index to fully qualified name of model to notify users
+  // about errors when certain parameters do not get gradient.
+  std::unordered_map<size_t, std::string> param_names_;
+  // Variable indices stored sequentially in order of when the gradient is ready
+  // for the current backwards pass.
+  std::vector<int64_t> grad_ready_order_indices_;
+  // Bytes capacity of first bucket, can be configured by user
+  int64_t first_bucket_bytes_cap_;
+  // Per iteration set of parameter indices that have been marked ready.
+  std::unordered_set<size_t> perIterationReadyParams_;
+  // Retrieves parameter names that have not been marked as ready as part of
+  // previous iteration.
+  std::vector<std::string> getUnmarkedParamsForIteration();
+  // Retrieves parameter indices that have not been marked as ready as part of
+  // previous iteration.
+  std::vector<size_t> getUnmarkedParamIndicesForIteration();
+  // Raises appropriate error if mark_variable_ready is called on the same
+  // variable twice, which is unexpected.
+  void checkAndRaiseMarkedTwiceError(size_t curVariableIndex);
+  // Retrieves parameter corresponding to the given VariableIndex.
+  at::Tensor& get_param_from_index(size_t index);
+  // Python reducer keeps C++ reducer initialized. To remove this flag,
+  // we need to refactor the DDP wrapper's initialization.
+  bool use_python_reducer_;
+
+  const std::vector<int64_t> bucket_bytes_cap_list_;
+
+  // Cached bucket index to model parameter mapping. Populated after buckets
+  // are rebuilt after which this mapping is static.
+  mutable std::unordered_map<size_t, std::vector<at::Tensor>>
+      cached_variables_for_bucket_;
+
+  bool optim_in_backward_{false};
+  friend class Logger;
+};
+
+// This is equivalent to take_tensors but returns indices into the
+// tensor list argument for bucket assignment. Also, it is aware
+// of device placement and will not allow buckets to span devices.
+// The index of tensors[i] assigned to bucket is tensor_indices[i],
+// when tensor_indices is empty, the index of tensors[i] assigned to
+// bucket is i.
+TORCH_API std::tuple<std::vector<std::vector<size_t>>, std::vector<size_t>>
+compute_bucket_assignment_by_size(
+    const std::vector<at::Tensor>& tensors,
+    const std::vector<size_t>& bucket_size,
+    const std::vector<bool>& expect_sparse_gradient = {},
+    const std::vector<int64_t>& tensor_indices = {},
+    const std::optional<std::weak_ptr<c10d::Logger>>& logger = {});
+
+// Verify models across all processes are the same as model on rank 0 with
+// respect to no. of params and matching dtype/size/layout.
+TORCH_API void verify_params_across_processes(
+    const c10::intrusive_ptr<c10d::ProcessGroup>& process_group,
+    const std::vector<at::Tensor>& params,
+    const std::optional<std::weak_ptr<c10d::Logger>>& logger);
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/reducer_timer.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/reducer_timer.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..7b4320325c5ef027f4ad47a789b67e5168ee7dd7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/reducer_timer.hpp
@@ -0,0 +1,86 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/ApproximateClock.h>
+#include <torch/csrc/autograd/profiler.h>
+
+namespace c10d {
+constexpr int kUnsetTime = -1;
+
+inline int64_t current_time_in_nanos() {
+  return c10::getTime();
+}
+
+class TORCH_API Timer {
+ private:
+  // The timestamp of forward call start time in each iteration.
+  int64_t forward_start_time = kUnsetTime;
+  // The timestamp of backward computation start and end time in each
+  // iteration.
+  int64_t backward_compute_start_time = kUnsetTime;
+  int64_t backward_compute_end_time = kUnsetTime;
+  // The timestamp of first communication call start time in each iteration.
+  int64_t backward_comm_start_time = kUnsetTime;
+  // The timestamp of last communication call end time in each iteration.
+  int64_t backward_comm_end_time = kUnsetTime;
+
+ public:
+  enum class Event : uint8_t {
+    kForwardStart,
+    kBackwardComputeStart,
+    kBackwardComputeEnd,
+    kBackwardCommStart,
+    kBackwardCommEnd,
+  };
+
+  // Record the current event, i.e., mark it as having occurred now. Default
+  // CPU implementation.
+  virtual void record(Event event) {
+    getTimeRef(event) = current_time_in_nanos();
+  }
+
+  // Return the difference between when two events occurred, in nanoseconds.
+  // Or nullopt if one of them hasn't been recorded.
+  virtual std::optional<int64_t> measureDifference(Event start, Event end) = 0;
+
+  virtual ~Timer() = default;
+
+  // Return host-side timestamp, or nullopt if it has not yet been recorded.
+  std::optional<int64_t> getTimestamp(Event event) {
+    auto time = getTimeRef(event);
+    if (time == kUnsetTime) {
+      return std::nullopt;
+    } else {
+      return time;
+    }
+  }
+
+  // Return host-side time member variable corresponding to the given event.
+  int64_t& getTimeRef(Event event) {
+    switch (event) {
+      case Event::kForwardStart:
+        return forward_start_time;
+      case Event::kBackwardComputeStart:
+        return backward_compute_start_time;
+      case Event::kBackwardComputeEnd:
+        return backward_compute_end_time;
+      case Event::kBackwardCommStart:
+        return backward_comm_start_time;
+      case Event::kBackwardCommEnd:
+        return backward_comm_end_time;
+      default:
+        TORCH_INTERNAL_ASSERT(false);
+    }
+  }
+};
+
+TORCH_DECLARE_TYPED_REGISTRY(
+    TimerRegistry,
+    c10::DeviceType,
+    Timer,
+    std::unique_ptr,
+    c10::Device);
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/sequence_num.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/sequence_num.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..fa88322a6fa02adc0285afecf0fa128048f310f4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/sequence_num.hpp
@@ -0,0 +1,71 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Macros.h>
+#include <c10/util/irange.h>
+#include <mutex>
+#include <optional>
+#include <vector>
+
+namespace c10d {
+constexpr int kUnsetSeqNum = 0;
+
+namespace {
+constexpr int kByteOffset = 8;
+} // namespace
+
+// Converts from int to char vec to write in store
+template <typename T>
+inline std::vector<T> toVec(uint64_t num, int numBytes) {
+  std::vector<T> values;
+  // Read off bytes from right to left, pushing them into
+  // char array.
+  for (const auto i : c10::irange(numBytes)) {
+    uint8_t x = (num >> (kByteOffset * i)) & 0xff;
+    values.push_back(static_cast<T>(x));
+  }
+  return values;
+}
+
+// Converts from char vec (such as from store read) to int.
+template <typename T>
+inline uint64_t fromVec(const std::vector<T>& values) {
+  uint64_t num = 0;
+  // Set each byte at the correct location on num
+  for (const auto i : c10::irange(values.size())) {
+    uint8_t x = static_cast<uint8_t>(values[i]);
+    num |= (static_cast<int64_t>(x) << (kByteOffset * i));
+  }
+  return num;
+}
+
+class TORCH_API SequenceNum {
+ public:
+  SequenceNum();
+  explicit SequenceNum(const uint64_t num);
+  // Retrieve num_. Will throw if not set.
+  uint64_t get() const;
+  // Increment num_. Will throw if not set.
+  void increment();
+  // Increment num_ and return the old value. Will throw if not set.
+  uint64_t getAndIncrement();
+  // Sets num_
+  void set(const uint64_t num);
+  // Returns true if this SequenceNum is properly initialized with a value, else
+  // false.
+  bool isSet() const;
+
+  SequenceNum& operator=(const SequenceNum& other);
+
+  SequenceNum(const SequenceNum& other);
+
+ private:
+  std::optional<uint64_t> num_;
+  mutable std::mutex lock_;
+};
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/socket.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/socket.h
new file mode 100644
index 0000000000000000000000000000000000000000..8d2ec93016c028f4d7c79410ab5c3a11cfa5264f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/socket.h
@@ -0,0 +1,110 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <chrono>
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include <c10/macros/Macros.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/distributed/c10d/Backoff.hpp>
+#include <torch/csrc/distributed/c10d/exception.h>
+
+namespace c10d::detail {
+
+class SocketOptions {
+ public:
+  SocketOptions& prefer_ipv6(bool value) noexcept {
+    prefer_ipv6_ = value;
+
+    return *this;
+  }
+
+  bool prefer_ipv6() const noexcept {
+    return prefer_ipv6_;
+  }
+
+  SocketOptions& connect_timeout(std::chrono::milliseconds value) noexcept {
+    connect_timeout_ = value;
+
+    return *this;
+  }
+
+  std::chrono::milliseconds connect_timeout() const noexcept {
+    return connect_timeout_;
+  }
+
+  // Sets the backoff policy to use for socket connect ops.
+  SocketOptions& connect_backoff(std::shared_ptr<Backoff> value) noexcept {
+    connect_backoff_ = std::move(value);
+
+    return *this;
+  }
+
+  const std::shared_ptr<Backoff>& connect_backoff() const noexcept {
+    return connect_backoff_;
+  }
+
+ private:
+  bool prefer_ipv6_ = true;
+  std::chrono::milliseconds connect_timeout_{std::chrono::seconds{30}};
+  std::shared_ptr<Backoff> connect_backoff_{
+      std::make_shared<FixedBackoff>(std::chrono::milliseconds(1000))};
+};
+
+class SocketImpl;
+
+class Socket {
+ public:
+  // This function initializes the underlying socket library and must be called
+  // before any other socket function.
+  static void initialize();
+
+  static Socket listen(std::uint16_t port, const SocketOptions& opts = {});
+
+  static Socket listenFromFd(int fd, std::uint16_t expected_port);
+
+  static Socket connect(
+      const std::string& host,
+      std::uint16_t port,
+      const SocketOptions& opts = {});
+
+  Socket() noexcept = default;
+
+  Socket(const Socket& other) = delete;
+
+  Socket& operator=(const Socket& other) = delete;
+
+  Socket(Socket&& other) noexcept;
+
+  Socket& operator=(Socket&& other) noexcept;
+
+  ~Socket();
+
+  Socket accept() const;
+
+  int handle() const noexcept;
+
+  std::uint16_t port() const;
+
+  bool waitForInput(std::chrono::milliseconds timeout);
+
+  std::string repr() const;
+
+ private:
+  explicit Socket(std::unique_ptr<SocketImpl>&& impl) noexcept;
+
+  std::unique_ptr<SocketImpl> impl_;
+};
+} // namespace c10d::detail
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/socket_fmt.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/socket_fmt.h
new file mode 100644
index 0000000000000000000000000000000000000000..f333c8fa12d8ad4f321072b9379c24529c7d3507
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/socket_fmt.h
@@ -0,0 +1,35 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+/*
+This file should not be included from other .h files and only used in cpp files
+as it exposes the underlying platform specific socket headers.
+*/
+
+#include <string>
+
+#ifdef _WIN32
+#include <mutex>
+
+#include <winsock2.h>
+#include <ws2tcpip.h>
+#else
+#include <netinet/in.h>
+#endif
+
+namespace c10d::detail {
+
+// Returns a human-readable representation of the given socket address.
+std::string formatSockAddr(const struct ::sockaddr* addr, socklen_t len);
+
+} // namespace c10d::detail
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..813594d55b02f174457d598a8d1b672995f87248
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp
@@ -0,0 +1,160 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAAllocatorConfig.h>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+
+namespace c10d::symmetric_memory {
+
+// Resource wrapper that owns a (vaddr, allocation handle) pair. Upon
+// destruction, it unmaps the vaddr and releases the allocation handle.
+struct AllocationRef : public c10::intrusive_ptr_target {
+  void* ptr;
+  HandleType handle;
+  size_t block_size;
+  int device_idx;
+  bool is_multicast;
+
+  AllocationRef(
+      void* ptr,
+      HandleType handle,
+      size_t block_size,
+      int device_idx,
+      bool is_multicast = false);
+
+  ~AllocationRef();
+};
+
+// Forward declaration of CUDAPeerAllocInfo
+class CUDAPeerAllocInfo;
+
+class CUDASymmetricMemory : public SymmetricMemory {
+ public:
+  // This is mostly a shallow copy that shares the pointer to
+  // `CUDAPeerAllocInfo` which corresponds to the base Block. The
+  // CUDASymmetricMemory handle is specified by the offset to the base ptr.
+  CUDASymmetricMemory(
+      const c10::intrusive_ptr<CUDAPeerAllocInfo>& pai,
+      size_t offset);
+
+  ~CUDASymmetricMemory() override {};
+
+  std::vector<void*> get_buffer_ptrs() override;
+  std::vector<void*> get_signal_pad_ptrs() override;
+  void** get_buffer_ptrs_dev() override;
+  void** get_signal_pad_ptrs_dev() override;
+  size_t get_buffer_size() override;
+  size_t get_offset() override;
+
+  bool has_multicast_support() override;
+  void* get_multicast_ptr() override;
+
+  void barrier(int channel, size_t timeout_ms) override;
+  void put_signal(int dst_rank, int channel, size_t timeout_ms) override;
+  void wait_signal(int src_rank, int channel, size_t timeout_ms) override;
+
+  int get_rank() override;
+  int get_world_size() override;
+  c10::Device get_device() override;
+  bool world_within_direct_access() override;
+
+ private:
+  int local_device_idx_;
+  int rank_;
+  int world_size_;
+  c10::intrusive_ptr<CUDAPeerAllocInfo> pai_;
+  size_t offset_{0}; // in byte
+};
+
+// A class to hold the base pointers and signal pad pointers for a group of
+// peers. One `CUDAPeerAllocInfo` object can be shared by multiple
+// `CUDASymmetricMemory` objects when latter reside on the same allocation
+// and rendezvous over the same group. (The `CUDASymmetricMemory` objects may
+// have different offsets compared to the base address.)
+class CUDAPeerAllocInfo : public c10::intrusive_ptr_target {
+ public:
+  CUDAPeerAllocInfo(
+      std::vector<c10::intrusive_ptr<AllocationRef>> alloc_refs,
+      std::vector<void*> buffers,
+      std::vector<void*> signal_pads,
+      HandleType mc_handle,
+      void* mc_addr,
+      size_t buffer_size,
+      int local_device_idx,
+      int rank,
+      int world_size,
+      std::string group_name);
+
+ private:
+  std::vector<c10::intrusive_ptr<AllocationRef>> alloc_refs_;
+  std::vector<void*> buffers_;
+  std::vector<void*> signal_pads_;
+  HandleType mc_handle_;
+  void* mc_addr_;
+  size_t buffer_size_;
+  int local_device_idx_;
+  int rank_;
+  int world_size_;
+  void** buffers_dev_;
+  void** signal_pads_dev_;
+  std::string group_name_;
+
+  friend class CUDASymmetricMemory;
+};
+
+// Metadata associated with each allocation performed by
+// `CUDASymmetricMemoryAllocator`.
+struct Block : public c10::intrusive_ptr_target {
+  c10::intrusive_ptr<AllocationRef> alloc_ref;
+  int device_idx;
+  size_t block_size;
+  size_t buffer_size;
+  size_t signal_pad_offset;
+  std::optional<std::string> default_group_name;
+  std::map<std::string, c10::intrusive_ptr<CUDAPeerAllocInfo>> symm_mems;
+
+  Block(
+      c10::intrusive_ptr<AllocationRef> alloc_ref,
+      int device_idx,
+      size_t block_size,
+      size_t buffer_size,
+      size_t signal_pad_offset,
+      const std::optional<std::string>& group_name);
+};
+
+class CUDASymmetricMemoryAllocator : public SymmetricMemoryAllocator {
+ public:
+  void* alloc(
+      size_t size,
+      int device_idx,
+      const std::optional<std::string>& group_name) override;
+
+  void free(void* ptr) override;
+  size_t get_alloc_size(void* ptr) override;
+  c10::intrusive_ptr<SymmetricMemory> rendezvous(
+      void* ptr,
+      const std::optional<std::string>& group_name) override;
+  bool has_multicast_support(int device_idx) override;
+  bool has_allocation(void* ptr) override;
+  c10::DeviceType supported_device_type() override;
+  std::string name() override;
+
+ private:
+  c10::intrusive_ptr<Block> find_block(void* ptr);
+  c10::intrusive_ptr<Block> find_block_covering(void* ptr, size_t& offset);
+
+  std::shared_mutex mutex_;
+  std::unordered_map<void*, c10::intrusive_ptr<Block>> ptr_to_block_;
+  c10::cuda::CUDACachingAllocator::Expandable_Segments_Handle_Type
+      handle_type_ = c10::cuda::CUDACachingAllocator::
+          Expandable_Segments_Handle_Type::UNSPECIFIED;
+};
+
+} // namespace c10d::symmetric_memory
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..b12c53f4410bea2b4301d33fb91dca6f1616f405
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp
@@ -0,0 +1,48 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include <c10/util/hash.h>
+
+#if defined(USE_ROCM)
+#include <hip/hip_runtime_api.h>
+#endif
+
+namespace c10d::symmetric_memory {
+
+// Key type for the symmetric memory map. `void*` for tensor storage ptr,
+// `std::string` for group name.
+using SymmMemKey = std::pair<void*, std::string>;
+// Hash function for the symmetric memory map. c10::hash has a std::pair
+// specialization (line 323-329 of hash.h) that delegates to the tuple hasher
+// which combines hashes of each element.
+using SymmMemKeyHash = c10::hash<SymmMemKey>;
+
+// Covers NVL72
+constexpr int max_cuda_p2p_domain_size = 72;
+// Maximum number of channels
+constexpr int symm_max_nblocks = 32;
+
+// Maximally, a rank will need to sync with all other ranks, over all
+// channels. Each signal is 32 bits, which is the minimum unit for atomic cas.
+// Default signal pad size, can be overridden via set_signal_pad_size().
+constexpr size_t default_signal_pad_size =
+    symm_max_nblocks * max_cuda_p2p_domain_size * sizeof(uint32_t);
+
+#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
+using HandleType = CUmemGenericAllocationHandle;
+#elif defined(USE_ROCM)
+using HandleType = hipMemGenericAllocationHandle_t;
+#else
+using HandleType = void*;
+#endif
+
+} // namespace c10d::symmetric_memory
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..5cd60c7368f80b16841848b00085bd488724fa78
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp
@@ -0,0 +1,121 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+#include <utility>
+
+namespace c10d {
+namespace symmetric_memory {
+
+bool device_has_multicast_support(int device_idx);
+
+bool allow_overlapping_devices();
+
+// Query environment variable to get the backend used for CUDA Symmetric Memory.
+std::string getSymmMemBackendCUDA();
+
+class IpcChannel {
+ public:
+  IpcChannel();
+  ~IpcChannel();
+
+  void send_fd(int dst_pid, int fd);
+  int recv_fd();
+
+  std::vector<int> all_gather_fds(
+      int rank,
+      const std::vector<int>& pids,
+      int fd);
+
+  int broadcast_fds(
+      int rank,
+      int src_rank,
+      const std::vector<int>& pids,
+      int fd);
+
+ private:
+  static std::string get_socket_name(int pid);
+
+  std::string socket_name_;
+  int socket_;
+};
+
+// A set of store-based exchange methods with a preset prefix typically type of
+// the SymmetricMemory.  Most used as static instances at respective
+// SymmetricMemory implementation files.
+class StoreExchange {
+ public:
+  StoreExchange(std::string store_prefix)
+      : store_prefix_(std::move(store_prefix)) {}
+
+  // Put template function in header file so that compiler can easily access it.
+  template <typename T>
+  std::vector<T> all_gather(
+      const c10::intrusive_ptr<c10d::Store>& store,
+      int rank,
+      int world_size,
+      T val) {
+    static_assert(std::is_trivially_copyable_v<T>);
+
+    std::vector<std::string> peer_keys;
+    peer_keys.reserve(world_size);
+    for (int r = 0; r < world_size; ++r) {
+      std::ostringstream oss;
+      oss << store_prefix_ << '/' << seq_id_ << '/' << r;
+      peer_keys.push_back(oss.str());
+    }
+    ++seq_id_;
+
+    {
+      std::vector<uint8_t> payload(
+          reinterpret_cast<uint8_t*>(&val),
+          reinterpret_cast<uint8_t*>(&val) + sizeof(T));
+      store->set(peer_keys[rank], payload);
+    }
+
+    std::vector<T> peer_vals;
+    peer_vals.reserve(world_size);
+    for (int r = 0; r < world_size; ++r) {
+      if (r == rank) {
+        peer_vals.push_back(val);
+        continue;
+      }
+      store->wait({peer_keys[r]});
+      auto payload = store->get(peer_keys[r]);
+      TORCH_CHECK(payload.size() == sizeof(T));
+      T peer_val{};
+      std::memcpy(&peer_val, payload.data(), sizeof(T));
+      peer_vals.push_back(peer_val);
+    }
+    return peer_vals;
+  }
+
+  void barrier(
+      const c10::intrusive_ptr<c10d::Store>& store,
+      int rank,
+      int world_size) {
+    // TODO: implement an efficient one?
+    all_gather(store, rank, world_size, 0);
+  }
+
+ private:
+  const std::string store_prefix_;
+  size_t seq_id_ = 0;
+};
+
+// Returns a pointer of virtual address that is mapped to the physical memory
+// held by the handle.
+void map_block(
+    void** ptr,
+    c10d::symmetric_memory::HandleType handle,
+    size_t size,
+    int device_idx);
+
+} // namespace symmetric_memory
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/DMAConnectivity.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/DMAConnectivity.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..1fbf2d774c5543fcf496bb87f2599fbace5c46b7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/DMAConnectivity.hpp
@@ -0,0 +1,43 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+
+namespace c10d {
+
+struct TORCH_API DMAConnectivity : c10::intrusive_ptr_target {
+  c10::DeviceType device_type;
+  std::string connection_type;
+
+  // This is an NxN matrix representing the connectivity between N devices,
+  // where each element matrix[i][j] indicates the connectivity between device
+  // i and device j. A value of 0 denotes that there is no connection between
+  // device i and j. The meaning of non-zero values are specific to the
+  // connection type (e.g., for NVLink it represents the number of NVLinks).
+  std::vector<std::vector<int>> matrix;
+
+  explicit DMAConnectivity(
+      c10::DeviceType device_type,
+      std::string connection_type,
+      std::vector<std::vector<int>> matrix);
+};
+
+struct DMAConnectivityDetector : c10::intrusive_ptr_target {
+  virtual c10::intrusive_ptr<DMAConnectivity> detect() = 0;
+  ~DMAConnectivityDetector() override = default;
+};
+
+C10_EXPORT void register_dma_connectivity_detector(
+    c10::DeviceType device_type,
+    const std::string& connection_type,
+    c10::intrusive_ptr<DMAConnectivityDetector> detector);
+
+TORCH_API c10::intrusive_ptr<DMAConnectivity> detect_dma_connectivity(
+    c10::DeviceType device_type,
+    const std::string& connection_type);
+
+} // namespace c10d
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..423606d09f3b5a32a31dbdf176cf2d16927bf607
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp
@@ -0,0 +1,67 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
+
+#ifdef NCCL_HAS_SYMMEM_SUPPORT
+
+namespace c10d {
+namespace symmetric_memory {
+
+class NCCLPeerAllocInfo;
+
+class NCCLSymmetricMemory : public SymmetricMemory {
+ public:
+  NCCLSymmetricMemory(c10::intrusive_ptr<NCCLPeerAllocInfo> pai, size_t offset);
+
+  ~NCCLSymmetricMemory() override = default;
+
+  std::vector<void*> get_buffer_ptrs() override;
+
+  std::vector<void*> get_signal_pad_ptrs() override;
+
+  void** get_buffer_ptrs_dev() override;
+
+  void** get_signal_pad_ptrs_dev() override;
+
+  size_t get_buffer_size() override;
+
+  std::string get_group_name();
+
+  bool has_multicast_support() override;
+
+  void* get_multicast_ptr() override;
+
+  void barrier(int channel, size_t timeout_ms) override;
+
+  void put_signal(int dst_rank, int channel, size_t timeout_ms) override;
+
+  void wait_signal(int src_rank, int channel, size_t timeout_ms) override;
+
+  int get_rank() override;
+
+  int get_world_size() override;
+
+  c10::Device get_device() override;
+
+  ncclWindow_t get_window();
+
+  ncclWindow_t get_signal_pad_handle();
+
+  size_t get_offset() override;
+
+ private:
+  c10::intrusive_ptr<NCCLPeerAllocInfo> pai_;
+  size_t offset_;
+  int rank_;
+  int world_size_;
+  int device_idx_;
+};
+
+} // namespace symmetric_memory
+} // namespace c10d
+#endif // NCCL_HAS_SYMMEM_SUPPORT
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..38c98774bcf563c4990b1bb6e90f1acf6657e4c3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp
@@ -0,0 +1,225 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+
+namespace c10d::symmetric_memory {
+
+// SymmetricMemory represents symmetric allocations across a group of devices.
+// The allocations represented by a SymmetricMemory object are accessible by
+// all devices in the group. The class can be used for op-level custom
+// communication patterns (via the get_buffer APIs and the synchronization
+// primitives), as well as custom communication kernels (via the buffer and
+// signal_pad device pointers).
+//
+// To acquire a SymmetricMemory object, each rank first allocates
+// identical-sized memory via SymmetricMemoryAllocator::alloc(), then invokes
+// SymmetricMemoryAllocator::rendezvous() on the memory to establish the
+// association across peer buffers. The rendezvous is a one-time process, and
+// the mapping between a local memory memory and the associated SymmetricMemory
+// object is unique.
+//
+// NOTE [symmetric memory signal pad]
+// Signal pads are P2P-accessible memory regions designated for
+// synchronization. SymmetricMemory offers built-in synchronization primitives
+// such as barriers, put_signal, and wait_signal, which are all based on signal
+// pads. Users may utilize signal pads for their own synchronization logic,
+// provided that the signal pads remain zero-filled following successful
+// synchronization.
+//
+// NOTE [symmetric memory synchronization channel]
+// Synchronization channels allow users to use a single SymmetricMemory object
+// to perform isolated synchronizations on different streams. For example,
+// consider the case in which two barriers are issued on two streams for
+// different purposes. Without the concept of channels, we cannot guarantee the
+// correctness of the barriers since signals issued from barrier on stream A
+// can be received by the barrier on stream B. By specifying different channels
+// for these two barriers, they can operate correctly in parallel.
+class TORCH_API SymmetricMemory : public torch::CustomClassHolder {
+ public:
+  ~SymmetricMemory() override = default;
+
+  virtual std::vector<void*> get_buffer_ptrs() = 0;
+  virtual std::vector<void*> get_signal_pad_ptrs() = 0;
+
+  // get_buffer_ptrs_dev() and get_signal_pad_ptrs_dev() each return a pointer
+  // to a device array of size world_size, containing buffer pointers and
+  // signal pad pointers, respectively.
+  virtual void** get_buffer_ptrs_dev() = 0;
+  virtual void** get_signal_pad_ptrs_dev() = 0;
+  virtual size_t get_buffer_size() = 0;
+  size_t get_signal_pad_size();
+
+  virtual size_t get_offset() = 0;
+
+  virtual bool has_multicast_support() = 0;
+  virtual void* get_multicast_ptr() = 0;
+
+  at::Tensor get_buffer(
+      int rank,
+      c10::IntArrayRef sizes,
+      c10::ScalarType dtype,
+      int64_t storage_offset);
+
+  at::Tensor get_signal_pad(
+      int rank,
+      c10::IntArrayRef sizes,
+      std::optional<c10::ScalarType> dtype = std::nullopt,
+      int64_t storage_offset = 0);
+
+  at::Tensor get_remote_tensor(
+      int peer,
+      c10::IntArrayRef sizes,
+      c10::ScalarType dtype);
+
+  virtual void barrier(int channel, size_t timeout_ms) = 0;
+  virtual void put_signal(int dst_rank, int channel, size_t timeout_ms) = 0;
+  virtual void wait_signal(int src_rank, int channel, size_t timeout_ms) = 0;
+
+  virtual int get_rank() = 0;
+  virtual int get_world_size() = 0;
+  virtual c10::Device get_device() = 0;
+
+  virtual const std::vector<int>& get_rank_to_global_rank() {
+    TORCH_CHECK(false, "NYI");
+  }
+
+  virtual int* get_rank_to_global_rank_dev() {
+    TORCH_CHECK(false, "NYI");
+  }
+
+  // Returns true if *all* peers within the group are accessible via direct
+  // memory load and store.
+  virtual bool world_within_direct_access() {
+    TORCH_CHECK(false, "NYI");
+  }
+};
+
+class SymmetricMemoryAllocator : public c10::intrusive_ptr_target {
+ public:
+  ~SymmetricMemoryAllocator() override = default;
+
+  virtual void* alloc(
+      size_t size,
+      int device_idx,
+      const std::optional<std::string>& group_name) = 0;
+
+  virtual void free(void* ptr) = 0;
+  virtual size_t get_alloc_size(void* ptr) = 0;
+  virtual c10::intrusive_ptr<SymmetricMemory> rendezvous(
+      void* ptr,
+      const std::optional<std::string>& group_name) = 0;
+  virtual bool has_multicast_support(int device_idx) = 0;
+  virtual c10::DeviceType supported_device_type() = 0;
+  virtual std::string name() = 0;
+  virtual bool has_allocation(void* ptr) {
+    return false;
+  }
+};
+
+C10_EXPORT bool is_finalizing();
+
+C10_EXPORT void register_allocator(
+    c10::DeviceType device_type,
+    c10::intrusive_ptr<SymmetricMemoryAllocator> allocator);
+
+C10_EXPORT void register_availability(
+    const std::string& name,
+    c10::intrusive_ptr<SymmetricMemoryAllocator> allocator);
+
+C10_EXPORT bool has_allocator(c10::DeviceType device_type);
+
+C10_EXPORT c10::intrusive_ptr<SymmetricMemoryAllocator> get_allocator(
+    c10::DeviceType device_type);
+
+// Set a store for rendezvousing symmetric allocations on a group of devices
+// identified by `group_name`. The concept of groups is logical; users can
+// utilize predefined groups (e.g., a group of device identified by a
+// ProcessGroup) or create custom ones. Note that a SymmetricMemoryAllocator
+// backends might employ a more efficient communication channel for the actual
+// rendezvous process and only use the store for bootstrapping purposes.
+TORCH_API void set_group_info(
+    const std::string& group_name,
+    int rank,
+    int world_size,
+    c10::intrusive_ptr<Store> store);
+
+struct GroupInfo {
+  int rank;
+  int world_size;
+  c10::intrusive_ptr<c10d::Store> store;
+};
+
+C10_EXPORT GroupInfo& get_group_info(const std::string& group_name);
+
+// Identical to empty_strided, but allows symmetric memory access to be
+// established for the allocated tensor via SymmetricMemory::rendezvous(). This
+// function itself is not a collective operation. It invokes
+// SymmetricMemoryAllocator::alloc() for the requested device under the hood.
+//
+// NOTE [symmetric memory persistent allocation]
+// If an `alloc_id` is supplied, empty_strided_p2p will perform persistent
+// allocation. This makes the function cache allocated memory and ensure that
+// invocations with the same `alloc_id` receive tensors backed by the same
+// memory address. For safety, if a previous persistent allocation is still
+// active (i.e., the storage of the returned tensor is still alive), persistent
+// allocations with the same `alloc_id` will fail. This determinism coupled
+// with memory planning of communication buffers (e.g., by Inductor) allows
+// communication algorithms to reliably reuse previously established remote
+// memory access.
+TORCH_API at::Tensor empty_strided_p2p(
+    c10::IntArrayRef size,
+    c10::IntArrayRef stride,
+    c10::ScalarType dtype,
+    c10::Device device,
+    const std::optional<std::string>& group_name,
+    std::optional<uint64_t> alloc_id);
+
+// Establishes symmetric memory access on tensors allocated via
+// empty_strided_p2p() and empty_strided_p2p_persistent(). rendezvous() is a
+// one-time process, and the mapping between a local memory region and the
+// associated SymmetricMemory object is unique. Subsequent calls to
+// rendezvous() with the same tensor, or tensors allocated with
+// empty_strided_p2p_persistent() using the same alloc_id, will receive the
+// cached SymmetricMemory object.
+//
+// The function has a collective semantic and must be invoked simultaneously
+// from all rendezvous participants.
+TORCH_API c10::intrusive_ptr<SymmetricMemory> rendezvous(
+    const at::Tensor& tensor,
+    const std::optional<std::string>& group_name = std::nullopt);
+
+TORCH_API bool has_multicast_support(
+    c10::DeviceType device_type,
+    int device_idx);
+
+TORCH_API bool is_symm_mem_tensor(const at::Tensor& tensor);
+
+TORCH_API void set_backend(const std::string& name);
+
+TORCH_API std::optional<std::string> get_backend(c10::Device device);
+
+// Get the current signal pad size for symmetric memory allocations.
+// Returns the user-configured size if set, otherwise returns the default size.
+TORCH_API size_t get_signal_pad_size();
+
+// Set the signal pad size for future symmetric memory allocations.
+// This must be called before any symmetric memory allocations are made.
+// The size should be proportional to the number of blocks the user launches
+// and the world size.
+TORCH_API void set_signal_pad_size(size_t size);
+
+C10_EXPORT void register_mempool_allocator(
+    c10::DeviceType device_type,
+    std::shared_ptr<c10::Allocator> allocator);
+
+TORCH_API std::shared_ptr<c10::Allocator> get_mempool_allocator(
+    c10::Device device);
+
+} // namespace c10d::symmetric_memory
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/env.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/env.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..ec10b3d7f4a84d66f1bb54ce582c56155d7f9257
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/env.hpp
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <c10/util/env.h>
+
+namespace c10d::symmetric_memory {
+
+static int getenv_nblocks() {
+  static int num_blocks = -1; // Uninitialized
+  if (num_blocks == -1) {
+    auto str = c10::utils::get_env("TORCH_SYMMMEM_NBLOCKS");
+    if (str.has_value()) {
+      num_blocks = std::stoi(str.value());
+    } else {
+      num_blocks = -2; // Not set
+    }
+  }
+  return num_blocks;
+}
+
+} // namespace c10d::symmetric_memory
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..661d205e79e86760b0f03f5013fac343c3775241
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
@@ -0,0 +1,98 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/Work.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+
+namespace c10d::intra_node_comm {
+
+using namespace c10d::symmetric_memory;
+
+constexpr size_t kMaxDevices = 8;
+constexpr size_t kDefaultBufferSize = 10ull * 1024 * 1024;
+
+using NvlMesh = std::array<std::array<size_t, kMaxDevices>, kMaxDevices>;
+
+enum class Topology : uint8_t {
+  UNKNOWN = 0,
+  FULLY_CONNECTED = 1,
+};
+
+enum class AllReduceAlgo : uint8_t {
+  NONE = 0,
+  ONE_SHOT = 1,
+  TWO_SHOT = 2,
+};
+
+// NOTE: this class will be be removed soon in favor of SymmetricMemory
+class TORCH_API IntraNodeComm : public c10::intrusive_ptr_target {
+ public:
+  IntraNodeComm(
+      c10::intrusive_ptr<c10d::Store> store,
+      size_t rank,
+      size_t worldSize,
+      std::optional<size_t> bufferSize = std::nullopt,
+      std::string groupName = "");
+
+  ~IntraNodeComm() override;
+
+  static bool isEnabled();
+
+  /**
+   * Performs rendezvous.
+   * If rendezvous fails, the IntraNodeComm object will be in an invalid
+   * state and it is the caller's responsibility to dispose it.
+   */
+  bool rendezvous();
+
+  /**
+   * Selects a AllReduceAlgo that we think will outperform nccl.
+   * Returns AllReduceAlgo::NONE if we don't think we can outperform nccl.
+   */
+  AllReduceAlgo selectAllReduceAlgo(const at::Tensor& input);
+
+  at::Tensor allReduce(const at::Tensor& input, AllReduceAlgo algo);
+
+ private:
+  at::Tensor oneShotAllReduce(
+      const at::Tensor& input,
+      at::cuda::CUDAStream& stream);
+
+  at::Tensor twoShotAllReduce(
+      const at::Tensor& input,
+      at::cuda::CUDAStream& stream);
+
+  c10::intrusive_ptr<Store> store_;
+  size_t rank_;
+  size_t worldSize_;
+  size_t bufferSize_;
+  std::string groupName_;
+
+  /**
+   * Members initialized after rendezvous
+   */
+  bool isInitialized_ = false;
+  int deviceIdx_{0};
+  Topology topology_ = Topology::UNKNOWN;
+  void* symmetricMemoryPtr_ = nullptr;
+  c10::intrusive_ptr<SymmetricMemory> symmetricMemory_ = nullptr;
+};
+
+class IntraNodeCommWork : public c10d::Work {
+ public:
+  bool wait(std::chrono::milliseconds timeout = kNoTimeout) override {
+    return true;
+  }
+};
+
+TORCH_API int64_t getIntraNodeCommUsageCounter();
+
+bool isIntraNodeCommSupported();
+} // namespace c10d::intra_node_comm
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/macros.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/macros.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..84b176a9dae3df9c7cec8f6feeb42534705f1694
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/macros.hpp
@@ -0,0 +1,29 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Macros for type dispatch and common utilities for symmetric memory
+#pragma once
+
+#include <ATen/ATen.h>
+
+// Convert ATen floating point types to NV floating point types
+// at::kBFloat16 -> __nv_bfloat16
+// at::kHalf -> __half
+// Float is the same.
+
+#define AT_DISPATCH_CASE_CONVERT(enum_type, scalar_type, ...) \
+  case enum_type: {                                           \
+    AT_PRIVATE_CHECK_SELECTIVE_BUILD(enum_type);              \
+    using scalar_t = scalar_type;                             \
+    return __VA_ARGS__();                                     \
+  }
+
+#define AT_DISPATCH_NV_FLOATS(scalar_type, name, ...)                      \
+  AT_DISPATCH_SWITCH(                                                      \
+      scalar_type,                                                         \
+      name,                                                                \
+      AT_DISPATCH_CASE_CONVERT(at::kBFloat16, __nv_bfloat16, __VA_ARGS__); \
+      AT_DISPATCH_CASE_CONVERT(at::kHalf, __half, __VA_ARGS__);            \
+      AT_DISPATCH_CASE(at::kFloat, __VA_ARGS__));
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..10674d46201a1e26f404eac979e4a0261d6a4e55
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#if USE_NCCL
+
+#include <nccl.h>
+#include <torch/csrc/cuda/nccl.h>
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 27, 0)
+#define NCCL_HAS_SYMMEM_SUPPORT
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)
+#if !defined(USE_ROCM)
+#define NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+#include <nccl_device.h>
+#endif
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 0)
+#define NCCL_HAS_ONE_SIDED_API
+#endif
+
+#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 29, 7)
+#define NCCL_DEVICE_HAS_REDUCE_COPY
+#endif
+#endif // USE_NCCL
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..57d34924e9c0fc6abb36b9bc2883690216f0c910
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp
@@ -0,0 +1,248 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/util/Exception.h>
+#include <c10/util/Logging.h>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
+#include <functional>
+#include <mutex>
+#include <optional>
+#include <string>
+#include <unordered_map>
+
+#ifdef NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+
+namespace c10d::symmetric_memory {
+
+// Manages NCCL device communicators for symmetric memory operations.
+// This is a singleton class that maintains a registry of device communicators
+// organized by process group name and an optional key (typically the caller
+// function name). This allows different functions within the same process group
+// to use different device communicators, which is useful for concurrent
+// collective operations.
+//
+// The registry uses a two-level map structure:
+// - First level: keyed by process group name
+// - Second level: keyed by an optional key (defaults to caller function name)
+//
+// Device communicators are stored by value in the registry, but methods return
+// references wrapped in std::optional for safe access.
+class NCCLDevCommManager {
+ public:
+  // Constructor
+  // @param device The CUDA device this manager is associated with
+  explicit NCCLDevCommManager(const c10::Device device) : device_(device) {}
+
+  // Get the singleton instance for the given device.
+  // This ensures there's only one manager per device. If called with a
+  // different device than the one used to create the singleton, it will throw.
+  // @param device The CUDA device to get the manager for
+  // @return Reference to the singleton manager instance
+  static NCCLDevCommManager& get(const c10::Device device) {
+    static NCCLDevCommManager manager(device);
+    TORCH_CHECK_VALUE(
+        manager.device_ == device,
+        "Detected use of NCCLDevCommManager on multiple devices. This is not supported.");
+    return manager;
+  }
+
+  // Get an NCCL device communicator for a group, for the caller function.  By
+  // default, we search for the device communicator using the caller function
+  // name as the key.  If you previously registered a device communicator with a
+  // different key, you should provide that key instead.
+  // Returns std::nullopt if the device communicator is not found.
+  // Example:
+  // void foo(const std::string& group_name) {
+  //   // Try to get first.
+  //   auto devcomm_opt = get_devcomm(group_name);
+  //   if (!devcomm_opt) {
+  //     // Not found, create then register.
+  //     ncclDevComm devcomm = ncclDevCommCreate(...);
+  //     devcomm_opt = register_devcomm(group_name, devcomm);
+  //   }
+  //   ncclDevComm& devcomm_ref = *devcomm_opt;
+  //   // Use devcomm_ref
+  // }
+  std::optional<std::reference_wrapper<ncclDevComm>> get_devcomm(
+      const std::string& group_name,
+      const std::string& key = __builtin_FUNCTION()) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    // First, look up the group in the registry
+    auto group_it = devcomm_registry_.find(group_name);
+    if (group_it == devcomm_registry_.end()) {
+      return std::nullopt;
+    }
+    // Then, look up the key within that group's map
+    auto key_it = group_it->second.find(key);
+    if (key_it == group_it->second.end()) {
+      return std::nullopt;
+    }
+    // Return a reference wrapper to the device communicator
+    // Using reference_wrapper because std::optional cannot hold references
+    // directly
+    return std::make_optional(std::ref(key_it->second));
+  }
+
+  // Get a host-side NCCL communicator for a group.
+  // This is the regular host-side communicator, not the device communicator.
+  // @param group_name The process group name
+  // @return The host-side NCCL communicator
+  // @throws TORCH_CHECK if the communicator is not found
+  ncclComm_t get_comm(const std::string& group_name) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto it = group_to_comm_.find(group_name);
+    if (it == group_to_comm_.end()) {
+      TORCH_CHECK(
+          false,
+          "NCCL host communicator for group ",
+          group_name,
+          " not found. Have you rendezvoused any tensor with this group?");
+    }
+    return it->second;
+  }
+
+  // Register a device communicator for a group. If `key` is not
+  // specified, we use the caller function name as the default `key`, to
+  // distinguish between different collective functions within the same group.
+  // You can provide your own `key` if your function uses two different
+  // device communicators on the same group at the same time, for example,
+  // when concurrent collective operations are used.
+  // Returns a reference to the newly registered device communicator.
+  // @throws TORCH_CHECK if the device communicator is already registered for
+  //         the given group and key combination.
+  // Example:
+  // void foo(const std::string& group_name) {
+  //   // Try to get first.
+  //   auto devcomm_opt = get_devcomm(group_name);
+  //   if (!devcomm_opt) {
+  //     // Not found, create then register.
+  //     ncclDevComm devcomm = ncclDevCommCreate(...);
+  //     devcomm_opt = register_devcomm(group_name, devcomm);
+  //   }
+  //   ncclDevComm& devcomm_ref = *devcomm_opt;
+  //   // Use devcomm_ref
+  // }
+  // void bar(const std::string& group_name) {
+  //   ncclDevComm devcomm0 = ncclDevCommCreate(...);
+  //   ncclDevComm devcomm1 = ncclDevCommCreate(...);
+  //   // You can provide your own `key` if you want to, for example, to
+  //   // distinguish between concurrent collective operations.
+  //   register_devcomm(group_name, devcomm0, "bar0");
+  //   register_devcomm(group_name, devcomm1, "bar1");
+  // }
+  std::optional<std::reference_wrapper<ncclDevComm>> register_devcomm(
+      const std::string& group_name,
+      ncclDevComm devcomm,
+      const std::string& key = __builtin_FUNCTION()) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    // Ensure the group exists in the registry, creating an empty map if needed
+    auto [group_it, inserted] = devcomm_registry_.try_emplace(
+        group_name, std::unordered_map<std::string, ncclDevComm>());
+    auto& group_map = group_it->second;
+    // Try to insert the device communicator with the given key
+    // Use std::move to avoid copying the device communicator
+    auto [key_it, key_inserted] =
+        group_map.try_emplace(key, std::move(devcomm));
+    if (!key_inserted) {
+      // Already registered - this is a programming error, so throw
+      TORCH_CHECK(
+          false,
+          "NCCL device communicator for group ",
+          group_name,
+          " with key ",
+          key,
+          " already registered.");
+    }
+    // Return a reference to the newly registered device communicator
+    return std::make_optional(std::ref(key_it->second));
+  }
+
+  // Register a host-side NCCL communicator for a group.
+  // This should be called before registering any device communicators for the
+  // same group, as device communicators need the host communicator for cleanup.
+  // @param group_name The process group name
+  // @param comm The host-side NCCL communicator to register
+  // @throws TORCH_CHECK if the group is already registered with a different
+  // communicator.
+  void register_comm(const std::string& group_name, ncclComm_t comm) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto [it, inserted] = group_to_comm_.try_emplace(group_name, comm);
+    // If the communicator is already registered, check if it is the same one.
+    // If not, throw an error.
+    TORCH_CHECK(
+        inserted || it->second == comm, // this is just a pointer comparison
+        "NCCL host communicator for group ",
+        group_name,
+        " already registered.");
+  }
+
+  // Destructor: Clean up all registered device communicators.
+  // This is a best-effort cleanup. If the CUDA context has already been
+  // destroyed, the cleanup will be skipped. All errors are caught and ignored
+  // to prevent exceptions from propagating during destruction.
+  ~NCCLDevCommManager() noexcept {
+    // Best effort to destroy the device communicators. Skip if CUDA context has
+    // exited.
+    try {
+      c10::cuda::CUDAGuard guard(device_);
+      // Make sure all kernels have completed before destroying the device
+      // communicator. This is important to ensure no kernels are still using
+      // the device communicator when we destroy it.
+      C10_CUDA_CHECK(cudaDeviceSynchronize());
+      // Iterate through all groups and their device communicators
+      for (auto& [group_name, group_map] : devcomm_registry_) {
+        // Find the host communicator for the group.
+        // Device communicators need the host communicator for destruction.
+        auto comm_it = group_to_comm_.find(group_name);
+        if (comm_it != group_to_comm_.end()) {
+          // Destroy each device communicator in this group
+          for (auto& [_, devcomm] : group_map) {
+            // Destroy the device communicator using the host communicator
+            ncclDevCommDestroy(comm_it->second, &devcomm);
+          }
+        }
+      }
+    } catch (...) {
+      // Ignore the error - we're in a destructor and can't throw
+      // Log a warning for debugging purposes
+      LOG(WARNING)
+          << "Failed to destroy the NCCL device communicator, skipping";
+    }
+  }
+
+ private:
+  // Device where the NCCL device communicator manager is created.
+  // The manager is device-specific and cannot be used across multiple devices.
+  const c10::Device device_;
+
+  // Mutex to protect the registry maps.
+  std::mutex mutex_;
+
+  // A map from process group name to the host-side NCCL communicator.
+  // The host communicator is required for creating and destroying device
+  // communicators. It should be registered before any device communicators
+  // for the same group.
+  std::unordered_map<std::string, ncclComm_t> group_to_comm_;
+
+  // A two-level map for device communicators:
+  // - First level: keyed by process group name
+  // - Second level: keyed by an optional key (defaults to caller function name
+  //   via __builtin_FUNCTION())
+  //
+  // This structure allows multiple device communicators per process group,
+  // which is useful when different functions need separate device communicators
+  // for concurrent operations. The key defaults to the caller's function name,
+  // but can be customized for cases where a single function needs multiple
+  // device communicators.
+  std::unordered_map<std::string, std::unordered_map<std::string, ncclDevComm>>
+      devcomm_registry_;
+};
+
+} // namespace c10d::symmetric_memory
+#endif // NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..c256e2b25116ea92d4d80361eca4e9e7dd438f0b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp
@@ -0,0 +1,38 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/macros/Macros.h>
+
+namespace c10d::nccl_extension {
+
+TORCH_API bool is_nccl_symmem_available();
+
+TORCH_API void nccl_put(at::Tensor& tensor, const int64_t peer);
+
+TORCH_API void nccl_get(at::Tensor& tensor, const int64_t peer);
+
+TORCH_API void nccl_wait_for_signal(at::Tensor& sigpad, int64_t signal);
+
+TORCH_API void nccl_put_with_signal(
+    at::Tensor& tensor,
+    int64_t signal,
+    int64_t peer);
+
+// Simultaneously reduce N blocks of a 2-D input tensor from a shared symmetric
+// memory buffer, routing each to a specific destination rank. Blocks are
+// described by inclusive-prefix-sum offsets along `dim` (0 or 1); all blocks
+// must have equal size.
+TORCH_API void nccl_reduce_scatter_offset(
+    const at::Tensor& input,
+    at::TensorList out,
+    const std::string& group_name,
+    int64_t dim,
+    std::optional<at::IntArrayRef> offsets,
+    std::optional<at::IntArrayRef> dst_ranks,
+    const std::string& red_op);
+} // namespace c10d::nccl_extension
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..6a524af7b6f85cca05fa65c9f4738af65d961f3d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp
@@ -0,0 +1,90 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/macros/Macros.h>
+
+#define NVSHMEM_CHECK(stmt, msg)                                             \
+  do {                                                                       \
+    int result = (stmt);                                                     \
+    TORCH_CHECK(                                                             \
+        result == 0,                                                         \
+        std::string(__FILE__) + ":" + std::to_string(__LINE__) + " " + msg + \
+            ". Error code: " + std::to_string(result));                      \
+  } while (0)
+
+namespace c10d::nvshmem_extension {
+
+// Check if NVSHMEM is available
+TORCH_API bool is_nvshmem_available();
+
+// Initializes the device state in CUmodule so that it’s able to perform NVSHMEM
+// operations.
+TORCH_API void nvshmemx_cumodule_init(uintptr_t module);
+
+TORCH_API void nvshmem_put(at::Tensor& tensor, const int64_t peer);
+
+TORCH_API void nvshmem_get(at::Tensor& tensor, const int64_t peer);
+
+at::Tensor nvshmem_broadcast(
+    at::Tensor& input,
+    const int64_t root,
+    const std::string& group_name);
+
+TORCH_API void nvshmem_wait_for_signal(
+    at::Tensor& sigpad,
+    int64_t signal,
+    int64_t peer);
+
+TORCH_API void nvshmem_put_with_signal(
+    at::Tensor& tensor,
+    at::Tensor& sigpad,
+    int64_t signal,
+    int64_t peer);
+
+at::Tensor nvshmem_all_to_all(
+    at::Tensor& input,
+    at::Tensor& out,
+    std::string group_name);
+
+void all_to_all_vdev(
+    at::Tensor& input,
+    at::Tensor& out,
+    at::Tensor& in_splits,
+    at::Tensor& out_splits_offsets,
+    std::string group_name);
+
+void all_to_all_vdev_2d(
+    at::Tensor& input,
+    at::Tensor& out,
+    at::Tensor& in_splits,
+    at::Tensor& out_splits_offsets,
+    std::string group_name,
+    std::optional<int64_t> major_align = std::nullopt);
+
+void all_to_all_vdev_2d_offset(
+    at::Tensor& input,
+    at::Tensor& out,
+    at::Tensor& in_splits_offsets,
+    at::Tensor& out_splits_offsets,
+    std::string group_name);
+
+void tile_reduce(
+    at::Tensor& in_tile,
+    at::Tensor& out_tile,
+    int64_t root,
+    std::string group_name,
+    std::string reduce_op = "sum");
+
+void multi_root_tile_reduce(
+    at::ArrayRef<at::Tensor> in_tiles,
+    at::Tensor& out_tile,
+    at::ArrayRef<int64_t> roots,
+    std::string group_name,
+    std::string reduce_op = "sum");
+
+} // namespace c10d::nvshmem_extension
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..606ef22400b1fd9829967ba6bc425f871a4261bf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp
@@ -0,0 +1,179 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/util/Exception.h>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+// Starting from NVSHMEM 3.3.9, nvshmem_host.h exists so that we can cleanly
+// include only the nvshmem host library headers:
+// #include <nvshmem_host.h>
+// It translates into the following two lines:
+#if !defined(USE_ROCM)
+#include <host/nvshmem_api.h>
+#include <host/nvshmemx_api.h>
+#else
+#include <rocshmem/rocshmem.hpp>
+#endif
+// For maximum compatibility, we use the "host/" style for now.
+
+namespace c10d::nvshmem_extension {
+
+// This corresponds to max nblocks
+constexpr int MAX_N_TEAMS = 128;
+
+// A pool of teams for each group. These are duplicate teams.
+using TeamPool = std::vector<nvshmem_team_t>;
+
+// Manage all the team business. Singleton.
+class TeamManager {
+ public:
+  // Constructor
+  explicit TeamManager(const c10::Device device) : device_(device) {}
+
+  // Get single, global manager.
+  static TeamManager& get(const c10::Device device) {
+    static TeamManager manager(device);
+    TORCH_CHECK(
+        manager.device_ == device,
+        "Detected use of TeamManager on multiple devices. This is not supported.");
+    return manager;
+  }
+
+  // Get a team for a group.
+  nvshmem_team_t get_team(
+      const std::string& group_name,
+      const std::vector<int>& global_ranks) {
+    auto [team_pool, pool_updated] =
+        group_to_team_pool(group_name, global_ranks, 1);
+    // Return the fist available team
+    return team_pool[0];
+  }
+
+  // Get n teams for a group.
+  // The first element of the returned pair is the team pool on host side.
+  // The second element of the returned pair is the team pool on device side.
+  // This API must be call with a device guard.
+  std::pair<const TeamPool&, nvshmem_team_t*> get_n_teams(
+      const std::string& group_name,
+      const std::vector<int>& global_ranks,
+      const int need_n) {
+    // A device guard is required for malloc and memcpy below
+    c10::cuda::CUDAGuard guard(device_);
+    // Get the team pool with the requested number of teams
+    auto [team_pool, pool_updated] =
+        group_to_team_pool(group_name, global_ranks, need_n);
+    // Check if the pool already exists in device memory
+    nvshmem_team_t* team_pool_dev = nullptr;
+    constexpr auto pool_bytes = sizeof(nvshmem_team_t) * MAX_N_TEAMS;
+    auto it = team_pool_devptrs_.find(group_name);
+    if (it == team_pool_devptrs_.end()) {
+      // If not, allocate a new pool in device memory
+      team_pool_dev = reinterpret_cast<nvshmem_team_t*>(
+          c10::cuda::CUDACachingAllocator::raw_alloc(pool_bytes));
+      team_pool_devptrs_[group_name] = team_pool_dev;
+    } else {
+      team_pool_dev = it->second;
+    }
+    // Update the pool in device memory if host side pool is updated
+    if (pool_updated) {
+      TORCH_INTERNAL_ASSERT(team_pool.size() == MAX_N_TEAMS);
+      auto stream = at::cuda::getCurrentCUDAStream();
+      C10_CUDA_CHECK(cudaMemcpyAsync(
+          team_pool_dev,
+          team_pool.data(),
+          pool_bytes,
+          cudaMemcpyHostToDevice,
+          stream));
+    }
+    return std::make_pair(std::cref(team_pool), team_pool_dev);
+  }
+
+  ~TeamManager() noexcept {
+    // Free the team pools in device memory
+    // Note that we do it in a best effort manner because the team pool is
+    // managed by a static TeamManager and the destruction order of static
+    // objects is undetermined. If the destructor is called after the CUDA
+    // context is destroyed, cudaFree would fail.
+    try {
+      // cudaFree generally implies a device synchronization, meaning it will
+      // block until all preceding CUDA operations on the device have completed
+      // before freeing the memory. Thus we don't need to worry about freeing
+      // the memory before CUDA kernels complete.
+      for (auto& [_, team_pool_dev] : team_pool_devptrs_) {
+        c10::cuda::CUDACachingAllocator::raw_delete(team_pool_dev);
+      }
+    } catch (...) {
+      // Ignore the error
+      std::cerr << "Failed to free the team pool in device memory, skipping\n";
+    }
+  }
+
+ private:
+  // Get the team pool for a group. If the pool doesn't exist, create it. If the
+  // pool exists but is not large enough, create more teams.
+  // The first element of the returned pair is the team pool on host side.
+  // The second element of the returned pair is a boolean indicating if the pool
+  // is updated.
+  std::pair<const TeamPool&, bool> group_to_team_pool(
+      const std::string& group_name,
+      const std::vector<int>& global_ranks,
+      const int need_n) {
+    TORCH_CHECK(need_n < MAX_N_TEAMS, "Too many teams requested");
+    // Guarding the NVSHMEM API calls below just to be safe
+    c10::cuda::CUDAGuard guard(device_);
+
+    // Insert a new team pool if not exists
+    auto [it, inserted] = group_name_to_team_pool_.emplace(
+        group_name, TeamPool(MAX_N_TEAMS, NVSHMEM_TEAM_INVALID));
+    auto& team_pool = it->second;
+    bool pool_updated = inserted;
+
+    // Create new teams if what's requested is more than what we have
+    int stride = 0; // stride in globe, uninitialized
+    for (int i = 0; i < need_n; ++i) {
+      if (team_pool[i] != NVSHMEM_TEAM_INVALID) {
+        continue;
+      }
+      // Some checks before we create new teams
+      if (stride == 0) { // Check only once
+        TORCH_CHECK(global_ranks.size() > 1);
+        stride = global_ranks[1] - global_ranks[0];
+        for (size_t r = 1; r < global_ranks.size(); ++r) {
+          TORCH_CHECK(global_ranks[r] - global_ranks[r - 1] == stride);
+        }
+      }
+      nvshmem_team_t team = NVSHMEM_TEAM_INVALID;
+      nvshmem_team_split_strided(
+          NVSHMEM_TEAM_WORLD,
+          global_ranks[0],
+          stride,
+          global_ranks.size(),
+          nullptr,
+          0,
+          &team);
+      TORCH_CHECK(team != NVSHMEM_TEAM_INVALID, "Failed to create a new team");
+      team_pool[i] = team;
+      pool_updated = true;
+    }
+    return std::make_pair(std::cref(team_pool), pool_updated);
+  }
+
+ private:
+  // Device where the team manager is created
+  const c10::Device device_;
+  // A map from group name to team pool for that group.
+  std::unordered_map<std::string, TeamPool> group_name_to_team_pool_;
+  // A map from group name to team pool array in device memory.
+  std::unordered_map<std::string, nvshmem_team_t*> team_pool_devptrs_;
+};
+
+} // namespace c10d::nvshmem_extension
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/python_placement.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/python_placement.h
new file mode 100644
index 0000000000000000000000000000000000000000..a15d330768b9788f18b96b6620c750ab0abdf9da
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/python_placement.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/python_stub.h>
+
+namespace torch::distributed {
+void initPlacementBindings(PyObject* module);
+} // namespace torch::distributed
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/agent_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/agent_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..25f915f21a7cc914b352ba4b8208e3ce4617d9a7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/agent_utils.h
@@ -0,0 +1,47 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/c10d/PrefixStore.hpp>
+#include <torch/csrc/distributed/rpc/utils.h>
+
+namespace torch::distributed::rpc {
+
+// All RPC peers should call into this function at the same time. Each peer
+// provides its own id and name, and this function uses the given Store to
+// gather global name-to-id mapping on all peers.
+TORCH_API std::unordered_map<std::string, worker_id_t> collectNames(
+    ::c10d::PrefixStore store,
+    const worker_id_t selfId,
+    const std::string& selfName,
+    const int worldSize);
+
+// Ranks in dynamic RPC groups will initially call into this to establish the
+// name-to-id mapping for the current peers in the group. The current rank will
+// put its own worker info in the store and discover all the ranks that came
+// before it. NOTE: This needs to be called with the Dynamic RPC group
+// membership management token held.
+TORCH_API std::unordered_map<std::string, worker_id_t> collectCurrentNames(
+    ::c10d::PrefixStore store,
+    const worker_id_t selfId,
+    const std::string& selfName);
+
+// Remove name from Store, used in dynamic RPC groups.
+// NOTE: This needs to be called with the Dynamic RPC group
+// membership management token held.
+TORCH_API void removeCurrentName(
+    ::c10d::PrefixStore store,
+    const worker_id_t selfId,
+    const std::string& selfName);
+
+// This performs a synchronization of all call counts by using store.
+// All RPC peers wait for others to join to exit at the same time.
+TORCH_API int syncCallCount(
+    ::c10d::PrefixStore store,
+    const int worldSize,
+    int activeCalls = 0);
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/message.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/message.h
new file mode 100644
index 0000000000000000000000000000000000000000..f6e71490f27f885d1a2e4e897a29fe50a1cb0ea5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/message.h
@@ -0,0 +1,198 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/types.h>
+#include <vector>
+
+namespace torch::distributed::rpc {
+
+// An enum denoting common RPC errors to allow specific error handling for them.
+// NOLINTNEXTLINE(performance-enum-size)
+enum RPCErrorType {
+  UNKNOWN_ERROR = 0, /* Indicates that error type could not be parsed */
+  TIMEOUT = 1, /* Indicates that the RPC has timed out */
+  INTENTIONAL_FAILURE = 2 /* Deliberate failure, such as those injected by
+                             FaultyAgent for testing */
+};
+
+// The enum values are bitwise ORed with MessageType
+// They are bit flags starting from 0x100 and should have
+// value such as 0x100, 0x200, 0x400, 0x800, 0xF00, etc.
+// NOLINTNEXTLINE(performance-enum-size)
+enum MessageTypeFlags {
+  REQUEST_TYPE = 0x100,
+  RESPONSE_TYPE = 0x200,
+};
+
+// Message types must have values between 0x00 to 0xff
+// NOLINTNEXTLINE(performance-enum-size)
+enum MessageType {
+  // messages for dist.rpc on builtin operators
+  SCRIPT_CALL = 0x00 | MessageTypeFlags::REQUEST_TYPE,
+  SCRIPT_RET = 0x01 | MessageTypeFlags::RESPONSE_TYPE,
+
+  // messages for dist.rpc on Python UDF
+  PYTHON_CALL = 0x02 | MessageTypeFlags::REQUEST_TYPE,
+  PYTHON_RET = 0x03 | MessageTypeFlags::RESPONSE_TYPE,
+
+  // messages for dist.remote on builtin operators and Python UDF
+  SCRIPT_REMOTE_CALL = 0x04 |
+      MessageTypeFlags::REQUEST_TYPE, // A remote call on a builtin operator
+  PYTHON_REMOTE_CALL =
+      0x05 | MessageTypeFlags::REQUEST_TYPE, // A remote call on a Python UDF
+  REMOTE_RET =
+      0x06 | MessageTypeFlags::RESPONSE_TYPE, // Response for remote calls for
+                                              // UDF, builtin, or script
+
+  // RRef related internal messages
+  SCRIPT_RREF_FETCH_CALL =
+      0x07 | MessageTypeFlags::REQUEST_TYPE, // A UserRRef<IValue> fetches value
+                                             // from owner
+  PYTHON_RREF_FETCH_CALL =
+      0x08 | MessageTypeFlags::REQUEST_TYPE, // A UserRRef<py::object> fetches
+                                             // value from owner
+  SCRIPT_RREF_FETCH_RET = 0x09 |
+      MessageTypeFlags::RESPONSE_TYPE, // An OwnerRRef sends ivalue to user
+  PYTHON_RREF_FETCH_RET = 0x0a |
+      MessageTypeFlags::RESPONSE_TYPE, // An OwnerRRef sends py::object to user
+  RREF_USER_DELETE = 0x0b |
+      MessageTypeFlags::REQUEST_TYPE, // A UserRRef tells the owner to deref
+  RREF_FORK_REQUEST =
+      0x0c | MessageTypeFlags::REQUEST_TYPE, // A child UserRRef tells the owner
+                                             // about itself
+  RREF_CHILD_ACCEPT =
+      0x0d | MessageTypeFlags::REQUEST_TYPE, // A child UserRRef tells parent
+                                             // that owner knows it
+  RREF_ACK =
+      0x0e | MessageTypeFlags::RESPONSE_TYPE, // ACK to internal RRef messages
+
+  // Messages with autograd info
+  FORWARD_AUTOGRAD_REQ = 0x0f | MessageTypeFlags::REQUEST_TYPE,
+  FORWARD_AUTOGRAD_RESP = 0x10 | MessageTypeFlags::RESPONSE_TYPE,
+
+  // Messages to propagate gradients on the backward pass.
+  BACKWARD_AUTOGRAD_REQ = 0x11 | MessageTypeFlags::REQUEST_TYPE,
+  BACKWARD_AUTOGRAD_RESP = 0x12 | MessageTypeFlags::RESPONSE_TYPE,
+
+  // Messages to tell workers to clean up their autograd context.
+  CLEANUP_AUTOGRAD_CONTEXT_REQ = 0x13 | MessageTypeFlags::REQUEST_TYPE,
+  CLEANUP_AUTOGRAD_CONTEXT_RESP = 0x14 | MessageTypeFlags::RESPONSE_TYPE,
+
+  // Messages that tell workers to run requests with profiling enabled.
+  RUN_WITH_PROFILING_REQ = 0x15 | MessageTypeFlags::REQUEST_TYPE,
+  RUN_WITH_PROFILING_RESP = 0x16 | MessageTypeFlags::RESPONSE_TYPE,
+
+  // Messages to support RRef.backward().
+  RREF_BACKWARD_REQ = 0x17 | MessageTypeFlags::REQUEST_TYPE,
+  RREF_BACKWARD_RESP = 0x18 | MessageTypeFlags::RESPONSE_TYPE,
+
+  // Other internal message types
+  EXCEPTION = 0x37 | MessageTypeFlags::RESPONSE_TYPE,
+  UNKNOWN = 0x3c
+};
+
+// A message to be sent/received by an RpcAgent.
+//
+// A Message object contains 4 fields:
+//    payload (std::vector<char>): a binary chunk of data.
+//    tensors (std::vector<torch::Tensor>): all tensors. Tensor data are not
+//        included in the payload, and it is up to the RpcAgent implementation
+//        to determine how to serialize them. This design is helpful for
+//        communicating super large tensors where serializing all the data at
+//        once leads to excessively large memory footprint. An implementation
+//        can then serialize and send tensors chunk-by-chunk, in the streaming
+//        fashion.
+//    type (MessageType): type of the message.
+//    id (int64_t): message id, this is used to match request and response.
+//               Other implementation can ignore it if they have their own
+//               ways to do matching.
+//
+// Layers above ``RpcAgent`` only converts ScriptCall, ScriptResp, PythonCall,
+// and PythonResp into a Message, and it is up to the RpcAgent
+// implementation to determine how to serialize a message.
+class TORCH_API Message final : public torch::CustomClassHolder {
+ private:
+  // Keep these private in order to force users to go through make_intrusive and
+  // thus prevent creating a Message that's not held by an intrusive_ptr.
+  Message();
+
+  Message(
+      std::vector<char>&& payload,
+      std::vector<torch::Tensor>&& tensors,
+      MessageType type);
+
+  Message(
+      std::vector<char>&& payload,
+      std::vector<torch::Tensor>&& tensors,
+      MessageType type,
+      int64_t id);
+
+  friend c10::intrusive_ptr<Message>;
+
+ public:
+  Message(const Message& other) = delete;
+  Message(Message&& other) = delete;
+  Message& operator=(Message const& rhs) = delete;
+  Message& operator=(Message&& rhs) = delete;
+  ~Message() override = default;
+
+  // Destructively retrieves the payload.
+  std::vector<char>&& movePayload() &&;
+  std::vector<torch::Tensor>&& moveTensors() &&;
+
+  std::vector<char>& payload();
+  const std::vector<char>& payload() const;
+  std::vector<torch::Tensor>& tensors();
+  const std::vector<torch::Tensor>& tensors() const;
+  MessageType type() const;
+
+  bool isRequest() const;
+  bool isResponse() const;
+  bool isShutdown() const;
+
+  // id is an optional field to match request/response. If an RpcAgent
+  // implementation is able to do the matching without using this id, it can be
+  // dropped during message serialization.
+  int64_t id() const;
+  void setId(int64_t id);
+
+  std::vector<c10::weak_intrusive_ptr<c10::StorageImpl>> getStorages() const;
+
+ private:
+  std::vector<char> payload_;
+  std::vector<torch::Tensor> tensors_;
+  MessageType type_ = MessageType::UNKNOWN;
+  int64_t id_ = -1;
+};
+
+// Create a response Message of type Exception.
+// The exception string representation will be used as the message's payload.
+// A message ID corresponding to the request that resulted in this response can
+// be provided for matching requests/responses.
+TORCH_API c10::intrusive_ptr<Message> createExceptionResponse(
+    const std::exception& e,
+    int64_t id);
+
+// Create a response Message of type Exception.
+// The passed in string representation will be used as the message's payload.
+// A message ID corresponding to the request that resulted in this response can
+// be provided for matching requests/responses.
+TORCH_API c10::intrusive_ptr<Message> createExceptionResponse(
+    const std::string& exceptionStr,
+    int64_t id);
+
+inline std::tuple<
+    c10::intrusive_ptr<Message>,
+    std::vector<c10::weak_intrusive_ptr<c10::StorageImpl>>>
+withStorages(c10::intrusive_ptr<Message> message) {
+  auto storages = message->getStorages();
+  return std::make_tuple(std::move(message), std::move(storages));
+}
+
+using JitFuture = c10::ivalue::Future;
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/metrics/RpcMetricsHandler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/metrics/RpcMetricsHandler.h
new file mode 100644
index 0000000000000000000000000000000000000000..d46a0e9be89a5957fd498fe2fbd17c7e02767e4b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/metrics/RpcMetricsHandler.h
@@ -0,0 +1,45 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Registry.h>
+#include <string>
+
+namespace torch::distributed::rpc {
+// All metrics are prefixed with the following key.
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+constexpr char kRpcMetricsKeyPrefix[] = "torch.distributed.rpc.";
+// APIs for logging time-series metrics for RPC-based distributed
+// training. Implementations of this class should provide thread safety so that
+// metrics can be logged from multiple threads without the user needing to
+// coordinate serialization.
+class RpcMetricsHandler {
+ public:
+  // Accumulates the metric value specified by the name for purposes of
+  // computing aggregate statistics over time.
+  virtual void accumulateMetric(const std::string& name, double value) = 0;
+  // Increment a count for the metric given by the name.
+  virtual void incrementMetric(const std::string& name) = 0;
+  virtual ~RpcMetricsHandler() = default;
+};
+
+// Configuration struct for metrics handling.
+struct RpcMetricsConfig {
+  explicit RpcMetricsConfig(std::string handlerName, bool enabled)
+      : handlerName_(std::move(handlerName)), enabled_(enabled) {}
+
+  // Handler name
+  std::string handlerName_;
+  // Whether metrics exporting should be enabled or not.
+  bool enabled_;
+};
+
+// A registry for different implementations of RpcMetricsHandler. Classes
+// implementing the above interface should use this to register implementations.
+TORCH_DECLARE_REGISTRY(
+    RpcMetricsHandlerRegistry,
+    torch::distributed::rpc::RpcMetricsHandler);
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/profiler/remote_profiler_manager.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/profiler/remote_profiler_manager.h
new file mode 100644
index 0000000000000000000000000000000000000000..ba39ea8a02305f64f642b6927ee43b35459d10b8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/profiler/remote_profiler_manager.h
@@ -0,0 +1,60 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <mutex>
+#include <optional>
+#include <unordered_map>
+
+namespace torch::distributed::rpc {
+extern const std::string REMOTE_PROFILING_KEY_PREFIX;
+
+class TORCH_API RemoteProfilerManager {
+ public:
+  // Retrieves the lazily-initialized RemoteProfilerManager singleton instance.
+  static RemoteProfilerManager& getInstance();
+  // Sets the current, thread-local profiling key.
+  void setCurrentKey(std::string key);
+  // Returns whether the current profiling key is set.
+  bool isCurrentKeySet() const;
+  // Unsets the current, thread-local profiling key to allow other RPCs to reset
+  // it.
+  void unsetCurrentKey();
+  // inserts a pair (globallyUniqueId, key) to an in-memory map. The
+  // corresponding ID is used in RPC deserialization to prefix remotely profiled
+  // events with the right key.
+  void saveRPCKey(
+      ProfilingId globallyUniqueId,
+      const std::string& rpcProfilingKey);
+  // Retrieves the profiling key corresponding to the given globallyUniqueId.
+  // Throws if it is not found.
+  std::string retrieveRPCProfilingKey(const ProfilingId& globallyUniqueId);
+  // Generates the next globally unique ID for profiling.
+  ProfilingId getNextProfilerId();
+  // Retrieves the currently set thread-local profiling key. Throws if it is not
+  // set.
+  std::string& getCurrentProfilingKey();
+  // erases the globallyUniqueId from the map. This can help save memory in the
+  // case that many RPCs are being profiled.
+  void eraseKey(const ProfilingId& globallyUniqueId);
+
+  RemoteProfilerManager(const RemoteProfilerManager& other) = delete;
+  RemoteProfilerManager operator=(const RemoteProfilerManager& other) = delete;
+  RemoteProfilerManager(RemoteProfilerManager&&) = delete;
+  RemoteProfilerManager& operator=(RemoteProfilerManager&&) = delete;
+
+ private:
+  RemoteProfilerManager();
+  ~RemoteProfilerManager() = default;
+  local_id_t getNextLocalId();
+  std::unordered_map<ProfilingId, std::string, ProfilingId::Hash>
+      profiledRpcKeys_;
+  static thread_local std::optional<std::string> currentThreadLocalKey_;
+  std::mutex mutex_;
+  local_id_t currentLocalId_;
+};
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/profiler/server_process_global_profiler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/profiler/server_process_global_profiler.h
new file mode 100644
index 0000000000000000000000000000000000000000..f86461f6f895bee8a46b54b9f57dd88b66362ad6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/profiler/server_process_global_profiler.h
@@ -0,0 +1,134 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <shared_mutex>
+#include <utility>
+
+#include <torch/csrc/autograd/profiler.h>
+
+namespace torch::distributed::rpc::profiler::processglobal {
+
+using namespace torch::autograd::profiler;
+
+// Process global profiler state.
+//
+// This class holds information about a profiling range, from "enable" to
+// "disable".
+// An instance of this ``State`` will be
+// pushed into a global stack, so nested profiling range is supported.
+//
+// It has 2 members.
+// One is ``autograd::profiler::ProfilerConfig``. It's set by user and
+// will be copied to thread-local profiler state of RPC threads.
+// The other is a container that aggregates recorded
+// ``autograd::profiler::Event``s from all thread-local profilers on RPC
+// threads.
+class State {
+ public:
+  explicit State(ProfilerConfig config) : config_(std::move(config)) {}
+  ~State() = default;
+
+  const ProfilerConfig& config() const {
+    return config_;
+  }
+
+  void pushResult(thread_event_lists result) {
+    std::unique_lock<std::mutex> lock(resultsMutex_);
+
+    // NB: When a thread wants to push an entry into the this container,
+    // main control logic might have exited the process-global profile range.
+    results_.emplace_back(std::move(result));
+  }
+
+  std::vector<thread_event_lists> results();
+
+ private:
+  // Each result comes from a profile range. In each profile range, there is a
+  // "__profiler_start" marker event that all following events calculate time
+  // relative to it, so it's required to call
+  // parse_cpu_trace(result) for results of all profile range.
+  std::mutex resultsMutex_;
+  std::vector<thread_event_lists> results_;
+  const ProfilerConfig config_ = ProfilerConfig(ProfilerState::Disabled);
+};
+
+class StateStackEntry;
+
+#if defined(__MACH__)
+// Compiler error: 'shared_timed_mutex' is unavailable: introduced in
+// macOS 10.12
+using mutexType = std::mutex;
+// Compiler error: 'shared_lock' is unavailable: introduced in
+// macOS 10.12
+using rLockType = std::unique_lock<std::mutex>;
+using wLockType = std::unique_lock<std::mutex>;
+#else
+using mutexType = std::shared_timed_mutex;
+using rLockType = std::shared_lock<std::shared_timed_mutex>;
+using wLockType = std::unique_lock<std::shared_timed_mutex>;
+#endif
+
+// This is the global stack of ``State``s.
+TORCH_API extern std::shared_ptr<StateStackEntry> currentStateStackEntryPtr;
+TORCH_API extern mutexType currentStateStackEntryMutex;
+
+// This class is used to implement a stack of ``State``s.
+// It has 2 members.
+// One is `prevPtr`, a shared_ptr pointing to previous element in the
+// stack.
+// The other is ``statePtr``, a shared_ptr pointing to ``State``.
+class StateStackEntry {
+ public:
+  StateStackEntry(
+      std::shared_ptr<StateStackEntry> prevPtr,
+      std::shared_ptr<State> statePtr)
+      : prevPtr_(std::move(prevPtr)), statePtr_(std::move(statePtr)) {}
+
+  static void pushRange(std::shared_ptr<State> profilerProcessGlobalStatePtr);
+  static std::shared_ptr<State> popRange();
+
+  static std::shared_ptr<StateStackEntry> current() {
+    rLockType rlock(currentStateStackEntryMutex);
+
+    return currentStateStackEntryPtr;
+  }
+
+  std::shared_ptr<StateStackEntry> prevPtr() const {
+    return prevPtr_;
+  }
+
+  std::shared_ptr<State> statePtr() const {
+    return statePtr_;
+  }
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::shared_ptr<StateStackEntry> prevPtr_{nullptr};
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::shared_ptr<State> statePtr_{nullptr};
+};
+
+// Push the result to ``State``s of current profile range and recursively outer
+// profile ranges.
+TORCH_API void pushResultRecursive(
+    std::shared_ptr<StateStackEntry> stateStackEntryPtr,
+    const thread_event_lists& result);
+
+// User-facing API.
+//
+// Enter a server-side process-global profiling range.
+// Profiling range can be neste, so it's ok to call this API for multiple
+// times. This enables all RPC threads running server-side request callbacks.
+TORCH_API void enableServer(const ProfilerConfig& new_config);
+//
+// Exit a server-side process-global profiling range.
+// Profiling range can be neste, so it's possible that profiler is still on
+// after calling this API.
+// This enables all RPC threads running server-side request callbacks.
+TORCH_API std::vector<thread_event_lists> disableServer();
+
+} // namespace torch::distributed::rpc::profiler::processglobal
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/py_rref.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/py_rref.h
new file mode 100644
index 0000000000000000000000000000000000000000..c5735e844b6e1f62dbfc779a76db78857e111593
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/py_rref.h
@@ -0,0 +1,86 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/rref_impl.h>
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::distributed::rpc {
+
+// NOLINTNEXTLINE(performance-enum-size)
+enum RRefProxyType { RPC_SYNC, RPC_ASYNC, REMOTE };
+
+// Python wrapper of an RRef shared_ptr that supports Python
+// pickle and unpickle.
+class PYBIND11_EXPORT PyRRef {
+ public:
+  // The first ctor can only be called while holding GIL. See its implementation
+  // for more explanations.
+  explicit PyRRef(const py::object& value, const py::object& type_hint);
+  explicit PyRRef(c10::intrusive_ptr<RRef> rref);
+  PyRRef(const PyRRef&) = default;
+  ~PyRRef();
+
+  bool isOwner() const;
+  bool confirmedByOwner() const;
+  WorkerInfo owner() const;
+  std::string ownerName() const;
+  py::object toHere(
+      const float timeoutSeconds =
+          torch::distributed::rpc::kUnsetRpcTimeout) const;
+  py::object localValue() const;
+  std::string str() const;
+  py::tuple pickle() const;
+  static PyRRef unpickle(const py::tuple& t);
+  c10::IValue toIValue() const;
+  // Future that is associated with the creation of this RRef on the remote end.
+  // This is only used to get the future corresponding to the rref for profiling
+  // use cases.
+  c10::intrusive_ptr<JitFuture> getFuture() const;
+  // Keeps track of the future responsible for profiling owner creation
+  // acknowledgement
+  c10::intrusive_ptr<JitFuture> getProfilingFuture() const;
+  // Sets the future responsible for profiling owner creation acknowledgement.
+  // This future is set from python to be a future that returns when profiling
+  // callbacks have been run.
+  void setProfilingFuture(c10::intrusive_ptr<JitFuture> profilingFuture);
+
+  // create a proxy on this RRef, which can be used to launch RPC on the owner
+  // of this RRef to run functions on the object referenced by this RRef.
+  py::object createRRefProxy(
+      const RRefProxyType& mode,
+      float timeoutSeconds = rpc::kUnsetRpcTimeout) const;
+
+  // get the type of the data object referenced by this RRef. Timeout argument
+  // is only used in the first invocation of this function as an argument to the
+  // RPC to the owner node of the RRef.
+  py::object getRRefType(
+      float timeout = rpc::kUnsetRpcTimeout,
+      bool blocking = true);
+
+  // Run the backward pass with the RRef as the root.
+  void backward(int64_t autogradContextId, bool retainGraph);
+
+  // Helper static function to run backward on a given rref.
+  static void backward(
+      int64_t autogradContextId,
+      bool retainGraph,
+      const c10::intrusive_ptr<RRef>& rref);
+
+  // Specialization of backward if the rref is an OwnerRRef.
+  static void backwardOwnerRRef(
+      int64_t autogradContextId,
+      bool retainGraph,
+      IValue value);
+
+ private:
+  c10::intrusive_ptr<RRef> rref_;
+  std::optional<c10::intrusive_ptr<JitFuture>> profilingFuture_;
+  std::optional<py::object> type_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_call.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_call.h
new file mode 100644
index 0000000000000000000000000000000000000000..ea339cae11b4dedb5a2aefa2b702443bc53708a4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_call.h
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+
+namespace torch::distributed::rpc {
+
+// RPC call representing calling a Python function over RPC.
+class TORCH_API PythonCall final : public RpcCommandBase {
+ public:
+  PythonCall(SerializedPyObj&& serializedPyObj, bool isAsyncExecution);
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+
+  static std::unique_ptr<PythonCall> fromMessage(const Message& message);
+
+  const SerializedPyObj& serializedPyObj() const;
+
+  inline bool isAsyncExecution() const {
+    return isAsyncExecution_;
+  }
+
+ private:
+  SerializedPyObj serializedPyObj_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const bool isAsyncExecution_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..ba69781bdcb0a200b52d30da80cabeb52ff8608d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_functions.h
@@ -0,0 +1,71 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/py_rref.h>
+#include <torch/csrc/distributed/rpc/rpc_agent.h>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::distributed::rpc {
+
+// Converts an internal ivalue::Future of Message into a user-facing
+// ivalue::Future of py::object type by creating a new ivalue::Future and call
+// its  markCompleted as a callback in the given ivalue::Future.
+// If hasValue is true, the Message will be converted into a py::object and then
+// wrap it with an IValue. If hasValue is false, this ivalue::Future is only
+// used for signaling and launching callbacks. In this case, the message will be
+// discarded and then set the ivalue::Future using an empty IValue or the given
+// FutureError if there is an error.
+c10::intrusive_ptr<JitFuture> toPyJitFuture(
+    const c10::intrusive_ptr<JitFuture>& messageJitFuture,
+    bool hasValue = true);
+
+c10::intrusive_ptr<JitFuture> pyRpcBuiltin(
+    const WorkerInfo& dst,
+    const std::string& opName,
+    const py::args& args,
+    const py::kwargs& kwargs,
+    const float rpcTimeoutSeconds);
+
+c10::intrusive_ptr<JitFuture> pyRpcPythonUdf(
+    const WorkerInfo& dst,
+    std::string& pickledPythonUDF,
+    std::vector<torch::Tensor>& tensors,
+    const float rpcTimeoutSeconds,
+    const bool isAsyncExecution);
+
+c10::intrusive_ptr<JitFuture> pyRpcTorchscript(
+    const std::string& dstWorkerName,
+    const std::string& qualifiedNameStr,
+    const py::tuple& argsTuple,
+    const py::dict& kwargsDict,
+    const float rpcTimeoutSeconds,
+    const bool isAsyncExecution);
+
+PyRRef pyRemoteBuiltin(
+    const WorkerInfo& dst,
+    const std::string& opName,
+    const float rpcTimeoutSeconds,
+    const py::args& args,
+    const py::kwargs& kwargs);
+
+PyRRef pyRemotePythonUdf(
+    const WorkerInfo& dst,
+    std::string& pickledPythonUDF,
+    std::vector<torch::Tensor>& tensors,
+    const float rpcTimeoutSeconds,
+    const bool isAsyncExecution);
+
+PyRRef pyRemoteTorchscript(
+    const std::string& dstWorkerName,
+    const std::string& qualifiedNameStr,
+    const float rpcTimeoutSeconds,
+    const bool isAsyncExecution,
+    const py::args& args,
+    const py::kwargs& kwargs);
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_remote_call.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_remote_call.h
new file mode 100644
index 0000000000000000000000000000000000000000..b34cb40349850c4c643e9e5d899b53293f220a2b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_remote_call.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+namespace torch::distributed::rpc {
+
+class TORCH_API PythonRemoteCall : public RpcCommandBase {
+ public:
+  PythonRemoteCall(
+      SerializedPyObj&& serializedPyObj,
+      at::IValue retRRefId,
+      at::IValue retForkId,
+      const bool isAsyncExecution);
+
+  inline const SerializedPyObj& serializedPyObj() const {
+    return serializedPyObj_;
+  }
+
+  inline const at::IValue& retRRefId() const {
+    return retRRefId_;
+  }
+
+  inline const at::IValue& retForkId() const {
+    return retForkId_;
+  }
+
+  inline bool isAsyncExecution() const {
+    return isAsyncExecution_;
+  }
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<PythonRemoteCall> fromMessage(const Message& message);
+
+ private:
+  SerializedPyObj serializedPyObj_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const at::IValue retRRefId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const at::IValue retForkId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const bool isAsyncExecution_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_resp.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_resp.h
new file mode 100644
index 0000000000000000000000000000000000000000..cc47fbe631a219bcd16a847d1bb849fa812068cf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_resp.h
@@ -0,0 +1,28 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+
+namespace torch::distributed::rpc {
+
+// RPC call representing the response of a Python UDF over RPC.
+class TORCH_API PythonResp final : public RpcCommandBase {
+ public:
+  explicit PythonResp(SerializedPyObj&& serializedPyObj);
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+
+  static std::unique_ptr<PythonResp> fromMessage(const Message& message);
+
+  const SerializedPyObj& serializedPyObj() const;
+
+ private:
+  SerializedPyObj serializedPyObj_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_rpc_handler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_rpc_handler.h
new file mode 100644
index 0000000000000000000000000000000000000000..8cd486389d7b0efe1fe438478a9d3313d15c1e2d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/python_rpc_handler.h
@@ -0,0 +1,134 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <torch/csrc/jit/frontend/script_type_parser.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::distributed::rpc {
+
+// Singleton class provides interface to execute python UDF remote call
+// and deserialize the returned results by running python function
+// in internal_rpc_utilities.
+// The singleton object is constructed at first when RPC agent is
+// constructed, where the python function in
+// torch/distributed/internal_rpc_utils.py are imported only once.
+class PYBIND11_EXPORT PythonRpcHandler {
+ public:
+  struct RRefProxyFunctions {
+    py::object rrefProxyCtor_;
+    py::object rpcSync_;
+    py::object rpcAsync_;
+    py::object remote_;
+  };
+
+  struct RRefTypeFunctions {
+    py::object onOwner_;
+    py::object onUser_;
+  };
+
+  static PythonRpcHandler& getInstance();
+
+  // Run a pickled Python UDF and return the result py::object
+  py::object runPythonUdf(const py::object& pythonUdf);
+
+  // Serialized a py::object into a string
+  SerializedPyObj serialize(const py::object& obj);
+
+  // Deserialize a string into a py::object
+  py::object deserialize(const SerializedPyObj& serializedObj);
+
+  // Check if obj is RemoteException, then throw it
+  void handleException(const py::object& obj);
+  // Alternative if the caller is already holding the GIL.
+  void handleExceptionGILHeld(const py::object& obj);
+  // Check if obj is an RemoteException instance.
+  bool isRemoteException(const py::object& obj);
+
+  // Explicitly clean up py::objects to avoid segment faults when
+  // py::objects with CPython are cleaned up later at program exit
+  // See similar issues reported https://github.com/pybind/pybind11/issues/1598
+  // and https://github.com/pybind/pybind11/issues/1493
+  // Our local tests also caught this segment faults if py::objects are cleaned
+  // up at program exit. The explanation is: CPython cleans up most critical
+  // utilities before cleaning up PythonRpcHandler singleton, so when
+  // PythonRpcHandler singleton cleans up py::objects and call dec_ref(), it
+  // will crash.
+  // The solution is to clean up py::objects earlier when Rpc agent join().
+  // Be note that py::objects can not be cleaned up when Rpc agent is destroyed
+  // as well, as Rpc agent is global variable and it will have same issue as
+  // PythonRpcHandler.
+  void cleanup();
+
+  std::shared_ptr<torch::jit::CompilationUnit> jitCompilationUnit();
+
+  // Parse the string to recover the jit_type, this is used for RRef python
+  // pickling/unpickling type recovery. The type string inference rule is as
+  // follows:
+  // 1. first try to parse if this is primitive types.
+  //    i.e. TensorType, IntType, PyObjectType, etc.
+  // 2. if not primitive type, we query the python_cu to see if it is a
+  //    class type or interface type registered in python
+  // We use a ScriptTypeParser instance with custom PythonTypeResolver
+  // to resolve types according to the above rules.
+  TypePtr parseTypeFromStr(const std::string& typeStr);
+
+  // Return a set of Python functions for RRef helpers.
+  const RRefProxyFunctions& getRRefProxyFunctions() const;
+
+  // Return a set of Python functions to retrieve the type of the object
+  // referenced by a given RRef.
+  const RRefTypeFunctions& getRRefTypeFunctions() const;
+
+  PythonRpcHandler(const PythonRpcHandler&) = delete;
+  PythonRpcHandler& operator=(const PythonRpcHandler&) = delete;
+  PythonRpcHandler(PythonRpcHandler&&) = delete;
+  PythonRpcHandler& operator=(PythonRpcHandler&&) = delete;
+
+ private:
+  void init();
+  PythonRpcHandler();
+  ~PythonRpcHandler() = default;
+
+  // Ref to `torch.distributed.rpc.internal._run_function`.
+  py::object pyRunFunction_;
+
+  // Ref to `torch.distributed.rpc.internal.serialize`.
+  py::object pySerialize_;
+
+  // Ref to `torch.distributed.rpc.internal.deserialize`.
+  py::object pyDeserialize_;
+
+  // Ref to 'torch.distributed.rpc.internal._handle_exception'
+  py::object pyHandleException_;
+
+  // Python functions for RRef proxy
+  RRefProxyFunctions rrefProxyFunctions_;
+
+  // Ref to 'torch.distributed.rpc.api._rref_typeof_on_'
+  RRefTypeFunctions rrefTypeFunctions_;
+
+  // Shared ptr to python compilation unit in jit, it is constructed in python
+  // side (see _python_cu = torch._C.CompilationUnit() in jit/__init__.py)
+  // and imported in C++ (see get_python_cu() in
+  // csrc/jit/python/pybind_utils.h). We import the compilation unit here only
+  // once for less cost and thread safety.
+  std::shared_ptr<torch::jit::CompilationUnit> jitCompilationUnit_;
+
+  // jit type parser to parse type_str back to TypePtr for RRef type
+  // recovery when pickling and unpickling RRef
+  std::shared_ptr<jit::ScriptTypeParser> typeParser_;
+
+  // Indicates whether or not we have properly initialized the handler.
+  bool initialized_{false};
+
+  // Lock to protect initialization.
+  std::mutex init_lock_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback.h
new file mode 100644
index 0000000000000000000000000000000000000000..f0b0a52371b27e14b530df19a47491337c856d63
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+
+namespace torch::distributed::rpc {
+
+// Functor which is invoked to process an RPC message. This is an abstract class
+// with some common functionality across all request handlers. Users need to
+// implement this interface to perform the actual business logic.
+class TORCH_API RequestCallback {
+ public:
+  // Invoke the callback.
+  c10::intrusive_ptr<JitFuture> operator()(
+      Message& request,
+      std::vector<c10::Stream> streams) const;
+
+  virtual ~RequestCallback() = default;
+
+ protected:
+  // RpcAgent implementation should invoke ``RequestCallback`` to process
+  // received requests. There is no restriction on the implementation's
+  // threading model. This function takes an rvalue reference of the Message
+  // object. It is expected to return the future to a response message or
+  // message containing an exception. Different rpc agent implementations are
+  // expected to ensure delivery of the response/exception based on their
+  // implementation specific mechanisms.
+  virtual c10::intrusive_ptr<JitFuture> processMessage(
+      Message& request,
+      std::vector<c10::Stream> streams) const = 0;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..7bea64f4b19c05dec4201eead386978cf6a32c64
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback_impl.h
@@ -0,0 +1,66 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/request_callback_no_python.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/jit/python/pybind.h>
+
+namespace torch::distributed::rpc {
+
+class TORCH_API RequestCallbackImpl : public RequestCallbackNoPython {
+ public:
+  std::unique_ptr<RpcCommandBase> deserializePythonRpcCommand(
+      std::unique_ptr<RpcCommandBase> rpc,
+      const MessageType& messageType) const override;
+
+  c10::intrusive_ptr<JitFuture> processPythonCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const override;
+
+  c10::intrusive_ptr<JitFuture> processScriptCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const override;
+
+  c10::intrusive_ptr<JitFuture> processScriptRemoteCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const override;
+
+  c10::intrusive_ptr<JitFuture> processPythonRemoteCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const override;
+
+  c10::intrusive_ptr<JitFuture> processPythonRRefFetchCall(
+      RpcCommandBase& rpc) const override;
+
+  void handleRRefDelete(c10::intrusive_ptr<RRef>& rref) const override;
+
+  c10::intrusive_ptr<JitFuture> processRpcWithErrors(
+      RpcCommandBase& rpc,
+      const MessageType& messageType,
+      const std::vector<c10::Stream>& streams) const override;
+
+  bool cudaAvailable() const override;
+
+  c10::intrusive_ptr<JitFuture> processRRefBackward(
+      RpcCommandBase& rpc) const override;
+
+  // Helpers to run user-defined functions, operators and other computations.
+
+  c10::intrusive_ptr<JitFuture> runJitFunction(
+      const c10::QualifiedName& name,
+      std::vector<at::IValue>& stack,
+      const std::vector<c10::Stream>& streams,
+      bool isAsyncExecution) const;
+
+  c10::intrusive_ptr<JitFuture> runPythonFunction(
+      const py::object& function,
+      const std::vector<c10::Stream>& streams,
+      bool isAsyncExecution) const;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback_no_python.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback_no_python.h
new file mode 100644
index 0000000000000000000000000000000000000000..e8632437b14fefed1c3de5c8ea390059c98c87a1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/request_callback_no_python.h
@@ -0,0 +1,120 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/request_callback.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/rref_impl.h>
+#include <torch/csrc/distributed/rpc/script_call.h>
+#include <torch/csrc/distributed/rpc/script_remote_call.h>
+
+namespace torch::distributed::rpc {
+
+// RequestCallback implementation with no Python dependencies.
+class TORCH_API RequestCallbackNoPython : public RequestCallback {
+ public:
+  c10::intrusive_ptr<JitFuture> processMessage(
+      Message& request,
+      std::vector<c10::Stream> streams) const override;
+
+ protected:
+  virtual std::unique_ptr<RpcCommandBase> deserializePythonRpcCommand(
+      std::unique_ptr<RpcCommandBase> rpc,
+      const MessageType& messageType) const;
+
+  virtual c10::intrusive_ptr<JitFuture> processScriptCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const;
+
+  virtual c10::intrusive_ptr<JitFuture> processPythonCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const;
+
+  c10::intrusive_ptr<JitFuture> assignOwnerRRef(
+      const RRefId& rrefId,
+      const RRefId& forkId,
+      const c10::intrusive_ptr<JitFuture>& valueFuture) const;
+
+  virtual c10::intrusive_ptr<JitFuture> processScriptRemoteCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const;
+
+  virtual c10::intrusive_ptr<JitFuture> processPythonRemoteCall(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const;
+
+  c10::intrusive_ptr<JitFuture> retrieveOwnerRRef(const RRefId& rrefId) const;
+
+  c10::intrusive_ptr<JitFuture> processScriptRRefFetchCall(
+      RpcCommandBase& rpc) const;
+
+  virtual c10::intrusive_ptr<JitFuture> processPythonRRefFetchCall(
+      RpcCommandBase& rpc) const;
+
+  c10::intrusive_ptr<JitFuture> processRRefUserDelete(
+      RpcCommandBase& rpc) const;
+
+  c10::intrusive_ptr<JitFuture> processRRefChildAccept(
+      RpcCommandBase& rpc) const;
+
+  c10::intrusive_ptr<JitFuture> processRRefForkRequest(
+      RpcCommandBase& rpc) const;
+
+  c10::intrusive_ptr<JitFuture> processForwardAutogradReq(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const;
+
+  c10::intrusive_ptr<JitFuture> processBackwardAutogradReq(
+      RpcCommandBase& rpc,
+      const std::vector<c10::Stream>& streams) const;
+
+  c10::intrusive_ptr<JitFuture> processCleanupAutogradContextReq(
+      RpcCommandBase& rpc) const;
+
+  c10::intrusive_ptr<JitFuture> processRunWithProfilingReq(
+      RpcCommandBase& rpc) const;
+
+  virtual void handleRRefDelete(c10::intrusive_ptr<RRef>& rref) const;
+
+  c10::intrusive_ptr<JitFuture> processRpc(
+      RpcCommandBase& rpc,
+      const MessageType& messageType,
+      const std::vector<c10::Stream>& streams) const;
+
+  virtual c10::intrusive_ptr<JitFuture> processRpcWithErrors(
+      RpcCommandBase& rpc,
+      const MessageType& messageType,
+      const std::vector<c10::Stream>& streams) const;
+
+  c10::intrusive_ptr<Message> handleError(
+      const std::exception& e,
+      const MessageType messageType,
+      int64_t messageId) const;
+
+  virtual bool cudaAvailable() const;
+
+  virtual c10::intrusive_ptr<JitFuture> processRRefBackward(
+      RpcCommandBase& rpc) const;
+
+  // Helpers to run user-defined functions, operators and other computations.
+
+  c10::intrusive_ptr<JitFuture> runJitOperator(
+      const jit::Operator& op,
+      std::vector<at::IValue>& stack,
+      const std::vector<c10::Stream>& streams) const;
+
+  // Helpers to convert various kinds of objects into already-completed futures.
+
+  c10::intrusive_ptr<JitFuture> asFuture(IValue value, TypePtr type) const;
+
+  c10::intrusive_ptr<JitFuture> asFuture(
+      c10::intrusive_ptr<Message> message) const;
+
+  c10::intrusive_ptr<JitFuture> asFuture(std::exception_ptr err) const;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc.h
new file mode 100644
index 0000000000000000000000000000000000000000..9e0ced778117bcf6dc19a32b2a7fc90749172e88
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::distributed::rpc {
+
+PyMethodDef* python_functions();
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc_agent.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc_agent.h
new file mode 100644
index 0000000000000000000000000000000000000000..7ce5590994cca2b262f3015e183d36c616f69137
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc_agent.h
@@ -0,0 +1,345 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/request_callback.h>
+#include <torch/csrc/distributed/rpc/types.h>
+
+#include <cctype>
+#include <chrono>
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace torch::distributed::rpc {
+
+using DeviceMap = std::unordered_map<c10::Device, c10::Device>;
+
+// Default RPC timeout
+constexpr float kDefaultRpcTimeoutSeconds = 60;
+// Unset RPC timeout. This is the value agent::send() will have if user does not
+// pass in a specific timeout, and indicates that we must use the default
+// timeout for RPCs.
+constexpr float kUnsetRpcTimeout = -1;
+constexpr auto kDefaultInitMethod = "env://";
+constexpr float kSecToMsConversion = 1000;
+constexpr auto kRpcTimeoutErrorStr =
+    "RPC ran for more than set timeout ({} ms) and will now be marked with an error";
+constexpr auto kDefaultNumWorkerThreads = 16;
+
+using steady_clock_time_point =
+    std::chrono::time_point<std::chrono::steady_clock>;
+// Input is qualified name string, output is JIT StrongTypePtr
+// Same as jit::TypeResolver, did not import jit::TypeResolver to here
+// because it could introduce cyclic dependencies.
+using TypeResolver =
+    std::function<c10::StrongTypePtr(const c10::QualifiedName&)>;
+
+struct TORCH_API RpcBackendOptions {
+  RpcBackendOptions()
+      : RpcBackendOptions(kDefaultRpcTimeoutSeconds, kDefaultInitMethod) {}
+
+  RpcBackendOptions(float rpcTimeoutSeconds, std::string initMethod)
+      : rpcTimeoutSeconds(rpcTimeoutSeconds),
+        initMethod(std::move(initMethod)) {
+    TORCH_CHECK(rpcTimeoutSeconds >= 0, "RPC Timeout must be non-negative");
+  }
+
+  float rpcTimeoutSeconds;
+  std::string initMethod;
+};
+
+// A globally unique ID to identify an RpcAgent
+struct TORCH_API WorkerInfo : torch::CustomClassHolder {
+  WorkerInfo(std::string name, int64_t id);
+
+  WorkerInfo(std::string name, worker_id_t id);
+
+  bool operator==(const WorkerInfo& rhs) {
+    return (id_ == rhs.id_) && (name_ == rhs.name_);
+  }
+
+  static constexpr size_t MAX_NAME_LEN = 128;
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::string name_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const worker_id_t id_;
+};
+
+struct TORCH_API RegisterWorkerInfoOnce {
+  RegisterWorkerInfoOnce();
+};
+
+TORCH_API std::ostream& operator<<(
+    std::ostream& os,
+    const WorkerInfo& workerInfo);
+
+// Struct for options to configure the RPC Retry protocol.
+struct TORCH_API RpcRetryOptions {
+  // Using a default constructor like all other Options structs in the RPC
+  // codebase. TORCH_CHECKs for input validation are done in the
+  // sendWithRetries function.
+  RpcRetryOptions() = default;
+  // Maximum number of times we will retry the RPC
+  int maxRetries{5};
+  // Initial duration between consecutive RPC send attempts
+  std::chrono::milliseconds rpcRetryDuration{std::chrono::milliseconds(1000)};
+  // Constant for exponential backoff used while calculating future wait
+  // durations
+  float retryBackoff{1.5};
+};
+
+// Struct that stores all the metadata needed to retry a given RPC.
+struct TORCH_API RpcRetryInfo {
+  RpcRetryInfo(
+      const WorkerInfo& to,
+      c10::intrusive_ptr<Message> message,
+      c10::intrusive_ptr<JitFuture> originalFuture,
+      int retryCount,
+      RpcRetryOptions options)
+      : to_(to),
+        message_(std::move(message)),
+        originalFuture_(std::move(originalFuture)),
+        retryCount_(retryCount),
+        options_(options) {}
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const WorkerInfo& to_;
+  c10::intrusive_ptr<Message> message_;
+  // Future that is returned to the caller of sendWithRetries().
+  c10::intrusive_ptr<JitFuture> originalFuture_;
+  // Number of send attempts completed so far.
+  int retryCount_;
+  RpcRetryOptions options_;
+};
+
+// ``RpcAgent`` is the base class for sending and receiving RPC messages. It
+// provides a unified ``send`` API for both request and response messages, and
+// will invoke the given ``RequestCallback`` to process received requests. It
+// should immediately become ready to serve request and accept response after
+// construction.
+class TORCH_API RpcAgent {
+ public:
+  // `WorkerInfo` is the globally unique identifier for this RpcAgent instance.
+  // It contains a ``name_`` field and an ``id_`` field. ``name_`` is the
+  // globally unique name for this ``RpcAgent``. It is up to the ``RpcAgent``
+  // implementation to determine how to resolve names. ``id_`` is the globally
+  // unique ID for this ``RpcAgent``. This should be determined by the
+  // ``RpcAgent`` implementation.
+  // The ``RequestCallback`` will be invoked to handle received requests. This
+  // ``RpcAgent`` base class makes no assumption on the thread-safeness of the
+  // ``RequestCallback``. ``RpcAgent`` implementations need to make sure that
+  // its threading model conform to ``RequestCallback``'s requirement.
+  // NB: RpcAgent implementations should not start serving requests until
+  // ``start()`` is called, as there could be other contexts that have not been
+  // initialized yet at this time.
+  RpcAgent(
+      WorkerInfo id,
+      std::unique_ptr<RequestCallback> cb,
+      std::chrono::milliseconds rpcTimeout);
+
+  virtual ~RpcAgent();
+
+  // Send a message to the ``RpcAgent`` of id ``to`` and returns a
+  // ``JitFuture`` ptr. The implementation must be asynchronous, i.e., it
+  // cannot block until it receives the response.
+  //
+  // If ``message.isRequest()`` is true, the ``JitFuture`` will be
+  // completed when the response arrives. For other message types, the Future
+  // should be ignored by the caller.
+  virtual c10::intrusive_ptr<JitFuture> send(
+      const WorkerInfo& to,
+      c10::intrusive_ptr<Message> message,
+      const float rpcTimeoutSeconds = kUnsetRpcTimeout,
+      const DeviceMap& deviceMap = {}) = 0;
+
+  // Retries sending the message up to maxRetries times until an ACK is
+  // received. The duration between consecutive sends is increased over
+  // time using an exponential backoff algorithm.
+  //
+  // Sends ``message`` to the ``RpcAgent`` of id ``to`` and returns a
+  // ``JitFuture`` ptr, just like send(). Caller can specify the maximum
+  // number of retries for this RPC (default is 5), initial duration between
+  // sends (default is 1000ms), and backoff constant (default is 1.5) by
+  // passing in the RpcRetryOptions struct. This API might end up
+  // executing a method twice on the remote end (it does not guarantee
+  // exactly-once semantics). Therefore, the user must ensure their requests
+  // are idempotent.
+  c10::intrusive_ptr<JitFuture> sendWithRetries(
+      const WorkerInfo& to,
+      c10::intrusive_ptr<Message> message,
+      RpcRetryOptions retryOptions = RpcRetryOptions());
+
+  // Return a reference to the ``WorkerInfo`` of this RpcAgent.
+  // NB: not using ``std::optional<const std::string&>`` here because we might
+  // need to create a separate RPC API lib and avoid forcing all ``RpcAgent``
+  // implementations to depend on libtorch.
+  const WorkerInfo& getWorkerInfo() const;
+
+  // Return a reference to the ``WorkerInfo`` of the given ``workerName``.
+  virtual const WorkerInfo& getWorkerInfo(
+      const std::string& workerName) const = 0;
+
+  virtual const WorkerInfo& getWorkerInfo(worker_id_t id) const = 0;
+
+  virtual std::vector<WorkerInfo> getWorkerInfos() const = 0;
+
+  // Retrieve the timeout for all RPCs.
+  inline std::chrono::milliseconds getRpcTimeout() const {
+    return rpcTimeout_.load();
+  }
+
+  // Set the timeout for all RPCs
+  inline void setRpcTimeout(const std::chrono::milliseconds& rpcTimeout) {
+    rpcTimeout_.store(rpcTimeout);
+  }
+
+  // Call sync and join all internal threads. This method should be called
+  // before every RPC process exits.
+  virtual void join(bool shutdown = false, float timeout = 0) = 0;
+
+  // Synchronize the this process with other ``RpcAgent`` processes. Block until
+  // all ``RpcAgent``s reach this method and send all pending messages.
+  virtual void sync() = 0;
+
+  // Sets up backend-agnostic state for accepting requests. Currently, this
+  // entails setting rpcAgentRunning_ to true, creating the retry thread, and
+  // calling the backend's startImpl.
+  void start();
+
+  // Derived classes must override this function to start accepting requests.
+  // This is used to initialize any backend-specific state. Users must call
+  // start, not startImpl, to initialize the RPC Agent.
+  virtual void startImpl() = 0;
+
+  // Stop accepting requests and shutdown the RPC framework as soon as possible
+  // by terminating all RPC threads.
+  void shutdown();
+
+  // Derived classes must override this function to start accepting requests.
+  // THis is used to clean up any backend-specific state. Users must call
+  // shutdown, not shutdownImpl, to shutdown the RPC Agent.
+  virtual void shutdownImpl() = 0;
+
+  // Check if current RPC agent is set.
+  static bool isCurrentRpcAgentSet();
+
+  // Retrieve the valid current RPC agent.
+  static std::shared_ptr<RpcAgent> getCurrentRpcAgent();
+
+  // Set the current RPC agent.
+  static void setCurrentRpcAgent(std::shared_ptr<RpcAgent> rpcAgent);
+
+  // Retrieve metrics as KV map
+  virtual std::unordered_map<std::string, std::string> getMetrics() = 0;
+
+  // Retrieve debug info in addition to metrics as KV map
+  virtual std::unordered_map<std::string, std::string> getDebugInfo();
+
+  // Flag to control whether GIL wait times
+  // should be profiled or not.
+  void enableGILProfiling(bool flag);
+
+  // Retrieve whether we should profile GIL wait times or not.
+  bool isGILProfilingEnabled();
+
+  // Set type resolver that will be passed to JIT pickler to resolver type Ptr
+  // based on type str.
+  void setTypeResolver(std::shared_ptr<TypeResolver> typeResolver);
+
+  // Get the type resolver
+  std::shared_ptr<TypeResolver> getTypeResolver();
+
+  // Retrieves the device map for the provided destination worker.
+  virtual DeviceMap getDeviceMap(const WorkerInfo& dst) const;
+
+  // Retrieve the (non-CPU) devices that are supported by the agent.
+  virtual const std::vector<c10::Device>& getDevices() const;
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const WorkerInfo workerInfo_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::unique_ptr<RequestCallback> cb_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::atomic<std::chrono::milliseconds> rpcTimeout_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::atomic<bool> profilingEnabled_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::shared_ptr<TypeResolver> typeResolver_;
+  // Atomic boolean indicating whether this agent is running. It controls
+  // whether several background threads should be running. It is set in
+  // RpcAgent::start() and unset in the derived class shutdown().
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::atomic<bool> rpcAgentRunning_;
+
+ private:
+  static std::shared_ptr<RpcAgent> currentRpcAgent_;
+  // Add GIL wait time data point to metrics
+  virtual void addGilWaitTime(const std::chrono::microseconds gilWaitTime) = 0;
+  friend class PythonRpcHandler;
+
+  // Map that stores metadata for RPC's that may need to be re-tried as well as
+  // the timepoint at which we should re-try them.
+  std::map<
+      steady_clock_time_point,
+      std::unordered_set<std::shared_ptr<RpcRetryInfo>>>
+      rpcRetryMap_;
+
+  // Thread that checks for retryable RPC's in the rpcRetryMap_ and sleeps until
+  // the next unACKed RPC's timeout has expired.
+  std::thread rpcRetryThread_;
+
+  // Function that rpcRetryThread_ calls in a loop as long as RpcAgent is
+  // running.
+  void retryExpiredRpcs();
+
+  // This is the callback attached to futures corresponding to send retries.
+  // This handles 3 cases: 1). send was completed, 2). send failed with an
+  // error and we've done maxRetries failed send attempts, and 3). send
+  // failed with an error and we have more retries to go. In case 1, we mark
+  // the original future as complete. In case 2, we mark the future with an
+  // error and do not retry again. In case 3, we move the RpcRetryInfo struct
+  // to another time point in the map to schedule the RPC for a future send.
+  void rpcRetryCallback(
+      JitFuture& message,
+      steady_clock_time_point newTime,
+      std::shared_ptr<RpcRetryInfo> earliestRpc);
+
+  // Function that uses the exponential backoff algorithm to compute the next
+  // time point to retry a given RPC.
+  inline steady_clock_time_point computeNewRpcRetryTime(
+      RpcRetryOptions& options,
+      int retryCount) {
+    // The exponential backoff algorithm being used here is:
+    // newTime = timeNow + (retryDuration * (backoffConstant ^ retryCount)).
+    std::chrono::milliseconds timedelta =
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            options.rpcRetryDuration * pow(options.retryBackoff, retryCount));
+    return std::chrono::time_point_cast<std::chrono::milliseconds>(
+        std::chrono::steady_clock::now() + timedelta);
+  }
+
+  // Condition Variable to signal when the rpcRetryMap_ has been populated.
+  std::condition_variable rpcRetryMapCV_;
+
+  // Mutex to protect RpcRetryMap_.
+  std::mutex rpcRetryMutex_;
+};
+
+} // namespace torch::distributed::rpc
+
+namespace std {
+template <>
+struct hash<torch::distributed::rpc::WorkerInfo> {
+  std::size_t operator()(
+      const torch::distributed::rpc::WorkerInfo& worker_info) const noexcept {
+    return worker_info.id_;
+  }
+};
+} // namespace std
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc_command_base.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc_command_base.h
new file mode 100644
index 0000000000000000000000000000000000000000..2ea338813b121c913e1cd0d78fdf5dfe22c03bb1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rpc_command_base.h
@@ -0,0 +1,28 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/types.h>
+
+namespace torch::distributed::rpc {
+
+// Base class for all RPC request and responses.
+class RpcCommandBase {
+ public:
+  // Need to override this to serialize the RPC. This should destructively
+  // create a message for the RPC (Hence the &&).
+  c10::intrusive_ptr<Message> toMessage() && {
+    JitRRefPickleGuard jitPickleGuard;
+    return std::move(*this).toMessageImpl();
+  }
+  virtual c10::intrusive_ptr<Message> toMessageImpl() && = 0;
+  virtual ~RpcCommandBase() = 0;
+};
+
+inline RpcCommandBase::~RpcCommandBase() = default;
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_context.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..6f1703a51e6f67be61b2591e3728eda4da5d5566
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_context.h
@@ -0,0 +1,340 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_agent.h>
+#include <torch/csrc/distributed/rpc/rref_impl.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <torch/csrc/distributed/rpc/utils.h>
+
+#include <atomic>
+#include <optional>
+
+namespace torch::distributed::rpc {
+
+namespace callback {
+// It's the callback for RemoteCall.
+void TORCH_API
+confirmPendingUser(const JitFuture& jitFuture, const ForkId& expectedForkId);
+
+// It's the callback for finishing creating owner rref, it returned deletedRRef,
+// so that the deletedRRef can be handled under GIL in python_functions.cpp if
+// deletedRRef contains python object.
+c10::intrusive_ptr<RRef> TORCH_API
+finishCreatingOwnerRRef(const JitFuture& jitFuture, const RRefId& rrefId);
+} // namespace callback
+
+// Manages RRef lifetime and keeps track of RRef forks.
+class TORCH_API RRefContext {
+ public:
+  static RRefContext& getInstance();
+  // NB: This method must be called before destructing RRefContext singleton.
+  // Similar to delForkOfOwner, this method returns a vector of OwnerRRefs that
+  // hold py::object. The call-site is also responsible for resetting those
+  // shared_ptr objects with a GIL. See comments at delForkOfOwner() for more
+  // details.
+  static std::vector<c10::intrusive_ptr<RRef>> destroyInstance(
+      bool ignoreRRefLeak = true);
+
+  static void handleException(const JitFuture& jitFuture);
+
+  // handle exception without throw ::c10::Error again
+  static void handleExceptionSilent(const JitFuture& jitFuture);
+
+  RRefContext(const RRefContext&) = delete;
+  RRefContext(RRefContext&& other) = delete;
+  void operator=(const RRefContext&) = delete;
+  RRefContext& operator=(RRefContext&& other) = delete;
+
+  ~RRefContext();
+
+  // get the worker id of the current worker
+  inline worker_id_t getWorkerId() const {
+    return agent_->getWorkerInfo().id_;
+  }
+
+  // get the worker name of the current worker
+  inline const std::string& getWorkerName() const {
+    return agent_->getWorkerInfo().name_;
+  }
+
+  //  generate a globally unique ID
+  inline GloballyUniqueId genGloballyUniqueId() {
+    return GloballyUniqueId(getWorkerId(), nextLocalId_++);
+  }
+
+  inline const std::shared_ptr<RpcAgent>& agent() const {
+    return agent_;
+  }
+
+  // create a ``UserRRef`` owned by the worker ``ownerId``
+  c10::intrusive_ptr<UserRRef> createUserRRef(
+      worker_id_t ownerId,
+      const TypePtr& type);
+
+  // Convert an RRefForkData into an RRef. This RRef could be user or owner.
+  // This RRef could have already existed before, or could be created in this
+  // method, we pass type here to validate or help the rref creation.
+  c10::intrusive_ptr<RRef> getOrCreateRRef(
+      const RRefForkData& rfd,
+      const TypePtr& type);
+
+  // Get the ``OwnerRRef`` of id ``rrefId``. If it does not exist, create a new
+  // one. This function is called in two places:
+  // 1. when processing ``rpc.remote()``, i.e., ``SCRIPT_REMOTE_CALL``
+  //    ``PYTHON_REMOTE_CALL``.
+  // 2. when unpickling ``OwnerRRef``.
+  // What's common in these two cases are, 1) the RRefId is already generated
+  // 2) the TypePtr is presented. So it can always create the ``OwnerRRef`` if
+  // it is not yet available.
+  c10::intrusive_ptr<OwnerRRef> getOrCreateOwnerRRef(
+      const RRefId& rrefId,
+      const TypePtr& type);
+
+  // Create an empty owner rref of type.
+  // This method is called to first time generate an ``OwnerRRef``, e.g.,
+  // 1) ``rpc.RRef(obj)``
+  // 2) create the ``OwnerRRef`` on `rpc.remote()` caller side.
+  // What's common in these two cases are, 1) the RRefId hasn't been generated
+  // 2) the TypePtr is presented.
+  c10::intrusive_ptr<OwnerRRef> createOwnerRRef(const TypePtr& type);
+
+  // Returns a Future of the OwnerRRef, which will be marked completed when
+  // ``OwnerRRef`` is created. This method is used when the TypePtr is not
+  // available, e.g., when processing to_here(). The forceCreated flag can be
+  // used to ensure that the rref is created on the owner, otherwise throw in
+  // cases where the user of this API expects this to return a completed future.
+  // Note that the return value is a intrusive_ptr to a c10::ivalue::Future that
+  // holds the RRef.
+  c10::intrusive_ptr<JitFuture> getOwnerRRef(
+      const RRefId& rrefId,
+      bool forceCreated = false);
+
+  // Adding the RRefId of an OwnerRRef into the forks_ map. This is useful when
+  // making a remote call to self, which as for now, still goes through serde
+  // and invokes request callback. In this case, the OwnerRRef has already been
+  // created on the send side, and we need to pass it to the receive side,
+  // instead of creating a new OwnerRRef. This is done by adding the OwnerRRef
+  // into owners_. However, that alone is not enough, as it could be deleted
+  // when all UserRRef die, which would then remove the OwnerRRef from owners_
+  // and this could happen before the self remote call finishes. To prevent
+  // that, this API adds the RRefId as a ForkId, which will then delete the
+  // ForkId when the self remote is done.
+  void addSelfAsFork(c10::intrusive_ptr<OwnerRRef>& rref);
+
+  // Register a fork of the ``OwnerRRef``, and inserts a intrusive_ptr of the
+  // ``OwnerRRef`` in a map to keep it alive.
+  void addForkOfOwner(const RRefId& rrefId, const ForkId& forkId);
+  // Performs the same function as addForkOfOwner but ignores duplicate
+  // requests. This idempotent function is used with RREF_FORK_REQUEST calls,
+  // whereas all other message types use the non-idempotent variant.
+  void addForkOfOwnerIfNotPresent(const RRefId& rrefId, const ForkId& forkId);
+  // Delete a fork of the ``OwnerRRef``. NB: this could trigger deletion on the
+  // IValue or py::object. For the later, this method will acquire GIL.
+  // NB: If this fork deletion triggered deleting OwnerRRef, this method will
+  // return a shared_ptr to the OwnerRRef, which is likely to be the last
+  // shared_ptr instance for it. Therefore, deleting this shared_ptr<OwnerRRef>
+  // will also trigger deleting the object it points to. If OwnerRRef holds a
+  // py::object, deleting it require GIL. The call site should guarded it with
+  // a GIL and reset the shared_ptr. The GIL-guarded deletion is intentionally
+  // left out of this function to avoid creating dependency on pybind.
+  c10::intrusive_ptr<RRef> delForkOfOwner(
+      const RRefId& rrefId,
+      const ForkId& forkId);
+
+  // Invoked when pickling an RRef to setup child/fork properly
+  RRefForkData prepareChildFork(const c10::intrusive_ptr<RRef>& rref);
+  // Invoked when unpickling an RRef to send RREF_FORK_REQUEST to owner and
+  // send RREF_CHILD_ACCEPT to the parent.
+  // NB: forkId is necessary here as the rref could be an OwnerRRef
+  void notifyOwnerAndParentOfFork(
+      const ForkId& forkId,
+      worker_id_t parent,
+      const c10::intrusive_ptr<RRef>& rref);
+
+  // When a UserRRef is forked to another worker (user or owner), it is added
+  // into pendingChildren_ to be held alive until it receives RREF_CHILD_ACCEPT
+  // from the child.
+  // NB: This is necessary for both user and owner child. As we do not have FIFO
+  // communication between workers, we need this strategy to make sure that all
+  // previously submitted rpc/remote calls are acked before sending out the
+  // RREF_USER_DELETE message. Otherwise, the OwnerRRef could be deleted too
+  // soon.
+  void addPendingChild(
+      const ForkId& forkId,
+      const c10::intrusive_ptr<RRef>& rref);
+  void delPendingChild(const ForkId& forkId);
+
+  // When a UserRRef is created, it is added into pendingUsers_ to be held alive
+  // until it receives RREF_USER_ACCEPT from the owner.
+  void addPendingUser(
+      const ForkId& forkId,
+      const c10::intrusive_ptr<RRef>& rref);
+  void delPendingUser(const ForkId& forkId);
+  void addConfirmedUser(
+      const ForkId& forkId,
+      const c10::intrusive_ptr<RRef>& rref);
+
+  // Retrieve a pending user given the fork ID. Throws if the user has already
+  // been confirmed (i.e. is no longer in the pendingUsers_ map).
+  c10::intrusive_ptr<RRef> getPendingUser(const ForkId& forkId);
+
+  // Start recording new pending UserRRefs. All pending UserRRefs introduced
+  // after this point will be put into the thread_local userTable_, which will
+  // then be consumed and cleared in waitForThreadLocalPendingRRefs().
+  void recordThreadLocalPendingRRefs();
+  // End recording new pending UserRRefs, and clear the thread_local userTable_.
+  // Returns a Future which will be marked as completed when all pending
+  // UserRRefs in the current userTable_ are confirmed by their owners. The bool
+  // value in the Future is unused.
+  // This method is useful to make sure RRefs in user function arguments are
+  // confirmed before launching user code.
+  // NB: Callers of this method does not need to keep the returned Future alive,
+  // because this Future is already captured in callbacks of the
+  // PendingUserState. If there is no pending UserRRefs, this method returns a
+  // completed future.
+  c10::intrusive_ptr<JitFuture> waitForThreadLocalPendingRRefs();
+  // Only call this function when there are errors during a recording session,
+  // and it is likely that waitForThreadLocalPendingRRefs() cannot be invoked
+  // properly.
+  // TODO: make this a context guard
+  void clearRecordedPendingRRefsOnError();
+
+  void delUser(
+      const worker_id_t owner,
+      const RRefId& rrefId,
+      const ForkId& forkId);
+  void delAllUsersAndUnforkedOwners(std::chrono::milliseconds timeoutMillis);
+
+  std::unordered_map<std::string, std::string> getDebugInfo();
+
+ private:
+  struct PendingUserState {
+    PendingUserState(c10::intrusive_ptr<RRef> rref)
+        : rref_(std::move(rref)),
+          confirmationFuture_(c10::make_intrusive<JitFuture>(BoolType::get())) {
+    }
+
+    inline void confirm() {
+      c10::static_intrusive_pointer_cast<UserRRef>(rref_)->confirm();
+      confirmationFuture_->markCompleted();
+    }
+
+    c10::intrusive_ptr<RRef> rref_;
+    // Use Future.wait() and Future.markCompleted() to block and unblock user
+    // functions. The bool value wrapped by the future_ is not used.
+    c10::intrusive_ptr<JitFuture> confirmationFuture_;
+  };
+
+  RRefContext(std::shared_ptr<RpcAgent> /*agent*/);
+
+  c10::intrusive_ptr<UserRRef> createUserRRef(
+      worker_id_t ownerId,
+      const RRefId& rrefId,
+      const ForkId& forkId,
+      const TypePtr& type);
+
+  void finishForkRequest(const ForkId& forkId, worker_id_t parent);
+
+  // If there is any leak on any RRef, this method will throw an error.
+  void checkRRefLeaks(bool ignoreRRefLeak);
+
+  static std::atomic<local_id_t> nextLocalId_;
+
+  const std::shared_ptr<RpcAgent> agent_;
+  mutable std::mutex mutex_;
+  // Keep OwnerRRefs alive until there is no living UserRRefs.
+  std::unordered_map<RRefId, c10::intrusive_ptr<RRef>, RRefId::Hash> owners_;
+  // A map to track OwnerRRefs that are requested but not yet created. This can
+  // happen if the to_here() message is processed on the owner before the
+  // corresponding creator rpc.remote() message. If this happens, instead of
+  // to_here() RPC thread to block waiting for the OwnerRRef creation, the
+  // RRefContext returns a Future, so that the RPC request processing logic can
+  // attach subsequent code as a callback to that Future.
+  // NB: the OwnerRRefs in this map must be cleared when the corresponding
+  // OwnerRRef is created. Note that the values in this map are intrusive_ptrs
+  // to c10::ivalue::Future that will be marked completed with the owner RRef.
+  std::unordered_map<RRefId, c10::intrusive_ptr<JitFuture>, RRefId::Hash>
+      pendingOwners_;
+  // Tracks known living UserRRefs of an OwnerRRef
+  std::unordered_map<
+      RRefId,
+      std::unordered_set<ForkId, ForkId::Hash>,
+      RRefId::Hash>
+      forks_;
+
+  // This cond var is used by deleteAllUsers(), a event notification is sent if
+  // number of pending UserRRef or UserRRef children is reduced, or
+  // number of owned OwnerRRef is reduced.
+  std::condition_variable deleteAllUsersCV_;
+  // The follow 3 maps keep UserRRefs alive by holding a intrusive_ptr to the
+  // RRef instances. A UserRRef must be added into this map if any of the
+  // following two conditions is true:
+  //
+  // (1) A UserRRef has not been accepted by owner yet.
+  //
+  //     It can be used or shared, but cannot be deleted, and hence kept alive
+  //     in this map. A message of type RREF_USER_ACCEPT will move the
+  //     corresponding RRef from pendingUsers_ map to confirmedUsers_ map.
+  std::unordered_map<ForkId, std::shared_ptr<PendingUserState>, ForkId::Hash>
+      pendingUsers_;
+  //     UserRRefs are added into this map when it is confirmed by the owner.
+  //     When destroying RRefContext this map helps to find local UserRRefs
+  //     and send delete messages if they are still not deleted by Python
+  //     garbage collection.
+  std::unordered_map<ForkId, c10::weak_intrusive_ptr<RRef>, ForkId::Hash>
+      confirmedUsers_;
+
+  // (2) A UserRRef has forked a child UserRRef which has not been accepted by
+  //     the owner yet.
+  //
+  //     In this case, this UserRRef cannot send out RREF_USER_DELETE message,
+  //     as it could potentially trigger the OwnerRRef been deleted before the
+  //     owner learns about the forked child.
+  std::unordered_map<ForkId, c10::intrusive_ptr<RRef>, ForkId::Hash>
+      pendingChildren_;
+
+  // The RRef context performs its operations through async RPC requests, in
+  // order to not block the user code. Therefore the RRef context's state may be
+  // lagging a bit behind what it is intended to be, while it waits for these
+  // requests to complete. To allow syncing when needed, we store the count of
+  // these pending requests, so that users can wait for it to reach zero.
+  std::atomic<int64_t> numPendingFutures_{0};
+
+  std::mutex destroyedMutex_;
+  bool destroyed_{false};
+
+  // Thread local states to keep UserRRefs deserialized from user function
+  // arguments.
+  static thread_local std::vector<std::shared_ptr<PendingUserState>> userTable_;
+  // A flag indicating whether subsequently created UserRRefs should be added to
+  // the thread_local userTable_. The flag is set to true before serializing
+  // RPC arguments and then set to false before running the corresponding
+  // user code. See addPendingUser and delPendingUser for more details.
+  // NB: The reason for having this flag is because addPendingUser are called in
+  // two cases, and we only want to track the 2nd case.
+  // (1) RRef as the return value: when calling rpc.remote, the UserRRef on the
+  //     caller side is added to the context using addPendingUser.
+  // (2) RRef as an argument: When running an RPC using RRefs as arguments, the
+  //     RRef is forwarded to the callee as new UserRRefs (if the callee is not
+  //     the owner). In this case, we block running the user function until all
+  //     UserRRefs are confirmed by the owner.
+  // This contract guarantees that no UserRRefs can be used remotely without
+  // confirmation. Note that, however, the UserRRef created by rpc.remote can
+  // still be passed to local functions as arguments and used there. This is by
+  // design, because this feature is especially useful when, say a master node
+  // creates multiple UserRRefs in a loop and then shares them with other nodes.
+  // Blocking every iteration in the loop until RRefs are confirmed will slow
+  // this down. This nuance on UserRRef can be interpreted as we only make
+  // exceptions for UserRRef creators. And using the UserRRef on its creator
+  // without confirmation is OK, because the creator would either call to_here
+  // or forward the UserRRef, and both would then require confirmations from the
+  // owner.
+  static thread_local bool recording_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..ae01733c5d9dab82ad33d47164120d1037ecb3f8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_impl.h
@@ -0,0 +1,426 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/jit_type.h>
+#include <ATen/core/rref_interface.h>
+#include <c10/core/Event.h>
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_agent.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <optional>
+
+#include <atomic>
+
+namespace torch::distributed::rpc {
+
+class RRef;
+class RRefContext;
+class UserRRef;
+
+constexpr int OWNER_IDX = 0; // index of ownerId in the tuple
+constexpr int RREFID_ON_IDX = 1; // index of RRefId.createdOn_ in the tuple
+constexpr int RREFID_ID_IDX = 2; // index of RRefId.localId_ in the tuple
+constexpr int FORKID_ON_IDX = 3; // index of ForkId.createdOn_ in the tuple
+constexpr int FORKID_ID_IDX = 4; // index of ForkId.localId_ in the tuple
+constexpr int PARENT_IDX = 5; // index of parent in the tuple
+constexpr int TYPE_IDX = 6; // index of parent in the tuple
+
+// NB: if more fields are added, make sure this field is also bumped
+constexpr int RFD_TUPLE_SIZE = 7; // number of RRefForkData fields in py::tuple
+
+// Represents fork of an RRef to be sent over the wire.
+struct TORCH_API RRefForkData {
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const worker_id_t ownerId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const RRefId rrefId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const ForkId forkId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const worker_id_t parent_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::string typeStr_;
+
+  RRefForkData(
+      worker_id_t ownerId,
+      const RRefId& rrefId,
+      const ForkId& forkId,
+      worker_id_t parent,
+      std::string typeStr);
+};
+
+// Note [RRef Protocol]
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~
+//
+// [Background]
+//
+// RRef stands for Remote REFerence. Each RRef is owned by a single worker
+// (i.e., owner) and can be used by multiple users. The owner stores the real
+// data referenced by its RRefs. RRef needs to support fast and scalable RPC.
+// Hence, in the design, we avoid using a single global master to keep RRef
+// states, instead owners will keep track of the global reference counts
+// for its RRefs. Every RRef can be uniquely identified by a global RRefId,
+// which is assigned at the time it is first created either on a user or on the
+// owner.
+//
+// On the owner worker, there is only one OwnerRRef instance, which contains the
+// real data, while on user workers, there can be as many UserRRefs as
+// necessary, and UserRRef does not hold the data. All usage on the OwnerRRef
+// should retrieve the unique OwnerRRef instance using the globally unique
+// RRefId. //A UserRRef will be created when it is used as an argument or return
+// value in dist.rpc or dist.remote call, but RRef forking and reference
+// counting (RC) are completely transparent to applications. Every UserRRef will
+// also have its globally unique ForkId.
+//
+// [Assumptions]
+//
+// 1. Transient Network Failures
+//
+// TODO: current RRef implementation does not tolerate failures
+//
+// The RRef design handles transient network failures by retrying
+// messages. Node crashes or permanent network partition is beyond the scope.
+// When those incidents occur, the application may take down all workers, revert
+// to the previous checkpoint, and resume training.
+//
+// 2. Non-idempotent UDFs
+//
+// We assume UDFs are not idempotent and therefore cannot be retried. However,
+// internal RRef control messages are idempotent and retried upon message
+// failure.
+//
+// TODO: RRef internal messages are not yet idempotent
+//
+// 3. Out of Order Message Delivery
+//
+// We do not assume message delivery order between any pair of nodes, because
+// both sender and receiver are using multiple threads. There is no guarantee on
+// which message will be processed first.
+//
+// [RRef Lifetime]
+//
+// The goal of the protocol is to delete an OwnerRRef at an appropriate time.
+// The right time to delete an OwnerRRef is when there are no living UserRRefs
+// and Python GC also agrees to delete the OwnerRRef instance on the owner. The
+// tricky part is to determine if there are any living UserRRefs.
+//
+// A user can get a UserRRef in three situations:
+//
+// (1). Receiving a UserRRef from the owner.
+// (2). Receiving a UserRRef from another user.
+// (3). Creating a new UserRRef owned by another worker.
+//
+// (1) is the simplest case where the owner initiates the fork, and hence it can
+// easily increment local RC. The only requirement is that any UserRRef must
+// notify the owner before destruction. Hence, we need the first guarantee:
+//
+// G1. The owner will be notified when any UserRRef is deleted.
+//
+// As messages might come delayed or out-of-order, we need more one guarantee to
+// make sure the delete message is not sent out too soon. Let us first introduce
+// a new concept. If A sends an RPC to B that involves an RRef, we call the RRef
+// on A the parent RRef and the RRef on B the child RRef.
+//
+// G2. Parent RRef cannot be deleted until the child RRef is confirmed by the
+//     owner.
+//
+// Under (1), where the caller is UserRRef and callee is OwnerRRef, it simply
+// means that the user will not send out the delete message until all previous
+// messages are ACKed. Note that ACKed does not mean the owner finishes
+// executing the function, instead, it only means the owner has retrieved its
+// local OwnerRRef and about to pass it to the function, which is sufficient to
+// keep the OwnerRRef alive even if the delete message from the user arrives at
+// the owner before the function finishes execution.
+//
+// With (2) and (3), it is possible that the owner only partially knows the RRef
+// fork graph or not even knowing it at all. For example, the RRef could be
+// constructed on a user, and before the owner receives the RPC call, the
+// creator user might have already shared the RRef with other users, and those
+// users could further share the RRef. One invariant is that the fork graph of
+// any RRef is always a tree rooted at the owner, because forking an RRef always
+// creates a new RRef instance, and hence every RRef has a single parent. One
+// nasty detail is that when an RRef is created on a user, technically the owner
+// is not its parent but we still consider it that way and it does not break the
+// argument below.
+//
+// The owner's view on any node (fork) in the tree has three stages:
+//
+//       1) unknown -> 2) known -> 3) deleted.
+//
+// The owner's view on the entire tree keeps changing. The owner deletes its
+// OwnerRRef instance when it thinks there are no living UserRRefs, i.e., when
+// OwnerRRef is deleted, all UserRRefs could be either indeed deleted or
+// unknown. The dangerous case is when some forks are unknown and others are
+// deleted.
+//
+// G2 trivially guarantees that no parent UserRRef Y can be deleted before the
+// owner knows all of Y's children UserRRefs.
+//
+// However, it is possible that the child UserRRef Z may be deleted before the
+// owner knows its parent Y. More specifically, this can happen when all of Z's
+// messages are processed by the owner before all messages from Y, including the
+// delete message. Nevertheless, this does not cause any problem. Because, at
+// least one of Y's ancestor will be alive, and it will prevent the owner from
+// deleting the OwnerRRef. Consider the following example: (NB: this scenario
+// will no longer relevant when we block UDF until all RRefs are confirmed by
+// the owner)
+//
+//     OwnerRRef -> A -> Y -> Z
+//
+// OwnerRRef forks to A, then A forks to Y, and Y forks to Z. Z can be deleted
+// without OwnerRRef knowing Y. However, the OwnerRRef will at least know A, as
+// the owner directly forks the RRef to A. A won't die before the owner knows Y.
+//
+// Things get a little trickier if the RRef is created on a user:
+//
+//  OwnerRRef
+//      ^
+//      |
+//      A -> Y -> Z
+//
+// If Z calls to_here on the UserRRef, the owner at least knows A when Z is
+// deleted, because otherwise to_here wouldn't finish. If Z does not call
+// to_here, it is possible that the owner receives all messages from Z before
+// any message from A and Y. In this case, as the real data of the OwnerRRef has
+// not been created yet, there is nothing to be deleted either. It is the same
+// as Z does not exist at all Hence, it's still OK.
+//
+// See #26759 for more details and discussions.
+//
+// TODO: make RRef an IValue, and edit createStackForSchema accordingly
+// TODO: make RRef system messages idempotent and retry on failures.
+//
+// ``RRef`` is the base type for both ``UserRRef`` and ``OwnerRRef``.
+// Each ``RRef`` has a globally unique ``RRefId``.
+class TORCH_API RRef : public RRefInterface {
+ public:
+  // RRef is made NOT copyable NOT movable to prevent messing up reference
+  // counting.
+  explicit RRef(const RRef& other) = delete;
+  explicit RRef(RRef&& other) = delete;
+  RRef& operator=(RRef&& other) = delete;
+
+  ~RRef() override = default;
+
+  // returns the worker id of the owner
+  inline worker_id_t owner() const override {
+    return ownerId_;
+  }
+
+  // returns the worker name of the owner
+  inline std::string ownerName() const override {
+    return RpcAgent::getCurrentRpcAgent()->getWorkerInfo(ownerId_).name_;
+  }
+
+  // returns the worker info of the owner
+  inline WorkerInfo ownerWorkerInfo() const {
+    return RpcAgent::getCurrentRpcAgent()->getWorkerInfo(ownerId_);
+  }
+
+  // Returns the globally unique RRefId of this RRef
+  inline const RRefId& rrefId() const {
+    return rrefId_;
+  }
+
+  inline bool isPyObj() const {
+    return type_ == PyObjectType::get();
+  }
+  inline const TypePtr type() const override {
+    return type_;
+  }
+
+  // Save the future corresponding to the creation of this RRef on a remote
+  // node. Note that this is only set when processing requests invoked with
+  // rpc.remote. This is only used to get the future corresponding to the rref
+  // for profiling use cases.
+  inline void registerOwnerCreationFuture(c10::intrusive_ptr<JitFuture> fut) {
+    ownerCreationFuture_ = std::move(fut);
+  }
+
+  // Get the future corresponding to the creation of this rref.
+  inline c10::intrusive_ptr<JitFuture> getOwnerCreationFuture() const {
+    return ownerCreationFuture_;
+  }
+
+  // Check if creation of this RRef on owner node has timed out.
+  inline bool getTimedOut() const {
+    return timedOut_.load();
+  }
+
+  // Dispatches an error to the correct handler based on its RPCErrorType.
+  void handleError(RPCErrorType errorType, const JitFuture& JitFuture);
+
+  // Send delete UserRRef request to Owner,
+  // if the request hasn't been sent yet.
+  // There are 2 cases to call it,
+  // 1, Python GC decides end of UserRRef lifetime, calling destructor.
+  // 2, RPC module graceful shutdown calls it on all UserRRefs tracked
+  //    in the RRefContext.
+  virtual void tryDel() {}
+
+ protected:
+  // Indicates that the creation of this RRef on owner node has timed out.
+  inline void setTimedOut() {
+    timedOut_ = true;
+  }
+  friend class RRefContext;
+
+  RRef(worker_id_t ownerId, const RRefId& rrefId, TypePtr type);
+
+  virtual RRefForkData fork() const;
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const worker_id_t ownerId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const RRefId rrefId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::atomic<bool> timedOut_{false};
+
+  // type field to denote the type of the element that the RRef is holding
+  // it could be any TypePtr that JIT support, including PyObjectType
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const TypePtr type_;
+  // Future corresponding to request to create RRef on remote node.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  c10::intrusive_ptr<JitFuture> ownerCreationFuture_;
+};
+
+// ``UserRRef`` represents a user of an RRef. Besides the ``RRefId``, each user
+// also has a globally unique ``ForkId`` to identify this user. ``UserRRef``
+// never owns the real value, the only way to get the value of the ``RRef`` is
+// to call ``to_here()`` and get a copy..
+class TORCH_API UserRRef final : public RRef {
+ public:
+  UserRRef(const UserRRef& other) = delete;
+  UserRRef(UserRRef&& other) = delete;
+  UserRRef& operator=(const UserRRef& other) = delete;
+  UserRRef& operator=(UserRRef&& other) = delete;
+
+  UserRRef(
+      worker_id_t ownerId,
+      const RRefId& rrefId,
+      const ForkId& forkId,
+      TypePtr type);
+
+  inline bool isOwner() const override {
+    return false;
+  }
+
+  inline bool confirmedByOwner() const override {
+    return confirmedByOwner_;
+  }
+
+  // Returns the globally unique ForkId of this RRef
+  const ForkId& forkId() const;
+
+  // Get of copy of the value from the ``OwnerRRef``. If the value is not ready
+  // yet, this call will block.
+  IValue toHere(
+      const float timeoutSeconds =
+          torch::distributed::rpc::kUnsetRpcTimeout) const;
+
+  void tryDel() override;
+
+  // Will be called when refcount reaches 0.
+  // Upon destruction, this ``UserRRef`` will tell the owner to deref.
+  void release_resources() override;
+
+  // Will be called when both refcount and weakcount reach 0. See
+  // https://github.com/pytorch/pytorch/blob/9116f02bebf3a5260feef5732d36c54ecb3b4033/c10/util/intrusive_ptr.h#L204
+  // This is called on destructing the wrapping intrusive_ptr_target instance
+  // and it's data members.
+  ~UserRRef() override;
+
+ private:
+  friend class RRefContext;
+
+  RRefForkData fork() const override;
+  inline void confirm() {
+    confirmedByOwner_ = true;
+  }
+
+  const ForkId forkId_;
+
+  // Indicates if this user has sent delete message to it's owner.
+  // Note, thread safety is needed because delete message could be sent by
+  // either the destructor called by Python garbage collection or RRefContext
+  // proactive cleanup on RPC graceful shutdown.
+  std::mutex deletedOnOwnerMutex_;
+  bool deletedOnOwner_{false};
+  // Indicating whether this UserRRef has been confirmed by its owner.
+  std::atomic<bool> confirmedByOwner_;
+};
+
+// Keep the template only on the derived class because ``RRefContext`` needs to
+// erase the type on ``RRef`` and keep them in one map.
+class TORCH_API OwnerRRef final : public RRef {
+ public:
+  OwnerRRef(const OwnerRRef& other) = delete;
+  OwnerRRef(OwnerRRef&& other) = delete;
+  OwnerRRef& operator=(const OwnerRRef& other) = delete;
+  OwnerRRef& operator=(OwnerRRef&& other) = delete;
+
+  OwnerRRef(
+      worker_id_t ownerId,
+      const RRefId& rrefId,
+      TypePtr type,
+      std::vector<c10::Device> devices);
+
+  OwnerRRef(
+      worker_id_t ownerId,
+      const RRefId& rrefId,
+      TypePtr type,
+      std::optional<IValue> value,
+      std::vector<c10::Device> devices);
+
+  inline bool isOwner() const override {
+    return true;
+  }
+
+  // OwnerRRef is always confirmed, while UserRRef is only confirmed when the
+  // owner knows about it.
+  inline bool confirmedByOwner() const override {
+    return true;
+  }
+
+  // Get a constant reference of the real value. This method will block if the
+  // value is not ready. This method does not need GIL as it does not create
+  // any new py::object. It will throw if there is an error.
+  const IValue& getValue() const;
+
+  // Set the value of this ``OwnerRRef``. This method does not need GIL as it
+  // does not create any new py::object.
+  void setValue(IValue&& value);
+  // Sets the value of this ``OwnerRRef`` to contain an exception.
+  void setError(std::exception_ptr eptr);
+
+  // Has a value or error been set?
+  bool hasValue() const;
+  // Gets a future that is satisfied when the value or error is set.
+  c10::intrusive_ptr<JitFuture> getFuture();
+
+ private:
+  friend class RRefContext;
+
+  c10::intrusive_ptr<JitFuture> future_;
+};
+
+TORCH_API std::ostream& operator<<(std::ostream& os, const RRef& rref);
+
+// Helper function that casts from c10::RRefInterface to OwnerRRef
+inline TORCH_API c10::intrusive_ptr<OwnerRRef> fromRRefInterface(
+    const c10::intrusive_ptr<c10::RRefInterface>& rrefInterface) {
+  return c10::static_intrusive_pointer_cast<OwnerRRef>(rrefInterface);
+}
+
+// Helper function that casts from OwnerRRef to c10::RRefInterface
+inline TORCH_API c10::intrusive_ptr<c10::RRefInterface> fromOwnerRRef(
+    const c10::intrusive_ptr<RRef>& ownerRRef) {
+  return c10::static_intrusive_pointer_cast<c10::RRefInterface>(ownerRRef);
+}
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_proto.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_proto.h
new file mode 100644
index 0000000000000000000000000000000000000000..b96713516e679436aee835898be9d554283e0a3e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/rref_proto.h
@@ -0,0 +1,168 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <torch/csrc/jit/runtime/operator.h>
+#include <vector>
+
+namespace torch::distributed::rpc {
+
+// Temporary solution of RRef operations.
+// TODO: Remove all these messages and use rpc + registered functions instead.
+class TORCH_API RRefMessageBase : public RpcCommandBase {
+ public:
+  RRefMessageBase(const RRefId& rrefId, MessageType type)
+      : rrefId_(rrefId), type_(type) {}
+
+  const RRefId& rrefId();
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines*)
+  const RRefId rrefId_;
+  // NOLINTNEXTLINE(cppcoreguidelines*)
+  const MessageType type_;
+};
+
+class TORCH_API ForkMessageBase : public RRefMessageBase {
+ public:
+  ForkMessageBase(const RRefId& rrefId, const ForkId& forkId, MessageType type)
+      : RRefMessageBase(rrefId, type), forkId_(forkId) {}
+
+  const ForkId& forkId();
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::pair<RRefId, ForkId> fromMessage(
+      const Message& message,
+      MessageType type);
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines*)
+  const ForkId forkId_;
+};
+
+// UserRRef uses this message to fetch the remote RRef value from the owner.
+class TORCH_API ScriptRRefFetchCall final : public RRefMessageBase {
+ public:
+  ScriptRRefFetchCall(worker_id_t fromWorkerId, const RRefId& rrefId)
+      : RRefMessageBase(rrefId, MessageType::SCRIPT_RREF_FETCH_CALL),
+        fromWorkerId_(fromWorkerId) {}
+
+  inline worker_id_t fromWorkerId() const {
+    return fromWorkerId_;
+  }
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<ScriptRRefFetchCall> fromMessage(
+      const Message& message);
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const worker_id_t fromWorkerId_;
+};
+
+class TORCH_API PythonRRefFetchCall final : public RRefMessageBase {
+ public:
+  PythonRRefFetchCall(worker_id_t fromWorkerId, const RRefId& rrefId)
+      : RRefMessageBase(rrefId, MessageType::PYTHON_RREF_FETCH_CALL),
+        fromWorkerId_(fromWorkerId) {}
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<PythonRRefFetchCall> fromMessage(
+      const Message& message);
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const worker_id_t fromWorkerId_;
+};
+
+// OwnerRRef uses this message to send the RRef value to a remote UserRRef
+class TORCH_API RRefFetchRet : public RpcCommandBase {
+ public:
+  RRefFetchRet(std::vector<at::IValue> values, MessageType type)
+      : values_(std::move(values)), type_(type) {}
+
+  const std::vector<at::IValue>& values();
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+
+ private:
+  std::vector<at::IValue> values_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const MessageType type_;
+};
+
+class TORCH_API ScriptRRefFetchRet final : public RRefFetchRet {
+ public:
+  explicit ScriptRRefFetchRet(std::vector<at::IValue> values)
+      : RRefFetchRet(std::move(values), MessageType::SCRIPT_RREF_FETCH_RET) {}
+
+  static std::unique_ptr<ScriptRRefFetchRet> fromMessage(
+      const Message& message);
+};
+
+class TORCH_API PythonRRefFetchRet final : public RRefFetchRet {
+ public:
+  explicit PythonRRefFetchRet(std::vector<at::IValue> values)
+      : RRefFetchRet(std::move(values), MessageType::PYTHON_RREF_FETCH_RET) {}
+
+  static std::unique_ptr<PythonRRefFetchRet> fromMessage(
+      const Message& message);
+};
+
+// UserRRef (regardless it's the creator or not) uses this message to notify
+// OwnerRRef on delete.
+class TORCH_API RRefUserDelete final : public ForkMessageBase {
+ public:
+  RRefUserDelete(const RRefId& rrefId, const ForkId& forkId)
+      : ForkMessageBase(rrefId, forkId, MessageType::RREF_USER_DELETE) {}
+
+  static std::unique_ptr<RRefUserDelete> fromMessage(const Message& message);
+};
+
+class TORCH_API RemoteRet final : public ForkMessageBase {
+ public:
+  RemoteRet(const RRefId& rrefId, const ForkId& forkId)
+      : ForkMessageBase(rrefId, forkId, MessageType::REMOTE_RET) {}
+
+  static std::unique_ptr<RemoteRet> fromMessage(const Message& message);
+};
+
+// A child RRef uses this message to notify its parent that the child has been
+// confirmed by the owner.
+class TORCH_API RRefChildAccept final : public RpcCommandBase {
+ public:
+  explicit RRefChildAccept(const ForkId& forkId) : forkId_(forkId) {}
+
+  const ForkId& forkId() const;
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<RRefChildAccept> fromMessage(const Message& message);
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const ForkId forkId_;
+};
+
+// A child RRef uses this message to send a fork request to the owner.
+class TORCH_API RRefForkRequest final : public ForkMessageBase {
+ public:
+  RRefForkRequest(const RRefId& rrefId, const ForkId& forkId)
+      : ForkMessageBase(rrefId, forkId, MessageType::RREF_FORK_REQUEST) {}
+
+  static std::unique_ptr<RRefForkRequest> fromMessage(const Message& message);
+};
+
+class TORCH_API RRefAck final : public RpcCommandBase {
+ public:
+  RRefAck() = default;
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<RRefAck> fromMessage(const Message& message);
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_call.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_call.h
new file mode 100644
index 0000000000000000000000000000000000000000..b4073693ec762921a3816b558a8c76913b940357
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_call.h
@@ -0,0 +1,72 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/jit/runtime/operator.h>
+#include <optional>
+#include <vector>
+
+namespace torch::distributed::rpc {
+
+using torch::jit::Operator;
+
+// A ScriptCall instance represents an invocation of a builtin operator for a
+// TorchScript function. If it is a builtin operator, it
+// contains a shared ptr to the `Operator` and a list of arguments.
+// If it is a TorchScript function, it contains a non empty qualifiedName string
+// to the TorchScript function schema name and a list of arguments.
+class TORCH_API ScriptCall : public RpcCommandBase {
+ public:
+  // Constructor for builtin operator call.
+  ScriptCall(std::shared_ptr<Operator> op, std::vector<at::IValue>&& stack);
+  // Constructor for TorchScript function call.
+  ScriptCall(
+      const c10::QualifiedName& qualifiedName,
+      std::vector<at::IValue>&& stack,
+      const bool isAsyncExecution = false);
+
+  bool hasOp() const;
+  std::shared_ptr<Operator> op() const;
+  bool hasQualifiedName() const;
+  const c10::QualifiedName& qualifiedName() const;
+  // return the argument stack of this builtin operator
+  const std::vector<at::IValue>& stack() const;
+  std::vector<at::IValue>& stackRef();
+  inline bool isAsyncExecution() const {
+    return isAsyncExecution_;
+  }
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<ScriptCall> fromMessage(const Message& message);
+
+  ~ScriptCall() override = default;
+
+ protected:
+  virtual void toIValues(std::vector<at::IValue>& ivalues) const;
+  static std::unique_ptr<ScriptCall> fromIValues(
+      std::vector<at::IValue>& ivalues);
+
+ private:
+  // Given an operator symbol and a string schema, return the matched operator.
+  static std::shared_ptr<Operator> matchOperator(const std::string& str_schema);
+
+  static const std::string BUILTIN_OP_NAMESPACE_;
+  static const std::string ATEN_PREFIX_;
+
+  // This field has value if this ScriptCall represents invocation of a builtin
+  // operator.
+  std::optional<std::shared_ptr<Operator>> op_;
+  // This field has non empty string if this ScriptCall represents invocation of
+  // an annotated torchscript function defined by users.
+  std::optional<const c10::QualifiedName> qualifiedName_;
+  std::vector<at::IValue> stack_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const bool isAsyncExecution_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_remote_call.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_remote_call.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ae72a328d457a150ae78f30c8c4c1d18c4b2664
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_remote_call.h
@@ -0,0 +1,59 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/script_call.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <torch/csrc/jit/runtime/operator.h>
+#include <vector>
+
+namespace torch::distributed::rpc {
+
+using torch::jit::Operator;
+
+// A ScriptRemoteCall instance represents an invocation of `dist.remote` on a
+// builtin operator. Currently, it does not support using RRef as arguments yet.
+// Besides the operator and a vector of arguments, ScriptRemoteCall also
+// contains the RRefId and the ForkId of the return value RRef.
+class TORCH_API ScriptRemoteCall final : public ScriptCall {
+ public:
+  // Constructor for builtin operator call.
+  ScriptRemoteCall(
+      std::shared_ptr<Operator> op,
+      std::vector<at::IValue>&& stack,
+      const RRefId& retRRefId,
+      const ForkId& retForkId);
+
+  // Constructor for TorchScript function call.
+  ScriptRemoteCall(
+      const c10::QualifiedName& qualifiedName,
+      std::vector<at::IValue>&& stack,
+      const RRefId& retRRefId,
+      const ForkId& retForkId,
+      const bool isAsyncExecution);
+
+  inline const RRefId& retRRefId() const {
+    return retRRefId_;
+  }
+
+  inline const ForkId& retForkId() const {
+    return retForkId_;
+  }
+
+  static std::unique_ptr<ScriptRemoteCall> fromIValues(
+      std::vector<at::IValue>& ivalues);
+
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<ScriptRemoteCall> fromMessage(const Message& message);
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const RRefId retRRefId_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const ForkId retForkId_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_resp.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_resp.h
new file mode 100644
index 0000000000000000000000000000000000000000..e4fb8e7ca92d1389f1501908ad7062b0a223a204
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/script_resp.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+
+namespace torch::distributed::rpc {
+
+// Return value of a builtin operator or a TorchScript function.
+class TORCH_API ScriptResp final : public RpcCommandBase {
+ public:
+  explicit ScriptResp(at::IValue&& values);
+
+  const at::IValue& value();
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  static std::unique_ptr<ScriptResp> fromMessage(const Message& message);
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const at::IValue value_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/tensorpipe_agent.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/tensorpipe_agent.h
new file mode 100644
index 0000000000000000000000000000000000000000..b5f3a788d0cb86657effd94171062c84f0efc0e9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/tensorpipe_agent.h
@@ -0,0 +1,498 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_TENSORPIPE
+
+#include <atomic>
+#include <thread>
+
+#include <c10/core/thread_pool.h>
+#include <torch/csrc/distributed/c10d/PrefixStore.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/rpc/rpc_agent.h>
+#include <utility>
+
+// Forward-declare the TensorPipe classes we need, to avoid including its
+// headers in PyTorch's ones and thus have it become a public dependency.
+
+namespace tensorpipe {
+
+class Context;
+class Error;
+class Listener;
+class Message;
+class Pipe;
+
+namespace transport {
+class Context;
+} // namespace transport
+
+namespace channel {
+class Context;
+} // namespace channel
+
+} // namespace tensorpipe
+
+namespace torch::distributed::rpc {
+
+// These priorities instruct TensorPipe on which transport/channel to pick
+// during handshake. Higher priorities will take precedence over lower ones.
+// The transport with lowest priority will be the one used to bootstrap pipes.
+
+constexpr int64_t kShmTransportPriority = 200;
+constexpr int64_t kIbvTransportPriority = 100;
+// The UV transport just uses TCP and should work everywhere, thus keep it last.
+constexpr int64_t kUvTransportPriority = 0;
+
+constexpr int64_t kCmaChannelPriority = 1200;
+constexpr int64_t kMultiplexedUvChannelPriority = 1100;
+// The basic channel reuses a transport as a channel, and is thus our fallback.
+constexpr int64_t kBasicChannelPriority = 1000;
+
+// CPU channel have higher priority than CUDA channels, since the latter might
+// handle CPU-to-CPU transfers, but will always be less efficient than their
+// CPU-only counterparts.
+constexpr int64_t kCudaIpcChannelPriority = 300;
+constexpr int64_t kCudaGdrChannelPriority = 200;
+constexpr int64_t kCudaXthChannelPriority = 400;
+constexpr int64_t kCudaBasicChannelPriority = 0;
+
+using steady_clock_time_point =
+    std::chrono::time_point<std::chrono::steady_clock>;
+
+struct TORCH_API TransportRegistration {
+  std::shared_ptr<tensorpipe::transport::Context> transport;
+  int64_t priority;
+  std::string address;
+};
+
+TORCH_DECLARE_REGISTRY(TensorPipeTransportRegistry, TransportRegistration);
+
+struct TORCH_API ChannelRegistration {
+  std::shared_ptr<tensorpipe::channel::Context> channel;
+  int64_t priority;
+};
+
+TORCH_DECLARE_REGISTRY(TensorPipeChannelRegistry, ChannelRegistration);
+
+struct TORCH_API TensorPipeRpcBackendOptions : public RpcBackendOptions {
+  TensorPipeRpcBackendOptions(
+      int numWorkerThreads,
+      std::optional<std::vector<std::string>> transports,
+      std::optional<std::vector<std::string>> channels,
+      float rpc_timeout,
+      std::string init_method,
+      std::unordered_map<std::string, DeviceMap> device_maps = {},
+      std::vector<c10::Device> devices = {})
+      : RpcBackendOptions(rpc_timeout, std::move(init_method)),
+        numWorkerThreads(numWorkerThreads),
+        transports(std::move(transports)),
+        channels(std::move(channels)),
+        deviceMaps(std::move(device_maps)),
+        devices(std::move(devices)) {
+    TORCH_CHECK(
+        numWorkerThreads > 0,
+        "num_worker_threads must be positive, got ",
+        numWorkerThreads);
+
+    if (this->transports.has_value()) {
+      for (const std::string& transportName : this->transports.value()) {
+        TORCH_CHECK(
+            TensorPipeTransportRegistry()->Has(transportName),
+            "Unknown transport: ",
+            transportName);
+      }
+    }
+
+    if (this->channels.has_value()) {
+      for (const std::string& channelName : this->channels.value()) {
+        TORCH_CHECK(
+            TensorPipeChannelRegistry()->Has(channelName),
+            "Unknown channel: ",
+            channelName);
+      }
+    }
+  }
+
+  void setDeviceMap(const std::string& workerName, const DeviceMap& deviceMap) {
+    auto iter = deviceMaps.find(workerName);
+    if (iter == deviceMaps.end()) {
+      deviceMaps[workerName] = deviceMap;
+    } else {
+      for (auto& entry : deviceMap) {
+        // c10::Device has no default constructor, hence map[device] doesn't
+        // work In C++-17 we can use insert_or_assign.
+        auto entryIter = iter->second.find(entry.first);
+        if (entryIter == iter->second.end()) {
+          iter->second.emplace(entry.first, entry.second);
+        } else {
+          entryIter->second = entry.second;
+        }
+      }
+    }
+  }
+
+  int numWorkerThreads;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::optional<std::vector<std::string>> transports;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::optional<std::vector<std::string>> channels;
+  std::unordered_map<std::string, DeviceMap> deviceMaps;
+  std::vector<c10::Device> devices;
+};
+
+// Struct to track the network source metrics
+struct TORCH_API NetworkSourceInfo {
+  worker_id_t srcRank;
+  std::vector<uint8_t> srcMachineAddr;
+};
+
+// Struct to track aggregated network metrics
+struct TORCH_API AggregatedNetworkData {
+  uint64_t numCalls{0};
+  uint64_t totalSentBytes{0};
+  uint64_t totalRecvBytes{0};
+  uint64_t totalErrors{0};
+};
+
+// TensorPipeAgent leverages TensorPipe (https://github.com/pytorch/tensorpipe)
+// to transparently move tensors and payloads through the fastest available
+// transport or channel. It acts like a hybrid RPC transport, providing shared
+// memory (linux) and TCP (linux & mac) support. CUDA support is in progress.
+class TORCH_API TensorPipeAgent : public RpcAgent {
+ public:
+  TensorPipeAgent(
+      const c10::intrusive_ptr<::c10d::Store>& store,
+      std::string selfName,
+      worker_id_t selfId,
+      std::optional<int> worldSize,
+      TensorPipeRpcBackendOptions opts,
+      std::unordered_map<std::string, DeviceMap> reverseDeviceMaps,
+      std::vector<c10::Device> devices,
+      std::unique_ptr<RequestCallback> cb);
+
+  TensorPipeAgent(const TensorPipeAgent&) = delete;
+  TensorPipeAgent& operator=(const TensorPipeAgent&) = delete;
+
+  c10::intrusive_ptr<JitFuture> send(
+      const WorkerInfo& to,
+      c10::intrusive_ptr<Message> message,
+      const float rpcTimeoutSeconds = kUnsetRpcTimeout,
+      const DeviceMap& deviceMap = {}) override;
+
+  // join() and sync() would be deprecated -
+  // https://github.com/pytorch/pytorch/issues/27647
+  void join(bool shutdown = false, float timeout = 0) override;
+  void sync() override {}
+  void startImpl() override;
+  void shutdownImpl() override;
+
+  ~TensorPipeAgent() override;
+
+  const WorkerInfo& getWorkerInfo(const std::string& workerName) const override;
+  const WorkerInfo& getWorkerInfo(worker_id_t workerId) const override;
+  std::vector<WorkerInfo> getWorkerInfos() const override;
+  void updateGroupMembership(
+      const WorkerInfo& workerInfo,
+      const std::vector<c10::Device>& devices,
+      const std::unordered_map<std::string, DeviceMap>& reverseDeviceMaps,
+      bool isJoin);
+
+  std::unordered_map<std::string, std::string> getMetrics() override;
+
+  void addGilWaitTime(const std::chrono::microseconds gilWaitTime) override;
+
+  TensorPipeRpcBackendOptions getBackendOptions() const;
+
+  const c10::intrusive_ptr<::c10d::Store> getStore() const;
+
+  DeviceMap getDeviceMap(const WorkerInfo& dest) const override;
+
+  const std::vector<c10::Device>& getDevices() const override;
+
+  using NetworkDataDict =
+      std::unordered_map<std::string, AggregatedNetworkData>;
+
+  // Returns metrics tracked by the NetworkDataDict
+  NetworkDataDict getNetworkData();
+  // Returns NetworkSourceInfo struct
+  NetworkSourceInfo getNetworkSourceInfo();
+
+  static const std::string& guessAddress();
+
+  // For testing purposes.
+  size_t timeoutMapSize();
+  size_t numPendingResponses();
+  size_t messageIdToTimeoutMapSize();
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const bool isStaticGroup_;
+
+ protected:
+  // TensorPipe write function that could be used to write response
+  // messages by server, and write request messages by client. This
+  // is a protected method since it is overwritten by FaultyTensorPipeAgent
+  virtual void pipeWrite(
+      const std::shared_ptr<tensorpipe::Pipe>& /*pipe*/,
+      const c10::intrusive_ptr<Message>& message,
+      std::vector<c10::Device>&& devices,
+      std::vector<c10::Stream> streams,
+      std::function<void(const tensorpipe::Error&)> /*fn*/) noexcept;
+
+ private:
+  // Removes the given messageId with the given expirationTime from the
+  // timeoutMap_.
+  void removeFromTimeoutMap(uint64_t messageId);
+
+  // Populates workerIdToInfo_ and workerNameToInfo_ using addressStore_
+  void prepareNames(bool isStaticGroup);
+
+  // Check the static group attribute with the value set in store
+  void checkAndSetStaticGroup(const c10::intrusive_ptr<::c10d::Store>& store);
+
+  const std::string& findWorkerURL(const WorkerInfo& worker) const;
+
+  // Only use for Dynamic RPC groups, method to have worker leave group
+  void leaveGroup();
+
+  // TensorPipe read function that could be used to read response messages
+  // by client, and read request messages by server.
+  void pipeRead(
+      const std::shared_ptr<tensorpipe::Pipe>& /*pipe*/,
+      std::function<void(
+          const tensorpipe::Error&,
+          c10::intrusive_ptr<Message>,
+          std::vector<c10::Stream>)> /*fn*/) noexcept;
+
+  // Callback of listener accept()
+  void onListenerAccepted(
+      const tensorpipe::Error& error,
+      std::shared_ptr<tensorpipe::Pipe>& pipe);
+
+  // Respond to a call from a peer
+  void respond(std::shared_ptr<tensorpipe::Pipe>& pipe);
+
+  void sendCompletedResponseMessage(
+      std::shared_ptr<tensorpipe::Pipe>& pipe,
+      JitFuture& futureResponseMessage,
+      uint64_t messageId,
+      std::vector<c10::Stream> stream);
+
+  // Collects metrics from successful RPC calls
+  void trackNetworkData(
+      uint64_t requestSize,
+      uint64_t responseSize,
+      const std::string& destWorkerName);
+
+  // Collects metrics from failed RPC calls
+  void trackNetworkError(
+      uint64_t requestSize,
+      const std::string& destWorkerName);
+
+  inline std::vector<c10::Device> getDevicesForRemote(
+      const std::string& remoteName,
+      const Message& message) const;
+
+  // When a request+response completes, we need to mark the future message as
+  // complete. However, if its timeout has already expired, it already has an
+  // error set. There is no atomic "test-and-set" way to mark a future complete
+  // only if it isn't yet. It does exist for errors (setErrorIfNeeded) but, even
+  // then, it ends up printing a log message, which may worry the user. To solve
+  // both issues we use a separate atomic flag to know the status of the future.
+  struct AtomicJitFuture {
+    explicit AtomicJitFuture(const std::vector<c10::Device>& devices) {
+      jitFuture = c10::make_intrusive<at::ivalue::Future>(
+          at::AnyClassType::get(), devices);
+    }
+
+    std::atomic_flag isComplete = ATOMIC_FLAG_INIT;
+    c10::intrusive_ptr<JitFuture> jitFuture;
+  };
+
+  // Maintains state per client pipe to track pending response messages and
+  // error states. pendingResponseMessage_ should be protected by a mutex since
+  // it can be raced with user send() call.
+  // TODO: To achieve better performance we can have a pipe pool per
+  // client that can be configured using RpcBackendOptions.
+  struct ClientPipe {
+    explicit ClientPipe(std::shared_ptr<tensorpipe::Pipe> pipe)
+        : pipe_(std::move(pipe)) {}
+    std::shared_ptr<tensorpipe::Pipe> pipe_;
+    mutable std::mutex mutex_;
+    bool inError_{false};
+    // Map from Message Request ID's to corresponding futures.
+    std::unordered_map<uint64_t, std::shared_ptr<AtomicJitFuture>>
+        pendingResponseMessage_;
+  };
+
+  const c10::intrusive_ptr<::c10d::Store> store_;
+
+  const TensorPipeRpcBackendOptions opts_;
+  // For dynamic RPC, the reverse device maps are updated whenever a new rank
+  // joins or leaves the group
+  std::unordered_map<std::string, DeviceMap> reverseDeviceMaps_;
+  // Local devices used by this agent. If application didn't specify this
+  // field, it will be initialized using corresponding local devices in
+  // opts_.deviceMaps and reverseDeviceMaps_;
+  std::vector<c10::Device> devices_;
+
+  ThreadPool threadPool_;
+  std::shared_ptr<tensorpipe::Context> context_;
+  std::shared_ptr<tensorpipe::Listener> listener_;
+
+  mutable std::mutex connectedPipesMutex_;
+  std::unordered_map<worker_id_t, ClientPipe> connectedPipes_;
+
+  // Maps keyed on name and id for easy WorkerInfo lookup.
+  std::unordered_map<worker_id_t, WorkerInfo> workerIdToInfo_;
+  std::unordered_map<std::string, WorkerInfo> workerNameToInfo_;
+  std::unordered_map<std::string, std::string> workerNameToURL_;
+
+  ::c10d::PrefixStore rankToNameStore_;
+  ::c10d::PrefixStore nameToAddressStore_;
+  // Store keys that will used to count joined processes and active calls during
+  // the shutdown process
+  ::c10d::PrefixStore shutdownStore_;
+  int worldSize_ = 0;
+  std::atomic<uint64_t> nextMessageID_{0};
+
+  // Metadata used for tracking of whether certain RPCs have timed out or not.
+  struct TimeoutMessageMetadata {
+    TimeoutMessageMetadata(
+        uint64_t messageId_,
+        std::shared_ptr<AtomicJitFuture> responseFuture_,
+        std::chrono::milliseconds timeout_)
+        : messageId(messageId_),
+          responseFuture(std::move(responseFuture_)),
+          timeout(timeout_) {}
+    uint64_t messageId;
+    std::shared_ptr<AtomicJitFuture> responseFuture;
+    std::chrono::milliseconds timeout;
+  };
+
+  // Map to store the expiration times for each message.
+  std::map<steady_clock_time_point, std::vector<TimeoutMessageMetadata>>
+      timeoutMap_;
+
+  // Map to store the messageId to expiry time.
+  std::unordered_map<uint64_t, steady_clock_time_point> messageIdToTimeout_;
+
+  // Thread that will poll the timeoutMap_ for timed out messages and mark them
+  // with an error accordingly
+  std::thread timeoutThread_;
+
+  // Function run by the timeoutThread_ to check for timed out RPCs
+  void pollTimeoutRpcs();
+
+  // Mutex to guard the timeoutMap_
+  std::mutex timeoutMapMutex_;
+
+  // Condition Variable to signal population of the timeoutMap_
+  std::condition_variable timeoutThreadCV_;
+
+  // Returns the expiration time for an RPC by adding the current time to the
+  // passed in timeout.
+  inline steady_clock_time_point computeRpcMessageExpiryTime(
+      std::chrono::milliseconds timeout) const {
+    return std::chrono::time_point_cast<std::chrono::milliseconds>(
+        std::chrono::steady_clock::now() + timeout);
+  }
+
+  // Handle error on an outgoing pipe
+  void handleClientError(
+      ClientPipe& clientPipe,
+      const tensorpipe::Error& error);
+
+  // This is a generic struct for capturing Time-Series Metrics. It keeps a
+  // running sum and count of data points (observations), and can return an
+  // average of the data points seen so far. This is currently only used for
+  // tracking the GIL Wait Time in RPC Agents, but can be used for other metrics
+  // as well.
+  struct TimeSeriesMetricsTracker {
+    // Running sum of the data points seen so far
+    uint64_t currentSum_;
+    // Running count of the data points seen so far
+    uint64_t currentCount_;
+
+    explicit TimeSeriesMetricsTracker(
+        uint64_t currentSum = 0,
+        uint64_t currentCount = 0);
+
+    // Adds a data point (which is basically one observation for the metric
+    // being tracked) to the running sum and count.
+    void addData(uint64_t dataPoint);
+    // Returns the average of all the data points seen so far.
+    float computeAverage() const;
+  };
+
+  // Map of Time-Series metrics tracked by the RPC Agent
+  std::unordered_map<std::string, TimeSeriesMetricsTracker> timeSeriesMetrics_;
+  // Mutex to guard timeSeriesMetrics_
+  std::mutex metricsMutex_;
+
+  // Custom lock guard used to check if the RPC group is dynamic and lock the
+  // mutex if so
+  struct GroupMembershipLockGuard {
+    GroupMembershipLockGuard(std::mutex& mutex, bool isStaticGroup)
+        : ref_(mutex), isStaticGroup_(isStaticGroup) {
+      if (isStaticGroup_) {
+        ref_.lock();
+      }
+    }
+
+    ~GroupMembershipLockGuard() {
+      if (isStaticGroup_) {
+        ref_.unlock();
+      }
+    }
+
+    GroupMembershipLockGuard(const GroupMembershipLockGuard&) = delete;
+
+   private:
+    std::mutex& ref_;
+    bool isStaticGroup_;
+  };
+  // Mutex to guard access to group membership data
+  // e.g. updates to (workerIdToInfo_, workerNameToInfo_, workerNameToURL_)
+  mutable std::mutex groupMembershipMutex_;
+
+  // Map to Track Network Data
+  NetworkDataDict networkData_;
+  // Mutex to guard networkData_
+  std::mutex networkDataMutex_;
+
+  // A mutex and a cv to guard access to the call counts and watch for changes.
+  std::mutex callCountMutex_;
+  std::condition_variable callCountCV_;
+  // Running total of un-processed, un-errored RPC calls sent
+  int32_t clientActiveCalls_{0};
+  // Running total of un-processed RPC requests received
+  int32_t serverActiveCalls_{0};
+  // Running total of RPC requests that will be completed asynchronously
+  int32_t serverActiveAsyncCalls_{0};
+
+  // Whether a global graceful shutdown has begun, in which case we'll silence
+  // error messages due to remote workers closing their pipes.
+  std::atomic<bool> shuttingDown_{false};
+
+  // Helpers to modify the counts while correctly dealing with the mutex and cv.
+  void increaseCallCount(int32_t& count);
+  void decreaseCallCount(int32_t& count);
+
+  // Helpers to set the state of the requests.
+  void markFutureAsComplete(
+      std::shared_ptr<AtomicJitFuture> atomicFuture,
+      c10::intrusive_ptr<Message> message,
+      std::vector<c10::Stream> streams);
+  void markFutureWithError(
+      std::shared_ptr<AtomicJitFuture> atomicFuture,
+      std::string errorMsg);
+};
+
+} // namespace torch::distributed::rpc
+
+#endif // USE_TENSORPIPE
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/tensorpipe_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/tensorpipe_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..025e143190c2df7c2898f03187e70064619bbf3c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/tensorpipe_utils.h
@@ -0,0 +1,124 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_TENSORPIPE
+
+#include <torch/csrc/distributed/rpc/utils.h>
+
+namespace tensorpipe {
+class Message;
+class Allocation;
+class Descriptor;
+} // namespace tensorpipe
+
+namespace torch::distributed::rpc {
+
+TORCH_API const c10::Stream& getStreamForDevice(
+    const std::vector<c10::Stream>& streams,
+    const c10::Device& device);
+
+// Inspired by c10/core/impl/DeviceGuardImplInterface.h.
+
+class TensorpipeDeviceTypeConverter {
+ public:
+  // Ideally we'd want this to also return a tensorpipe::Message::Tensor object
+  // but we cannot forward-declare that class (because it's nested), and we
+  // cannot include the TensorPipe headers because it's a private dependency.
+  // Thus we bend over backwards and entrust this method with appending that
+  // object to the `tensors` field of the tensorpipe::Message object we pass.
+  virtual std::optional<std::vector<char>> prepareTensorForSending(
+      const c10::Storage& storage,
+      const std::vector<c10::Stream>& streams,
+      tensorpipe::Message& message) const = 0;
+
+  // Same as above: this method cannot return a tensorpipe::Allocation::Tensor,
+  // thus it appends it to the `tensors` field of the tensorpipe::Allocation.
+  virtual at::DataPtr allocateTensorForReceiving(
+      c10::DeviceIndex deviceIndex,
+      size_t length,
+      const std::vector<c10::Stream>& streams,
+      tensorpipe::Allocation& allocation) const = 0;
+
+  virtual ~TensorpipeDeviceTypeConverter() = default;
+};
+
+extern TORCH_API std::array<
+    std::atomic<const TensorpipeDeviceTypeConverter*>,
+    static_cast<size_t>(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)>
+    device_type_converter_registry;
+
+class TORCH_API TensorpipeDeviceTypeConverterRegistrar {
+ public:
+  TensorpipeDeviceTypeConverterRegistrar(
+      DeviceType /*type*/,
+      const TensorpipeDeviceTypeConverter* /*impl*/);
+};
+
+#define C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(                     \
+    DevType, TensorpipeDeviceTypeConverter)                                \
+  static ::torch::distributed::rpc::TensorpipeDeviceTypeConverterRegistrar \
+      C10_ANONYMOUS_VARIABLE(g_##DeviceType)(                              \
+          ::c10::DeviceType::DevType, new TensorpipeDeviceTypeConverter());
+
+inline const TensorpipeDeviceTypeConverter* getDeviceTypeConverter(
+    DeviceType type) {
+  return device_type_converter_registry[static_cast<size_t>(type)].load();
+}
+
+// A struct that holds pointers that keep alive all the memory that will be
+// accessed by TensorPipe during a write operation.
+struct TensorpipeWriteBuffers {
+  // Allocate on heap so pointers stay valid as we move the holder.
+  std::unique_ptr<MessageType> type;
+  std::unique_ptr<int64_t> id;
+  std::vector<char> payload;
+  std::vector<char> pickle;
+  // This contains the original tensors and the clones of the sparse tensors.
+  std::vector<torch::Tensor> tensors;
+  // This contains the copies of the data of the tensors that didn't own their
+  // memory, e.g., the ones created from torch::from_blob() with no deleter.
+  std::vector<std::vector<char>> copiedTensors;
+};
+
+// A struct that holds pointers that keep alive all the memory that will be
+// accessed by TensorPipe during a read operation.
+struct TensorpipeReadBuffers {
+  // Allocate on heap so pointers stay valid as we move the holder.
+  std::unique_ptr<MessageType> type;
+  std::unique_ptr<int64_t> id;
+  std::vector<char> payload;
+  std::vector<char> pickle;
+  std::vector<c10::DataPtr> tensors;
+};
+
+// Convert an RPC message into a TensorPipe message, plus a holder to all the
+// data that must be kept alive while the write is performed asynchronously.
+TORCH_API std::tuple<tensorpipe::Message, TensorpipeWriteBuffers>
+tensorpipeSerialize(
+    const c10::intrusive_ptr<Message>& rpcMessage,
+    std::vector<c10::Device> devices,
+    const std::vector<c10::Stream>& streams);
+
+// Allocate the buffers that will hold the incoming data. They will be managed
+// by the returned holder, which must be kept alive until the asynchronous read
+// has finished. Pointers to these buffers will be stored in the returned
+// tensorpipe::Allocation struct.
+TORCH_API std::pair<tensorpipe::Allocation, TensorpipeReadBuffers>
+tensorpipeAllocate(
+    const tensorpipe::Descriptor& tpDescriptor,
+    const std::vector<c10::Stream>& streams);
+
+// Convert a TensorPipe message back into an RPC message. This requires the data
+// to be available and can thus only be performed once the asynchronous read has
+// completed. The holder can be destroyed once this function returns.
+TORCH_API c10::intrusive_ptr<Message> tensorpipeDeserialize(
+    const tensorpipe::Descriptor& tpDescriptor,
+    TensorpipeReadBuffers&& holder);
+
+} // namespace torch::distributed::rpc
+
+#endif // USE_TENSORPIPE
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.h
new file mode 100644
index 0000000000000000000000000000000000000000..c0adb349f2095a721970b5bfdd9acd7141abe827
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.h
@@ -0,0 +1,109 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_TENSORPIPE
+
+#include <torch/csrc/distributed/rpc/message.h>
+#include <torch/csrc/distributed/rpc/tensorpipe_agent.h>
+
+namespace torch::distributed::rpc {
+
+struct TORCH_API FaultyTensorPipeRpcBackendOptions
+    : public TensorPipeRpcBackendOptions {
+  FaultyTensorPipeRpcBackendOptions(
+      int num_worker_threads,
+      float rpc_timeout,
+      std::string init_method,
+      std::vector<std::string> messages_to_fail,
+      std::unordered_map<std::string, float> messages_to_delay,
+      int num_fail_sends = 0)
+      : TensorPipeRpcBackendOptions(
+            num_worker_threads,
+            std::optional<std::vector<std::string>>(),
+            std::optional<std::vector<std::string>>(),
+            rpc_timeout,
+            std::move(init_method)),
+        messagesToFail(std::move(messages_to_fail)),
+        messagesToDelay(std::move(messages_to_delay)),
+        numFailSends(num_fail_sends) {
+    TORCH_CHECK(numFailSends >= 0, "numFailSends should be non-negative");
+  }
+
+  std::vector<std::string> messagesToFail;
+  std::unordered_map<std::string, float> messagesToDelay;
+  int numFailSends;
+};
+
+class TORCH_API FaultyTensorPipeAgent : public TensorPipeAgent {
+ public:
+  FaultyTensorPipeAgent(
+      const c10::intrusive_ptr<::c10d::Store>& store,
+      std::string selfName,
+      worker_id_t selfId,
+      int worldSize,
+      FaultyTensorPipeRpcBackendOptions opts,
+      std::unordered_map<std::string, DeviceMap> reverseDeviceMaps,
+      std::vector<c10::Device> devices,
+      std::unique_ptr<RequestCallback> callback);
+
+  // Faulty send function for this class.
+  c10::intrusive_ptr<JitFuture> send(
+      const WorkerInfo& to,
+      c10::intrusive_ptr<Message> message,
+      const float rpcTimeoutSeconds = torch::distributed::rpc::kUnsetRpcTimeout,
+      const DeviceMap& deviceMap = {}) override;
+
+  // Add delay to writes
+  void pipeWrite(
+      const std::shared_ptr<tensorpipe::Pipe>& pipe,
+      const c10::intrusive_ptr<Message>& rpcMessage,
+      std::vector<c10::Device>&& devices,
+      std::vector<c10::Stream> streams,
+      std::function<void(const tensorpipe::Error&)> fn) noexcept override;
+
+ protected:
+  // This function checks the messageTypesToFail_ to determine whether to use
+  // the faulty send or not.
+  bool shouldFailMessage(MessageType type) const;
+
+ private:
+  // This function parses the list of strings passed in by the python tests and
+  // resolves the Message Types that must use the faulty send.
+  std::vector<MessageType> parseMessagesToFailInput(
+      const std::vector<std::string>& messagesToFail) const;
+
+  // Returns amount of time in seconds to delay sending of the given message
+  // type.
+  float getDelayForMessage(MessageType type) const;
+
+  // Parse message types that we should inject arbitrary delays for.
+  std::unordered_map<MessageType, float, std::hash<int>> parseMessagesToDelay(
+      const std::unordered_map<std::string, float>& messageTypesToDelay) const;
+
+  // Number of sends to intentionally fail before allowing one to succeed.
+  const int numFailSends_;
+
+  // Vector of the MessageTypes that we must use the faulty send for. This is
+  // parsed based on a list of strings passed in by the python tests.
+  const std::vector<MessageType> messageTypesToFail_;
+
+  // Mapping of message types to amount we should delay send for in the ::send()
+  // function.
+  std::unordered_map<MessageType, float, std::hash<int>> messageTypesToDelay_;
+
+  // Map to track the number of sends we've failed for each RPC.
+  std::unordered_map<std::string, int> failMessageCountMap_;
+
+  // Mutex to guard failMessageCountMap_
+  std::mutex failMapMutex_;
+
+  MessageType messageStringToType(const std::string& messageString) const;
+};
+
+} // namespace torch::distributed::rpc
+
+#endif // USE_TENSORPIPE
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/testing/testing.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/testing/testing.h
new file mode 100644
index 0000000000000000000000000000000000000000..baf94a5397fe21394cf9ab024800a29cbc4fe9d6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/testing/testing.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::distributed::rpc::testing {
+
+PyMethodDef* python_functions();
+
+} // namespace torch::distributed::rpc::testing
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/torchscript_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/torchscript_functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..37c6975d05559d13de07463d279ea11cab121f79
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/torchscript_functions.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/autograd/profiler.h>
+#include <torch/csrc/distributed/autograd/utils.h>
+#include <torch/csrc/distributed/rpc/rref_context.h>
+#include <torch/csrc/distributed/rpc/script_remote_call.h>
+
+namespace torch::distributed::rpc {
+
+// This function sends an rpc call to run torchscript function, currently the
+// torchscript function could only be a user defined python function with
+// "@torch.jit.script" annotation. The torchscript function could not be
+// a class constructor, class method, instance method or a script module.
+//   dst: destination worker name
+//   qualifiedName: torchscript function qualified name string like
+//                  "moduleName::torchscriptFunctionName", e.g,
+//                  "dist_autograd_test::my_py_add"
+//   stack: a bag of IValue args passed to torchscriptFunctionName
+// It returns c10::intrusive_ptr<ivalue::Future>
+c10::intrusive_ptr<c10::ivalue::Future> TORCH_API rpcTorchscript(
+    const std::string& dstWorkerName,
+    const c10::QualifiedName& qualifiedName,
+    const c10::FunctionSchema& functionSchema,
+    std::vector<c10::IValue> stack,
+    const float rpcTimeoutSeconds = torch::distributed::rpc::kUnsetRpcTimeout,
+    const bool isAsyncExecution = false);
+
+c10::intrusive_ptr<RRef> TORCH_API remoteTorchscript(
+    const std::string& dstWorkerName,
+    const c10::QualifiedName& qualifiedName,
+    const c10::FunctionSchema& functionSchema,
+    std::vector<c10::IValue>& stack,
+    const float rpcTimeoutSeconds = torch::distributed::rpc::kUnsetRpcTimeout,
+    const bool isAsyncExecution = false);
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/types.h
new file mode 100644
index 0000000000000000000000000000000000000000..f8ec54b86e986a64fd2f55d2c35e1d5c594f30f4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/types.h
@@ -0,0 +1,75 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+
+namespace torch::distributed::rpc {
+
+using worker_id_t = int16_t;
+using local_id_t = int64_t;
+
+bool getAllowJitRRefPickle();
+TORCH_API void enableJitRRefPickle();
+TORCH_API void disableJitRRefPickle();
+
+struct TORCH_API JitRRefPickleGuard {
+  JitRRefPickleGuard();
+  JitRRefPickleGuard(JitRRefPickleGuard&& other) = delete;
+  JitRRefPickleGuard(const JitRRefPickleGuard&) = delete;
+  JitRRefPickleGuard& operator=(const JitRRefPickleGuard&) = delete;
+  JitRRefPickleGuard& operator=(JitRRefPickleGuard&&) = delete;
+  ~JitRRefPickleGuard();
+};
+
+struct TORCH_API GloballyUniqueId final {
+  GloballyUniqueId(worker_id_t createdOn, local_id_t localId);
+  GloballyUniqueId(const GloballyUniqueId& other) = default;
+  GloballyUniqueId& operator=(const GloballyUniqueId& other) = delete;
+  GloballyUniqueId(GloballyUniqueId&& other) = default;
+  GloballyUniqueId& operator=(GloballyUniqueId&& other) = delete;
+  ~GloballyUniqueId() = default;
+
+  bool operator==(const GloballyUniqueId& other) const;
+  bool operator!=(const GloballyUniqueId& other) const;
+
+  at::IValue toIValue() const;
+  static GloballyUniqueId fromIValue(const at::IValue& /*ivalue*/);
+
+  struct Hash {
+    size_t operator()(const GloballyUniqueId& key) const {
+      return (uint64_t(key.createdOn_) << kLocalIdBits) | key.localId_;
+    }
+  };
+
+  static constexpr int kLocalIdBits = 48;
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const worker_id_t createdOn_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const local_id_t localId_;
+};
+
+TORCH_API std::ostream& operator<<(
+    std::ostream& os,
+    const GloballyUniqueId& globalId);
+
+using RRefId = GloballyUniqueId;
+using ForkId = GloballyUniqueId;
+using ProfilingId = GloballyUniqueId;
+
+struct TORCH_API SerializedPyObj final {
+  SerializedPyObj(std::string&& payload, std::vector<at::Tensor>&& tensors)
+      : payload_(std::move(payload)), tensors_(std::move(tensors)) {}
+
+  std::vector<at::IValue> toIValues() &&;
+  static SerializedPyObj fromIValues(std::vector<at::IValue> value);
+
+  std::string payload_;
+  std::vector<at::Tensor> tensors_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/unpickled_python_call.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/unpickled_python_call.h
new file mode 100644
index 0000000000000000000000000000000000000000..da76292342019059afd7d1c868c7cac1e375fd88
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/unpickled_python_call.h
@@ -0,0 +1,43 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::distributed::rpc {
+
+// This class converts the content in a PythonCall into py::object. This is a
+// helper class to make sure that all arguments deserialization is done before
+// entering RequestCallbackImpl::processRpc(...), so that the deserialization
+// related logic can be carried out in one spot instead of scattered in multiple
+// places for different message types.
+// NB: The reason for not consolidating class into PythonCall is because
+// PythonCall is a libtorch type which should not depend on Python types.
+class TORCH_API UnpickledPythonCall : public RpcCommandBase {
+ public:
+  UnpickledPythonCall(
+      const SerializedPyObj& serializedPyObj,
+      bool isAsyncExecution);
+  ~UnpickledPythonCall() override;
+
+  // toMessage() method is not implemented, as objects of this class should
+  // never be directly converted into a Message object.
+  c10::intrusive_ptr<Message> toMessageImpl() && override;
+  const py::object& pythonUdf() const;
+
+  inline bool isAsyncExecution() const {
+    return isAsyncExecution_;
+  }
+
+ private:
+  py::object pythonUdf_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const bool isAsyncExecution_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/unpickled_python_remote_call.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/unpickled_python_remote_call.h
new file mode 100644
index 0000000000000000000000000000000000000000..afe8a977a615e59b2f77e180275e9fc0e6adc92b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/unpickled_python_remote_call.h
@@ -0,0 +1,38 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/distributed/rpc/types.h>
+#include <torch/csrc/distributed/rpc/unpickled_python_call.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::distributed::rpc {
+
+// This class converts the content in a PythonRemoteCall into py::object. This
+// is a helper class to make sure that all arguments deserialization is done
+// before entering RequestCallbackImpl::processRpc(...), so that the
+// deserialization related logic can be carried out in one spot instead of
+// scattered in multiple places for different message types.
+// NB: The reason for not consolidating class into PythonRemoteCall is because
+// PythonRemoteCall is a libtorch type which should not depend on Python types.
+class TORCH_API UnpickledPythonRemoteCall final : public UnpickledPythonCall {
+ public:
+  explicit UnpickledPythonRemoteCall(
+      const SerializedPyObj& serializedPyObj,
+      const at::IValue& retRRefId,
+      const at::IValue& retForkId,
+      const bool isAsyncExecution);
+
+  const RRefId& rrefId() const;
+  const ForkId& forkId() const;
+
+ private:
+  RRefId rrefId_;
+  ForkId forkId_;
+};
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..324d76b2e4dc48b6b12b593fa58c198daa723ad9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/distributed/rpc/utils.h
@@ -0,0 +1,91 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/Device.h>
+#include <c10/core/Event.h>
+#include <c10/core/Stream.h>
+#include <torch/csrc/autograd/profiler.h>
+#include <torch/csrc/distributed/rpc/rpc_command_base.h>
+#include <torch/csrc/jit/serialization/pickle.h>
+#include <torch/csrc/utils/byte_order.h>
+
+namespace torch::distributed::rpc {
+
+// Parse error message and return RPCErrorType based on the message.
+TORCH_API RPCErrorType getRPCErrorType(const JitFuture& jitFuture);
+// Create an error string given the error description and error type
+TORCH_API std::string makeRPCError(
+    const std::string& rpcErrorStr,
+    RPCErrorType errorType);
+
+// Given an RPC message received as a request over the wire, deserialize it into
+// the appropriate 'RpcCommandBase' type.
+TORCH_API std::unique_ptr<RpcCommandBase> deserializeRequest(
+    const Message& request);
+
+// Given an RPC message received as a response over the wire, deserialize it
+// into the appropriate 'RpcCommandBase' type, if the response is
+// FORWARD_AUTOGRAD_RESP type, unwrap it, attach recvBackward() functions
+// to received tensors and set the wrappedMsgType to its wrapped message type.
+TORCH_API std::unique_ptr<RpcCommandBase> deserializeResponse(
+    const Message& response,
+    MessageType& wrappedMsgType);
+
+// Given an RPC message received as a response over the wire, deserialize it
+// into the valid IValue if the message is for a script rpc result,
+// otherwise deserialize it into dummy none ivalue that will never be used.
+// In this deserialization, we also attach recv rpc backward functions if
+// needed.
+IValue deserializeResptoIValueInternal(
+    RpcCommandBase& rpc,
+    MessageType messageType);
+TORCH_API IValue deserializeRespToIValue(const Message& message);
+
+// Note: format is subject to change and intended for RPCs.
+// For saving persistently to disk, use torch::save().
+TORCH_API std::string wireSerialize(
+    const std::vector<char>& payload,
+    const std::vector<at::Tensor>& tensors);
+
+TORCH_API std::pair<std::vector<char>, std::vector<at::Tensor>> wireDeserialize(
+    const void* data,
+    size_t data_size);
+
+// We use vector<char> as the type of blobs because it's what rpc::Message uses
+// for its payload, even though it has the disadvantage that it cannot be
+// allocated with uninitialized memory: it is always zeroed out.
+
+// Some Tensors are effectively views of larger Tensors, where only a small
+// subset of the Storage data is referenced. This normally is good and avoids
+// copies when kept locally, but if we naively push the whole Storage over the
+// wire, we'll end up with excess network traffic. This change clones tensors if
+// we'd save at least half the data, and over a minimum hurdle.
+TORCH_API c10::List<at::Tensor> cloneSparseTensors(
+    const std::vector<at::Tensor>& tensors);
+
+// Combines an original payload and wrapped payload into the original payload.
+// Used to generate the overall payload for the wrapped RPC.
+TORCH_API void writeWrappedPayload(
+    std::vector<char>& originalPayload,
+    std::vector<char>& additionalPayload);
+
+// Reads the additional, wrapped payload from a wrapped RPC off of the input
+// payload. After this, payload will contain the payload of the original,
+// un-wrapped RPC.
+TORCH_API std::vector<at::IValue> readWrappedPayload(
+    std::vector<char>& payload,
+    const rpc::Message& message);
+
+// Takes a list of events from autograd profiler and populates them into
+// profiledEvents to be carried over RPC.
+TORCH_API void populateRemoteProfiledEvents(
+    std::vector<torch::autograd::profiler::LegacyEvent>& profiledEvents,
+    const torch::autograd::profiler::ProfilerConfig& profilerConfig,
+    const std::vector<std::vector<torch::autograd::profiler::LegacyEvent>>&
+        eventLists);
+
+} // namespace torch::distributed::rpc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cache_entry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cache_entry.h
new file mode 100644
index 0000000000000000000000000000000000000000..c178ce57cf456cbd3e7bc3913364e37d80b5c592
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cache_entry.h
@@ -0,0 +1,100 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <Python.h>
+
+#ifdef __cplusplus
+
+#include <torch/csrc/dynamo/utils.h>
+#include <torch/csrc/utils/pybind.h>
+#include <list>
+
+extern "C" {
+
+#endif
+
+/*
+Our cache resides on the extra scratch space of the code object. The structure
+of the cache is as follows:
+
+-> ExtraState
+  -> CacheEntry (list)
+    -> guard_manager (a wrapper that contains the actual guard manager at its
+attr named root)
+    -> code
+  -> FrameState
+
+CacheEntry is a linked list node containing the guard_manager for guards
+and the optimized code.
+
+The FrameState is a PyDict that enables sharing between different frames. This
+is used to detect dynamism in automatic dynamic shapes.
+
+These two are encapsulated into a ExtraState.
+*/
+
+typedef struct CacheEntry CacheEntry;
+typedef struct ExtraState ExtraState;
+
+#ifdef __cplusplus
+
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED(
+    "-Wdeprecated-copy-with-user-provided-dtor")
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wdeprecated-copy-dtor")
+// NOLINTNEXTLINE(cppcoreguidelines-special-member-functions)
+typedef struct VISIBILITY_HIDDEN CacheEntry {
+  // check the guards: lambda: <locals of user function>: bool
+  py::object guard_manager;
+  // modified user bytecode (protected by guard_manager's guards)
+  py::object code;
+  // CompileId corresponding to this compilation
+  py::object compile_id;
+  // root guard manager if exists
+  void* root_mgr{nullptr};
+  // diff guard root guard manager if exists
+  void* diff_guard_root_mgr{nullptr};
+  // backend used to create this cache entry
+  py::object backend;
+  // Reference to owning ExtraState
+  ExtraState* _owner{nullptr};
+  // Reference to this CacheEntry's location in owner's linked list
+  std::list<CacheEntry>::iterator _owner_loc;
+  // Reference to string representation of the CompileContext
+  std::string trace_annotation;
+
+  CacheEntry(const py::handle& guarded_code, PyObject* backend);
+  CacheEntry(const CacheEntry&) = default;
+  CacheEntry(CacheEntry&&) = default;
+  CacheEntry& operator=(const CacheEntry&) = default;
+  CacheEntry& operator=(CacheEntry&&) = default;
+  ~CacheEntry();
+
+  // Warning: returns a reference whose lifetime is controlled by C++
+  py::object next();
+
+  void invalidate(py::object deleted_guard_manager);
+  // Called from the python side to update the diff guard root manager
+  void update_diff_guard_root_manager();
+} CacheEntry;
+C10_DIAGNOSTIC_POP()
+C10_DIAGNOSTIC_POP()
+
+#endif
+
+// Returns borrowed reference
+PyCodeObject* CacheEntry_get_code(CacheEntry* e);
+
+// Returns borrowed string representation of CompileContext
+const char* CacheEntry_get_trace_annotation(CacheEntry* e);
+
+// Returns a borrowed reference to CacheEntry as a PyObject
+// Warning: lifetime is controlled by C++
+PyObject* CacheEntry_to_obj(CacheEntry* e);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/compiled_autograd.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/compiled_autograd.h
new file mode 100644
index 0000000000000000000000000000000000000000..5120939dfd51c04c87739e98ae83cd091e5708d4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/compiled_autograd.h
@@ -0,0 +1,1574 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/TensorGeometry.h>
+#include <ATen/core/ivalue.h>
+#include <c10/core/impl/TorchDispatchModeTLS.h>
+#include <c10/util/flat_hash_map.h>
+#include <torch/csrc/autograd/function.h>
+#include <torch/csrc/autograd/input_metadata.h>
+#include <torch/csrc/autograd/saved_variable.h>
+#include <torch/csrc/autograd/variable_info.h>
+#include <torch/csrc/utils/python_stub.h>
+#include <torch/csrc/utils/torch_dispatch_mode.h>
+#include <typeindex>
+#include <vector>
+
+// see [Note: Compiled Autograd]
+
+namespace torch::dynamo::autograd {
+using namespace torch::autograd;
+
+// This is a layer of indirection for calling methods on the Python
+// AutogradCompilerInstance (referred to as the "py_compiler") from
+// libtorch_cpu (where Python is not available).
+// A PyCompilerInterfaceImpl in libtorch_python subclasses it and
+// overrides the methods to do the actual calls back to Python.
+struct TORCH_API PyCompilerInterface {
+  PyCompilerInterface() = default;
+  PyCompilerInterface(const PyCompilerInterface&) = delete;
+  PyCompilerInterface& operator=(const PyCompilerInterface&) = delete;
+  PyCompilerInterface(PyCompilerInterface&&) = delete;
+  PyCompilerInterface& operator=(PyCompilerInterface&&) = delete;
+  virtual ~PyCompilerInterface() = default;
+
+  // Invokes py_compiler.bind_function
+  virtual std::string bind_function(
+      PyObject* py_compiler,
+      const std::string& fn_name,
+      // NOLINTNEXTLINE(performance-unnecessary-value-param)
+      functional_apply_t fn,
+      // NOLINTNEXTLINE(performance-unnecessary-value-param)
+      std::vector<at::TypePtr> packed_args_schema,
+      bool is_custom_function = false,
+      bool is_traceable = true) const {
+    TORCH_INTERNAL_ASSERT(false, "Needs to be overridden");
+  }
+
+  // Invokes py_compiler.method_name(fn_name, inputs, packed_args,
+  // output_metadata)
+  virtual variable_list call_function(
+      PyObject* py_compiler,
+      const char* method_name,
+      const std::string& fn_name,
+      const variable_list& inputs,
+      const ivalue_list& packed_args,
+      const c10::IValue& output_metadata) const {
+    TORCH_INTERNAL_ASSERT(false, "Needs to be overridden");
+  }
+  virtual variable_list call_copy_slices_prologue(
+      PyObject* py_compiler,
+      const variable_list& inputs,
+      const at::TensorGeometry& base,
+      const at::TensorGeometry& view) const {
+    TORCH_INTERNAL_ASSERT(false, "Needs to be overridden");
+  }
+  virtual variable_list call_copy_slices_epilogue(
+      PyObject* py_compiler,
+      const std::vector<bool>& needs_input_grad,
+      const at::Tensor& result,
+      const variable_list& res,
+      const at::Tensor& grad_slice) const {
+    TORCH_INTERNAL_ASSERT(false, "Needs to be overridden");
+  }
+  virtual at::Tensor call_unpack(
+      PyObject* py_compiler,
+      std::optional<size_t> hook_id,
+      size_t hook_input_id) const {
+    TORCH_INTERNAL_ASSERT(false, "Needs to be overridden");
+  }
+  virtual void call_accumulate_grad(
+      PyObject* py_compiler,
+      const at::Tensor& variable,
+      const at::Tensor& grad,
+      bool has_post_hooks) const {
+    TORCH_INTERNAL_ASSERT(false, "Needs to be overridden");
+  }
+};
+
+TORCH_API const std::unique_ptr<PyCompilerInterface>& getPyCompilerInterface();
+struct TORCH_API PyCompilerGuard {
+  explicit PyCompilerGuard(std::unique_ptr<PyCompilerInterface>&& impl);
+  PyCompilerGuard(const PyCompilerGuard&) = delete;
+  PyCompilerGuard& operator=(const PyCompilerGuard&) = delete;
+  PyCompilerGuard(PyCompilerGuard&&) = delete;
+  PyCompilerGuard& operator=(PyCompilerGuard&&) = delete;
+
+  ~PyCompilerGuard();
+};
+
+// including torch/csrc/autograd/engine.h breaks BC by somehow introducing
+// symbol resolution issues. Instead requiring downstream users to include
+// engine.h to access collect_input_metadata, we provide it here (with a
+// different name to avoid ambiguous symbols...)
+TORCH_API std::vector<std::optional<InputMetadata>> get_input_metadata(
+    const edge_list& edges);
+
+struct SizeInput {
+  // Note: int value is still needed when dynamic to pass as an arg
+  enum DynType : uint8_t { STATIC = 0, DYNAMIC = 1 };
+  SizeInput(DynType dt, int64_t v) : dyn_type(dt), value(v) {}
+  DynType dyn_type;
+  int64_t value;
+};
+
+struct CacheKeyBuffer {
+  CacheKeyBuffer(const uint8_t* key, uint16_t len) : data(new uint8_t[len]) {
+    std::memcpy(data.get(), key, len);
+  }
+  const uint8_t* get() const {
+    return data.get();
+  }
+
+ private:
+  // NOLINTNEXTLINE(*c-array*)
+  std::unique_ptr<uint8_t[]> data;
+};
+
+struct CacheKey {
+  // Key to find the next node in the shadow graph.  We use C++ RTTI for the
+  // type of the node (ntype), then a key generated with a visitor pattern.
+  CacheKey(const std::type_index& ntype, const uint8_t* key, uint16_t len)
+      : node_type(ntype), key_size(len), key(key) {}
+
+  bool operator<(const CacheKey& other) const {
+    if (node_type != other.node_type) {
+      return node_type < other.node_type;
+    }
+    if (key_size != other.key_size) {
+      return key_size < other.key_size;
+    }
+    return std::memcmp(key, other.key, key_size) < 0;
+  }
+
+  bool operator==(const CacheKey& other) const {
+    return node_type == other.node_type && key_size == other.key_size &&
+        std::memcmp(key, other.key, key_size) == 0;
+  }
+
+  size_t hash() const noexcept {
+    // don't bother hashing the key data, common case 1 cache entry per node
+    return std::hash<std::type_index>()(node_type) ^ key_size;
+  }
+
+  std::type_index node_type;
+  uint16_t key_size;
+  const uint8_t* key;
+};
+
+struct NodeCall {
+  NodeCall(uint32_t id_, std::shared_ptr<Node> node_)
+      : id(id_), node(std::move(node_)) {}
+
+  void mark_output(int input_nr, int output_idx) {
+    graph_output.emplace_back(input_nr, output_idx);
+  }
+
+  uint32_t id;
+  std::shared_ptr<Node> node;
+  std::vector<std::pair<int, int>> tensor_pre_hooks;
+  std::vector<std::pair<int, int>> cpp_tensor_pre_hooks;
+  std::vector<int> pre_hooks;
+  std::vector<int> post_hooks;
+  std::vector<int> post_acc_grad_hooks;
+  std::vector<std::pair<int, int>> graph_output;
+  bool needed = true;
+};
+
+struct NodeCalls : public std::unordered_map<Node*, NodeCall> {
+  NodeCall& lookup(const std::shared_ptr<Node>& function) {
+    auto it = find(function.get());
+    if (it == end()) {
+      it = emplace(function.get(), NodeCall(_next_id++, function)).first;
+      nodes.emplace_back(function.get());
+    }
+    return it->second;
+  }
+
+  const NodeCall& lookup(uint32_t id) const {
+    TORCH_INTERNAL_ASSERT(id < nodes.size());
+    auto it = find(nodes[id]);
+    TORCH_INTERNAL_ASSERT(it != end());
+    return it->second;
+  }
+
+  void clear() {
+    _next_id = 0;
+    std::unordered_map<Node*, NodeCall>::clear();
+    nodes.clear();
+  }
+
+ private:
+  uint32_t _next_id = 0;
+  std::vector<Node*> nodes;
+};
+
+struct TensorArg {
+  // Represents a de-duplicated tensor that will be passed into the graph
+  TensorArg(uint32_t i = 0) : id(i) {}
+  uint32_t index() const {
+    TORCH_INTERNAL_ASSERT(defined());
+    return id - 1;
+  }
+  bool defined() const {
+    return id != 0;
+  }
+  uint32_t id;
+  at::Tensor proxy_tensor;
+};
+
+struct TensorArgs {
+  // Manages a collection of TensorArgs and mappings from Tensors/SavedVariables
+  // to them.  This also allows us to unpack SavedVariable exactly once and
+  // store the unpacked Tensor.
+  TensorArgs(const std::optional<size_t>& active_node_call_idx)
+      : active_node_call_idx(active_node_call_idx) {}
+
+  TensorArg& lookup(const at::Tensor& tensor, bool create = false) {
+    if (!tensor.defined()) {
+      return _undefined;
+    }
+    auto impl = tensor.unsafeGetTensorImpl();
+    auto it = _args.find(impl);
+    if (it == _args.end()) {
+      TORCH_INTERNAL_ASSERT(create && inputs.size() == _next_id - 1);
+      it = _args.emplace(impl, TensorArg(_next_id++)).first;
+      inputs.emplace_back(tensor);
+      if (active_node_call_idx.has_value()) {
+        input_origins.emplace_back(active_node_call_idx.value());
+      }
+    }
+    return it->second;
+  }
+
+  TensorArg& lookup(const SavedVariable& sv) {
+    if (auto it = _saved_variables.find(&sv); it != _saved_variables.end()) {
+      // unpacked before graph
+      return *it->second;
+    }
+    // unpacked in graph
+    auto it2 = _saved_variables_proxies.find(&sv);
+    TORCH_INTERNAL_ASSERT(it2 != _saved_variables_proxies.end());
+    return *it2->second;
+  }
+
+  TensorArg& add(const at::Tensor& tensor) {
+    return lookup(tensor, true);
+  }
+
+  TensorArg& add(const SavedVariable& sv, const std::shared_ptr<Node>& node) {
+    // no unpack hooks in this codepath
+    at::Tensor tensor = sv.unpack(node);
+    TensorArg& arg = add(tensor);
+    _saved_variables.emplace(&sv, &arg);
+    return arg;
+  }
+
+  // the concrete tensors that will get passed into the graph as inputs
+  std::vector<at::Tensor> inputs;
+  // NodeCall id of each input, only when verbose logging is enabled
+  std::vector<uint32_t> input_origins;
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::optional<size_t>& active_node_call_idx;
+  std::unordered_map<const c10::TensorImpl*, TensorArg> _args;
+  // Every TensorArg from this is actually owned by _args (or _undefined) and
+  // that's why we have an un-owned pointer here.
+  std::unordered_map<const SavedVariable*, TensorArg*> _saved_variables;
+  std::unordered_map<const SavedVariable*, TensorArg*> _saved_variables_proxies;
+  TensorArg _undefined;
+  uint32_t _next_id = 1; // id=0 used by _undefined
+};
+
+struct LiftedIValueArg {
+  LiftedIValueArg() = delete;
+  LiftedIValueArg(const at::IValue* ptr)
+      : actual_ptr(ptr), proxy(at::IValue::uninitialized()) {}
+
+  const at::IValue* actual_ptr; // lifetime handled by autograd node
+  at::IValue proxy;
+};
+
+struct LiftedIValueArgs {
+  LiftedIValueArgs(const std::optional<size_t>& active_node_call_idx)
+      : active_node_call_idx(active_node_call_idx) {}
+
+  at::IValue& next_proxy(const at::IValue* actual_ptr) {
+    TORCH_INTERNAL_ASSERT(next < args.size());
+    auto& iv_arg = args.at(next++);
+    TORCH_INTERNAL_ASSERT(iv_arg.actual_ptr == actual_ptr);
+    return iv_arg.proxy;
+  }
+
+  void add(const at::IValue* iv) {
+    args.emplace_back(iv);
+    if (active_node_call_idx.has_value()) {
+      args_origins.emplace_back(active_node_call_idx.value());
+    }
+  }
+
+  std::vector<LiftedIValueArg> args;
+  size_t next = 0;
+  // NodeCall id of each arg, only when verbose logging is enabled
+  std::vector<uint32_t> args_origins;
+
+ private:
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::optional<size_t>& active_node_call_idx;
+};
+
+struct AutogradCompilerCall {
+  AutogradCompilerCall(SizeInput::DynType default_dyn_type)
+      : active_node_call_idx(std::nullopt),
+        tensor_args(active_node_call_idx),
+        lifted_ivalue_args(active_node_call_idx),
+        default_dyn_type(default_dyn_type) {}
+  void add_size_input(const c10::SymInt& s) {
+    all_size_inputs.emplace_back(
+        default_dyn_type, s.guard_int(__FILE__, __LINE__));
+    if (active_node_call_idx.has_value()) {
+      size_input_origins.emplace_back(active_node_call_idx.value());
+    }
+  }
+
+  size_t emplace_hook(c10::SafePyObject&& fn) {
+    hooks.emplace_back(std::move(fn));
+    return hooks.size() - 1;
+  }
+
+  size_t emplace_cpp_tensor_pre_hook(
+      std::function<at::TensorBase(const at::TensorBase&)>&& fn) {
+    cpp_tensor_pre_hooks.emplace_back(std::move(fn));
+    return cpp_tensor_pre_hooks.size() - 1;
+  }
+
+  size_t emplace_packed_input(c10::SafePyObject&& input) {
+    packed_inputs.emplace_back(std::move(input));
+    return packed_inputs.size() - 1;
+  }
+
+  void set_active_node_call_idx(size_t node_call_idx) {
+    active_node_call_idx = node_call_idx;
+  }
+
+  std::optional<size_t> active_node_call_idx;
+  TensorArgs tensor_args;
+  std::vector<SizeInput> all_size_inputs;
+  LiftedIValueArgs lifted_ivalue_args;
+  std::vector<int64_t> dyn_size_inputs;
+  std::vector<c10::SafePyObject> hooks;
+  std::vector<std::function<at::TensorBase(const at::TensorBase&)>>
+      cpp_tensor_pre_hooks;
+  std::vector<c10::SafePyObject> packed_inputs;
+  NodeCalls node_calls;
+  SizeInput::DynType default_dyn_type;
+  // NodeCall id of each size, only when verbose logging is enabled
+  std::vector<uint32_t> size_input_origins;
+  std::unordered_map<const SavedVariable*, std::pair<size_t, size_t>>
+      sv_to_hooks;
+  // pynode -> backward idx, backward state idx, opaque object indices
+  std::unordered_map<
+      const Node*,
+      std::tuple<size_t, std::optional<size_t>, std::vector<size_t>>>
+      pynode_objs;
+};
+
+class CompiledNodeArgs {
+  // CompiledNodeArgs builds a representation of the constant values found
+  // across all the nodes in the compiled graph, via 'collect' overloads. The
+  // collected constants are specialized on by concatenation into a cache key.
+  // Tensor, symint arguments (which are lifted to become graph inputs rather
+  // than specialized on) are forwarded to the compiler and not included in the
+  // key.
+ public:
+  void collect(const TensorArg& t) {
+    collect_size(t.id);
+    if (t.defined()) {
+      const at::Tensor& tensor = _compiler.tensor_args.inputs[t.index()];
+      // including these in the cache key means dynamo-level tensor guards can
+      // be skipped
+      collect(tensor.device());
+      collect(tensor.dtype());
+      collect(tensor.requires_grad());
+    }
+  }
+
+  void collect(const at::Tensor& t) {
+    collect(_compiler.tensor_args.add(t));
+  }
+  void collect(const SavedVariable& sv, bool is_output) {
+    if (auto hook_data = sv.retrieve_unpack_hook_data();
+        hook_data.has_value()) {
+      // hooks, unpack in graph
+      auto& [hook, packed_input] = hook_data.value();
+      size_t hook_id = _compiler.emplace_hook(std::move(hook));
+      // rely on dynamo to dedup packed tensors from unpacked tensors
+      size_t input_id = _compiler.emplace_packed_input(std::move(packed_input));
+      _compiler.sv_to_hooks.emplace(&sv, std::make_pair(hook_id, input_id));
+    } else {
+      // no hooks, unpack now
+      collect(
+          _compiler.tensor_args.add(sv, is_output ? _node_call.node : nullptr));
+    }
+  }
+  void collect(const c10::SymInt& t) {
+    _compiler.add_size_input(t);
+  }
+  void collect(const std::vector<SavedVariable>& t, bool is_output) {
+    collect_size(t.size());
+    for (const SavedVariable& i : t) {
+      collect(i, is_output);
+    }
+  }
+  template <typename T>
+  void collect(const std::vector<T>& t) {
+    collect_size(t.size());
+    for (const T& i : t) {
+      collect(i);
+    }
+  }
+  void collect(const c10::ArrayRef<SavedVariable>& t, bool is_output) {
+    collect_size(t.size());
+    for (const SavedVariable& i : t) {
+      collect(i, is_output);
+    }
+  }
+  template <typename T>
+  void collect(const c10::ArrayRef<T>& t) {
+    collect_size(t.size());
+    for (const T& i : t) {
+      collect(i);
+    }
+  }
+  template <typename T>
+  void collect(const c10::OptionalArray<T>& t) {
+    collect(t.list);
+  }
+  template <typename T>
+  void collect(const std::optional<T>& t) {
+    if (cond(t.has_value())) {
+      // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
+      collect(*t);
+    }
+  }
+  template <typename A, typename B>
+  void collect(const std::pair<A, B>& t) {
+    collect(t.first);
+    collect(t.second);
+  }
+  template <typename V>
+  void collect(const ska::flat_hash_map<std::string, V>& m) {
+    collect_size(m.size());
+
+    std::vector<std::string> keys;
+    keys.reserve(m.size());
+    std::transform(
+        m.begin(), m.end(), std::back_inserter(keys), [](const auto& entry) {
+          return entry.first;
+        });
+    std::sort(keys.begin(), keys.end());
+    for (const auto& k : keys) {
+      collect(k);
+      collect(m.at(k));
+    }
+  }
+  void collect(const at::IValue& iv, bool nested = false) {
+    // used by AutogradContext::saved_data from CppNode
+    if (iv.isList()) {
+      c10::List<at::IValue> list = iv.toList();
+      collect_size(list.size());
+      for (auto&& value : list) {
+        collect(value, true);
+      }
+    } else if (iv.isGenericDict()) {
+      c10::Dict<at::IValue, at::IValue> ordered_dict = iv.toGenericDict();
+      collect_size(ordered_dict.size());
+      // NOLINTNEXTLINE(modernize-loop-convert)
+      for (auto it = ordered_dict.begin(); it != ordered_dict.end(); it++) {
+        collect(it->key());
+        collect(it->value(), true);
+      }
+    } else if (iv.isTensor()) {
+      collect(iv.toTensor());
+    } else if (
+        !nested &&
+        (iv.isInt() || iv.isSymInt() || iv.isDouble() || iv.isSymFloat())) {
+      // can't lift ivalues nested in collections
+      _compiler.lifted_ivalue_args.add(&iv);
+    } else {
+      try {
+        collect(static_cast<uint64_t>(at::IValue::hash(iv)));
+      } catch (const std::runtime_error& e) {
+        std::string msg =
+            "Compiled autograd can not trace unhashable IValues, error: " +
+            std::string(e.what());
+        TORCH_CHECK_NOT_IMPLEMENTED(false, msg);
+      }
+    }
+  }
+  void collect(const c10::Scalar& t) {
+    auto type = t.type();
+    specialize_on_bytes(type);
+    if (type == c10::ScalarType::Double) {
+      collect(t.toDouble());
+    } else if (type == c10::ScalarType::Long) {
+      collect(t.toLong());
+    } else if (type == c10::ScalarType::Bool) {
+      collect(t.toBool());
+    } else if (type == c10::ScalarType::ComplexDouble) {
+      auto c = t.toComplexDouble();
+      collect(c.real());
+      collect(c.imag());
+    } else {
+      TORCH_INTERNAL_ASSERT(false);
+    }
+  }
+  void collect(const c10::TensorOptions& t) {
+    collect(t.device());
+    collect(t.dtype());
+    collect(t.layout());
+    collect(t.requires_grad());
+    collect(t.pinned_memory());
+    collect(t.memory_format_opt());
+  }
+  void collect(const at::TensorGeometry& t) {
+    collect(t.sym_sizes());
+    collect(t.sym_strides());
+    collect(t.sym_storage_offset());
+  }
+  void collect(const torch::autograd::TypeAndSize& t) {
+    collect(t.sym_sizes);
+    collect(t.options);
+  }
+  void collect(const c10::Device& t) {
+    collect(t.type());
+    collect(t.index());
+  }
+  void collect(const std::string& t) {
+    collect_size(t.size());
+    for (char c : t) {
+      collect(c);
+    }
+  }
+  void collect(const caffe2::TypeMeta& t) {
+    specialize_on_bytes(t.id());
+  }
+  void collect(const std::shared_ptr<Node>& t) {
+    // Note: this is only capturing the ID of the node not everything
+    // contained inside it.  This is used for tracking connections between
+    // nodes and the actual details of the node itself must be handled by
+    // a separate call to `node->compiled_args()`.
+    if (cond((bool)t)) {
+      collect(_compiler.node_calls.lookup(t));
+    }
+  }
+  void collect(const NodeCall& t) {
+    collect_size(t.id);
+    collect(t.graph_output);
+    collect_hooks_from(t.node.get());
+  }
+  void collect(const Edge& t) {
+    if (cond(t.is_valid())) {
+      collect_size(_compiler.node_calls.lookup(t.function).id);
+      collect_size(t.input_nr);
+      collect(t.function->input_metadata(t.input_nr)); // for validate_outputs
+    }
+  }
+  void collect(const InputMetadata& t) {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        !t.is_nested_tensor(), "NestedTensor support not implemented. ");
+    collect(t.options());
+    collect(t.is_tensor_subclass());
+    collect(t.shape_as_dim_vector());
+  }
+  void collect(const VariableInfo& t) {
+    collect(t.layout);
+    collect(t.device);
+    collect(t.scalar_type);
+    collect(t.size);
+    collect(t.requires_grad);
+    collect(t.is_empty);
+  }
+  bool cond(bool cond) {
+    collect(cond);
+    return cond;
+  }
+
+#define COLLECT_AS_BYTES(T) \
+  void collect(T t) {       \
+    specialize_on_bytes(t); \
+  }
+  COLLECT_AS_BYTES(c10::ScalarType)
+  COLLECT_AS_BYTES(c10::DeviceType)
+  COLLECT_AS_BYTES(c10::Layout)
+  COLLECT_AS_BYTES(c10::MemoryFormat)
+  COLLECT_AS_BYTES(int8_t)
+  COLLECT_AS_BYTES(int16_t)
+  COLLECT_AS_BYTES(int32_t)
+  COLLECT_AS_BYTES(int64_t)
+  COLLECT_AS_BYTES(uint8_t)
+  COLLECT_AS_BYTES(uint16_t)
+  COLLECT_AS_BYTES(uint32_t)
+  COLLECT_AS_BYTES(uint64_t)
+  COLLECT_AS_BYTES(bool)
+  COLLECT_AS_BYTES(float)
+  COLLECT_AS_BYTES(double)
+#undef COLLECT_AS_BYTES
+
+  void collect_hooks_from(Node* fn) {
+    for (auto& i : fn->tensor_pre_hooks()) {
+      i->compiled_args(*this);
+    }
+    for (auto& [_, i] : fn->retains_grad_hooks()) {
+      i->compiled_args(*this);
+    }
+    for (auto& i : fn->pre_hooks()) {
+      i->compiled_args(*this);
+    }
+    for (auto& i : fn->post_hooks()) {
+      i->compiled_args(*this);
+    }
+    collect_size(_node_call.tensor_pre_hooks.size());
+    collect_size(_node_call.pre_hooks.size());
+    collect_size(_node_call.post_hooks.size());
+    for (const auto& h : _node_call.tensor_pre_hooks) {
+      collect_size(static_cast<size_t>(h.second));
+    }
+  }
+
+  CacheKey key() const {
+    Node* node = _node_call.node.get();
+    return CacheKey(
+        typeid(*node), _specialization_key, _specialization_key_size);
+  }
+
+  void collect_pynode_objs(
+      const Node* pynode,
+      c10::SafePyObject&& bwd,
+      std::optional<c10::SafePyObject>&& bwd_state,
+      std::vector<c10::SafePyObject>&& opaque_objs) {
+    size_t bwd_idx = _compiler.emplace_hook(std::move(bwd));
+    std::optional<size_t> bwd_state_idx;
+    if (auto state = std::move(bwd_state); state.has_value()) {
+      bwd_state_idx = _compiler.emplace_hook(std::move(state.value()));
+    }
+    std::vector<size_t> opaque_indices(opaque_objs.size());
+    for (size_t i = 0; i < opaque_objs.size(); i += 1) {
+      opaque_indices[i] = _compiler.emplace_hook(std::move(opaque_objs[i]));
+    }
+    _compiler.pynode_objs.emplace(
+        pynode,
+        std::make_tuple(bwd_idx, bwd_state_idx, std::move(opaque_indices)));
+  }
+
+  void add_tensor_pre_hook(c10::SafePyObject&& obj, int index) {
+    auto fn_id = _compiler.emplace_hook(std::move(obj));
+    collect_size(fn_id);
+    _node_call.tensor_pre_hooks.emplace_back(fn_id, index);
+  }
+
+  void add_cpp_single_tensor_pre_hook(
+      const std::function<at::TensorBase(const at::TensorBase&)>& hook,
+      size_t idx) {
+    auto wrapper = [hook](const at::TensorBase& grad) {
+      // handle when hook returns nothing
+      auto out = hook(grad);
+      if (!out.defined()) {
+        return grad;
+      }
+      return out;
+    };
+
+    auto hook_id = _compiler.emplace_cpp_tensor_pre_hook(std::move(wrapper));
+    collect_size(hook_id);
+    _node_call.cpp_tensor_pre_hooks.emplace_back(hook_id, idx);
+  }
+
+  void add_pre_hook(c10::SafePyObject&& obj) {
+    auto fn_id = _compiler.emplace_hook(std::move(obj));
+    collect_size(fn_id);
+    _node_call.pre_hooks.emplace_back(fn_id);
+  }
+
+  void add_post_hook(c10::SafePyObject&& obj) {
+    auto fn_id = _compiler.emplace_hook(std::move(obj));
+    collect_size(fn_id);
+    _node_call.post_hooks.emplace_back(fn_id);
+  }
+
+  void add_post_acc_grad_hook(c10::SafePyObject&& obj) {
+    auto fn_id = _compiler.emplace_hook(std::move(obj));
+    collect_size(fn_id);
+    _node_call.post_acc_grad_hooks.emplace_back(fn_id);
+  }
+
+  // Need to template the size_t to silence internal 32-bit build errors due to
+  // a mix of -Werror, -Wtautological-type-limit-compare and
+  // -Wunknown-pragmas
+  template <typename T>
+  std::enable_if_t<std::is_unsigned_v<T>, void> collect_size(T s) {
+    // we expect sizes to be small, so try to cram them into a single byte
+    constexpr uint8_t encode_as_u64 = std::numeric_limits<uint8_t>::max();
+    constexpr uint8_t encode_as_u32 = encode_as_u64 - 1;
+    constexpr uint8_t encode_as_u16 = encode_as_u64 - 2;
+    if (C10_UNLIKELY(s >= encode_as_u16)) {
+      // first write a byte indicating the path we followed, then the data
+      if (s <= std::numeric_limits<uint16_t>::max()) {
+        // 3 bytes
+        specialize_on_bytes(encode_as_u16);
+        specialize_on_bytes(static_cast<uint16_t>(s));
+      } else if (s <= std::numeric_limits<uint32_t>::max()) {
+        // 5 bytes
+        specialize_on_bytes(encode_as_u32);
+        specialize_on_bytes(static_cast<uint32_t>(s));
+      } else {
+        // 9 bytes
+        specialize_on_bytes(encode_as_u64);
+        specialize_on_bytes(s);
+      }
+    } else {
+      // happy case, 1 byte
+      specialize_on_bytes(static_cast<uint8_t>(s));
+    }
+  }
+
+  SizeInput::DynType set_default_dyn_type(SizeInput::DynType default_dyn_type) {
+    return std::exchange(_compiler.default_dyn_type, default_dyn_type);
+  }
+
+  CompiledNodeArgs(AutogradCompilerCall& compiler, NodeCall& node_call)
+      : _compiler(compiler),
+        _node_call(node_call),
+        _specialization_key(
+            // NOLINTNEXTLINE(cppcoreguidelines-no-malloc)
+            (uint8_t*)std::malloc(_specialization_key_storage)) {}
+  CompiledNodeArgs(const CompiledNodeArgs&) = delete;
+  CompiledNodeArgs(CompiledNodeArgs&&) = delete;
+  CompiledNodeArgs& operator=(const CompiledNodeArgs&) = delete;
+  CompiledNodeArgs& operator=(CompiledNodeArgs&&) = delete;
+  ~CompiledNodeArgs() {
+    // NOLINTNEXTLINE(cppcoreguidelines-no-malloc)
+    std::free(_specialization_key);
+  }
+
+ private:
+  template <typename T>
+  void specialize_on_bytes(const T& t) {
+    while (C10_UNLIKELY(
+        _specialization_key_size + sizeof(T) > _specialization_key_storage)) {
+      _specialization_key_storage *= 2;
+      // NOLINTNEXTLINE(cppcoreguidelines-no-malloc)
+      _specialization_key = (uint8_t*)std::realloc(
+          _specialization_key, _specialization_key_storage);
+    }
+    std::memcpy(_specialization_key + _specialization_key_size, &t, sizeof(T));
+    _specialization_key_size += sizeof(T);
+  }
+
+  AutogradCompilerCall& _compiler;
+  NodeCall& _node_call;
+  size_t _specialization_key_size{0};
+  size_t _specialization_key_storage{1024};
+  uint8_t* _specialization_key;
+};
+
+struct TraceState {
+  TraceState(std::vector<std::optional<c10::SymInt>>&& ss, size_t num_outputs)
+      : sym_sizes(std::move(ss)), outputs(num_outputs) {}
+
+  void debug_asserts() {
+    TORCH_INTERNAL_ASSERT(sym_sizes_index == sym_sizes.size());
+  }
+  std::optional<c10::SymInt> next_sym_size() {
+    TORCH_INTERNAL_ASSERT(sym_sizes_index < sym_sizes.size());
+    return sym_sizes[sym_sizes_index++];
+  }
+
+  size_t sym_sizes_index{0};
+  std::vector<std::optional<c10::SymInt>> sym_sizes;
+  variable_list outputs;
+};
+
+class SwapSavedVariables {
+  // SwapSavedVariables is used during the tracing/compilation phase after a
+  // cache-miss. It swaps any 'lifted' inputs (tensors, symints) to proxy nodes,
+  // allows tracing to happen, then swaps them back afterwards.
+ public:
+  std::tuple<size_t, std::optional<size_t>, std::vector<size_t>>
+  retrieve_pynode_objs(Node* pynode) const {
+    auto it = compiler.pynode_objs.find(pynode);
+    TORCH_INTERNAL_ASSERT(it != compiler.pynode_objs.end());
+    return it->second;
+  }
+
+  void before(at::Tensor& t) {
+    TensorArg& arg = compiler.tensor_args.lookup(t);
+    stashed_tensors.save(&t, std::move(t));
+    if (arg.defined()) {
+      TORCH_INTERNAL_ASSERT(arg.proxy_tensor.defined());
+      t = arg.proxy_tensor;
+    }
+  }
+  void after(at::Tensor& t) {
+    stashed_tensors.restore(&t);
+  }
+
+  void before(SavedVariable& t) {
+    if (auto it = compiler.sv_to_hooks.find(&t);
+        it != compiler.sv_to_hooks.end()) {
+      const auto& pyinterface =
+          torch::dynamo::autograd::getPyCompilerInterface();
+      auto proxy_tensor = pyinterface->call_unpack(
+          get_py_compiler(), it->second.first, it->second.second);
+      stashed_variables.save(&t, std::move(t));
+      bool prior = at::SavedTensorDefaultHooks::set_tracing(true);
+      t = SavedVariable(proxy_tensor, false);
+      at::SavedTensorDefaultHooks::set_tracing(prior);
+    } else {
+      // no hooks, was already unpacked
+      TensorArg& arg = compiler.tensor_args.lookup(t);
+      stashed_variables.save(&t, std::move(t));
+      if (arg.defined()) {
+        bool prior = at::SavedTensorDefaultHooks::set_tracing(true);
+        TORCH_INTERNAL_ASSERT(arg.proxy_tensor.defined());
+        t = SavedVariable(arg.proxy_tensor, false);
+        at::SavedTensorDefaultHooks::set_tracing(prior);
+      }
+    }
+  }
+  void after(SavedVariable& t) {
+    stashed_variables.restore(&t);
+  }
+
+  void before(c10::SymInt& t) {
+    stashed_symints.save(&t, c10::SymInt(t));
+    auto opt_value = state.next_sym_size();
+    if (opt_value.has_value()) {
+      t = *opt_value; // dynamic shape
+    }
+  }
+  void after(c10::SymInt& t) {
+    stashed_symints.restore(&t);
+  }
+
+  void before(at::IValue& iv) {
+    if (iv.isTensor()) {
+      before(iv.toTensor());
+    } else {
+      stashed_ivalues.save(&iv, at::IValue(iv));
+      if (iv.isInt() || iv.isSymInt() || iv.isDouble() || iv.isSymFloat()) {
+        iv = compiler.lifted_ivalue_args.next_proxy(&iv);
+      }
+    }
+  }
+
+  void after(at::IValue& t) {
+    if (t.isTensor()) {
+      after(t.toTensor());
+    } else {
+      stashed_ivalues.restore(&t);
+    }
+  }
+
+  void before(Edge& t) {
+    if (t.is_valid()) {
+      // need for symints used by validate_outputs
+      before(t.function->mutable_input_metadata(t.input_nr));
+    }
+  }
+  void after(Edge& t) {
+    if (t.is_valid()) {
+      after(t.function->mutable_input_metadata(t.input_nr));
+    }
+  }
+  void before(InputMetadata& t) {
+    before(t.mutable_shape_as_dim_vector());
+  }
+  void after(InputMetadata& t) {
+    after(t.mutable_shape_as_dim_vector());
+  }
+  void before(at::TensorGeometry& t) {
+    before(t.mutable_sizes());
+    before(t.mutable_strides());
+    before(t.mutable_storage_offset());
+    t.recompute();
+  }
+  void after(at::TensorGeometry& t) {
+    after(t.mutable_sizes());
+    after(t.mutable_strides());
+    after(t.mutable_storage_offset());
+    t.recompute();
+  }
+  void before(torch::autograd::TypeAndSize& t) {
+    before(t.sym_sizes);
+    before(t.options);
+  }
+  void after(torch::autograd::TypeAndSize& t) {
+    after(t.sym_sizes);
+    after(t.options);
+  }
+  void before(VariableInfo& t) {
+    before(t.size);
+  }
+  void after(VariableInfo& t) {
+    after(t.size);
+  }
+
+  template <typename T>
+  void before(std::vector<T>& t) {
+    for (T& i : t) {
+      before(i);
+    }
+  }
+  template <typename T>
+  void after(std::vector<T>& t) {
+    for (T& i : t) {
+      after(i);
+    }
+  }
+  template <typename T, unsigned N>
+  void before(c10::SmallVector<T, N>& t) {
+    for (T& i : t) {
+      before(i);
+    }
+  }
+  template <typename T, unsigned N>
+  void after(c10::SmallVector<T, N>& t) {
+    for (T& i : t) {
+      after(i);
+    }
+  }
+
+  template <typename T>
+  void before(c10::OptionalArray<T>& t) {
+    before(t.list);
+  }
+  template <typename T>
+  void after(c10::OptionalArray<T>& t) {
+    after(t.list);
+  }
+
+  template <typename T>
+  void before(std::optional<T>& t) {
+    if (t.has_value()) {
+      before(*t);
+    }
+  }
+  template <typename T>
+  void after(std::optional<T>& t) {
+    if (t.has_value()) {
+      after(*t);
+    }
+  }
+
+  template <typename V>
+  void before(ska::flat_hash_map<std::string, V>& m) {
+    std::vector<std::string> keys;
+    keys.reserve(m.size());
+    std::transform(
+        m.begin(), m.end(), std::back_inserter(keys), [](const auto& entry) {
+          return entry.first;
+        });
+    std::sort(keys.begin(), keys.end());
+    for (auto& k : keys) {
+      before(m.at(k));
+    }
+  }
+
+  template <typename V>
+  void after(ska::flat_hash_map<std::string, V>& m) {
+    for (auto& [_, v] : m) {
+      after(v);
+    }
+  }
+
+#define NO_OP_VISIT(T)     \
+  void before(const T&) {} \
+  void after(const T&) {}
+  NO_OP_VISIT(caffe2::TypeMeta)
+  NO_OP_VISIT(c10::Device)
+  NO_OP_VISIT(c10::DeviceType)
+  NO_OP_VISIT(c10::Layout)
+  NO_OP_VISIT(c10::MemoryFormat)
+  NO_OP_VISIT(c10::ScalarType)
+  NO_OP_VISIT(c10::Scalar)
+  NO_OP_VISIT(c10::TensorOptions)
+  NO_OP_VISIT(std::string)
+  NO_OP_VISIT(int64_t)
+  NO_OP_VISIT(bool)
+  NO_OP_VISIT(double)
+#undef NO_OP_VISIT
+
+  SwapSavedVariables(
+      AutogradCompilerCall& c,
+      TraceState& s,
+      PyObject* p,
+      const NodeCall& n)
+      : compiler(c), state(s), py_compiler(p), curr_node_call(n) {}
+
+  PyObject* get_py_compiler() const {
+    return py_compiler;
+  }
+
+  const NodeCall& get_curr_node_call() {
+    return curr_node_call;
+  }
+
+  void debug_asserts() {
+    stashed_variables.debug_assert();
+    stashed_tensors.debug_assert();
+    stashed_symints.debug_assert();
+  }
+
+ private:
+  template <typename T>
+  struct Stashed {
+    Stashed(T&& v) : prior_value(std::move(v)) {}
+    T prior_value;
+    // Note: we need count here to support duplicate calls to before()
+    // which happen when we have multiple autograd::Edge objects pointing
+    // to the same autograd::Node
+    int count = 1;
+  };
+
+  template <typename T>
+  struct StashedVars : public std::unordered_map<const T*, Stashed<T>> {
+    void save(const T* key, T&& value) {
+      auto [it, inserted] = this->try_emplace(key, std::move(value));
+      if (!inserted) {
+        // keep the value from the prior save()
+        it->second.count++;
+      }
+    }
+    void restore(T* var) {
+      auto it = this->find(var);
+      TORCH_INTERNAL_ASSERT(it != this->end(), "missing before())");
+      if (--it->second.count == 0) {
+        // restore the value on the last restore()
+        *var = std::move(it->second.prior_value);
+        this->erase(it);
+      }
+    }
+    void debug_assert() {
+      TORCH_INTERNAL_ASSERT(this->empty(), "missing call to after()");
+    }
+  };
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  AutogradCompilerCall& compiler;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  TraceState& state;
+  // This is a borrowed reference, we do not increment ownership, or lower it,
+  // it's lifecycle is entirely longer than this objects.
+  PyObject* py_compiler;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const NodeCall& curr_node_call;
+
+  // These mappings are used to save the prior values when we overwrite things
+  // in before(). In after(), we use these to cleanup after ourselves.
+  StashedVars<SavedVariable> stashed_variables;
+  StashedVars<at::Tensor> stashed_tensors;
+  StashedVars<c10::SymInt> stashed_symints;
+  StashedVars<at::IValue> stashed_ivalues;
+};
+
+// NOTE: [Compiled Autograd and backward functions]
+// Built-in autograd nodes have functional apply variants
+// (e.g. MulBackward0_apply_functional). Compiled Autograd's initial graph
+// capture wants to take a variant of this function and proxy it into the graph.
+// Every autograd node defines an apply_with_saved function, that when invoked,
+// proxies a call to a function into the Compiled Autograd graph.
+//
+// Some requirements that we have are:
+// - The proxy'ed function must have inputs that are FX-graphable types.
+// - Windows has a DLL symbol limit of 65536.
+// - Node::apply_with_saved is in libtorch_cpu which does not have direct access
+// to Python
+//
+// There were multiple ways to skin the cat, but what we end up doing is:
+// - for e.g. MulBackward0_apply_functional, we create a new C++ function
+// MulBackward0_apply_functional_ivalue that accepts vector<IValue>.
+// - We define how to pack and unpack arbitrary C++ types into IValues.
+// - apply_with_saved passes MulBackward0_apply_functional_ivalue and
+// the IValue arguments to Python via an indirection.
+// In Python, these get proxy'ed into a graph.
+
+// Helper struct for packing/unpacking an arbitrary C++ type into a single
+// IValue. There are various full and partial specializations for IValuePacker
+// to handle packing specific types (like TensorOptions) into an IValue.
+template <typename T>
+struct IValuePacker {
+  // Defines how to pack T into an IValue.
+  static at::IValue pack(const T& t) {
+    return t;
+  }
+  // Defines how to unpack an IValue into T.
+  static T unpack(const at::IValue& t) {
+    return t.to<T>();
+  }
+  // Returns the TypePtr for the IValue (this is like the "type" of the IValue).
+  // We use this when passing the packed IValue from Python to C++.
+  // In Python, the IValue is just a PyObject* with the native type.
+  // For example, it may be a Python int, a Python List[int], etc.
+  // When passing this PyObject* into C++, we need to know how to parse it
+  // into a C++ type that then gets put into an IValue.
+  // That's what the TypePtr is for: it contains the information to do the
+  // parsing. See torch::jit::toIValue for more information.
+  static at::TypePtr packed_type() {
+    // On windows CPU is support compiled autograd.
+#if defined(_WIN32) && (defined(USE_CUDA) || defined(USE_ROCM))
+    // NB: the if-constexpr usage triggers compilation errors on Windows
+    // with certain compiler settings
+    // (see https://github.com/pytorch/pytorch/pull/144707 for examples).
+    // It's not clear what the problem is, so we're going to ignore it for now.
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        false, "torch.compile not supported on Windows");
+#else
+    if constexpr (::std::is_same_v<T, at::Tensor>) {
+      return at::TensorType::get();
+    } else if constexpr (::std::is_same_v<T, int64_t>) {
+      return at::IntType::get();
+    } else if constexpr (::std::is_same_v<T, c10::SymInt>) {
+      return at::SymIntType::get();
+    } else if constexpr (::std::is_same_v<T, bool>) {
+      return at::BoolType::get();
+    } else if constexpr (::std::is_same_v<T, double>) {
+      return at::FloatType::get();
+    } else if constexpr (::std::is_same_v<T, c10::SymFloat>) {
+      return at::SymFloatType::get();
+    } else if constexpr (::std::is_same_v<T, c10::SymBool>) {
+      return at::SymBoolType::get();
+    } else if constexpr (::std::is_same_v<T, c10::Layout>) {
+      return at::LayoutType::get();
+    } else if constexpr (::std::is_same_v<T, ::std::string>) {
+      return at::StringType::get();
+    } else if constexpr (::std::is_same_v<T, at::Device>) {
+      return at::DeviceObjType::get();
+    } else if constexpr (::std::is_same_v<T, at::Scalar>) {
+      return at::NumberType::get();
+    } else if constexpr (::std::is_same_v<T, at::MemoryFormat>) {
+      return at::MemoryFormatType::get();
+    } else if constexpr (::std::is_same_v<T, at::ScalarType>) {
+      return at::ScalarTypeType::get();
+    } else {
+      // If you got here, you have probably added a member of a new type
+      // to a built-in C++ autograd node.
+      // Unfortunately, we don't know how to handle this type yet.
+      // To get this new type to work with Compiled Autograd, please
+      // either change it to be an IValue-constructible type, or
+      // define how to pack and unpack an object of this time into an IValue
+      // by creating a specialization of IValuePacker for this type.
+      // See NOTE: [Compiled Autograd and backward functions] for context.
+      TORCH_CHECK_NOT_IMPLEMENTED(
+          false, "IValuePacker not implemented for type");
+      return at::NoneType::get();
+    }
+#endif
+  }
+};
+
+template <>
+struct IValuePacker<size_t> {
+  static at::IValue pack(const size_t& t) {
+    // We generally use size_t as the size of a list of Tensors or number of
+    // dimensions. The number of dimensions generally do not exceed 64
+    // (TensorIterator has that limitation), and lists of Tensors generally do
+    // not exceed the int64_t max (you'd probably run out of RAM or run into
+    // significant Tensor overhead). If you run into this limitation the fix is
+    // to figure out how to pack size_t into int64_t. Note that size_t has some
+    // weird behavior on Mac OS.
+    uint64_t maximum_value = std::numeric_limits<int64_t>::max();
+    TORCH_INTERNAL_ASSERT(
+        static_cast<uint64_t>(t) <= maximum_value,
+        "size_t too large to pack into IValue");
+    return static_cast<int64_t>(t); // pack as int64_t
+  }
+  static size_t unpack(const at::IValue& t) {
+    return static_cast<size_t>(t.toInt());
+  }
+  static at::TypePtr packed_type() {
+    return IValuePacker<int64_t>::packed_type();
+  }
+};
+
+template <>
+struct IValuePacker<std::vector<at::SymInt>> {
+  static at::IValue pack(const std::vector<at::SymInt>& t) {
+    return t;
+  }
+  static std::vector<at::SymInt> unpack(const at::IValue& t) {
+    // We need this because there's no t.to<std::vector<at::SymInt>>() override?
+    return t.toSymIntVector();
+  }
+  static at::TypePtr packed_type() {
+    return at::ListType::create(at::SymIntType::get());
+  }
+};
+
+template <>
+struct IValuePacker<VariableInfo> {
+  static at::IValue pack(const VariableInfo& t) {
+    auto tuple = std::make_tuple(
+        t.layout, t.device, t.scalar_type, t.size, t.requires_grad, t.is_empty);
+    return tuple;
+  }
+  static VariableInfo unpack(const at::IValue& t) {
+    auto tuple = t.toTuple();
+    const auto& tuple_elements = tuple->elements();
+    const auto elements = tuple_elements.asArrayRef();
+    TORCH_INTERNAL_ASSERT(elements.size() == 6);
+    VariableInfo v;
+    v.layout = elements[0].toLayout();
+    v.device = elements[1].toDevice();
+    v.scalar_type = elements[2].toScalarType();
+    v.size = elements[3].toSymIntVector();
+    v.requires_grad = elements[4].toBool();
+    v.is_empty = elements[5].toBool();
+    return v;
+  }
+  static at::TypePtr packed_type() {
+    return at::TupleType::create({
+        at::LayoutType::get(),
+        at::DeviceObjType::get(),
+        at::ScalarTypeType::get(),
+        at::ListType::create(at::SymIntType::get()),
+        at::BoolType::get(),
+        at::BoolType::get(),
+    });
+  }
+};
+
+template <>
+struct IValuePacker<caffe2::TypeMeta> {
+  static at::IValue pack(const caffe2::TypeMeta& t) {
+    return at::typeMetaToScalarType(t); // pack as at::ScalarType
+  }
+  static caffe2::TypeMeta unpack(const at::IValue& t) {
+    return caffe2::TypeMeta::fromScalarType(t.to<at::ScalarType>());
+  }
+  static at::TypePtr packed_type() {
+    return IValuePacker<at::ScalarType>::packed_type();
+  }
+};
+
+inline std::optional<at::ScalarType> optTypeMetaToScalarType(
+    const std::optional<caffe2::TypeMeta>& t) {
+  if (t.has_value()) {
+    return at::typeMetaToScalarType(t.value());
+  } else {
+    return std::nullopt;
+  }
+}
+
+using packed_tensoroptions_t = std::tuple<
+    std::optional<bool>,
+    std::optional<at::MemoryFormat>,
+    std::optional<at::Device>,
+    std::optional<at::ScalarType>,
+    std::optional<at::Layout>,
+    std::optional<bool>>;
+
+inline packed_tensoroptions_t pack_TensorOptions(const at::TensorOptions& t) {
+  auto tuple = std::make_tuple(
+      t.requires_grad_opt(),
+      t.memory_format_opt(),
+      t.device_opt(),
+      optTypeMetaToScalarType(t.dtype_opt()),
+      t.layout_opt(),
+      t.pinned_memory_opt());
+  return tuple;
+}
+inline at::TensorOptions unpack_TensorOptions(
+    const packed_tensoroptions_t& tuple) {
+  at::TensorOptions result;
+  auto maybe_requires_grad = std::get<0>(tuple);
+  if (maybe_requires_grad.has_value()) {
+    result = result.requires_grad(maybe_requires_grad);
+  }
+  auto maybe_memory_format = std::get<1>(tuple);
+  if (maybe_memory_format.has_value()) {
+    result = result.memory_format(maybe_memory_format);
+  }
+  auto maybe_device = std::get<2>(tuple);
+  if (maybe_device.has_value()) {
+    result = result.device(maybe_device.value());
+  }
+  auto maybe_dtype = std::get<3>(tuple);
+  if (maybe_dtype.has_value()) {
+    result =
+        result.dtype(caffe2::TypeMeta::fromScalarType(maybe_dtype.value()));
+  }
+  auto maybe_layout = std::get<4>(tuple);
+  if (maybe_layout.has_value()) {
+    result = result.layout(maybe_layout);
+  }
+  auto maybe_pinned_memory = std::get<5>(tuple);
+  if (maybe_pinned_memory.has_value()) {
+    result = result.pinned_memory(maybe_pinned_memory);
+  }
+  return result;
+}
+
+template <>
+struct IValuePacker<at::TensorOptions> {
+  static at::IValue pack(const at::TensorOptions& t) {
+    return pack_TensorOptions(t);
+  }
+  static at::TensorOptions unpack(const at::IValue& t) {
+    auto tuple = t.to<packed_tensoroptions_t>();
+    return unpack_TensorOptions(tuple);
+  }
+  static at::TypePtr packed_type() {
+    return at::TupleType::create(
+        {at::OptionalType::create(at::BoolType::get()),
+         at::OptionalType::create(at::MemoryFormatType::get()),
+         at::OptionalType::create(at::DeviceObjType::get()),
+         at::OptionalType::create(at::ScalarTypeType::get()),
+         at::OptionalType::create(at::LayoutType::get()),
+         at::OptionalType::create(at::BoolType::get())});
+  }
+};
+
+template <>
+struct IValuePacker<TypeAndSize> {
+  static at::IValue pack(const TypeAndSize& t) {
+    auto tuple = std::make_tuple(t.sym_sizes, pack_TensorOptions(t.options));
+    return tuple;
+  }
+  static TypeAndSize unpack(const at::IValue& t) {
+    auto tuple =
+        t.to<std::tuple<std::vector<at::SymInt>, packed_tensoroptions_t>>();
+    TypeAndSize result;
+    result.sym_sizes = std::get<0>(tuple);
+    result.options = unpack_TensorOptions(std::get<1>(tuple));
+    return result;
+  }
+  static at::TypePtr packed_type() {
+    return at::TupleType::create(
+        {IValuePacker<std::vector<at::SymInt>>::packed_type(),
+         IValuePacker<at::TensorOptions>::packed_type()});
+  }
+};
+
+template <typename T>
+struct IValuePacker<std::optional<T>> {
+  static at::IValue pack(const std::optional<T>& t) {
+    if (t.has_value()) {
+      return IValuePacker<T>::pack(t.value());
+    } else {
+      return std::nullopt;
+    }
+  }
+  static std::optional<T> unpack(const at::IValue& t) {
+    if (t.isNone()) {
+      return std::nullopt;
+    } else {
+      return IValuePacker<T>::unpack(t);
+    }
+  }
+  static at::TypePtr packed_type() {
+    return at::OptionalType::create(IValuePacker<T>::packed_type());
+  }
+};
+
+template <typename T>
+struct IValuePacker<std::vector<T>> {
+  static at::IValue pack(const std::vector<T>& t) {
+    if constexpr (::std::is_constructible_v<at::IValue, T>) {
+      return t;
+    }
+    if (t.empty()) {
+      auto lst = c10::impl::GenericList(at::AnyType::get());
+      return lst;
+    }
+    auto type_ptr = IValuePacker<T>::pack(t[0]).type();
+    auto lst = c10::impl::GenericList(type_ptr);
+    for (const auto& elt : t) {
+      lst.emplace_back(IValuePacker<T>::pack(elt));
+    }
+    return lst;
+  }
+  static std::vector<T> unpack(const at::IValue& t) {
+    if constexpr (::std::is_constructible_v<at::IValue, T>) {
+      return t.to<::std::vector<T>>();
+    }
+    std::vector<T> result;
+    auto lst = t.toList();
+    for (size_t i = 0; i < lst.size(); ++i) {
+      const at::IValue& elt = lst.get(i);
+      result.emplace_back(IValuePacker<T>::unpack(elt));
+    }
+    return result;
+  }
+  static at::TypePtr packed_type() {
+    return at::ListType::create(IValuePacker<T>::packed_type());
+  }
+};
+
+template <typename T>
+struct IValuePacker<c10::List<T>> {
+  static at::IValue pack(const c10::List<T>& t) {
+    return IValuePacker<std::vector<T>>::pack(t.vec());
+  }
+  static c10::List<T> unpack(const at::IValue& t) {
+    return c10::List<T>(IValuePacker<std::vector<T>>::unpack(t));
+  }
+  static at::TypePtr packed_type() {
+    return IValuePacker<std::vector<T>>::packed_type();
+  }
+};
+
+template <size_t N>
+struct IValuePacker<std::array<bool, N>> {
+  static at::IValue pack(const std::array<bool, N>& t) {
+    std::vector<bool> result(t.begin(), t.end());
+    return IValuePacker<std::vector<bool>>::pack(result);
+  }
+  static std::array<bool, N> unpack(const at::IValue& t) {
+    std::array<bool, N> result;
+    auto packed = IValuePacker<std::vector<bool>>::unpack(t);
+    for (size_t i = 0; i < packed.size(); i++) {
+      result[i] = packed[i];
+    }
+    return result;
+  }
+  static at::TypePtr packed_type() {
+    return IValuePacker<std::vector<bool>>::packed_type();
+  }
+};
+
+template <>
+struct IValuePacker<at::TensorGeometry> {
+  static at::IValue pack(const at::TensorGeometry& t) {
+    auto tuple = std::make_tuple(
+        t.sym_sizes().vec(), t.sym_strides().vec(), t.sym_storage_offset());
+    return tuple;
+  }
+  static at::TensorGeometry unpack(const at::IValue& t) {
+    auto tuple = t.to<std::tuple<
+        std::vector<at::SymInt>,
+        std::vector<at::SymInt>,
+        at::SymInt>>();
+    return at::TensorGeometry(
+        std::get<0>(tuple), std::get<1>(tuple), std::get<2>(tuple));
+  }
+  static at::TypePtr packed_type() {
+    return at::TupleType::create(
+        {IValuePacker<std::vector<at::SymInt>>::packed_type(),
+         IValuePacker<std::vector<at::SymInt>>::packed_type(),
+         at::SymIntType::get()});
+  }
+};
+
+template <>
+struct IValuePacker<InputMetadata> {
+  static at::IValue pack(const InputMetadata& t) {
+    TORCH_INTERNAL_ASSERT(!t.is_nested_tensor());
+    auto tuple = std::make_tuple(
+        pack_TensorOptions(t.options()),
+        t.shape_as_dim_vector().vec(),
+        t.is_tensor_subclass(),
+        t.grad_dtype());
+    return tuple;
+  }
+  static InputMetadata unpack(const at::IValue& t) {
+    auto tuple = t.to<std::tuple<
+        packed_tensoroptions_t,
+        std::vector<at::SymInt>,
+        bool,
+        std::optional<c10::ScalarType>>>();
+
+    return InputMetadata(
+        unpack_TensorOptions(std::get<0>(tuple)),
+        SymIntSmallVec(std::get<1>(tuple)),
+        std::get<2>(tuple),
+        false,
+        std::get<3>(tuple));
+  }
+  static at::TypePtr packed_type() {
+    return at::TupleType::create(
+        {IValuePacker<at::TensorOptions>::packed_type(),
+         IValuePacker<std::vector<at::SymInt>>::packed_type(),
+         at::BoolType::get(),
+         IValuePacker<std::optional<at::ScalarType>>::packed_type()});
+  }
+};
+
+template <typename T>
+struct IValuePacker<at::OptionalArray<T>> {
+  static at::IValue pack(const at::OptionalArray<T>& t) {
+    return IValuePacker<std::optional<std::vector<T>>>::pack(t.list);
+  }
+  static at::OptionalArray<T> unpack(const at::IValue& t) {
+    auto result = IValuePacker<std::optional<std::vector<T>>>::unpack(t);
+    if (result.has_value()) {
+      return {result.value()};
+    } else {
+      return {};
+    }
+  }
+  static at::TypePtr packed_type() {
+    return IValuePacker<std::optional<std::vector<T>>>::packed_type();
+  }
+};
+
+// This is a helper struct for packing and unpacking multiple arguments into
+// an ivalue_list. It leverages IValuePacker<T>.
+struct PackedArgs {
+  PackedArgs() = default;
+
+  explicit PackedArgs(std::vector<at::IValue> stack_)
+      : stack(std::move(stack_)) {}
+
+  const std::vector<at::IValue>& vec() const {
+    return stack;
+  }
+
+  template <typename T>
+  void pack(const T& t) {
+    stack.emplace_back(IValuePacker<T>::pack(t));
+  }
+  template <typename T>
+  T unpack() {
+    return IValuePacker<T>::unpack(std::move(stack[idx++]));
+  }
+
+  void pack_saved_data(const ska::flat_hash_map<std::string, at::IValue>& dct) {
+    std::vector<std::string> keys;
+    std::vector<at::IValue> values;
+    for (const auto& [key, value] : dct) {
+      keys.emplace_back(key);
+      values.emplace_back(value);
+    }
+    pack(keys);
+    for (const auto& value : values) {
+      pack(value);
+    }
+  }
+
+  ska::flat_hash_map<std::string, at::IValue> unpack_saved_data() {
+    ska::flat_hash_map<std::string, at::IValue> dct;
+    auto keys = unpack<std::vector<std::string>>();
+    for (const auto& key : keys) {
+      dct.insert({key, std::move(stack[idx++])});
+    }
+    return dct;
+  }
+
+ private:
+  std::vector<at::IValue> stack;
+  int64_t idx = 0;
+};
+
+} // namespace torch::dynamo::autograd
+
+template <>
+struct std::hash<torch::dynamo::autograd::CacheKey> {
+  size_t operator()(const torch::dynamo::autograd::CacheKey& k) const noexcept {
+    return k.hash();
+  }
+};
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpp_shim.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpp_shim.h
new file mode 100644
index 0000000000000000000000000000000000000000..d764297efb4acb23aaadfea5a2658f8b5a3512f1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpp_shim.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+struct _PytorchRecordFunctionState;
+typedef struct _PytorchRecordFunctionState _PytorchRecordFunctionState;
+
+_PytorchRecordFunctionState* _pytorch_record_function_enter(const char* name);
+void _pytorch_record_function_exit(_PytorchRecordFunctionState* state);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpython_defs.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpython_defs.h
new file mode 100644
index 0000000000000000000000000000000000000000..d2d361c2b8ec4fd4b6331738ee01f9bc3d54356a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpython_defs.h
@@ -0,0 +1,45 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/python_compat.h>
+
+// Functions that need to be copied from the CPython source
+// should go in cpython_defs.c. Copying is required when, e.g.,
+// we need to call internal CPython functions that are not exposed.
+
+#if IS_PYTHON_3_11_PLUS
+
+typedef struct _PyInterpreterFrame _PyInterpreterFrame;
+
+PyFunctionObject* _PyFunction_CopyWithNewCode(
+    PyFunctionObject* o,
+    PyCodeObject* code);
+
+void THP_PyFrame_Clear(_PyInterpreterFrame* frame);
+
+_PyInterpreterFrame* THP_PyThreadState_BumpFramePointerSlow(
+    PyThreadState* tstate,
+    size_t size);
+
+void THP_PyThreadState_PopFrame(
+    PyThreadState* tstate,
+    _PyInterpreterFrame* frame);
+
+#endif
+
+// pointers to _PyOpcode_Caches for C++
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+extern const uint8_t* THP_PyOpcode_Caches;
+extern int THP_PyOpcode_Caches_size;
+void init_THPCaches();
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpython_includes.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpython_includes.h
new file mode 100644
index 0000000000000000000000000000000000000000..ea994e6fce119bcb14ed14239d80fe4faa4de1b9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/cpython_includes.h
@@ -0,0 +1,76 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/python_compat.h>
+
+// Problem in CPython includes when mixing core and non-core build
+// The fix was not backported to 3.12 so this is needed here
+// https://github.com/python/cpython/issues/105268
+#if IS_PYTHON_3_12_PLUS
+#undef _PyGC_FINALIZED
+#endif
+
+// see https://bugs.python.org/issue35886
+#define Py_BUILD_CORE
+
+#ifndef __cplusplus
+// C-only headers
+#include <internal/pycore_pystate.h>
+
+#endif // __cplusplus
+
+#if IS_PYTHON_3_11_PLUS
+#include <internal/pycore_frame.h>
+
+#include <torch/csrc/dynamo/stackref_bridge.h>
+#if IS_PYTHON_3_14_PLUS && !defined(_WIN32)
+#include <internal/pycore_code.h>
+#include <internal/pycore_genobject.h>
+#include <internal/pycore_interpframe.h>
+#include <internal/pycore_stackref.h>
+#elif IS_PYTHON_3_14_PLUS && defined(_WIN32)
+#include <internal/pycore_interpframe_structs.h> // _PyInterpreterFrame
+#endif
+
+#endif
+
+#undef Py_BUILD_CORE
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#if IS_PYTHON_3_14_PLUS
+
+#define F_CODE(x) \
+  ((PyCodeObject*)THP_PyStackRef_AsPyObjectBorrow(&(x)->f_executable))
+#define PREV_INSTR(x) (x)->instr_ptr
+
+#else
+
+#if IS_PYTHON_3_13_PLUS
+#define F_CODE(x) ((PyCodeObject*)(x)->f_executable)
+#define PREV_INSTR(x) (x)->instr_ptr
+#else
+#define F_CODE(x) ((PyCodeObject*)(x)->f_code)
+#define PREV_INSTR(x) (x)->prev_instr
+#endif // IS_PYTHON_3_13_PLUS
+
+#endif // IS_PYTHON_3_14_PLUS
+
+#if IS_PYTHON_3_14_PLUS
+#define FUNC(x) \
+  ((PyFunctionObject*)THP_PyStackRef_AsPyObjectBorrow(&(x)->f_funcobj))
+#elif IS_PYTHON_3_12_PLUS
+#define FUNC(x) ((PyFunctionObject*)(x)->f_funcobj)
+#else
+#define FUNC(x) ((PyFunctionObject*)(x)->f_func)
+#endif
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/debug_macros.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/debug_macros.h
new file mode 100644
index 0000000000000000000000000000000000000000..49ef7b69fe4ab8a5a6948644c98008c11518489d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/debug_macros.h
@@ -0,0 +1,107 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/python_compat.h>
+
+#ifdef __cplusplus
+#include <cstdio>
+#else
+#include <stdio.h>
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifdef _WIN32
+#define unlikely(x) (x)
+#else
+#define unlikely(x) __builtin_expect((x), 0)
+#endif
+
+#define NULL_CHECK(val)                                         \
+  if (unlikely((val) == NULL)) {                                \
+    fprintf(stderr, "NULL ERROR: %s:%d\n", __FILE__, __LINE__); \
+    PyErr_Print();                                              \
+    abort();                                                    \
+  } else {                                                      \
+  }
+
+// CHECK might be previously declared
+#undef CHECK
+#define CHECK(cond)                                                     \
+  if (unlikely(!(cond))) {                                              \
+    fprintf(stderr, "DEBUG CHECK FAILED: %s:%d\n", __FILE__, __LINE__); \
+    abort();                                                            \
+  } else {                                                              \
+  }
+
+// Uncomment next line to print debug message
+// #define TORCHDYNAMO_DEBUG 1
+#ifdef TORCHDYNAMO_DEBUG
+
+#define DEBUG_CHECK(cond) CHECK(cond)
+#define DEBUG_NULL_CHECK(val) NULL_CHECK(val)
+#define DEBUG_TRACE(msg, ...) \
+  fprintf(stderr, "TRACE[%s:%d] " msg "\n", __func__, __LINE__, __VA_ARGS__)
+#define DEBUG_TRACE0(msg) \
+  fprintf(stderr, "TRACE[%s:%d] " msg "\n", __func__, __LINE__)
+
+#else
+
+#define DEBUG_CHECK(cond)
+#define DEBUG_NULL_CHECK(val)
+#define DEBUG_TRACE(msg, ...)
+#define DEBUG_TRACE0(msg)
+
+#endif
+
+inline _PyFrameEvalFunction _debug_set_eval_frame(
+    PyThreadState* tstate,
+    _PyFrameEvalFunction eval_frame) {
+  _PyFrameEvalFunction prev =
+      _PyInterpreterState_GetEvalFrameFunc(tstate->interp);
+  _PyInterpreterState_SetEvalFrameFunc(tstate->interp, eval_frame);
+  return prev;
+}
+
+// Inspect PyObject*'s from C/C++ at the Python level, in pdb.
+// e.g.
+//
+// PyObject* obj1 = PyList_New(...);
+// PyObject* obj2 = PyObject_CallFunction(...);
+// INSPECT(obj1, obj2);
+// (pdb) p args[0]
+// # list
+// (pdb) p args[1]
+// # some object
+// (pdb) p args[1].some_attr
+// # etc.
+//
+// Implementation: set eval frame callback to default, call
+// torch._dynamo.utils._breakpoint_for_c_dynamo, reset eval frame callback.
+#define INSPECT(...)                                                  \
+  {                                                                   \
+    PyThreadState* cur_tstate = PyThreadState_Get();                  \
+    _PyFrameEvalFunction prev_eval_frame =                            \
+        _debug_set_eval_frame(cur_tstate, &_PyEval_EvalFrameDefault); \
+    PyObject* torch__dynamo_utils_module =                            \
+        PyImport_ImportModule("torch._dynamo.utils");                 \
+    NULL_CHECK(torch__dynamo_utils_module);                           \
+    PyObject* breakpoint_for_c_dynamo_fn = PyObject_GetAttrString(    \
+        torch__dynamo_utils_module, "_breakpoint_for_c_dynamo");      \
+    NULL_CHECK(breakpoint_for_c_dynamo_fn);                           \
+    PyObject_CallFunctionObjArgs(                                     \
+        breakpoint_for_c_dynamo_fn, __VA_ARGS__, NULL);               \
+    _debug_set_eval_frame(cur_tstate, prev_eval_frame);               \
+    Py_DECREF(breakpoint_for_c_dynamo_fn);                            \
+    Py_DECREF(torch__dynamo_utils_module);                            \
+  }
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/eval_frame.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/eval_frame.h
new file mode 100644
index 0000000000000000000000000000000000000000..cd0bd96ab976dc831cf1a7d40158a3c7aebccf7e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/eval_frame.h
@@ -0,0 +1,67 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <stdbool.h>
+
+#include <torch/csrc/dynamo/extra_state.h>
+#include <torch/csrc/utils/python_compat.h>
+#ifdef __cplusplus
+
+extern "C" {
+
+PyObject* torch_c_dynamo_eval_frame_init(void);
+
+#endif
+
+// All the eval APIs change in 3.11 so we need to decide which one to use on the
+// fly https://docs.python.org/3/c-api/init.html#c._PyFrameEvalFunction
+#if IS_PYTHON_3_11_PLUS
+#define THP_EVAL_API_FRAME_OBJECT _PyInterpreterFrame
+#else
+#define THP_EVAL_API_FRAME_OBJECT PyFrameObject
+#endif // IS_PYTHON_3_11_PLUS
+
+// We need to be able to return the _PyInterpreterFrame to python so create
+// a python binding for it
+
+typedef struct THPPyInterpreterFrame {
+  PyObject_HEAD
+  THP_EVAL_API_FRAME_OBJECT* frame; // Borrowed reference
+  PyObject* locals;
+} THPPyInterpreterFrame;
+
+THPPyInterpreterFrame* THPPyInterpreterFrame_New(
+    THP_EVAL_API_FRAME_OBJECT* frame);
+
+extern bool is_skip_guard_eval_unsafe;
+extern int fullgraph_compiled_frame_count;
+extern bool fullgraph_error_on_nested_compile;
+
+void clear_old_frame_if_python_312_plus(
+    PyThreadState* tstate,
+    THP_EVAL_API_FRAME_OBJECT* frame);
+
+void eval_frame_callback_set(PyObject* obj);
+
+const char* get_frame_name(THP_EVAL_API_FRAME_OBJECT* frame);
+
+PyObject* dynamo_eval_frame_default(
+    PyThreadState* tstate,
+    THP_EVAL_API_FRAME_OBJECT* frame,
+    int throw_flag);
+
+PyObject* dynamo_eval_custom_code(
+    PyThreadState* tstate,
+    THP_EVAL_API_FRAME_OBJECT* frame,
+    PyCodeObject* code,
+    const char* trace_annotation,
+    int throw_flag);
+
+#ifdef __cplusplus
+
+} // extern "C"
+
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/eval_frame_cpp.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/eval_frame_cpp.h
new file mode 100644
index 0000000000000000000000000000000000000000..348643b2d6094f98bcd2610b81bbfb6cb9494e40
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/eval_frame_cpp.h
@@ -0,0 +1,55 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <Python.h>
+
+#include <torch/csrc/dynamo/eval_frame.h>
+#include <torch/csrc/dynamo/extra_state.h>
+#include <torch/csrc/dynamo/framelocals_mapping.h>
+
+#ifdef __cplusplus
+
+#include <torch/csrc/utils/pybind.h>
+
+extern "C" {
+
+#endif
+
+PyObject* dynamo__custom_eval_frame(
+    PyThreadState* tstate,
+    THP_EVAL_API_FRAME_OBJECT* frame,
+    int throw_flag,
+    PyObject* callback);
+
+PyObject* dynamo_set_code_exec_strategy(PyObject* dummy, PyObject* obj);
+void dynamo_skip_code_recursive(PyCodeObject* code);
+
+void dynamo_set_c_recursion_limit(int32_t limit);
+int32_t dynamo_get_c_recursion_limit();
+
+#ifdef __cplusplus
+
+} // extern "C"
+
+// Bytecode debugger callback functions
+void set_bytecode_debugger_callback(py::object callback);
+py::object get_bytecode_debugger_callback();
+
+// Breakpoint code object tracking
+void register_breakpoint_code(py::object code);
+
+// NullStackValue - sentinel class for representing NULL values on Python stack
+class NullStackValue {
+ public:
+  static NullStackValue& get_singleton();
+};
+
+py::object get_null_stack_value();
+py::list _get_frame_value_stack_with_depth(
+    const py::handle& frame_obj,
+    int depth);
+
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/extra_state.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/extra_state.h
new file mode 100644
index 0000000000000000000000000000000000000000..04d8789fd3aaba74bda9291f7042b75152c89e2d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/extra_state.h
@@ -0,0 +1,213 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <Python.h>
+
+#include <torch/csrc/dynamo/framelocals_mapping.h>
+
+#ifdef __cplusplus
+
+#include <torch/csrc/dynamo/utils.h>
+#include <torch/csrc/utils/pybind.h>
+#include <list>
+
+namespace py = pybind11;
+
+extern "C" {
+
+#else
+
+#include <stdbool.h>
+
+#endif
+
+enum FrameAction {
+  DEFAULT, // look through the cache, compile if not found
+  SKIP, // eager
+  RUN_ONLY, // look through the cache, run eager if not found
+};
+
+typedef struct FrameExecStrategy {
+  enum FrameAction cur_action; // action to take for current frame
+  enum FrameAction recursive_action; // action to take for recursive frames
+} FrameExecStrategy;
+
+// Points to the extra scratch space on the code object
+extern Py_ssize_t extra_index;
+
+// function to call when cache lookup errors
+extern PyObject* guard_error_hook;
+
+typedef PyObject FrameState;
+typedef struct CacheEntry CacheEntry;
+
+// ExtraState encasulates CacheEntry and FrameState. ExtraState is the highest
+// level of abstraction of what is stored on the extra code object. Previously,
+// we saved different parts on different extra indexes.  We prefer this way
+// because of cleaner abstraction and faster SetExtra access.
+
+#ifdef __cplusplus
+
+typedef struct VISIBILITY_HIDDEN PrecompileEntry {
+  py::object guard_manager;
+  py::object code;
+  void* root_mgr;
+
+  PrecompileEntry(py::object gm, py::object c);
+} PrecompileEntry;
+
+typedef struct VISIBILITY_HIDDEN ExtraState {
+  // A pointer to the orig_code object to prevent race conditions in invalidate
+  // function.
+  PyCodeObject* orig_code;
+  std::list<PrecompileEntry> precompile_entries;
+  // List of cache entries for compiled code objects
+  std::list<CacheEntry> cache_entry_list;
+  // Frame state to detect dynamic shape dims
+  py::dict frame_state;
+  // Actions to apply to all frames with this code object
+  FrameExecStrategy strategy{DEFAULT, DEFAULT};
+
+  ExtraState(PyCodeObject* orig_code_arg);
+  CacheEntry* get_first_entry();
+  void move_to_front(CacheEntry* cache_entry);
+  void move_to_back(CacheEntry* cache_entry);
+  void invalidate(CacheEntry* cache_entry, py::object deleted_guard_manager);
+} ExtraState;
+
+#else
+
+typedef struct ExtraState ExtraState;
+typedef struct PrecompileEntry PrecompileEntry;
+
+#endif
+
+// Helper to extra the cache_entry from the extra state.
+// Ownership contract
+// args
+//  - extra_state: Borrowed
+// return
+//  - CacheEntry: Borrowed.
+CacheEntry* extract_cache_entry(ExtraState* extra_state);
+
+// Returns either the previously stored frame state or an empty dict.
+// Ownership contract
+// args
+//  - extra_state: Borrowed
+// return
+//  - extra_state->frame_state: Borrowed.
+FrameState* extract_frame_state(ExtraState* extra_state);
+
+// Returns the FrameExecStrategy stored in extra_state.
+// Ownership contract
+// args
+//  - extra_state: Borrowed
+FrameExecStrategy extra_state_get_exec_strategy(ExtraState* extra_state);
+
+// Set the FrameExecStrategy to be done to all frames with code object
+// corresponding to this extra_state. Ownership contract
+// - extra_state: Borrowed
+void extra_state_set_exec_strategy(
+    ExtraState* extra_state,
+    FrameExecStrategy strategy);
+
+// Ownership contract
+// args
+//  - code: Borrowed
+// return
+//  - extra_state: Borrowed.
+ExtraState* get_extra_state(PyCodeObject* code);
+
+// This is passed as freefunc to _PyEval_RequestCodeExtraIndex. This acts as a
+// deleter for the object on extra scratch space. This function is called
+// internally in _PyCode_SetExtra and also during the code deallocation.
+
+// Destroys the extra state by deleting cache_entry, frame state and finally
+// freeing the constructed extra state.
+
+// Developer note - You should not call this function directly. This is called
+// directly inside set_extra_state. If you are in a situation trying to call
+// this function, consider if set_extra_state should be called.
+void destroy_extra_state(void* obj);
+
+// Clears the existing object sitting on the extra scratch spance and sets it
+// up with the new state. Note that _PyCode_SetExtra calls the
+// destroy_extra_state deleter internally, and therefore we don't call it
+// explicitly here.
+
+// Ownership contract
+// args
+//  - extra_state: Stolen
+// return
+//  - there is no return, but the extra_state is stolen, so it becomes
+//  set_extra_state responsibility to clean it up. It will be deleted during
+//  the reset_code, when the set_extra_state is called with NULL.
+
+// Invariant - Dont set the extra state for the extra state that is already on
+// the code object. Otherwise, we will first free up the old extra state
+// (which is also the new extra state) and write something invalid on the
+// scratch space.
+void set_extra_state(PyCodeObject* code, ExtraState* extra_state);
+
+// Creates a new extra state and put it on the extra scratch space of the code
+// object.
+
+// Ownership contract
+// args
+//  - code: Borrowed
+// return:
+//   - extra_state: New reference.
+// These references are then further passed to set_extra_state which becomes
+// the final owner of these references.
+ExtraState* init_and_set_extra_state(PyCodeObject* code);
+
+// Lookup the cache held by extra_state.
+// Ownership contract
+// args
+//  - extra_state: Borrowed
+// return:
+//   - Py_None or PyCodeObject: Borrowed reference.
+//   - Py_None or PyObject: Trace id of the compiled code.
+void lookup(
+    ExtraState* extra_state,
+    FrameLocalsMapping* f_locals,
+    PyObject* backend,
+    PyObject** maybe_cached_code,
+    const char** trace_annotation,
+    bool is_skip_guard_eval_unsafe);
+
+// Create a new cache entry at extra_state holding on to guarded_code.
+// Ownership contract
+// args
+//  - extra_state: Borrowed
+//  - guarded_code: Borrowed
+// return:
+//  - cache_entry: Borrowed reference
+CacheEntry* create_cache_entry(
+    ExtraState* extra_state,
+    PyObject* guraded_code,
+    PyObject* callback);
+
+// Extracts the backend fn from the callback.
+PyObject* get_backend(PyObject* callback);
+
+#ifdef __cplusplus
+
+} // extern "C"
+
+// Returns the list of CacheEntry corresponding to code_obj.
+// Warning: returns references whose lifetimes are controlled by C++
+py::list _debug_get_cache_entry_list(const py::handle& code_obj);
+void _reset_precompile_entries(const py::handle& code_obj);
+void _load_precompile_entry(
+    const py::handle& code_obj,
+    py::object guard_manager,
+    py::object dynamo_code);
+py::list _debug_get_precompile_entries(const py::handle& code_obj);
+void _set_lru_cache(py::object boolean);
+
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/framelocals_mapping.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/framelocals_mapping.h
new file mode 100644
index 0000000000000000000000000000000000000000..a604a35ff5a988ddfd6fc65c61b6200ee1d1814f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/framelocals_mapping.h
@@ -0,0 +1,97 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/python_compat.h>
+
+#ifdef __cplusplus
+
+#include <string>
+#include <unordered_map>
+
+#include <torch/csrc/dynamo/utils.h>
+#include <torch/csrc/utils/pybind.h>
+
+extern "C" {
+
+#if IS_PYTHON_3_11_PLUS
+using FrameLocalsFrameType = _PyInterpreterFrame;
+#else
+using FrameLocalsFrameType = PyFrameObject;
+#endif // IS_PYTHON_3_11_PLUS
+
+/**
+ * Utility to view a frame's localsplus (locals + cells + freevars)
+ * in C/C++ and Python, without changing the state of the frame.
+ *
+ * Notes on usage:
+ *  - C/C++ can directly read the frame's localsplus using an index.
+ *  - Cell/free variables are unboxed.
+ *  - Can be converted into a dict for use in Python.
+ *    The dict is constructed once per FrameLocalsMapping, lazily.
+ *  - Lifetime should not exceed the lifetime of the frame
+ *
+ * How do guards use FrameLocalsMapping?
+ * - When a guard accesses a frame's localsplus, we find the index of the
+ *   variable name in the frame's code object and create a
+ *   FrameLocalsGuardAccessor.
+ * - We create a FrameLocalsMapping for the frame that we pass on to guard eval.
+ * - LeafGuards/GuardManagers/GuardAccessors now need to define how they
+ *   handle FrameLocalsMapping. By default, the FrameLocalsMapping is converted
+ *   to a Python dict and the guard check is performed on the resulting dict.
+ * - Some guard checks don't actually depend on the input arguments, e.g. they
+ *   only check global state. In this case, no dict conversion of
+ *   FrameLocalsMapping is done.
+ * - FrameLocalsGuardAccessor is like DictGetItemGuardAccessor, except it knows
+ *   how to handle FrameLocalsMapping - by using the framelocals variable name
+ *   index that it was given when it was built.
+ */
+typedef struct VISIBILITY_HIDDEN FrameLocalsMapping {
+ private:
+  py::object _code_obj;
+  // can't use localsplus directly due to closure variables:
+  // - in 3.11+, the closure vars in the frame's closure object and
+  //   the corresponding localsplus entry is nullptr
+  // - regardless of Python version, we need to unbox the cell variable
+  std::vector<py::handle> _framelocals;
+
+  py::object _dict{py::none()};
+
+  void _realize_dict();
+
+ public:
+  explicit FrameLocalsMapping(FrameLocalsFrameType* frame);
+
+  PyObject* get(int idx);
+
+  bool dict_realized() const {
+    return _dict.is_none();
+  }
+
+  // Borrowed reference
+  PyDictObject* to_dict() {
+    if (this->dict_realized()) {
+      _realize_dict();
+    }
+    return (PyDictObject*)_dict.ptr();
+  }
+} FrameLocalsMapping;
+
+#else
+
+// opaque type for C
+typedef struct FrameLocalsMapping FrameLocalsMapping;
+
+#endif
+
+// Borrowed reference
+PyDictObject* framelocals_mapping_to_dict(FrameLocalsMapping* map);
+
+#ifdef __cplusplus
+} // extern "C"
+
+py::tuple code_framelocals_names(py::handle code);
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/guards.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/guards.h
new file mode 100644
index 0000000000000000000000000000000000000000..7e98d64bbfb5a6871f53e643d40458b4c145a7f0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/guards.h
@@ -0,0 +1,120 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/GradMode.h>
+#include <torch/csrc/dynamo/framelocals_mapping.h>
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::dynamo {
+
+PyObject* torch_c_dynamo_guards_init();
+
+// interfaces for extra_state and eval_frame.c because RootGuardManager class is
+// not visible there.
+void* convert_to_root_guard_manager(py::object root);
+bool run_root_guard_manager(void* root, FrameLocalsMapping* f_locals);
+
+extern thread_local bool tls_is_in_mode_without_ignore_compile_internals;
+
+void set_is_in_mode_without_ignore_compile_internals(bool value);
+
+// If we're in a mode with ignore_compile_internals=False, we WON'T mask
+// Python keys from guard checking (they should be visible, so eager fallback is
+// possible). Otherwise (invisible mode or no mode), we WILL mask Python keys to
+// avoid guard failures on the dispatch keyset at runtime.
+bool get_is_in_mode_without_ignore_compile_internals();
+
+struct LocalState {
+  // TLS state that changes operators
+  c10::impl::LocalDispatchKeySet dispatch_modifier;
+  c10::DispatchKeySet override_dispatch_key_set;
+  bool grad_mode_enabled;
+  bool should_mask_python_keys;
+
+  at::DispatchKeySet apply(at::DispatchKeySet ks) const {
+    if (override_dispatch_key_set.empty()) {
+      auto result =
+          (ks | dispatch_modifier.included_) - dispatch_modifier.excluded_;
+
+      if (should_mask_python_keys) {
+        result = result -
+            c10::DispatchKeySet(
+                     {c10::DispatchKey::Python,
+                      c10::DispatchKey::PythonTLSSnapshot,
+                      c10::DispatchKey::PythonDispatcher});
+      }
+
+      return result;
+    } else {
+      return override_dispatch_key_set;
+    }
+  }
+
+  LocalState()
+      : dispatch_modifier(c10::impl::tls_local_dispatch_key_set()),
+        override_dispatch_key_set(c10::BackendComponent::InvalidBit),
+        grad_mode_enabled(at::GradMode::is_enabled()),
+        should_mask_python_keys(
+            !get_is_in_mode_without_ignore_compile_internals()) {}
+
+  void overrideDispatchKeySet(c10::DispatchKeySet ks) {
+    override_dispatch_key_set = ks;
+  }
+};
+
+class TensorCheck {
+ public:
+  TensorCheck(
+      const LocalState& state,
+      PyTypeObject* pt,
+      const at::Tensor& v,
+      c10::DispatchKeySet dispatch_key_set,
+      std::vector<std::optional<c10::SymInt>> dynamic_dims_sizes,
+      std::vector<std::optional<c10::SymInt>> dynamic_dims_strides);
+
+  TensorCheck(
+      const LocalState& state,
+      PyTypeObject* pt,
+      c10::DispatchKeySet dispatch_key_set,
+      at::ScalarType dtype,
+      at::DeviceIndex device_index,
+      bool requires_grad,
+      std::vector<std::optional<c10::SymInt>> dynamic_dims_sizes,
+      std::vector<std::optional<c10::SymInt>> dynamic_dims_strides);
+
+  bool check(const LocalState& state, const at::Tensor& v);
+  bool check(
+      const LocalState& state,
+      const c10::DispatchKeySet& dispatch_key_set,
+      const at::ScalarType& dtype,
+      const c10::Device& device,
+      const c10::SymIntArrayRef& dynamic_dims_sizes,
+      const c10::SymIntArrayRef& dynamic_dims_strides,
+      const bool& requires_grad);
+  std::string check_verbose(
+      const LocalState& state,
+      const at::Tensor& v,
+      const std::string& tensor_name);
+
+  PyTypeObject* pytype;
+
+ private:
+  uint64_t dispatch_key_; // DispatchKeySet includes device/layout
+  at::ScalarType dtype_;
+  // Note(voz): While dispatch_key_ is sufficiently representative of a device
+  // In that keys are more granular AND device specific - they do not
+  // necessarily capture device indices correctly.
+  at::DeviceIndex device_index_;
+  bool requires_grad_;
+  // NB: These are unset if dynamic shapes is enabled.
+  std::vector<std::optional<c10::SymInt>> sizes_;
+  std::vector<std::optional<c10::SymInt>> strides_;
+  // Not strictly required for dense tensors, but nested tensors need it.
+  int64_t dim_;
+};
+
+} // namespace torch::dynamo
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..74f5673ff3a012da01afa6b8386db56c48c9516a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/init.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// C2039 MSVC
+#include <pybind11/complex.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include <Python.h>
+
+namespace torch::dynamo {
+void initDynamoBindings(PyObject* torch);
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/python_compiled_autograd.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/python_compiled_autograd.h
new file mode 100644
index 0000000000000000000000000000000000000000..dd48feff5884dacc2b37cb32264a6579f81f8daf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/python_compiled_autograd.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/utils/python_stub.h>
+
+// see [Note: Compiled Autograd]
+namespace torch::dynamo::autograd {
+PyObject* torch_c_dynamo_compiled_autograd_init();
+} // namespace torch::dynamo::autograd
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/stackref_bridge.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/stackref_bridge.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ad3a6390e68e35fdbbce7ce71b86169113094da
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/stackref_bridge.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/python_compat.h>
+
+#if IS_PYTHON_3_14_PLUS
+
+#ifdef __cplusplus
+extern "C" {
+#endif // __cplusplus
+
+// Use a void* to avoid exposing the internal _PyStackRef union on this
+// translation unit
+PyObject* THP_PyStackRef_AsPyObjectBorrow(void* stackref);
+
+#ifdef __cplusplus
+}
+#endif // __cplusplus
+#endif // IS_PYTHON_3_14_PLUS
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..3a6063bf47bf724386c321c87fe1baad658fada4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/dynamo/utils.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/python_headers.h>
+// C2039 MSVC
+#include <pybind11/complex.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include <Python.h>
+// The visibility attribute is to avoid a warning about storing a field in the
+// struct that has a different visibility (from pybind) than the struct.
+#ifdef _WIN32
+#define VISIBILITY_HIDDEN
+#else
+#define VISIBILITY_HIDDEN __attribute__((visibility("hidden")))
+#endif
+
+namespace torch::dynamo {
+PyObject* torch_c_dynamo_utils_init();
+} // namespace torch::dynamo
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/example_upgraders.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/example_upgraders.h
new file mode 100644
index 0000000000000000000000000000000000000000..1fb15f1f81f0c047e26443519d168d853dea5751
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/example_upgraders.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::_export {
+
+/// Register example upgraders for the upgrader system for testing.
+/// This function demonstrates common upgrade patterns and is primarily
+/// used for testing and demonstration purposes.
+void registerExampleUpgraders();
+
+/// Deregister example upgraders for the upgrader system for testing.
+/// This function cleans up the example upgraders that were registered
+/// by registerExampleUpgraders().
+void deregisterExampleUpgraders();
+
+} // namespace torch::_export
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/pt2_archive_constants.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/pt2_archive_constants.h
new file mode 100644
index 0000000000000000000000000000000000000000..c7a3894e7bd7e894828b5738e77b0003b35cb448
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/pt2_archive_constants.h
@@ -0,0 +1,78 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <string_view>
+
+namespace torch::_export::archive_spec {
+
+#define FORALL_CONSTANTS(DO)                                                   \
+  DO(ARCHIVE_ROOT_NAME, "package")                                             \
+  /* Archive format */                                                         \
+  DO(ARCHIVE_FORMAT_PATH, "archive_format")                                    \
+  DO(ARCHIVE_FORMAT_VALUE, "pt2")                                              \
+  /* Archive version */                                                        \
+  DO(ARCHIVE_VERSION_PATH, "archive_version")                                  \
+  DO(ARCHIVE_VERSION_VALUE, "0") /* Sep.4.2024: This is the initial version of \
+                                    the PT2 Archive Spec */                    \
+  /*                                                                           \
+   * ######## Note on updating ARCHIVE_VERSION_VALUE ########                  \
+   * When there is a BC breaking change to the PT2 Archive Spec,               \
+   * e.g. deleting a folder, or changing the naming convention of the          \
+   * following fields it would require bumping the ARCHIVE_VERSION_VALUE       \
+   * Archive reader would need corresponding changes to support loading both   \
+   * the current and older versions of the PT2 Archive.                        \
+   */                                                                          \
+  /* Model definitions */                                                      \
+  DO(MODELS_DIR, "models/")                                                    \
+  DO(MODELS_FILENAME_FORMAT, "models/{}.json") /* {model_name} */              \
+  /* AOTInductor artifacts */                                                  \
+  DO(AOTINDUCTOR_DIR, "data/aotinductor/")                                     \
+  /* MTIA artifacts */                                                         \
+  DO(MTIA_DIR, "data/mtia")                                                    \
+  /* weights, including parameters and buffers */                              \
+  DO(WEIGHTS_DIR, "data/weights/")                                             \
+  DO(WEIGHT_FILENAME_PREFIX, "weight_")                                        \
+  DO(WEIGHTS_PARAM_CONFIG_FORMAT, "data/weights/{}_model_param_config.json")   \
+  DO(WEIGHTS_CONFIG_FILENAME_FORMAT, "data/weights/{}_weights_config.json")    \
+  /* constants, including tensor_constants, non-persistent buffers and script  \
+   * objects */                                                                \
+  DO(CONSTANTS_DIR, "data/constants/")                                         \
+  DO(CONSTANTS_PARAM_CONFIG_FORMAT,                                            \
+     "data/constants/{}_model_constants_config.json")                          \
+  DO(CONSTANTS_CONFIG_FILENAME_FORMAT,                                         \
+     "data/constants/{}_constants_config.json")                                \
+  DO(TENSOR_CONSTANT_FILENAME_PREFIX, "tensor_")                               \
+  DO(CUSTOM_OBJ_FILENAME_PREFIX, "custom_obj_")                                \
+  /* example inputs */                                                         \
+  DO(SAMPLE_INPUTS_DIR, "data/sample_inputs/")                                 \
+  DO(SAMPLE_INPUTS_FILENAME_FORMAT,                                            \
+     "data/sample_inputs/{}.pt") /* {model_name} */                            \
+  DO(TS_SAMPLE_INPUTS_FILENAME_FORMAT,                                         \
+     "extra/{}.forward.sample_input.pt") /* {model_name} */                    \
+  /* ExecuTorch artifacts, including PTE files */                              \
+  DO(EXECUTORCH_DIR, "data/executorch/")                                       \
+  /* extra folder */                                                           \
+  DO(EXTRA_DIR, "extra/")                                                      \
+  DO(MODULE_INFO_PATH, "extra/module_info.json")                               \
+  /* xl_model_weights, this folder is used for storing per-feature-weights for \
+   * remote net data in this folder is consume by Predictor, and is not        \
+   * intended to be used by Sigmoid */                                         \
+  DO(XL_MODEL_WEIGHTS_DIR, "xl_model_weights/")                                \
+  DO(XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH, "xl_model_weights/model_param_config")
+
+#define DEFINE_GLOBAL(NAME, VALUE) \
+  inline constexpr std::string_view NAME = VALUE;
+FORALL_CONSTANTS(DEFINE_GLOBAL)
+#undef DEFINE_GLOBAL
+
+#define DEFINE_ENTRY(NAME, VALUE) std::pair(#NAME, VALUE),
+inline constexpr std::array kAllConstants{FORALL_CONSTANTS(DEFINE_ENTRY)};
+#undef DEFINE_ENTRY
+
+#undef FORALL_CONSTANTS
+} // namespace torch::_export::archive_spec
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/pybind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/pybind.h
new file mode 100644
index 0000000000000000000000000000000000000000..a621d3cc80658866a78f6f44a0aecef59c6dabf2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/pybind.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/python_headers.h>
+
+namespace torch::_export {
+
+void initExportBindings(PyObject* module);
+
+} // namespace torch::_export
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/upgrader.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/upgrader.h
new file mode 100644
index 0000000000000000000000000000000000000000..d0c9d1f72fdb1c2a56014351d011f60fe6af2cec
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/export/upgrader.h
@@ -0,0 +1,124 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <nlohmann/json.hpp>
+#include <functional>
+#include <string>
+#include <vector>
+
+namespace torch::_export {
+
+/// Function type for upgrading JSON fields during schema version migration.
+/// Takes a JSON field and returns the upgraded version of that field.
+using UpgraderFunction = std::function<nlohmann::json(const nlohmann::json&)>;
+
+/// Structure containing upgrader information for a specific keypath.
+/// The version is stored as the map key in the registry, so it's not
+/// duplicated here.
+struct Upgrader {
+  /// Path to the field that should be upgraded (e.g., {"graph_module", "graph",
+  /// "nodes"}) Assuming top-level is a JSON object that represents
+  /// ExportedProgram
+  std::vector<std::string> keypath;
+
+  /// Function that performs the actual upgrade transformation
+  UpgraderFunction upgrade_func;
+
+  /// Constructor for creating an upgrader with keypath and function
+  Upgrader(std::vector<std::string> kp, UpgraderFunction func);
+
+  /// Comparator for maintaining bottom-up ordering in the registry.
+  /// Deeper keypaths are processed first to ensure safe upgrade application
+  /// without conflicts between parent and child field modifications.
+  bool operator<(const Upgrader& other) const;
+};
+
+/// Register an upgrader function for a specific schema version and keypath.
+///
+/// This function allows registration of custom upgrade logic that will be
+/// applied when upgrading artifacts from the specified version. Upgraders
+/// are applied in bottom-up order (deeper keypaths first) to prevent
+/// conflicts between parent and child field modifications.
+///
+/// @param version The schema version this upgrader applies to
+/// @param keypath The key path to the field that should be upgraded
+/// @param upgrade_func Function that performs the upgrade transformation
+void registerUpgrader(
+    int version,
+    const std::vector<std::string>& keypath,
+    const UpgraderFunction& upgrade_func);
+
+/// Register an upgrader function using dot-separated keypath notation.
+///
+/// Convenience overload that accepts dot-separated keypath strings for
+/// simpler syntax. For example: "graph_module.graph.nodes" instead of
+/// {"graph_module", "graph", "nodes"}.
+///
+/// @param version The schema version this upgrader applies to
+/// @param dot_keypath Dot-separated keypath string (e.g., "graph.nodes")
+/// @param upgrade_func Function that performs the upgrade transformation
+void registerUpgrader(
+    int version,
+    const std::string& dot_keypath,
+    const UpgraderFunction& upgrade_func);
+
+/// Deregister an upgrader function for a specific schema version and keypath.
+///
+/// This function allows removal of previously registered upgrade logic for
+/// the specified version and keypath. This is useful for testing scenarios
+/// where you need to clean up registered upgraders or modify upgrader
+/// behavior dynamically.
+///
+/// @param version The schema version to deregister the upgrader from
+/// @param keypath The key path to the field that should be deregistered
+/// @return true if an upgrader was found and removed, false otherwise
+bool deregisterUpgrader(int version, const std::vector<std::string>& keypath);
+
+/// Deregister an upgrader function using dot-separated keypath notation.
+///
+/// Convenience overload that accepts dot-separated keypath strings for
+/// simpler syntax. For example: "graph_module.graph.nodes" instead of
+/// {"graph_module", "graph", "nodes"}.
+///
+/// @param version The schema version to deregister the upgrader from
+/// @param dot_keypath Dot-separated keypath string (e.g., "graph.nodes")
+/// @return true if an upgrader was found and removed, false otherwise
+bool deregisterUpgrader(int version, const std::string& dot_keypath);
+
+/// Utility function for throwing consistent upgrader errors.
+///
+/// This function formats error messages in a standardized way for upgrader
+/// failures, including version information and optional problematic object
+/// details for debugging.
+///
+/// @param upgrader_name Name of the upgrader that failed
+/// @param from_version Source schema version being upgraded from
+/// @param error_message Descriptive error message
+/// @param problematic_object Optional JSON object that caused the error
+/// @throws std::runtime_error Always throws with formatted error message
+void throwUpgraderError(
+    const std::string& upgrader_name,
+    int from_version,
+    const std::string& error_message,
+    const nlohmann::json& problematic_object = nlohmann::json::object());
+
+/// Upgrade a JSON artifact to a specific target version with available
+/// upgraders until a target version is reached.
+///
+/// This handles major version upgrade only. For minor version upgrade,
+/// e.g. adding a new field with default value, it's automatically handled by
+/// the default constructor in generated_serialization_types.h.
+///
+/// @param artifact The JSON artifact to upgrade(passed by value: function
+/// operates on a local copy, original remains unmodified)
+/// @param target_version The target schema version to upgrade to
+/// @return The upgraded JSON artifact with updated schema version
+/// @throws std::runtime_error if artifact is missing schema_version field
+/// @throws std::runtime_error if final version doesn't match target version
+nlohmann::json upgrade(nlohmann::json artifact, int target_version);
+
+} // namespace torch::_export
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/functionalization/Module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/functionalization/Module.h
new file mode 100644
index 0000000000000000000000000000000000000000..8d9ef7eee21ea132de0b5ce75a25345bdfc2539f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/functionalization/Module.h
@@ -0,0 +1,41 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/FunctionalStorageImpl.h>
+
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::functionalization {
+
+// Creates the default bindings for `ViewMeta` specializations.
+//
+// Defines a constructor using the types in `SerializableTuple`, as well
+// as pickle methods.
+template <class T>
+void create_binding_with_pickle(py::module m) {
+  py::class_<T, std::shared_ptr<T>, at::functionalization::ViewMeta>(
+      m, T::name())
+      .def(py::init<typename T::SerializableTuple>())
+      .def(
+          "as_tuple",
+          [](const std::shared_ptr<T>& meta) {
+            return meta->to_serializable_tuple();
+          })
+      .def(py::pickle(
+          [](const std::shared_ptr<T>& meta) {
+            return meta->to_serializable_tuple();
+          },
+          [](const typename T::SerializableTuple& tpl) {
+            return std::make_shared<T>(tpl);
+          }));
+}
+
+void initModule(PyObject* module);
+void initGenerated(PyObject* module);
+
+} // namespace torch::functionalization
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/functorch/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/functorch/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..e92e68fc321237beb24195636ff4f689105e5733
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/functorch/init.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <Python.h>
+
+namespace torch::functorch::impl {
+
+void initFuncTorchBindings(PyObject* module);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/fx/node.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/fx/node.h
new file mode 100644
index 0000000000000000000000000000000000000000..a4d3f2d4fdcb780ca683a63b5a384c1f6aafc93c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/fx/node.h
@@ -0,0 +1,11 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+bool NodeBase_init(PyObject* module);
+bool NodeIter_init(PyObject* module);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_eager/kernel_holder.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_eager/kernel_holder.h
new file mode 100644
index 0000000000000000000000000000000000000000..fa5b31156b755524c0bc3d3ceb1bf4b136ea2145
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_eager/kernel_holder.h
@@ -0,0 +1,131 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/boxing/KernelFunction.h>
+#include <ATen/core/function_schema.h>
+
+#include <torch/csrc/dynamo/guards.h>
+#include <torch/csrc/inductor/aoti_eager/kernel_meta_info.h>
+#include <torch/csrc/inductor/aoti_runner/model_container_runner.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include <string>
+
+namespace torch::inductor {
+
+// Represent AOTI kernel. It contains all the parameter metadata of the kernel
+// and the AOTI model runner.
+struct AOTIKernelMetadata {
+  // Represent all the parameters of AOTI kernel
+  std::vector<ParameterMetadata> parameter_metadata_list_;
+  // AOTI model runner to run the AOTI kernel
+  std::shared_ptr<AOTIModelContainerRunner> kernel_runner_;
+  // Whether this kernel was compiled with dynamic shapes. When true, cache
+  // matching skips exact size/stride comparison on tensors and matches by
+  // dtype/device/rank instead, allowing a single compiled kernel to serve
+  // multiple input shapes.
+  bool is_dynamic_{false};
+  AOTIKernelMetadata() : kernel_runner_(nullptr) {}
+
+  // Check whether the given parameter metadata list is the same as the
+  // parameter metadata list of the AOTI kernel.
+  bool check(
+      const std::vector<ParameterMetadata>& parameter_metadata_list) const {
+    if (parameter_metadata_list_.size() != parameter_metadata_list.size()) {
+      return false;
+    }
+
+    for (size_t i = 0; i < parameter_metadata_list_.size(); ++i) {
+      if (is_dynamic_) {
+        // Dynamic shapes: match by type/dtype/device/rank, skip exact sizes
+        if (!parameter_metadata_list_[i].dynamic_check(
+                parameter_metadata_list[i])) {
+          return false;
+        }
+      } else {
+        if (!(parameter_metadata_list_[i] == parameter_metadata_list[i])) {
+          return false;
+        }
+      }
+    }
+
+    return true;
+  }
+};
+
+// The AOTIPythonKernelHolder class uses the AOT Inductor to generate a kernel
+// for a specified operation. To speed up this process, the generated kernel
+// library is cached on disk. Detailed information from the input tensors is
+// used as the key for caching the kernel library. On subsequent runs, these
+// input tensors are used to search the cache. If a cache hit occurs, the cached
+// kernel library is loaded and executed. If a cache miss occurs, the AOT
+// Inductor is called again to generate the kernel library.
+class AOTIPythonKernelHolder : public c10::OperatorKernel {
+  // A DispatchKey object that represents the dispatch key for the kernel.
+  c10::DispatchKey dispatch_key_;
+  // Namespace of the kernel.
+  std::string ns_;
+  // Name of the operation the kernel performs.
+  std::string op_name_with_overload_;
+  // The device on which the kernel is to be executed.
+  c10::Device device_;
+  // The Python interpreter to get OpOverload object with the given op_name and
+  // op_overload_name.
+  c10::impl::PyInterpreter* pyinterpreter_;
+  // Cache the produced kernels by AOTI and its metadata
+  std::vector<AOTIKernelMetadata> aoti_kernel_cache_;
+  // Whether to compile with dynamic shapes support
+  bool dynamic_;
+
+ public:
+  AOTIPythonKernelHolder(
+      c10::DispatchKey dispatch_key,
+      std::string_view ns,
+      std::string_view op_name_with_overload,
+      bool dynamic = false);
+
+  void operator()(
+      const c10::OperatorHandle& op,
+      c10::DispatchKeySet keyset,
+      torch::jit::Stack* stack);
+
+ private:
+  bool cache_lookup(
+      const c10::OperatorHandle& op,
+      const c10::DispatchKeySet& keyset,
+      const torch::jit::Stack* stack,
+      AOTIKernelMetadata& aoti_kernel_metadata);
+  void cache_miss(
+      const c10::OperatorHandle& op,
+      const c10::DispatchKeySet& keyset,
+      torch::jit::Stack* stack);
+  void cache_hit(
+      const AOTIKernelMetadata& aoti_kernel_metadata,
+      const c10::OperatorHandle& op,
+      const c10::DispatchKeySet& keyset,
+      torch::jit::Stack* stack);
+  // Invoke python utility function on the Inductor side to produce AOTI kernel
+  // for the given operation.
+  //   Inductor utility function -
+  //   torch._inductor.utils.aoti_compile_with_persistent_cache
+  std::string produce_aoti_kernel_lib(
+      const c10::OperatorHandle& op,
+      const c10::DispatchKeySet& keyset,
+      const torch::jit::Stack* stack);
+  // Invoke python utility function on the Inductor side to load AOTI kernel for
+  // the given operation.
+  //   Inductor utility function - torch._inductor.utils.load_aoti_eager_cache
+  void init_aoti_kernel_cache();
+  // Load the AOTIModelContainerRunner object from the given file path.
+  std::shared_ptr<AOTIModelContainerRunner> load_aoti_model_runner(
+      const std::string& /*so_path*/);
+};
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_eager/kernel_meta_info.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_eager/kernel_meta_info.h
new file mode 100644
index 0000000000000000000000000000000000000000..64fd695595c4e0f36c88972ae97f455a8b27be71
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_eager/kernel_meta_info.h
@@ -0,0 +1,156 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/core/SymIntArrayRef.h>
+#include <torch/csrc/dynamo/guards.h>
+
+#include <string>
+
+namespace torch::inductor {
+
+// Regarding a aten operation implemented by AOTI, the metadata of the input
+// tensors will be cached on the disk to accelerate next run. TensorMetadata
+// structure is to represent the metadata of each input tensor. It includes
+// whether the tensor is symbolic, the dtype, the device, the sizes and the
+// strides of the tensor. When the metadata of the input tensors is the same as
+// the cached metadata, the cached kernel library will be loaded and executed.
+// Otherwise, the AOT Inductor will be called again to generate the kernel
+// library.
+// Beyond the TensorMetadata, we build guard/TensorCheck for each input tensor
+// as well to support symbolic shape. We intend to utilize TensorCheck to find
+// out the proper kernel rather than TensorMetadata comparison. Suppose an
+// operation with a single input tensor and two kernels:
+//   kernel1: TensorMetadata(is_symbolic=false, dtype=Float, device=CPU,
+//   sizes=[s0, s1, s2], strides=[s1 * s2, s2, 1]) kernel2:
+//   TensorMetadata(is_symbolic=false, dtype=Float, device=CPU, sizes=[3, s1,
+//   s2], strides=[s1 * s2, s2, 1])
+// If a tensor with sizes=[3, 4, 5] is passed to the operation, both kernel1 and
+// kernel2 support the tensor shape. In this case, we need to use TensorCheck
+// plus some heuristic rules to find out the proper kernel.
+struct TensorMetadata {
+  // Indicate whether the tensor is symbolic and it may be concluded by sizes_
+  // and strides_ in the future.
+  bool is_symbolic_;
+  // Dtype of a tensor(For scalar, we will wrap it as a scalar tensor)
+  c10::ScalarType dtype_ = c10::ScalarType::Undefined;
+  // Device of a tensor.
+  c10::Device device_;
+  // Dispatch key set of a tensor
+  c10::DispatchKeySet dispatch_key_set_;
+  // Sizes of a tensor. Currently, we only support static shape and use int64_t
+  // to represent the sizes. In the future, we will create symbolic size and use
+  // SymInt to represent it to support symbolic shape.
+  std::vector<int64_t> sizes_;
+  // Strides of a tensor. For symbolic shape support, it is the same as sizes_
+  std::vector<int64_t> strides_;
+  // requires grad
+  bool requires_grad_ = false;
+  // TensorCheck for the tensor
+  std::optional<dynamo::TensorCheck> tensor_check_;
+
+  TensorMetadata()
+      : is_symbolic_(false),
+        device_(c10::DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES),
+        sizes_({}),
+        strides_({}) {}
+  TensorMetadata(const at::Tensor& src_tensor);
+  TensorMetadata(
+      bool is_symbolic,
+      c10::ScalarType dtype,
+      c10::Device device,
+      c10::DispatchKeySet dispatch_key_set,
+      std::vector<int64_t> sizes,
+      std::vector<int64_t> strides,
+      bool requires_grad = false);
+
+  // Build TensorCheck for the tensor by using the data fields in TensorMetadata
+  void build_guard(const dynamo::LocalState& local_state);
+
+  // Compare two TensorMetadata objects
+  bool operator==(const TensorMetadata& other) const;
+
+  // Dynamic-shape-aware comparison: matches by dtype/device/rank but
+  // skips exact sizes/strides comparison.
+  bool dynamic_check(const TensorMetadata& other) const;
+};
+
+// ParameterTag is to represent the type of the input parameters of a aten
+// operation. Currently, we support the following types:
+//   1. TENSOR: a single tensor
+//   2. TENSOR_OPTIONAL: a single optional tensor
+//   3. TENSOR_LIST: a list of tensors
+//   4. TENSOR_LIST_OPTIONAL: a list of optional tensors
+//   5. SCALAR: a scalar value
+// If we need to support more types in the future, we will add more types in the
+// ParameterTag enum. For example, we will extend the enum to support string,
+// Dimname and so on to support more types of input parameters of aten
+// operations.
+enum ParameterTag {
+  TENSOR,
+  TENSOR_OPTIONAL,
+  TENSOR_LIST,
+  TENSOR_LIST_OPTIONAL,
+  SCALAR,
+  STRING,
+  DEVICE,
+  INVALID,
+};
+
+// ParameterMetadataValue is to represent the value of the input parameters of a
+// aten operation.
+using ParameterMetadataValue = std::variant<
+    TensorMetadata,
+    std::vector<TensorMetadata>,
+    c10::Scalar,
+    std::string,
+    c10::Device>;
+
+// ParameterMetadata is to represent the metadata of the input parameters of a
+// aten operation. It includes the tag of the parameter, the value of the
+// parameter and the order of the parameter.
+struct ParameterMetadata {
+  // The tag of the parameter. It indicates the type of the parameter.
+  ParameterTag tag_;
+  // The value of the parameter. It can be a tensor, a list of tensors or a
+  // scalar.
+  ParameterMetadataValue value_;
+  // The order of the parameter is used to distinguish the parameters with the
+  // same tag. For example, an operation with two input tensors, the first
+  // tensor is a optional tensor and the second tensor is a tensor. The first
+  // tensor will have the order 0 and the second tensor will have the order 1.
+  uint64_t order_{};
+
+  ParameterMetadata() : tag_(INVALID) {}
+  ParameterMetadata(TensorMetadata tensor_metadata, uint64_t input_order);
+  ParameterMetadata(const at::Tensor& tensor, uint64_t input_order);
+  ParameterMetadata(
+      const std::vector<at::Tensor>& tensor_list,
+      uint64_t input_order);
+  ParameterMetadata(
+      const std::vector<TensorMetadata>& tensor_metadata_list,
+      uint64_t input_order);
+  ParameterMetadata(const c10::Scalar& scalar, uint64_t input_order);
+  ParameterMetadata(const std::string& string_value, uint64_t input_order);
+  ParameterMetadata(const c10::Device& device, uint64_t input_order);
+
+  bool operator==(const ParameterMetadata& other) const;
+
+  // Dynamic-shape-aware comparison: matches by type/dtype/device/rank but
+  // skips exact size/stride comparison for tensors, so a single compiled
+  // kernel can serve multiple input shapes.
+  bool dynamic_check(const ParameterMetadata& other) const;
+
+ private:
+  // Helper function to compare two ParameterMetadata objects with the same
+  // SCALAR tag.
+  bool equal_to(const c10::Scalar& scalar) const;
+};
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/array_ref.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/array_ref.h
new file mode 100644
index 0000000000000000000000000000000000000000..27cd706592b540eab7f1dab3c76afec0d2789440
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/array_ref.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_include/common.h>
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/thread_local.h>
+#include <torch/csrc/inductor/array_ref_impl.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/common.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/common.h
new file mode 100644
index 0000000000000000000000000000000000000000..a676e55c9d3b9844b801adbcba8ca3c7473fea95
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/common.h
@@ -0,0 +1,21 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <filesystem>
+#include <optional>
+
+#include <torch/csrc/inductor/aoti_runtime/interface.h>
+#include <torch/csrc/inductor/aoti_runtime/model.h>
+
+#include <c10/util/generic_math.h>
+#include <torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h>
+
+// Round up to the nearest multiple of 64
+[[maybe_unused]] inline int64_t align(int64_t nbytes) {
+  return (nbytes + 64 - 1) & -64;
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/cpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/cpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..08fe09380478cca15008bdbf611eb30606c4b908
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/cpu.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_include/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..b7e910bf59b308eb889ea81c56a60dbbad8359e3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/cuda.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_include/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/mps.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/mps.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ed9ce406262d8402fe0ebbdb16822256051896e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/mps.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_include/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/mps.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..59c681ea6a49947bb2b079d31a2669aec7b95ff6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_include/xpu.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_include/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_package/model_package_loader.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_package/model_package_loader.h
new file mode 100644
index 0000000000000000000000000000000000000000..07ea91b062aac3878b110440f58cd343b0d75b7d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_package/model_package_loader.h
@@ -0,0 +1,64 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+#pragma once
+
+#include <ATen/Tensor.h>
+#include <c10/core/Device.h>
+#include <torch/csrc/inductor/aoti_runner/model_container_runner.h>
+
+namespace torch::inductor {
+class TORCH_API AOTIModelPackageLoader {
+ public:
+  AOTIModelPackageLoader(
+      const std::string& model_package_path,
+      const std::string& model_name = "model",
+      const bool run_single_threaded = false,
+      const size_t num_runners = 1,
+      const c10::DeviceIndex device_index = -1);
+  ~AOTIModelPackageLoader();
+
+  AOTIModelContainerRunner* get_runner();
+  std::unordered_map<std::string, std::string> get_metadata();
+
+  std::vector<at::Tensor> run(
+      const std::vector<at::Tensor>& inputs,
+      void* stream_handle = nullptr);
+
+  // boxed_run will steal the ownership of the input tensors
+  std::vector<at::Tensor> boxed_run(
+      std::vector<at::Tensor>&& inputs,
+      void* stream_handle = nullptr);
+
+  std::vector<std::string> get_call_spec();
+  void load_constants(
+      std::unordered_map<std::string, at::Tensor>& constants_map,
+      bool use_inactive,
+      bool check_full_update,
+      bool user_managed = false);
+  std::vector<std::string> get_constant_fqns();
+
+  void update_constant_buffer(
+      std::unordered_map<std::string, at::Tensor>& tensor_map,
+      bool use_inactive,
+      bool validate_full_updates,
+      bool user_managed = false);
+
+  // Static function to load metadata directly from a model package
+  static std::unordered_map<std::string, std::string> load_metadata_from_package(
+      const std::string& model_package_path,
+      const std::string& model_name);
+
+ private:
+  std::string temp_dir_;
+  std::unique_ptr<AOTIModelContainerRunner> runner_;
+  std::unordered_map<std::string, std::string> metadata_;
+
+  void load_metadata(const std::string& cpp_filename);
+};
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_package/pybind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_package/pybind.h
new file mode 100644
index 0000000000000000000000000000000000000000..7c27556e79df2539098ffd878ea9a6aa44658fec
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_package/pybind.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/python_headers.h>
+
+namespace torch::inductor {
+
+void initAOTIPackageBindings(PyObject* module);
+
+} // namespace torch::inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner.h
new file mode 100644
index 0000000000000000000000000000000000000000..fb471fcc1c39891ed984d50a3ef8a35e52489a78
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner.h
@@ -0,0 +1,148 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+#pragma once
+
+#include <ATen/Tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/interface.h>
+#include <torch/csrc/inductor/aoti_torch/proxy_executor.h>
+
+// Forward declare DynamicLibrary
+namespace at {
+struct DynamicLibrary;
+}
+
+namespace torch::inductor {
+using TensorConstantMap = std::unordered_map<std::string, at::Tensor*>;
+
+class TORCH_API AOTIModelContainerRunner {
+ public:
+  AOTIModelContainerRunner(const AOTIModelContainerRunner& other) = delete;
+  AOTIModelContainerRunner(AOTIModelContainerRunner&& other) = delete;
+  AOTIModelContainerRunner& operator=(const AOTIModelContainerRunner& other) =
+      delete;
+  AOTIModelContainerRunner& operator=(AOTIModelContainerRunner&& other) =
+      delete;
+  virtual ~AOTIModelContainerRunner();
+
+  std::vector<at::Tensor> run(
+      const std::vector<at::Tensor>& inputs,
+      void* stream_handle = nullptr);
+
+  // boxed_run will steal the ownership of the input tensors
+  std::vector<at::Tensor> boxed_run(
+      std::vector<at::Tensor>&& inputs,
+      void* stream_handle = nullptr);
+
+  std::unordered_map<std::string, std::string> getConstantNamesToOriginalFQNs()
+      const;
+  std::unordered_map<std::string, int32_t> getConstantNamesToDtypes() const;
+
+  const std::unordered_map<std::string, at::Tensor> extract_constants_map(
+      bool use_inactive) const;
+  void update_inactive_constant_buffer(const TensorConstantMap& const_map);
+  void update_constant_buffer(
+      std::unordered_map<std::string, at::Tensor>& tensor_map,
+      bool use_inactive,
+      bool validate_full_updates,
+      bool user_managed = false);
+  void update_constant_buffer(
+      const TensorConstantMap& const_map,
+      bool use_inactive,
+      bool validate_full_updates,
+      bool user_managed = false);
+  void run_const_fold(
+      bool use_inactive,
+      AOTInductorStreamHandle cuda_stream_handle = nullptr);
+  void swap_constant_buffer();
+  void free_inactive_constant_buffer();
+  void update_constant_buffer_from_blob(const std::string& weights_path);
+
+  std::vector<std::string> get_call_spec();
+
+ protected:
+  AOTIModelContainerRunner(
+      const std::string& model_so_path,
+      size_t num_models,
+      const std::string& device_str,
+      const std::string& cubin_dir,
+      const bool run_single_threaded);
+
+  // Default constructor for custom device implementations that don't
+  // use .so files. Derived classes must override run_impl().
+  AOTIModelContainerRunner();
+
+  virtual std::vector<at::Tensor> run_impl(
+      std::vector<AtenTensorHandle>& input_handles,
+      void* stream_handle);
+
+  std::unique_ptr<at::DynamicLibrary> model_so_;
+  decltype(&AOTInductorModelContainerCreateWithDevice) create_func_{nullptr};
+  decltype(&AOTInductorModelContainerDelete) delete_func_{nullptr};
+  decltype(&AOTInductorModelContainerGetNumOutputs) get_num_outputs_func_{
+      nullptr};
+  decltype(&AOTInductorModelContainerRun) run_func_{nullptr};
+  decltype(&AOTInductorModelContainerGetNumConstants) get_num_constants_func_{
+      nullptr};
+  decltype(&AOTInductorModelContainerGetConstantName) get_constant_name_func_{
+      nullptr};
+  decltype(&AOTInductorModelContainerGetConstantOriginalFQN)
+      get_constant_original_fqn_func_{nullptr};
+  decltype(&AOTInductorModelContainerGetConstantDtype) get_constant_dtype_func_{
+      nullptr};
+  decltype(&AOTInductorModelContainerExtractConstantsMap)
+      extract_constants_map_func_{nullptr};
+  decltype(&AOTInductorModelContainerUpdateUserManagedConstantBuffer)
+      update_user_managed_constant_buffer_func_{nullptr};
+  decltype(&AOTInductorModelContainerUpdateConstantBuffer)
+      update_constant_buffer_func_{nullptr};
+  decltype(&AOTInductorModelContainerUpdateInactiveConstantBuffer)
+      update_inactive_constant_buffer_func_{nullptr};
+  decltype(&AOTInductorModelContainerRunConstantFolding) run_const_fold_func_{
+      nullptr};
+  decltype(&AOTInductorModelContainerSwapConstantBuffer)
+      swap_constant_buffer_func_{nullptr};
+  decltype(&AOTInductorModelContainerFreeInactiveConstantBuffer)
+      free_inactive_constant_buffer_func_{nullptr};
+  decltype(&AOTInductorModelContainerGetCallSpec) get_call_spec_func_{nullptr};
+  decltype(&AOTInductorModelContainerGetConstantsBlobSize)
+      get_constants_blob_size_func_{nullptr};
+  decltype(&AOTInductorModelUpdateConstantsFromBlob)
+      update_constants_from_blob_func_{nullptr};
+
+  AOTInductorModelContainerHandle container_handle_ = nullptr;
+
+  AOTIProxyExecutorHandle proxy_executor_handle_ = nullptr;
+
+ private:
+  std::unique_ptr<torch::aot_inductor::ProxyExecutor> proxy_executor_;
+};
+
+using CreateAOTIModelRunnerFunc = std::unique_ptr<AOTIModelContainerRunner> (*)(
+    const std::string& model_so_path,
+    size_t num_models,
+    const std::string& device_str,
+    const std::string& bin_dir,
+    const bool run_single_threaded);
+
+// Return a global map "device name" -> "aoti model runner create function" for
+// all registered in AOTI external backends
+TORCH_API std::unordered_map<std::string, CreateAOTIModelRunnerFunc>&
+getAOTIModelRunnerRegistry();
+
+// To register a new external backend in AOTI one needs to create an instance of
+// this struct. It is not thread-safe. Because it is expected to be called
+// during the initialization of the program.
+struct TORCH_API RegisterAOTIModelRunner{RegisterAOTIModelRunner(
+    const std::string& name,
+    CreateAOTIModelRunnerFunc create_aoti_model_runner_fn){
+    getAOTIModelRunnerRegistry()[name] = create_aoti_model_runner_fn;
+} // namespace torch::inductor
+}
+;
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..83eecc1c684973fc4c30e5409ee9b4d48f61cf9a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runner/model_container_runner.h>
+
+namespace torch::inductor {
+class TORCH_API AOTIModelContainerRunnerCpu : public AOTIModelContainerRunner {
+ public:
+  AOTIModelContainerRunnerCpu(
+      const std::string& model_so_path,
+      size_t num_models = 1,
+      const bool run_single_threaded = false);
+
+  ~AOTIModelContainerRunnerCpu() override;
+};
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..dfdfdaf735409d2f5901b95b13678cd89ef4881c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h
@@ -0,0 +1,40 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/csrc/inductor/aoti_runner/model_container_runner.h>
+
+namespace torch::inductor {
+
+// NOTICE: Following APIs are subject to change due to active development
+// We provide NO BC guarantee for these APIs
+// NOLINTNEXTLINE(cppcoreguidelines-special-member-functions)
+class TORCH_CUDA_CPP_API AOTIModelContainerRunnerCuda
+    : public AOTIModelContainerRunner {
+ public:
+  // @param device_str: cuda device string, e.g. "cuda", "cuda:0"
+  AOTIModelContainerRunnerCuda(
+      const std::string& model_so_path,
+      size_t num_models = 1,
+      const std::string& device_str = "cuda",
+      const std::string& cubin_dir = "",
+      const bool run_single_threaded = false);
+
+  ~AOTIModelContainerRunnerCuda() override;
+
+  std::vector<at::Tensor> run_impl(
+      std::vector<AtenTensorHandle>& input_handles,
+      void* stream_handle) override;
+
+  std::vector<at::Tensor> run_with_cuda_stream(
+      const std::vector<at::Tensor>& inputs,
+      const at::cuda::CUDAStream& cuda_stream);
+};
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_mps.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_mps.h
new file mode 100644
index 0000000000000000000000000000000000000000..2e36e600a6f575e2d8b05cbfdde89fd6837a87d2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_mps.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if defined(__APPLE__)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runner/model_container_runner.h>
+
+namespace torch::inductor {
+class TORCH_API AOTIModelContainerRunnerMps : public AOTIModelContainerRunner {
+ public:
+  AOTIModelContainerRunnerMps(
+      const std::string& model_so_path,
+      size_t num_models = 1,
+      const bool run_single_threaded = false);
+
+  ~AOTIModelContainerRunnerMps() override;
+};
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..2a55999f70b8be633b9d90237b18ad8c6502b349
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_xpu.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#if !defined(C10_MOBILE) && !defined(ANDROID)
+#pragma once
+
+#include <c10/xpu/XPUStream.h>
+#include <torch/csrc/inductor/aoti_runner/model_container_runner.h>
+
+namespace torch::inductor {
+
+// NOTICE: Following APIs are subject to change due to active development
+// We provide NO BC guarantee for these APIs
+
+// HERE we use C10_EXPORT because libtorch_python needs this Symbol be exported.
+// And `TORCH_API and `TORCH_XPU_API`` do not export the symbol in Windows
+// build.
+class C10_EXPORT AOTIModelContainerRunnerXpu : public AOTIModelContainerRunner {
+ public:
+  // @param device_str: xpu device string, e.g. "xpu", "xpu:0"
+  AOTIModelContainerRunnerXpu(
+      const std::string& model_so_path,
+      size_t num_models = 1,
+      const std::string& device_str = "xpu",
+      const std::string& kernel_bin_dir = "",
+      const bool run_single_threaded = false);
+
+  ~AOTIModelContainerRunnerXpu() override;
+
+  std::vector<at::Tensor> run_impl(
+      std::vector<AtenTensorHandle>& input_handles,
+      void* stream_handle) override;
+
+  std::vector<at::Tensor> run_with_xpu_stream(
+      const std::vector<at::Tensor>& inputs,
+      const at::xpu::XPUStream& xpu_stream);
+};
+
+} // namespace torch::inductor
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/pybind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/pybind.h
new file mode 100644
index 0000000000000000000000000000000000000000..5a0eef2af2edaccf962ce0fc92de97761c2becf1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runner/pybind.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/python_headers.h>
+
+namespace torch::inductor {
+
+void initAOTIRunnerBindings(PyObject* module);
+
+} // namespace torch::inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..5080b9ee61f2fc57f1408a26d482c902d41bca17
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h
@@ -0,0 +1,249 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runtime/mini_array_ref.h>
+#include <torch/csrc/inductor/aoti_runtime/utils.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#include <cassert>
+#include <cstdint>
+#include <cstring>
+
+namespace torch::aot_inductor {
+
+using MiniIntArrayRef = MiniArrayRef<int64_t>;
+
+static_assert(
+    sizeof(MiniIntArrayRef) == sizeof(void*) + sizeof(size_t),
+    "changing the size of MiniArrayRef breaks ABI compatibility!");
+
+inline bool is_contiguous_strides_for_shape(
+    int64_t ndim,
+    const int64_t* strides_ptr,
+    const int64_t* sizes_ptr) {
+  int64_t z = 1;
+  for (int64_t d = ndim - 1; d >= 0; d--) {
+    const auto& size_d = sizes_ptr[d];
+    if (size_d != 1) {
+      if (strides_ptr[d] == z) {
+        z *= size_d;
+      } else {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+// Shim for AOTI generated code to pretend a raw array works like an
+// AtenTensorHandle.
+template <typename T>
+class ArrayRefTensor {
+ public:
+  using value_type = T;
+
+  ArrayRefTensor() = default;
+
+  explicit ArrayRefTensor(
+      MiniArrayRef<T> arr,
+      MiniArrayRef<const int64_t> sizes,
+      MiniArrayRef<const int64_t> strides,
+      int32_t device_type,
+      int32_t device_idx)
+      : arrayRef_(arr),
+        sizes_(sizes),
+        strides_(strides),
+        device_type_(device_type),
+        device_idx_(device_idx) {
+    assert(sizes.size() == strides.size());
+    assert(is_contiguous_strides_for_shape(
+        sizes.size(), strides.data(), sizes.data()));
+  }
+
+  AtenTensorHandle expensiveCopyToTensor() const {
+    AtenTensorHandle result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_empty_strided(
+        sizes_.size(),
+        sizes_.data(),
+        strides_.data(),
+        aoti_torch_dtype<std::remove_const_t<T>>(),
+        device_type_,
+        device_idx_,
+        &result));
+    void* dataPtr = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(result, &dataPtr));
+    std::memcpy(dataPtr, data(), numel() * sizeof(T));
+    return result;
+  }
+
+  // We need to look the same as RAIIAtenTensorHandle, which returns
+  // an owning AtenTensorHandle from release(). So, we allocate one!
+  AtenTensorHandle release() {
+    return expensiveCopyToTensor();
+  }
+
+  AtenTensorHandle borrowAsTensor() const {
+    AtenTensorHandle result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_v2(
+        data(),
+        sizes_.size(),
+        sizes_.data(),
+        strides_.data(),
+        0,
+        aoti_torch_dtype<std::remove_const_t<T>>(),
+        device_type_,
+        device_idx_,
+        &result,
+        aoti_torch_layout_strided(),
+        nullptr,
+        0));
+    return result;
+  }
+
+  // We don't need to free any memory.
+  void reset() {}
+
+  auto sizes() const {
+    return sizes_;
+  }
+
+  auto strides() const {
+    return strides_;
+  }
+
+  auto device_type() const {
+    return device_type_;
+  }
+
+  auto device_idx() const {
+    return device_idx_;
+  }
+
+  T* data() const {
+    return arrayRef_.data();
+  }
+
+  auto numel() const {
+    return arrayRef_.size();
+  }
+
+  void set_arrayref(MiniArrayRef<T> new_arrayref) {
+    arrayRef_ = new_arrayref;
+  }
+
+ private:
+  MiniArrayRef<T> arrayRef_;
+  // We expect generated code to have statically available sizes &
+  // strides for us.
+  MiniArrayRef<const int64_t> sizes_;
+  MiniArrayRef<const int64_t> strides_;
+  int32_t device_type_ = 0;
+  int32_t device_idx_ = 0;
+  // We continue to zero-initialize this field in case we repurpose
+  // the space later; having predictable contents can only help.
+  int32_t unusedDoNotRemoveForABICompatibility_ = 0;
+};
+
+static_assert(
+    sizeof(ArrayRefTensor<int>) ==
+        3 * sizeof(MiniIntArrayRef) + 3 * sizeof(int32_t) +
+            (alignof(ArrayRefTensor<int>) > 4 ? sizeof(int32_t) : 0),
+    "changing the size of ArrayRefTensor breaks ABI compatibility!");
+
+template <typename T>
+inline ArrayRefTensor<T> reinterpret_tensor_wrapper(
+    const ArrayRefTensor<T>& self,
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int64_t storage_offset) {
+  // REVIEW: we should add a way to build the DSO in debug mode during
+  // tests so we can have checks like this!
+  assert(is_contiguous_strides_for_shape(ndim, strides_ptr, sizes_ptr));
+  return ArrayRefTensor<T>(
+      MiniArrayRef<T>(
+          self.data() + storage_offset, self.numel() - storage_offset),
+      MiniArrayRef<const int64_t>(sizes_ptr, ndim),
+      MiniArrayRef<const int64_t>(strides_ptr, ndim),
+      self.device_type(),
+      self.device_idx());
+}
+
+template <typename T>
+inline T* get_data_ptr_wrapper(ArrayRefTensor<T>& tensor) {
+  return tensor.data();
+}
+
+template <typename T>
+inline T* get_data_ptr_wrapper(const MiniArrayRef<T>& arr) {
+  return arr.data();
+}
+
+template <typename T>
+inline const ArrayRefTensor<T>& unwrap_raii_handle_if_needed(
+    const ArrayRefTensor<T>& tensor) {
+  return tensor;
+}
+
+template <typename T>
+inline ArrayRefTensor<T>& unwrap_raii_handle_if_needed(
+    ArrayRefTensor<T>& tensor) {
+  return tensor;
+}
+
+template <typename T>
+inline const ArrayRefTensor<T>& wrap_with_raii_handle_if_needed(
+    const ArrayRefTensor<T>& tensor) {
+  return tensor;
+}
+
+template <typename T>
+inline ArrayRefTensor<T>& wrap_with_raii_handle_if_needed(
+    ArrayRefTensor<T>& tensor) {
+  return tensor;
+}
+
+template <typename T>
+inline ArrayRefTensor<T> wrap_with_raii_handle_if_needed(
+    ArrayRefTensor<T>&& tensor) {
+  return std::move(tensor);
+}
+
+template <typename T>
+inline RAIIAtenTensorHandle expensive_copy_to_tensor_if_needed(
+    const ArrayRefTensor<T>& tensor) {
+  return tensor.expensiveCopyToTensor();
+}
+
+inline AtenTensorHandle expensive_copy_to_tensor_if_needed(
+    AtenTensorHandle handle) {
+  return handle;
+}
+
+template <typename T>
+const T& copy_arrayref_tensor_to_tensor(const T& t) {
+  return t;
+}
+
+template <typename T>
+RAIIAtenTensorHandle copy_arrayref_tensor_to_tensor(
+    const ArrayRefTensor<T>& art) {
+  return art.expensiveCopyToTensor();
+}
+
+template <typename T>
+const T& borrow_arrayref_tensor_as_tensor(const T& t) {
+  return t;
+}
+
+template <typename T>
+RAIIAtenTensorHandle borrow_arrayref_tensor_as_tensor(
+    const ArrayRefTensor<T>& art) {
+  return art.borrowAsTensor();
+}
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h
new file mode 100644
index 0000000000000000000000000000000000000000..fe9957ff1ef6a530a63f8ae865362f9ccee57d27
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h
@@ -0,0 +1,90 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// Zero-copy conversion utilities between ArrayRefTensor<T> (C++ template) and
+// AOTInductorArrayRefTensor (plain C struct).
+//
+// These helpers allow the host process to marshal ArrayRefTensor objects into
+// the C-compatible AOTInductorArrayRefTensor descriptors before calling into a
+// DSO, and to unmarshal the descriptors back after the call.  Because only
+// C types cross the DSO boundary, the host and DSO can be linked against
+// different C++ standard libraries (e.g. libc++ vs libstdc++) without ABI
+// conflicts.
+//
+// IMPORTANT: Both sides share the same underlying data buffers -- no copies
+// are made.  The caller must ensure the data remains valid for the lifetime
+// of the descriptor.
+
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/interface.h>
+
+#include <cassert>
+#include <cstring>
+#include <stdexcept>
+#include <type_traits>
+
+namespace torch::aot_inductor {
+
+inline void validate_arrayref_tensor_ndim(int32_t ndim) {
+  if (ndim < 0 || ndim > AOTI_ARRAYREF_TENSOR_MAX_DIMS) {
+    throw std::runtime_error(
+        "AOTInductorArrayRefTensor ndim exceeds AOTI_ARRAYREF_TENSOR_MAX_DIMS");
+  }
+}
+
+// -------------------------------------------------------------------------
+// ArrayRefTensor<T> --> AOTInductorArrayRefTensor  (zero-copy)
+// -------------------------------------------------------------------------
+template <typename T>
+inline void arrayref_tensor_to_c(
+    const ArrayRefTensor<T>& src,
+    AOTInductorArrayRefTensor& dst) {
+  const auto sizes = src.sizes();
+  const auto strides = src.strides();
+  dst.data = const_cast<void*>(static_cast<const void*>(src.data()));
+  dst.numel = static_cast<int64_t>(src.numel());
+  dst.ndim = static_cast<int32_t>(sizes.size());
+  dst.dtype = aoti_torch_dtype<std::remove_const_t<T>>();
+  dst.device_type = src.device_type();
+  dst.device_idx = src.device_idx();
+
+  validate_arrayref_tensor_ndim(dst.ndim);
+  assert(dst.ndim <= AOTI_ARRAYREF_TENSOR_MAX_DIMS);
+  std::memcpy(dst.sizes, sizes.data(), dst.ndim * sizeof(int64_t));
+  std::memcpy(dst.strides, strides.data(), dst.ndim * sizeof(int64_t));
+  const int32_t remaining = AOTI_ARRAYREF_TENSOR_MAX_DIMS - dst.ndim;
+  std::memset(dst.sizes + dst.ndim, 0, remaining * sizeof(int64_t));
+  std::memset(dst.strides + dst.ndim, 0, remaining * sizeof(int64_t));
+  std::memset(dst.reserved, 0, sizeof(dst.reserved));
+}
+
+template <typename T>
+inline AOTInductorArrayRefTensor arrayref_tensor_to_c(
+    const ArrayRefTensor<T>& src) {
+  AOTInductorArrayRefTensor dst;
+  arrayref_tensor_to_c(src, dst);
+  return dst;
+}
+
+// -------------------------------------------------------------------------
+// AOTInductorArrayRefTensor --> ArrayRefTensor<T>  (zero-copy)
+// -------------------------------------------------------------------------
+template <typename T>
+inline ArrayRefTensor<T> c_to_arrayref_tensor(
+    const AOTInductorArrayRefTensor& src) {
+  validate_arrayref_tensor_ndim(src.ndim);
+  return ArrayRefTensor<T>(
+      MiniArrayRef<T>(
+          static_cast<T*>(const_cast<void*>(src.data)),
+          static_cast<size_t>(src.numel)),
+      MiniArrayRef<const int64_t>(src.sizes, static_cast<size_t>(src.ndim)),
+      MiniArrayRef<const int64_t>(src.strides, static_cast<size_t>(src.ndim)),
+      src.device_type,
+      src.device_idx);
+}
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/constant_type.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/constant_type.h
new file mode 100644
index 0000000000000000000000000000000000000000..8666e0556acbebc2e798e02d05b725be57a167ea
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/constant_type.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstdint>
+
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+
+namespace torch::aot_inductor {
+
+enum ConstantType : uint8_t {
+  Unknown = 0,
+  Parameter = 1,
+  Buffer = 2,
+  TensorConstant = 3,
+  FoldedConstant = 4,
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/device_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/device_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..1c57f319b594df4f4ad0135aab86eff3424244a9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/device_utils.h
@@ -0,0 +1,72 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+
+#ifdef USE_CUDA
+
+// FIXME: Currently, CPU and CUDA backend are mutually exclusive.
+// This is a temporary workaround. We need a better way to support
+// multi devices.
+
+#include <cuda.h>
+#include <cuda_runtime_api.h>
+
+#define AOTI_RUNTIME_CUDA_CHECK(EXPR)                      \
+  do {                                                     \
+    const cudaError_t code = EXPR;                         \
+    const char* msg = cudaGetErrorString(code);            \
+    if (code != cudaSuccess) {                             \
+      throw std::runtime_error(                            \
+          std::string("CUDA error: ") + std::string(msg)); \
+    }                                                      \
+  } while (0)
+
+namespace torch::aot_inductor {
+
+using DeviceStreamType = cudaStream_t;
+
+} // namespace torch::aot_inductor
+
+#elif defined(USE_XPU)
+#include <level_zero/ze_api.h>
+#include <sycl/sycl.hpp>
+#include <sstream>
+#define AOTI_RUNTIME_XPU_CHECK(EXPR)                                      \
+  do {                                                                    \
+    const ze_result_t status = EXPR;                                      \
+    if (status != ZE_RESULT_SUCCESS) {                                    \
+      std::stringstream ss;                                               \
+      ss << "L0 runtime error: " << std::hex << std::uppercase << status; \
+      throw std::runtime_error(ss.str());                                 \
+    }                                                                     \
+  } while (0)
+
+namespace torch::aot_inductor {
+
+using DeviceStreamType = sycl::queue*;
+
+} // namespace torch::aot_inductor
+
+#else
+
+#define AOTI_RUNTIME_CPU_CHECK(EXPR)               \
+  bool ok = EXPR;                                  \
+  if (!ok) {                                       \
+    throw std::runtime_error("CPU runtime error"); \
+  }
+
+namespace torch::aot_inductor {
+
+using DeviceStreamType = void*;
+
+} // namespace torch::aot_inductor
+
+#endif // USE_CUDA
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/interface.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/interface.h
new file mode 100644
index 0000000000000000000000000000000000000000..305261775da46f3225a59702dc2090f5fba40255
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/interface.h
@@ -0,0 +1,340 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+#include <torch/csrc/inductor/aoti_runtime/utils.h>
+
+#ifdef _WIN32
+/*
+On Windows, we need to explicit declaration for export APIs. And because the
+package loader call these API via GetProcAddress(ldsym on Linux), we can ignore
+the import case.
+*/
+#define AOTI_API __declspec(dllexport)
+#else
+#define AOTI_API __attribute__((__visibility__("default")))
+#endif
+
+extern "C" {
+struct AOTInductorModelOpaque;
+using AOTInductorModelHandle = AOTInductorModelOpaque*;
+
+struct AOTInductorModelContainerOpaque;
+using AOTInductorModelContainerHandle = AOTInductorModelContainerOpaque*;
+
+struct AOTInductorStreamOpaque;
+using AOTInductorStreamHandle = AOTInductorStreamOpaque*;
+
+struct AOTInductorConstantMap;
+using AOTInductorConstantMapHandle = AOTInductorConstantMap*;
+
+struct AOTInductorConstantMapEntry {
+  const char* name;
+  AtenTensorHandle handle;
+};
+
+// ---------------------------------------------------------------------------
+// C-compatible tensor descriptor for crossing the DSO boundary.
+//
+// This struct carries the same information as ArrayRefTensor<T> but uses only
+// C-compatible types so the host process and DSO can be built with different
+// C++ standard libraries (e.g. libc++ vs libstdc++).  All pointer fields
+// reference memory owned by the caller; no copies are made.
+//
+// Maximum supported number of dimensions.  8 covers all practical AOTI
+// models; tensors with more dims should fall back to the AtenTensorHandle
+// interface.
+// ---------------------------------------------------------------------------
+#define AOTI_ARRAYREF_TENSOR_MAX_DIMS 8
+
+struct AOTInductorArrayRefTensor {
+  // Pointer to the raw data buffer.  Not owned.
+  void* data;
+
+  // Number of elements in the data buffer (product of sizes for contiguous
+  // tensors).
+  int64_t numel;
+
+  // Static-size arrays for shape metadata.  Only the first `ndim` entries
+  // are meaningful.
+  int64_t sizes[AOTI_ARRAYREF_TENSOR_MAX_DIMS];
+  int64_t strides[AOTI_ARRAYREF_TENSOR_MAX_DIMS];
+
+  // Number of dimensions (0 <= ndim <= AOTI_ARRAYREF_TENSOR_MAX_DIMS).
+  int32_t ndim;
+
+  // Torch dtype encoded as int32_t (same encoding as aoti_torch_dtype_*()).
+  int32_t dtype;
+
+  // Device information.
+  int32_t device_type;
+  int32_t device_idx;
+
+  // Reserved for future extension.  Zero-initialize and do not read — a
+  // newer reader must tolerate zeros, and an older reader must ignore them.
+  int64_t reserved[4];
+};
+
+static_assert(
+    sizeof(AOTInductorArrayRefTensor) == 192,
+    "changing the size of AOTInductorArrayRefTensor breaks ABI compatibility!");
+
+// TODO: Deprecate this API. This was kept for BC compatibility.
+// Please use AOTInductorModelContainerCreateWithDevice instead.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerCreate(
+    AOTInductorModelContainerHandle* container_handle,
+    size_t num_models,
+    bool is_cpu,
+    const char* cubin_dir);
+
+// Creates an AOTInductor model container. The parameter num_models
+// specifies the number of model instances that may be run concurrently for
+// the same input model.
+// `device_str` MUST NOT be nullptr. It must be a valid device string, e.g.
+// "cpu", "cuda", "cuda:0", etc. If the device index is not specified for CUDA
+// device, runtime will use the device index returned by
+// "cudaGetDevice(&device_idx)"
+AOTI_API AOTIRuntimeError AOTInductorModelContainerCreateWithDevice(
+    AOTInductorModelContainerHandle* container_handle,
+    size_t num_models,
+    const char* device_str,
+    const char* cubin_dir);
+
+// Deletes the AOTInductor model container.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerDelete(
+    AOTInductorModelContainerHandle container_handle);
+
+// Runs the inference.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerRun(
+    AOTInductorModelContainerHandle container_handle,
+    AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
+                                     // are stolen; the array itself is borrowed
+    size_t num_inputs,
+    AtenTensorHandle*
+        output_handles, // array for writing output AtenTensorHandle; handles
+                        // will be stolen by the caller; the array itself is
+                        // borrowed
+    size_t num_outputs,
+    AOTInductorStreamHandle stream_handle,
+    AOTIProxyExecutorHandle proxy_executor_handle);
+
+// Single-threaded variant of previous.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerRunSingleThreaded(
+    AOTInductorModelContainerHandle container_handle,
+    AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
+                                     // are stolen; the array itself is borrowed
+    size_t num_inputs,
+    AtenTensorHandle*
+        output_handles, // array for writing output AtenTensorHandle; handles
+                        // will be stolen by the caller; the array itself is
+                        // borrowed
+    size_t num_outputs,
+    AOTInductorStreamHandle stream_handle,
+    AOTIProxyExecutorHandle proxy_executor_handle);
+
+// Retrieves the number of constants for the model.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetNumConstants(
+    AOTInductorModelContainerHandle container_handle,
+    size_t* num_constants);
+
+// Retrieves a constant's name.
+// idx is the index of the internal's constants.
+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetConstantName(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    const char** name);
+
+// Retrieves a constant's original FQN.
+// idx is the index of the internal's constants.
+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetConstantOriginalFQN(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    const char** original_fqn);
+
+// Retrieves whether a constant is from folded.
+// idx is the index of the internal's constants.
+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetConstantFromFolded(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    bool* from_folded);
+
+// Retrieves the inductor constant type.
+// idx is the index of the internal's constants.
+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetConstantType(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    int32_t* type);
+
+// Retrieves a constant's dtype.
+// idx is the index of the internal's constants.
+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetConstantDtype(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    int32_t* dtype);
+
+// Retrieves a constant's data size.
+// idx is the index of the internal's constants.
+// Need idx < num_constants from AOTInductorModelContainerGetNumConstants
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetConstantDataSize(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    size_t* data_size);
+
+// Extract the constants that is being used in the container.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerExtractConstantsMap(
+    AOTInductorModelContainerHandle container_handle,
+    AOTInductorConstantMapHandle constant_map_handle,
+    bool use_inactive);
+
+// Setup the constant buffer in model container with provided ConstantMap.
+// The ConstantMap is user managed, and the user would retain ownership.
+AOTI_API AOTIRuntimeError
+AOTInductorModelContainerUpdateUserManagedConstantBuffer(
+    AOTInductorModelContainerHandle container_handle,
+    AOTInductorConstantMapHandle constant_map_handle,
+    bool use_inactive,
+    bool validate_full_update);
+
+// Same as AOTInductorModelContainerUpdateUserManagedConstantBuffer,
+// but no std::unordered_map crosses DLL boundaries for cross-compilation.
+AOTI_API AOTIRuntimeError
+AOTInductorModelContainerUpdateUserManagedConstantBufferPairs(
+    AOTInductorModelContainerHandle container_handle,
+    const AOTInductorConstantMapEntry* pairs,
+    size_t num_pairs,
+    bool use_inactive,
+    bool validate_full_update);
+
+// Setup the constant buffer in model container with provided ConstantMap
+// use_inactive should be set as true if the inactive buffer is to be updated.
+// validate_full_update checks if all constants are included in the ConstantMap
+AOTI_API AOTIRuntimeError AOTInductorModelContainerUpdateConstantBuffer(
+    AOTInductorModelContainerHandle container_handle,
+    AOTInductorConstantMapHandle constant_map_handle,
+    bool use_inactive,
+    bool validate_full_update);
+
+// Setup the inactive constant buffer in model container with provided
+// ConstantMap
+AOTI_API AOTIRuntimeError AOTInductorModelContainerUpdateInactiveConstantBuffer(
+    AOTInductorModelContainerHandle container_handle,
+    AOTInductorConstantMapHandle constant_map_handle);
+
+// Free the inactive constant buffer in model container.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerFreeInactiveConstantBuffer(
+    AOTInductorModelContainerHandle container_handle);
+
+// Run constant folding on constant buffer.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerRunConstantFolding(
+    AOTInductorModelContainerHandle container_handle,
+    bool use_inactive,
+    AOTInductorStreamHandle stream_handle,
+    AOTIProxyExecutorHandle proxy_executor_handle);
+
+// Swap the constant buffer being used to the inactive one.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerSwapConstantBuffer(
+    AOTInductorModelContainerHandle container_handle);
+
+// Retrieves the number of inputs for the model.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetNumInputs(
+    AOTInductorModelContainerHandle container_handle,
+    size_t* ret_num_inputs);
+
+// Retrieves the input name at the given index.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetInputName(
+    AOTInductorModelContainerHandle container_handle,
+    size_t input_idx,
+    const char** ret_input_names);
+
+// Retrieves the number of outputs for the model.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetNumOutputs(
+    AOTInductorModelContainerHandle container_handle,
+    size_t* ret_num_outputs);
+
+// Retrieves the output name at the given index.
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetOutputName(
+    AOTInductorModelContainerHandle container_handle,
+    size_t output_idx,
+    const char** ret_output_names);
+
+// Creates an AOTInductorModel instance.  This is a thin and light wrapper
+// around the compiled model; it doesn't handle concurrency, queueing, device
+// management, etc.  Use this if bare-metal performance is needed and you are
+// willing to handle other "management" aspects yourself.
+//
+// constant_map_handle is an opaque type to satisfy the C ABI.  It should be a
+// std::unordered_map<std::string, at::Tensor*>*.
+AOTI_API AOTIRuntimeError AOTInductorModelCreate(
+    AOTInductorModelHandle* model_handle,
+    AOTInductorConstantMapHandle constant_map_handle);
+
+// Run an AOTInductorModel (see AOTInductorModelCreate for when one should use
+// this function versus AOTInductorModelContainerRun).
+AOTI_API AOTIRuntimeError AOTInductorModelRun(
+    AOTInductorModelHandle model_handle,
+    AtenTensorHandle* input_handles,
+    AtenTensorHandle* output_handles);
+
+// Replace AOTInductorModel's constant map. Note it doesn't handle concurrency
+// so be sure to handle ordering if AOTInductorModelRun is ran concurrently.
+AOTI_API AOTIRuntimeError AOTInductorModelUpdateConstantsMap(
+    AOTInductorModelHandle model_handle,
+    AOTInductorConstantMapHandle constant_map_handle);
+
+// Get the size of the constant blob
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetConstantsBlobSize(
+    AOTInductorModelContainerHandle container_handle,
+    uint64_t* ret_size);
+
+// Load weights from a single blob in weight_blob_ptr
+AOTI_API AOTIRuntimeError AOTInductorModelUpdateConstantsFromBlob(
+    AOTInductorModelContainerHandle container_handle,
+    const uint8_t* weight_blob_ptr);
+
+// Delete an AOTInductorModel created by AOTInductorModelCreate.
+AOTI_API AOTIRuntimeError
+AOTInductorModelDelete(AOTInductorModelHandle model_handle);
+
+AOTI_API AOTIRuntimeError AOTInductorModelGetNumOutputs(
+    AOTInductorModelHandle model_handle,
+    size_t* ret_num_outputs);
+
+AOTI_API AOTIRuntimeError AOTInductorModelContainerGetCallSpec(
+    AOTInductorModelContainerHandle container_handle,
+    const char** in_spec,
+    const char** out_spec);
+
+// ---------------------------------------------------------------------------
+// C-ABI-safe variant of AOTInductorModelRunMinimalArrayrefInterface.
+//
+// Instead of passing std::tuple<ArrayRefTensor<T>...>& (which encodes C++
+// standard library types into the ABI), this function accepts flat C arrays
+// of AOTInductorArrayRefTensor descriptors.  The descriptors reference the
+// same underlying data buffers -- no copies are made.
+//
+// The host process marshals its ArrayRefTensor<T> objects into
+// AOTInductorArrayRefTensor descriptors, calls into the DSO through this
+// pure-C interface, and then unmarshals the output descriptors back.
+// Because only C types cross the DSO boundary, the host and DSO can be
+// built with different C++ standard libraries (e.g. libc++ vs libstdc++).
+// ---------------------------------------------------------------------------
+AOTI_API AOTIRuntimeError AOTInductorModelRunMinimalArrayrefInterfaceV2(
+    AOTInductorModelHandle model_handle,
+    int32_t num_inputs,
+    const AOTInductorArrayRefTensor* inputs,
+    int32_t num_outputs,
+    AOTInductorArrayRefTensor* outputs);
+
+} // extern "C"
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h
new file mode 100644
index 0000000000000000000000000000000000000000..3bb92f6f84155ccb10fd19a18eb1af2a667e0c5c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h
@@ -0,0 +1,149 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstdio>
+#include <filesystem>
+#include <sstream>
+#include <string>
+#include <utility>
+
+namespace torch::aot_inductor {
+
+struct KernelContext {
+  std::string kernel_name;
+  std::string python_stack;
+  std::string compressed_python_stack;
+
+  KernelContext(std::string name, std::string stack)
+      : kernel_name(std::move(name)) {
+    python_stack = trim_stack(stack);
+    compressed_python_stack = compress_stack(python_stack);
+  }
+
+  KernelContext(const KernelContext&) = default;
+  KernelContext& operator=(const KernelContext&) = default;
+  KernelContext(KernelContext&&) = default;
+  KernelContext& operator=(KernelContext&&) = default;
+
+ private:
+  // Strip leading and trailing newlines from stack:
+  // - finds first and last non-newline characters
+  // - outputs substring between them (inclusive)
+  // - returns empty string if stack contains only newlines
+  static std::string trim_stack(const std::string& stack) {
+    std::string res_stack;
+    auto beg = stack.find_first_not_of('\n');
+    if (beg == std::string::npos) {
+      return res_stack;
+    }
+    auto end = stack.find_last_not_of('\n');
+    if (end == std::string::npos) {
+      res_stack = stack.substr(beg);
+    } else {
+      res_stack = stack.substr(beg, end - beg + 1);
+    }
+    return res_stack;
+  }
+
+  // Compress stack into compact format:
+  // - blank lines are treated as stack separators
+  // - lines with 0 or 2 spaces of indentation are parsed as filename, fileline
+  //   and function
+  // - lines with 4 spaces of indentation are parsed as the code snippet
+  // - outputs function[code], filename and fileline (at newline each)
+  // - returns empty string on any parse error
+  static std::string compress_stack(const std::string& stack) {
+    std::string res_stack;
+    namespace fs = std::filesystem;
+    char function[1025];
+    char filename[1025];
+    uint32_t fileline;
+    int ret, n, ws;
+    const char* p;
+    std::stringstream stream{stack};
+    std::string line;
+    std::string fmt = "File \"%1024[^\"]\", line %u, in %1024[^\n]\n%n";
+    while (std::getline(stream, line)) {
+      // check if new stack
+      if (line.empty()) {
+        res_stack += '\n';
+        continue;
+      }
+      p = line.c_str();
+      ws = 0;
+      while (*p == ' ') {
+        ++p;
+        ++ws;
+      }
+      // check if new file
+      if (ws != 0 && ws != 2) {
+        return {};
+      }
+      ret = sscanf(p, fmt.c_str(), filename, &fileline, function, &n);
+      if (ret != 3) {
+        return {};
+      }
+      if (!std::getline(stream, line)) {
+        return {};
+      }
+      p = line.c_str();
+      ws = 0;
+      while (*p == ' ') {
+        ++p;
+        ++ws;
+      }
+      // check if command
+      if (ws != 4) {
+        return {};
+      }
+      res_stack += std::string{function} + '[' + std::string{p} + ']';
+      res_stack += '\n';
+      res_stack += fs::path{filename}.filename();
+      res_stack += '\n';
+      res_stack += std::to_string(fileline);
+      res_stack += '\n';
+    }
+    return res_stack;
+  }
+};
+
+// Thread-local pointer
+extern thread_local KernelContext* tls_kernel_context;
+
+inline KernelContext* current_kernel_context() {
+  return tls_kernel_context;
+}
+
+inline void set_kernel_context(KernelContext* ctx) {
+  tls_kernel_context = ctx;
+}
+
+inline void clear_kernel_context() {
+  tls_kernel_context = nullptr;
+}
+
+struct KernelContextGuard {
+  KernelContextGuard(const std::string& name, const std::string& stack)
+      : owned_context_(name, stack) {
+    set_kernel_context(&owned_context_);
+  }
+  ~KernelContextGuard() {
+    clear_kernel_context();
+  }
+
+  // Delete copy constructor and copy assignment operator
+  KernelContextGuard(const KernelContextGuard&) = delete;
+  KernelContextGuard& operator=(const KernelContextGuard&) = delete;
+
+  KernelContextGuard(KernelContextGuard&&) = default;
+  KernelContextGuard& operator=(KernelContextGuard&&) = delete;
+
+ private:
+  KernelContext owned_context_;
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/mini_array_ref.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/mini_array_ref.h
new file mode 100644
index 0000000000000000000000000000000000000000..31cdf3063f928702619631d25997c1bd703e54db
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/mini_array_ref.h
@@ -0,0 +1,165 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <cassert>
+#include <cstdint>
+#include <cstring>
+#include <vector>
+
+namespace torch::aot_inductor {
+
+// Can't use c10::ArrayRef because it's not truly header-only and
+// pulls in other c10 headers. This is (sadly) copy-pasted and
+// adapted.
+template <typename T>
+class MiniArrayRef final {
+ public:
+  using iterator = T*;
+  using const_iterator = const T*;
+  using size_type = size_t;
+  using value_type = T;
+
+  using reverse_iterator = std::reverse_iterator<iterator>;
+
+ private:
+  /// The start of the array, in an external buffer.
+  T* Data;
+
+  /// The number of elements.
+  size_type Length;
+
+ public:
+  /// @name Constructors
+  /// @{
+
+  /// Construct an empty MiniArrayRef.
+  /* implicit */ constexpr MiniArrayRef() : Data(nullptr), Length(0) {}
+
+  /// Construct an MiniArrayRef from a single element.
+  // TODO Make this explicit
+  constexpr MiniArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {}
+
+  /// Construct an MiniArrayRef from a pointer and length.
+  constexpr MiniArrayRef(T* data, size_t length) : Data(data), Length(length) {}
+
+  /// Construct an MiniArrayRef from a range.
+  constexpr MiniArrayRef(T* begin, T* end) : Data(begin), Length(end - begin) {}
+
+  template <
+      typename Container,
+      typename = std::enable_if_t<std::is_same_v<
+          std::remove_const_t<decltype(std::declval<Container>().data())>,
+          T*>>>
+  /* implicit */ MiniArrayRef(Container& container)
+      : Data(container.data()), Length(container.size()) {}
+
+  /// Construct an MiniArrayRef from a std::vector.
+  // The enable_if stuff here makes sure that this isn't used for
+  // std::vector<bool>, because MiniArrayRef can't work on a std::vector<bool>
+  // bitfield.
+  template <typename A>
+  /* implicit */ MiniArrayRef(const std::vector<T, A>& Vec)
+      : Data(Vec.data()), Length(Vec.size()) {
+    static_assert(
+        !std::is_same_v<T, bool>,
+        "MiniArrayRef<bool> cannot be constructed from a std::vector<bool> bitfield.");
+  }
+
+  /// Construct an MiniArrayRef from a std::array
+  template <size_t N>
+  /* implicit */ constexpr MiniArrayRef(std::array<T, N>& Arr)
+      : Data(Arr.data()), Length(N) {}
+
+  /// Construct an MiniArrayRef from a C array.
+  template <size_t N>
+  // NOLINTNEXTLINE(*c-array*)
+  /* implicit */ constexpr MiniArrayRef(T (&Arr)[N]) : Data(Arr), Length(N) {}
+
+  // /// Construct an MiniArrayRef from an empty C array.
+  /* implicit */ constexpr MiniArrayRef(const volatile void* Arr)
+      : Data(nullptr), Length(0) {}
+
+  /// Construct an MiniArrayRef from a std::initializer_list.
+  /* implicit */ constexpr MiniArrayRef(const std::initializer_list<T>& Vec)
+      : Data(
+            std::begin(Vec) == std::end(Vec) ? static_cast<T*>(nullptr)
+                                             : std::begin(Vec)),
+        Length(Vec.size()) {}
+
+  /// @}
+  /// @name Simple Operations
+  /// @{
+
+  constexpr iterator begin() const {
+    return Data;
+  }
+  constexpr iterator end() const {
+    return Data + Length;
+  }
+
+  // These are actually the same as iterator, since MiniArrayRef only
+  // gives you const iterators.
+  constexpr const_iterator cbegin() const {
+    return Data;
+  }
+  constexpr const_iterator cend() const {
+    return Data + Length;
+  }
+
+  constexpr reverse_iterator rbegin() const {
+    return reverse_iterator(end());
+  }
+  constexpr reverse_iterator rend() const {
+    return reverse_iterator(begin());
+  }
+
+  /// empty - Check if the array is empty.
+  constexpr bool empty() const {
+    return Length == 0;
+  }
+
+  constexpr T* data() const {
+    return Data;
+  }
+
+  /// size - Get the array size.
+  constexpr size_t size() const {
+    return Length;
+  }
+
+  /// equals - Check for element-wise equality.
+  constexpr bool equals(MiniArrayRef RHS) const {
+    return Length == RHS.Length && std::equal(begin(), end(), RHS.begin());
+  }
+
+  /// @}
+  /// @name Operator Overloads
+  /// @{
+  constexpr const T& operator[](size_t Index) const {
+    return Data[Index];
+  }
+
+  /// Disallow accidental assignment from a temporary.
+  ///
+  /// The declaration here is extra complicated so that "arrayRef = {}"
+  /// continues to select the move assignment operator.
+  template <typename U>
+  std::enable_if_t<std::is_same_v<U, T>, MiniArrayRef<T>>& operator=(
+      // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
+      U&& Temporary) = delete;
+
+  /// Disallow accidental assignment from a temporary.
+  ///
+  /// The declaration here is extra complicated so that "arrayRef = {}"
+  /// continues to select the move assignment operator.
+  template <typename U>
+  std::enable_if_t<std::is_same_v<U, T>, MiniArrayRef<T>>& operator=(
+      std::initializer_list<U>) = delete;
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model.h
new file mode 100644
index 0000000000000000000000000000000000000000..eca53e62494e348a7c5f75684f650b71c45caa7f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model.h
@@ -0,0 +1,67 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+#include <torch/csrc/inductor/aoti_runtime/model_base.h>
+
+namespace torch::aot_inductor {
+
+class AOTInductorModel : public AOTInductorModelBase<AOTInductorModel> {
+ public:
+  AOTInductorModel(
+      std::shared_ptr<ConstantMap> constants_map,
+      std::shared_ptr<std::vector<ConstantHandle>> constants_array,
+      const std::string& device_str,
+      std::optional<std::string> cubin_dir);
+
+  std::unordered_map<std::string, AtenTensorHandle> const_run_impl(
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor,
+      bool initialization = false);
+
+  void _const_run_impl(
+      std::vector<AtenTensorHandle>& output_handles,
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor);
+
+  void run_impl(
+      AtenTensorHandle*
+          input_handles, // array of input AtenTensorHandle; handles
+                         // are stolen; the array itself is borrowed
+      AtenTensorHandle*
+          output_handles, // array for writing output AtenTensorHandle; handles
+                          // will be stolen by the caller; the array itself is
+                          // borrowed
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor);
+
+  template <typename Inputs, typename Outputs>
+  Outputs run_impl_minimal_arrayref_interface(
+      const Inputs& inputs,
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor);
+
+  static std::unique_ptr<AOTInductorModel> Create(
+      std::shared_ptr<ConstantMap> constants_map,
+      std::shared_ptr<std::vector<ConstantHandle>> constants_array,
+      const std::string& device_str,
+      std::optional<std::string> cubin_dir) {
+    return std::make_unique<AOTInductorModel>(
+        std::move(constants_map),
+        std::move(constants_array),
+        device_str,
+        std::move(cubin_dir));
+  }
+
+ private:
+  std::unique_ptr<AOTInductorModelKernelsBase> kernels_;
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model_base.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model_base.h
new file mode 100644
index 0000000000000000000000000000000000000000..12fef809d7eac0f1479206fb1bef92989bc90473
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model_base.h
@@ -0,0 +1,1126 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#ifdef _WIN32
+#include <windows.h>
+#include <functional> // std::function
+#ifdef USE_MMAP_SELF
+#include <errno.h>
+#include <fcntl.h>
+#include <io.h>
+#include <sys/stat.h>
+
+#define PROT_READ 0x1
+#define PROT_WRITE 0x2
+#define PROT_EXEC 0x4
+
+#define MAP_SHARED 0x01
+#define MAP_PRIVATE 0x02
+#define MAP_FAILED ((void*)-1)
+
+#define SEEK_SET 0
+#define SEEK_CUR 1
+#define SEEK_END 2
+
+struct Dl_info {
+  char dli_fname[MAX_PATH]; /**< Filename of defining object */
+  void* dli_fbase; /**< Load address of that object */
+  const char* dli_sname; /**< Name of nearest lower symbol */
+  void* dli_saddr; /**< Exact value of nearest symbol */
+};
+typedef struct Dl_info Dl_info;
+
+int dladdr(const void* addr, Dl_info* info) {
+  // only returns filename, FWIW.
+  CHAR tpath[MAX_PATH];
+  MEMORY_BASIC_INFORMATION mbi;
+  char* path;
+  char* tmp;
+  size_t length;
+  int ret = 0;
+
+  if (!info)
+    return 0;
+
+  HMODULE hModule;
+  if (!GetModuleHandleExA(
+          GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+              GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+          (LPCSTR)addr,
+          &hModule) ||
+      hModule == NULL)
+    return 0;
+
+  ret = GetModuleFileNameA(hModule, (LPSTR)&tpath, MAX_PATH);
+  if (!ret)
+    return 0;
+
+  path = tpath;
+
+  length = strlen(path);
+  if (length >= MAX_PATH) {
+    length = MAX_PATH - 1;
+    path[MAX_PATH - 1] = '\0';
+  }
+
+  tmp = path;
+  while (*tmp) {
+    if (*tmp == '\\')
+      *tmp = '/';
+    tmp++;
+  }
+
+  memcpy(info->dli_fname, path, length + 1);
+  info->dli_fbase = hModule;
+  info->dli_sname = NULL;
+  info->dli_saddr = NULL;
+  return 1;
+}
+
+static DWORD get_creation_disposition(int flags) {
+  if (flags & O_CREAT) {
+    if (flags & O_EXCL)
+      return CREATE_NEW;
+    if (flags & O_TRUNC)
+      return CREATE_ALWAYS;
+    return OPEN_ALWAYS;
+  }
+  if (flags & O_TRUNC)
+    return TRUNCATE_EXISTING;
+  return OPEN_EXISTING;
+}
+
+#define O_ACCMODE 03
+#define O_RDONLY 00
+#define O_WRONLY 01
+#define O_RDWR 02
+
+static DWORD get_access_mode(int flags) {
+  switch (flags & O_ACCMODE) {
+    case O_RDONLY:
+      return GENERIC_READ;
+    case O_WRONLY:
+      return GENERIC_WRITE;
+    case O_RDWR:
+      return GENERIC_READ | GENERIC_WRITE;
+    default:
+      return GENERIC_READ;
+  }
+}
+#ifndef O_DSYNC
+#define O_DSYNC 00010000 /* used to be O_SYNC, see below */
+#endif
+
+#ifndef O_SYNC
+#define __O_SYNC 04000000
+#define O_SYNC (__O_SYNC | O_DSYNC)
+#endif
+
+int open(char* pathname, int flags) {
+  DWORD dwDesiredAccess = get_access_mode(flags);
+  DWORD dwCreationDisposition = get_creation_disposition(flags);
+  DWORD dwShareMode = FILE_SHARE_READ | FILE_SHARE_WRITE;
+  DWORD dwFlagsAndAttributes = FILE_ATTRIBUTE_NORMAL;
+
+  if (flags & O_SYNC) {
+    dwFlagsAndAttributes |= FILE_FLAG_WRITE_THROUGH;
+  }
+
+  if (flags & O_SEQUENTIAL) {
+    dwFlagsAndAttributes |= FILE_FLAG_SEQUENTIAL_SCAN;
+  }
+
+  if (flags & O_RANDOM) {
+    dwFlagsAndAttributes |= FILE_FLAG_RANDOM_ACCESS;
+  }
+
+  HANDLE hFile = CreateFileA(
+      pathname,
+      dwDesiredAccess,
+      dwShareMode,
+      NULL,
+      dwCreationDisposition,
+      dwFlagsAndAttributes,
+      NULL);
+
+  if (hFile == INVALID_HANDLE_VALUE) {
+    switch (GetLastError()) {
+      case ERROR_FILE_NOT_FOUND:
+        errno = ENOENT;
+        break;
+      case ERROR_PATH_NOT_FOUND:
+        errno = ENOTDIR;
+        break;
+      case ERROR_ACCESS_DENIED:
+        errno = EACCES;
+        break;
+      case ERROR_FILE_EXISTS:
+        errno = EEXIST;
+        break;
+      case ERROR_TOO_MANY_OPEN_FILES:
+        errno = EMFILE;
+        break;
+      default:
+        errno = EIO;
+    }
+    return -1;
+  }
+
+  int fd = _open_osfhandle((intptr_t)hFile, flags);
+  if (fd == -1) {
+    CloseHandle(hFile);
+    errno = EMFILE;
+    return -1;
+  }
+
+  if (flags & O_APPEND) {
+    lseek(fd, 0, SEEK_END);
+  }
+
+  return fd;
+}
+
+int close(int fd) {
+  return _close(fd);
+}
+
+void* mmap(
+    void* addr,
+    size_t length,
+    int prot,
+    int flags,
+    int fd,
+    off_t offset) {
+  HANDLE hFile = (HANDLE)_get_osfhandle(fd);
+  if (hFile == INVALID_HANDLE_VALUE) {
+    errno = EBADF;
+    return MAP_FAILED;
+  }
+
+  DWORD flProtect;
+  if (prot & PROT_WRITE) {
+    flProtect = PAGE_READWRITE;
+  } else if (prot & PROT_READ) {
+    flProtect = PAGE_READONLY;
+  } else {
+    flProtect = PAGE_NOACCESS;
+  }
+
+  flProtect = PAGE_READONLY;
+
+  DWORD dwDesiredAccess = 0;
+  if (prot & PROT_READ)
+    dwDesiredAccess |= FILE_MAP_READ;
+  if (prot & PROT_WRITE)
+    dwDesiredAccess |= FILE_MAP_WRITE;
+  if (prot & PROT_EXEC)
+    dwDesiredAccess |= FILE_MAP_EXECUTE;
+
+  dwDesiredAccess = FILE_MAP_READ;
+
+  SYSTEM_INFO SysInfo;
+  GetSystemInfo(&SysInfo);
+  DWORD dwSysGran = SysInfo.dwAllocationGranularity;
+
+  DWORD dwFileMapStart = (offset / dwSysGran) * dwSysGran;
+  DWORD dwMapViewSize = (offset % dwSysGran) + length;
+  DWORD dwFileMapSize = offset + length;
+  int iViewDelta = offset - dwFileMapStart;
+
+  HANDLE hMapping =
+      CreateFileMapping(hFile, NULL, flProtect, 0, dwFileMapSize, NULL);
+
+  if (!hMapping) {
+    DWORD dwErrCode = GetLastError();
+    errno = EACCES;
+    return MAP_FAILED;
+  }
+
+  void* lpMapAddress = MapViewOfFileEx(
+      hMapping, dwDesiredAccess, 0, dwFileMapStart, dwMapViewSize, addr);
+  if (!lpMapAddress) {
+    DWORD dwErrCode = GetLastError();
+    errno = EINVAL;
+  }
+
+  void* pData = (char*)lpMapAddress + iViewDelta;
+
+  CloseHandle(hMapping);
+
+  if (!lpMapAddress) {
+    return MAP_FAILED;
+  }
+
+  return pData;
+}
+
+int munmap(void* addr, size_t length) {
+  if (!UnmapViewOfFile(addr)) {
+    errno = EINVAL;
+    return -1;
+  }
+  return 0;
+}
+#endif // USE_MMAP_SELF
+#else // !_WIN32
+#include <dlfcn.h>
+#include <sys/mman.h>
+#include <unistd.h>
+#endif // _WIN32
+
+#include <fcntl.h>
+#include <optional>
+#include <regex>
+#include <stdexcept>
+#include <unordered_map>
+#include <utility>
+
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+#include <torch/csrc/inductor/aoti_runtime/device_utils.h>
+#ifdef USE_MPS
+#include <torch/csrc/inductor/aoti_torch/c/shim_mps.h>
+#endif // USE_MPS
+#ifdef USE_XPU
+#include <torch/csrc/inductor/aoti_runtime/utils_xpu.h>
+#else
+#include <torch/csrc/inductor/aoti_runtime/utils.h>
+#endif // USE_XPU
+#include <torch/csrc/inductor/aoti_runtime/constant_type.h>
+
+#define AOTI_RUNTIME_CHECK(EXPR, MSG) \
+  do {                                \
+    bool ok = EXPR;                   \
+    if (!ok) {                        \
+      throw std::runtime_error(MSG);  \
+    }                                 \
+  } while (0)
+
+// At codegen time, we write out a binary file called constants.bin.
+// We then turn the raw binary to an object file that exposes this
+// symbol and link it into the final .so.
+// For information on the binary format, see `man objcopy`, under
+// the "binary-architecture" flag:
+// https://man7.org/linux/man-pages/man1/objcopy.1.html
+// todo: use #embed in C++ 23 once available
+// The constants are NOT readonly because they may be mutated.
+// NOLINTNEXTLINE(*array*)
+extern uint8_t _binary_constants_bin_start[];
+// NOLINTNEXTLINE(*array*)
+extern uint8_t _binary_constants_bin_end[];
+
+#if defined(USE_CUDA) || defined(USE_XPU)
+// Compute required blob size with 64-alignment if on GPU.
+#define AOTI_CONST_ALIGNMENT 64
+#else
+// Use 64-alignment (use something >=64)for better performance on CPU.
+#define AOTI_CONST_ALIGNMENT 64
+#endif
+
+namespace {
+
+using RAIIDataPtr = std::unique_ptr<void, std::function<void(void*)>>;
+
+#ifdef USE_CUDA
+
+// NOLINTNEXTLINE(clang-diagnostic-unneeded-internal-declaration)
+RAIIDataPtr RAII_gpuMalloc(size_t num_bytes) {
+#ifdef AOT_INDUCTOR_USE_CACHING_ALLOCATOR
+  // Use caching allocator for allocating GPU memory
+  void* data_ptr = nullptr;
+  AOTI_TORCH_ERROR_CODE_CHECK(
+      aoti_torch_cuda_caching_allocator_raw_alloc(num_bytes, &data_ptr));
+  auto deleter = [](void* ptr) {
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_cuda_caching_allocator_raw_delete(ptr));
+  };
+  return RAIIDataPtr(data_ptr, deleter);
+#else
+  // Use cudaMalloc directly for allocating GPU memory
+  void* data_ptr = nullptr;
+  AOTI_RUNTIME_CUDA_CHECK(cudaMalloc((void**)&data_ptr, num_bytes));
+  auto deleter = [](void* ptr) { AOTI_RUNTIME_CUDA_CHECK(cudaFree(ptr)); };
+  return RAIIDataPtr(data_ptr, deleter);
+#endif
+}
+
+#elif defined(USE_XPU)
+
+// NOLINTNEXTLINE(clang-diagnostic-unneeded-internal-declaration)
+RAIIDataPtr RAII_gpuMalloc(size_t num_bytes) {
+  sycl::queue* queue_ptr = nullptr;
+  aoti_torch_get_current_sycl_queue((void**)&queue_ptr);
+  void* data_ptr = sycl::malloc_device(num_bytes, *queue_ptr);
+  auto deleter = [queue_ptr](void* ptr) { sycl::free(ptr, *queue_ptr); };
+  return RAIIDataPtr(data_ptr, deleter);
+}
+
+#elif defined(USE_MPS)
+
+RAIIDataPtr RAII_gpuMalloc(size_t num_bytes) {
+  void* data_ptr = nullptr;
+  aoti_torch_mps_malloc(&data_ptr, num_bytes);
+  auto deleter = [](void* ptr) { aoti_torch_mps_free(ptr); };
+  return RAIIDataPtr(data_ptr, deleter);
+}
+
+#endif // USE_CUDA
+
+// NOLINTNEXTLINE(clang-diagnostic-unneeded-internal-declaration)
+RAIIDataPtr RAII_cpuMalloc(size_t num_bytes) {
+  void* data_ptr = std::malloc(num_bytes);
+  if (!data_ptr) {
+    throw std::bad_alloc();
+  }
+  auto deleter = [](void* ptr) { std::free(ptr); };
+  return RAIIDataPtr(data_ptr, deleter);
+}
+} // anonymous namespace
+
+namespace torch::aot_inductor {
+
+using ConstantMap =
+    std::unordered_map<std::string, MaybeOwningAtenTensorHandle>;
+
+// valid device strs are: cpu, cuda, cuda:0, cuda:1, ...
+// Update the list here if more devices are supported in the future
+inline void parse_device_str(
+    const std::string& device_str,
+    int32_t& device_type,
+    int32_t& device_idx) {
+  std::regex re("(cpu|cuda|xpu|mps)(:([0-9]+))?");
+  std::smatch sm;
+  bool matched = std::regex_match(device_str, sm, re);
+  AOTI_RUNTIME_CHECK(matched, "Invalid device: " + device_str);
+
+  if (sm[1].str() == "cpu") {
+    device_type = aoti_torch_device_type_cpu();
+  } else if (sm[1].str() == "cuda") {
+    device_type = aoti_torch_device_type_cuda();
+#ifdef USE_XPU
+  } else if (sm[1].str() == "xpu") {
+    device_type = aoti_torch_device_type_xpu();
+#endif
+#ifdef USE_MPS
+  } else if (sm[1].str() == "mps") {
+    device_type = aoti_torch_device_type_mps();
+#endif
+  } else {
+    AOTI_RUNTIME_CHECK(false, "Invalid device: " + device_str);
+  }
+
+  if (sm[3].matched) {
+    device_idx = stoi(sm[3].str());
+  } else {
+    device_idx = -1;
+  }
+}
+
+// Defines the base class for AOTInductorModel, which is generated by the
+// AOTInductor cpp codegen. Since we do not need dynamic dispatch, we rely
+// on curiously recurring template pattern (CRTP) to save some runtime
+// v-table overhead. The generated AOTInductorModel is specialized with
+// methods such as run_impl.
+template <typename Model>
+class AOTInductorModelBase {
+ public:
+  AOTInductorModelBase(
+      size_t num_inputs,
+      size_t num_outputs,
+      size_t num_constants,
+      const std::string& device_str,
+      std::optional<std::string> cubin_dir,
+      bool include_weights = true)
+      : inputs_info_(num_inputs),
+        outputs_info_(num_outputs),
+        constants_info_(num_constants),
+        cubin_dir_(std::move(cubin_dir)),
+        include_weights(include_weights) {
+    parse_device_str(device_str, device_type_, device_idx_);
+
+#ifdef USE_CUDA
+    if (device_idx_ == -1) {
+      AOTI_RUNTIME_CUDA_CHECK(cudaGetDevice(&device_idx_));
+    } else {
+      // If device_idx_ is passed in, we need to set the current device to it
+      AOTI_RUNTIME_CUDA_CHECK(cudaSetDevice(device_idx_));
+    }
+#endif // USE_CUDA
+#ifdef USE_XPU
+    if (device_idx_ == -1) {
+      aoti_torch_get_current_xpu_device(&device_idx_);
+    } else {
+      aoti_torch_set_current_xpu_device(device_idx_);
+    }
+#endif // USE_XPU
+#ifdef USE_MPS
+    if (device_idx_ == -1) {
+      device_idx_ = 0;
+    }
+#endif // USE_MPS
+  }
+
+  // NOLINTNEXTLINE(modernize-use-equals-default)
+  ~AOTInductorModelBase() {
+#ifdef USE_CUDA
+    if (run_finished_) {
+      auto code = cudaEventDestroy(*run_finished_);
+      if (code != cudaSuccess) {
+        std::cerr << "Failed to destroy CUDA event in AOTInductor model: "
+                  << cudaGetErrorString(code) << '\n';
+      }
+    }
+#endif // USE_CUDA
+#ifdef USE_XPU
+    if (run_finished_) {
+      (*run_finished_)->wait_and_throw();
+      delete *run_finished_;
+    }
+#endif // USE_XPU
+  }
+
+  AOTInductorModelBase(AOTInductorModelBase&&) = delete;
+  AOTInductorModelBase& operator=(AOTInductorModelBase&&) = delete;
+  AOTInductorModelBase(const AOTInductorModelBase&) = delete;
+  AOTInductorModelBase& operator=(const AOTInductorModelBase&) = delete;
+
+  void run(
+      AtenTensorHandle*
+          input_handles, // array of input AtenTensorHandle; handles
+                         // are stolen; the array itself is borrowed
+      AtenTensorHandle*
+          output_handles, // array for writing output AtenTensorHandle; handles
+                          // will be stolen by the caller; the array itself is
+                          // borrowed
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor) {
+#ifdef USE_CUDA
+    if (!run_finished_) {
+      cudaEvent_t run_finished = nullptr;
+      AOTI_RUNTIME_CUDA_CHECK(cudaEventCreate(&run_finished));
+      run_finished_.emplace(run_finished);
+    }
+#elif defined(USE_XPU)
+    if (run_finished_) {
+      (*run_finished_)->wait_and_throw();
+      delete *run_finished_;
+      run_finished_.reset();
+    }
+    if (stream == nullptr) {
+      aoti_torch_get_current_xpu_stream(this->device_idx_, (void**)&stream);
+    }
+#else // !USE_CUDA && !USE_XPU
+    run_finished_ = false;
+#endif
+
+    auto* model = static_cast<Model*>(this);
+    model->run_impl(input_handles, output_handles, stream, proxy_executor);
+
+#ifdef USE_CUDA
+    AOTI_RUNTIME_CUDA_CHECK(cudaEventRecord(*run_finished_, stream));
+#elif defined(USE_XPU)
+    run_finished_ = std::make_optional<sycl::event*>(new sycl::event(
+        static_cast<sycl::queue*>(stream)->ext_oneapi_submit_barrier()));
+#else // !USE_CUDA && !USE_XPU
+    run_finished_ = true;
+#endif // USE_CUDA
+  }
+
+  // Non-thread-aware variant of run(). Obviously unsafe to use in a threaded
+  // environment :)
+  void run_single_threaded(
+      AtenTensorHandle*
+          input_handles, // array of input AtenTensorHandle; handles
+                         // are stolen; the array itself is borrowed
+      AtenTensorHandle*
+          output_handles, // array for writing output AtenTensorHandle; handles
+                          // will be stolen by the caller; the array itself is
+                          // borrowed
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor) {
+    // don't bother with any of the run_finished stuff; this is unsafe to call
+    // in a threaded context
+    auto* model = static_cast<Model*>(this);
+    model->run_impl(input_handles, output_handles, stream, proxy_executor);
+  }
+
+  std::unordered_map<std::string, AtenTensorHandle> run_const_fold(
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor,
+      bool initialization = false) {
+#ifdef USE_CUDA
+    if (!run_finished_) {
+      cudaEvent_t run_finished = nullptr;
+      AOTI_RUNTIME_CUDA_CHECK(cudaEventCreate(&run_finished));
+      run_finished_.emplace(run_finished);
+    }
+#elif defined(USE_XPU)
+    if (run_finished_) {
+      (*run_finished_)->wait_and_throw();
+      delete *run_finished_;
+      run_finished_.reset();
+    }
+#else // !USE_CUDA && !USE_XPU
+    run_finished_ = false;
+#endif
+
+    auto* model = static_cast<Model*>(this);
+    auto folded_constants =
+        model->const_run_impl(stream, proxy_executor, initialization);
+
+#ifdef USE_CUDA
+    AOTI_RUNTIME_CUDA_CHECK(cudaEventRecord(*run_finished_, stream));
+#elif defined(USE_XPU)
+    // sycl::queue* queue_ptr = nullptr;
+    // aoti_torch_get_current_sycl_queue((void**)&queue_ptr);
+    run_finished_ = std::make_optional<sycl::event*>(new sycl::event(
+        static_cast<sycl::queue*>(stream)->ext_oneapi_submit_barrier()));
+
+#else // !USE_CUDA && !USE_XPU
+    run_finished_ = true;
+#endif // USE_CUDA
+
+    return folded_constants;
+  }
+
+  void update_constants_from_blob(const uint8_t* weight_blob_ptr) {
+#if defined(USE_MMAP_EXTERNAL)
+    user_managed_mmap = const_cast<uint8_t*>(weight_blob_ptr);
+    load_constants(true);
+#endif
+  }
+
+  void load_constants(bool force = false) {
+    size_t num_constants = this->num_constants();
+    size_t num_folded_constants = this->num_folded_constants();
+    constants_map_->reserve(num_constants);
+
+    // A CUDA model can still have constants on CPU,
+    // so we need a separate secondary blob for them.
+    std::vector<size_t> constants_internal_offset(
+        num_constants - num_folded_constants);
+    std::vector<size_t> secondary_cpu_constants_internal_offset(
+        num_constants - num_folded_constants);
+    size_t blob_size = 0;
+    size_t secondary_cpu_blob_size = 0;
+    compute_constant_blob(
+        blob_size,
+        constants_internal_offset,
+        secondary_cpu_blob_size,
+        secondary_cpu_constants_internal_offset);
+
+    if (!force && !include_weights) {
+      return;
+    }
+
+    // Allocate main blob
+    if (blob_size > 0) {
+#if defined(USE_CUDA) || defined(USE_XPU) || defined(USE_MPS)
+      constant_blob_ = RAII_gpuMalloc(blob_size);
+#else
+      constant_blob_ = RAII_cpuMalloc(blob_size);
+#endif
+    }
+
+    // Allocate secondary blob on CPU
+    if (secondary_cpu_blob_size > 0) {
+      secondary_cpu_constant_blob_ = RAII_cpuMalloc(secondary_cpu_blob_size);
+    }
+
+    size_t bytes_read = 0;
+    size_t main_blob_idx = 0;
+    size_t secondary_cpu_blob_idx = 0;
+
+    for (size_t i = 0; i < num_constants; i++) {
+      bool from_folded = this->constant_from_folded(i);
+      if (from_folded) {
+        continue;
+      }
+      std::string name = this->constant_name(i);
+      size_t data_size = this->constant_data_size(i);
+      int32_t const_device_type = this->constant_device_type(i);
+      bool device_type_matches = const_device_type == device_type_;
+
+      // Mixed-device constants are only supported when the secondary device is
+      // CPU. If a constant was compiled for a non-CPU device but we're loading
+      // on a different device, we cannot safely create the tensor.
+      AOTI_RUNTIME_CHECK(
+          device_type_matches ||
+              const_device_type == aoti_torch_device_type_cpu(),
+          "Mixed-device constants are only supported when the secondary "
+          "device is CPU. Constant '" +
+              name +
+              "' was compiled for a non-CPU device. "
+              "Hint: This can happen if you compiled on GPU but are loading "
+              "on CPU, which is not supported. In AOTI, you must compile and "
+              "load on the same device type.");
+
+      uint8_t* internal_ptr = nullptr;
+      if (data_size != 0) {
+        if (device_type_matches) {
+          internal_ptr = constant_ptr(
+              constants_internal_offset[main_blob_idx],
+              bytes_read,
+              data_size,
+              /* skip_copy = */ false);
+        } else {
+          auto* secondary_cpu_constants_ptr =
+              static_cast<uint8_t*>(secondary_cpu_constant_blob_.get());
+          internal_ptr = secondary_cpu_constants_ptr +
+              secondary_cpu_constants_internal_offset[secondary_cpu_blob_idx];
+          memcpy(internal_ptr, _get_constants_start() + bytes_read, data_size);
+        }
+      }
+
+      // Always increment blob indices to stay in sync with
+      // compute_constant_blob(), even for zero-size constants.
+      if (device_type_matches) {
+        main_blob_idx++;
+      } else {
+        secondary_cpu_blob_idx++;
+      }
+
+      bytes_read += data_size;
+
+      // Create at::Tensor from copied memory.
+      auto dtype = this->constant_dtype(i);
+      auto ndim = this->constant_ndim(i);
+      auto size = this->constant_shape(i);
+      auto stride = this->constant_stride(i);
+#ifdef USE_MPS
+      auto offset = this->constant_offset(i) +
+          (constants_internal_offset[i] / aoti_torch_dtype_element_size(dtype));
+#else
+      auto offset = this->constant_offset(i);
+#endif
+      auto layout = this->constant_layout(i);
+      auto opaque_metadata_ptr = this->opaque_metadata(i);
+      auto opaque_metadata_size = this->opaque_metadata_size(i);
+
+      AtenTensorHandle tensor_handle = nullptr;
+      AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_v2(
+          internal_ptr,
+          ndim,
+          size,
+          stride,
+          offset,
+          dtype,
+          const_device_type,
+          device_type_matches ? device_idx_ : 0,
+          &tensor_handle,
+          layout,
+          opaque_metadata_ptr,
+          opaque_metadata_size));
+      constants_map_->emplace(std::move(name), tensor_handle);
+    }
+    if (constants_map_) {
+      this->update_constants_array_from_map();
+    }
+  }
+
+  RAIIDataPtr&& release_constant_blob() {
+    return std::move(constant_blob_);
+  }
+
+  std::shared_ptr<std::vector<ConstantHandle>> get_constants_array() {
+    return constants_;
+  }
+
+  int32_t get_device_type() const {
+    return device_type_;
+  }
+
+  int32_t get_device_idx() const {
+    return device_idx_;
+  }
+
+  uint8_t* constant_ptr(
+      size_t constant_offset,
+      size_t bytes_read,
+      size_t data_size,
+      bool skip_copy) {
+    auto* constants_ptr = static_cast<uint8_t*>(constant_blob_.get());
+    uint8_t* internal_ptr = constants_ptr + constant_offset;
+    // TODO: Handle shared storage case.
+    if (!skip_copy) {
+#ifdef USE_XPU
+      sycl::queue* queue_ptr = nullptr;
+      aoti_torch_get_current_sycl_queue((void**)&queue_ptr);
+      queue_ptr
+          ->memcpy(internal_ptr, _get_constants_start() + bytes_read, data_size)
+          .wait();
+#elif USE_CUDA
+      AOTI_RUNTIME_CUDA_CHECK(cudaMemcpy(
+          internal_ptr,
+          _get_constants_start() + bytes_read,
+          data_size,
+          cudaMemcpyHostToDevice));
+#elif USE_MPS
+      aoti_torch_mps_memcpy(
+          constants_ptr,
+          constant_offset,
+          bytes_read,
+          data_size,
+          _get_constants_start());
+      return constants_ptr;
+#else
+      memcpy(internal_ptr, _get_constants_start() + bytes_read, data_size);
+#endif
+    }
+    return internal_ptr;
+  }
+
+  void compute_constant_blob(
+      size_t& blob_size,
+      std::vector<size_t>& constants_internal_offset,
+      size_t& secondary_cpu_blob_size,
+      std::vector<size_t>& secondary_cpu_constants_internal_offset) {
+    size_t num_constants = this->num_constants();
+    blob_size = 0;
+    secondary_cpu_blob_size = 0;
+    size_t main_idx = 0;
+    size_t secondary_idx = 0;
+
+    for (size_t i = 0; i < num_constants; i++) {
+      if (this->constant_from_folded(i)) {
+        continue;
+      }
+
+      size_t data_size = this->constant_data_size(i);
+      // ok to use same AOTI_CONST_ALIGNMENT for both main and secondary blobs
+      if (data_size % AOTI_CONST_ALIGNMENT) {
+        data_size = AOTI_CONST_ALIGNMENT +
+            (data_size / AOTI_CONST_ALIGNMENT) * AOTI_CONST_ALIGNMENT;
+      }
+
+      if (this->constant_device_type(i) == device_type_) {
+        constants_internal_offset[main_idx++] = blob_size;
+        blob_size += data_size;
+      } else {
+        secondary_cpu_constants_internal_offset[secondary_idx++] =
+            secondary_cpu_blob_size;
+        secondary_cpu_blob_size += data_size;
+      }
+    }
+  }
+
+  size_t num_inputs() const {
+    return inputs_info_.size();
+  }
+
+  size_t num_outputs() const {
+    return outputs_info_.size();
+  }
+
+  size_t num_constants() const {
+    return constants_info_.size();
+  }
+
+  size_t num_folded_constants() const {
+    size_t total_consts = this->num_constants();
+    size_t folded_consts = 0;
+    for (size_t i = 0; i < total_consts; i++) {
+      if (this->constant_from_folded(i)) {
+        folded_consts++;
+      }
+    }
+    return folded_consts;
+  }
+
+  const char* input_name(int64_t idx) const {
+    return inputs_info_.at(idx).name;
+  }
+
+  const char* output_name(int64_t idx) const {
+    return outputs_info_.at(idx).name;
+  }
+
+  const char* constant_name(int64_t idx) const {
+    return constants_info_.at(idx).name;
+  }
+
+  size_t constant_ndim(int64_t idx) {
+    return constants_info_.at(idx).shape.size();
+  }
+
+  const int64_t* constant_shape(int64_t idx) const {
+    return constants_info_.at(idx).shape.data();
+  }
+
+  const int64_t* constant_stride(int64_t idx) const {
+    return constants_info_.at(idx).stride.data();
+  }
+
+  int32_t constant_dtype(int64_t idx) const {
+    return constants_info_.at(idx).dtype;
+  }
+
+  int32_t constant_layout(int64_t idx) const {
+    return constants_info_.at(idx).layout;
+  }
+
+  size_t constant_offset(int64_t idx) const {
+    return constants_info_.at(idx).offset;
+  }
+
+  size_t constant_data_size(int64_t idx) const {
+    return constants_info_.at(idx).data_size;
+  }
+
+  const char* constant_original_fqn(int64_t idx) const {
+    return constants_info_.at(idx).original_fqn;
+  }
+
+  const uint8_t* opaque_metadata(int64_t idx) const {
+    return constants_info_.at(idx).opaque_metadata.data();
+  }
+
+  size_t opaque_metadata_size(int64_t idx) {
+    return constants_info_.at(idx).opaque_metadata.size();
+  }
+
+  bool constant_from_folded(int64_t idx) const {
+    return constants_info_.at(idx).from_folded;
+  }
+
+  int32_t constant_type(int64_t idx) const {
+    return constants_info_.at(idx).type;
+  }
+
+  int32_t constant_device_type(int64_t idx) const {
+    return constants_info_.at(idx).device_type;
+  }
+
+  const char* get_in_spec() const {
+    return in_spec_.c_str();
+  }
+
+  const char* get_out_spec() const {
+    return out_spec_.c_str();
+  }
+
+  uint64_t constant_blob_size() const {
+#if defined(USE_MMAP_SELF) || defined(USE_MMAP_EXTERNAL)
+    const uint64_t weights_size =
+        reinterpret_cast<const uint64_t*>(_binary_constants_bin_start)[0];
+    return weights_size;
+#else
+    throw std::runtime_error{
+        "constant blob size is only available for mmap'd weights"};
+#endif
+  }
+
+  void update_constants_array_from_map() {
+    if (!constants_map_) {
+      throw std::runtime_error{
+          "constants_map_ was not ready when constants_ is trying to be constructed from it!"};
+    }
+    if (!constants_) {
+      constants_ =
+          std::make_shared<std::vector<ConstantHandle>>(constants_info_.size());
+    } else {
+      constants_->resize(constants_info_.size());
+    }
+    int idx = 0;
+    for (const auto& info : constants_info_) {
+      const auto it = constants_map_->find(info.name);
+      if (it != constants_map_->end()) {
+        constants_->at(idx) = ConstantHandle(it->second);
+      }
+      idx++;
+    }
+  }
+
+  void update_constants_map(
+      std::shared_ptr<ConstantMap> constants_map,
+      bool remap_constants_array = true) {
+    constants_map_ = std::move(constants_map);
+    if (remap_constants_array) {
+      update_constants_array_from_map();
+    }
+  }
+
+  // This function allows us to update the constants_ that is used to look up
+  // the corresponding constant tensor during runtime.
+  void update_constants_array(
+      std::shared_ptr<std::vector<ConstantHandle>> constants_array) {
+    constants_ = std::move(constants_array);
+  }
+
+  /// Returns true if the model is complete.
+  bool is_finished() {
+#ifdef USE_CUDA
+    if (!run_finished_) {
+      throw std::runtime_error{"Model CUDA event was not initialized"};
+    }
+
+    auto event_status = cudaEventQuery(*run_finished_);
+    if (event_status == cudaSuccess) {
+      return true;
+    } else if (event_status == cudaErrorNotReady) {
+      return false;
+    }
+
+    throw std::runtime_error(
+        std::string("The model did not finish successfully. Error: ") +
+        cudaGetErrorString(cudaGetLastError()));
+#elif defined(USE_XPU)
+    if (!run_finished_) {
+      throw std::runtime_error{"Model XPU event was not initialized"};
+    }
+    using namespace sycl::info;
+    return (*run_finished_)->get_info<event::command_execution_status>() ==
+        event_command_status::complete;
+
+#else // !USE_CUDA && !USE_XPU
+    return run_finished_;
+#endif // USE_CUDA
+  }
+
+  /// Synchronizes completion event.
+  void wait_for_completion() {
+#ifdef USE_CUDA
+    if (!run_finished_) {
+      throw std::runtime_error{"Model event was not initialized"};
+    }
+
+    AOTI_RUNTIME_CUDA_CHECK(cudaEventSynchronize(*run_finished_));
+#endif // USE_CUDA
+#ifdef USE_XPU
+    if (!run_finished_) {
+      throw std::runtime_error{"Model event was not initialized"};
+    }
+    (*run_finished_)->wait_and_throw();
+#endif
+  }
+
+ protected:
+  uint8_t* _get_constants_start() {
+#if defined(USE_MMAP_EXTERNAL)
+    if (!user_managed_mmap) {
+      throw std::runtime_error{
+          "Constants are not mmap'd. Use AOTInductorModelUpdateConstantsBlob to initialize the constants first."};
+    }
+    // Mapped memory for weights
+    return user_managed_mmap;
+#endif
+
+#ifndef USE_MMAP_SELF
+    // NOLINTNEXTLINE(*const-cast*)
+    return const_cast<uint8_t*>(_binary_constants_bin_start);
+#else
+    if (self_mmap) {
+      return self_mmap;
+    }
+    Dl_info dl_info;
+    // get pointer to constant which are appended to the binary
+    AOTI_RUNTIME_CHECK(
+        dladdr(__func__, &dl_info), "Can't find shared library name");
+    int fd = open(dl_info.dli_fname, O_RDONLY);
+    AOTI_RUNTIME_CHECK(fd >= 0, "Shared library file cannot be opened");
+    auto fsize = lseek(fd, 0, SEEK_END);
+    auto weights_size =
+        reinterpret_cast<const uint64_t*>(_binary_constants_bin_start)[0];
+    auto magic_number =
+        reinterpret_cast<const uint64_t*>(_binary_constants_bin_start)[1];
+    auto weights_offset = fsize - weights_size;
+    AOTI_RUNTIME_CHECK(
+        (weights_offset & 0x3fff) == 0,
+        "weights_offset must be aligned to 16K boundary");
+    auto ptr = mmap(
+        NULL,
+        weights_size,
+        PROT_READ | PROT_WRITE,
+        MAP_PRIVATE,
+        fd,
+        weights_offset);
+    close(fd);
+    AOTI_RUNTIME_CHECK(ptr != MAP_FAILED, "mmap() failed");
+    self_mmap = static_cast<uint8_t*>(ptr);
+    AOTI_RUNTIME_CHECK(
+        reinterpret_cast<uint64_t*>(
+            self_mmap + weights_size - sizeof(uint64_t))[0] == magic_number,
+        "Weights data seems corrupt");
+    return self_mmap;
+#endif
+  }
+
+  struct ParamInfo {
+    const char* name = nullptr;
+  };
+
+  struct ConstInfo {
+    const char* name = nullptr;
+    std::vector<int64_t> shape;
+    std::vector<int64_t> stride;
+    int32_t dtype{};
+    int32_t device_type{};
+    int64_t offset{};
+    size_t data_size{};
+    int32_t layout{};
+    std::vector<uint8_t> opaque_metadata;
+    int64_t opaque_metadata_size{};
+    const char* original_fqn = nullptr;
+    bool from_folded{};
+    int32_t type{};
+  };
+
+  std::vector<ParamInfo> inputs_info_;
+  std::vector<ParamInfo> outputs_info_;
+  std::vector<ConstInfo> constants_info_;
+  std::string in_spec_;
+  std::string out_spec_;
+
+  std::shared_ptr<ConstantMap> constants_map_;
+  std::shared_ptr<std::vector<ConstantHandle>> constants_;
+
+  // Holds the blob storage for constants' at::Tensor.
+  RAIIDataPtr constant_blob_;
+  // For mixed-device models, secondary_cpu_constant_blob_ holds CPU constants
+  RAIIDataPtr secondary_cpu_constant_blob_;
+
+#if defined(USE_MMAP_SELF)
+  // Mapped memory for weights
+  uint8_t* self_mmap = NULL;
+#endif
+
+#if defined(USE_MMAP_EXTERNAL)
+  // Mapped memory for weights
+  uint8_t* user_managed_mmap = NULL;
+#endif
+
+  // A directory with CUDA binary files, e.g. compiled kernels, etc.
+  const std::optional<std::string> cubin_dir_;
+
+  // This is the flag that implies whether the weight is included in the model.
+  // If True, we would prepare the weight when loading the model, otherwise the
+  // model will be loaded without weights, and need to be provided by the user.
+  bool include_weights;
+
+  // Record if the model finishes an inference run so that its owning
+  // AOTModelContainer can reuse this instance.
+#ifdef USE_CUDA
+  std::optional<cudaEvent_t> run_finished_;
+#elif defined(USE_XPU)
+  std::optional<sycl::event*> run_finished_;
+#else // !USE_CUDA
+  bool run_finished_{};
+#endif
+
+  // Generated model uses this device index to create CUDA guards.
+  int32_t device_type_{};
+  int32_t device_idx_{};
+};
+
+// Codegen-ed classes can derive from this to keep pointers to loaded kernels.
+class AOTInductorModelKernelsBase {
+ public:
+  virtual ~AOTInductorModelKernelsBase() = default;
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model_container.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model_container.h
new file mode 100644
index 0000000000000000000000000000000000000000..742bc392fb313214e4689b096ce636f6f863df9d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/model_container.h
@@ -0,0 +1,832 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <algorithm>
+#include <condition_variable>
+#include <deque>
+#include <mutex>
+#include <shared_mutex>
+
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+#include <torch/csrc/inductor/aoti_runtime/model.h>
+
+namespace torch::aot_inductor {
+// The state transition is done by:
+// (1) NONE state: The default state when created. This state should only exist
+// when model_container is created and no constants are being loaded or updated.
+// (2) INITIALIZED state: This state get set whenever we load the constants into
+// the buffer. This could be done by load_constants or update_constants_buffer.
+// (3) FOLDED state: This state should transition from INITIALIZED after
+// const_fold is being invoked.
+enum class ConstantState : uint8_t { NONE, INITIALIZED, FOLDED, UNKNOWN };
+
+inline std::string toStringConstantState(ConstantState state) {
+  switch (state) {
+    case ConstantState::NONE:
+      return "ConstantState::NONE";
+    case ConstantState::INITIALIZED:
+      return "ConstantState::INITIALIZED";
+    case ConstantState::FOLDED:
+      return "ConstantState::FOLDED";
+    case ConstantState::UNKNOWN:
+      return "ConstantState::UNKNOWN";
+    default:
+      return "Unknown enum class state for ConstantState";
+  }
+}
+
+class AOTInductorModelContainer {
+ public:
+  AOTInductorModelContainer(
+      size_t num_models,
+      const std::string& device_str,
+      const std::optional<std::string>& cubin_dir = std::nullopt) {
+    constants_map_ = std::make_shared<ConstantMap>();
+    constants_array_ = std::make_shared<std::vector<ConstantHandle>>();
+
+    models_.reserve(num_models);
+    available_models_.reserve(num_models);
+    for (size_t i = 0; i < num_models; ++i) {
+      models_.push_back(AOTInductorModel::Create(
+          constants_map_, constants_array_, device_str, cubin_dir));
+      available_models_.push_back(models_.back().get());
+    }
+
+    // Note that the all following fields (input_names_, output_names,
+    // etc) can be filled in by the AOT
+    // codegen. However, we choose to query such information from
+    // the owned AOTInductorModel for a couple of reasons:
+    //   * simplify the codegen templates
+    //   * reduce information fragmentation and duplication
+    //   * the initialization process below is done only once when the container
+    //     is constructed, so it would have little performance impact
+    auto* model = available_models_[0];
+    size_t num_inputs = model->num_inputs();
+    input_names_.reserve(num_inputs);
+    for (size_t i = 0; i < num_inputs; i++) {
+      input_names_.emplace_back(model->input_name(static_cast<int64_t>(i)));
+    }
+
+    size_t num_outputs = model->num_outputs();
+    output_names_.reserve(num_outputs);
+    for (size_t i = 0; i < num_outputs; i++) {
+      output_names_.emplace_back(model->output_name(static_cast<int64_t>(i)));
+    }
+    model->load_constants();
+    constant_blob_ = model->release_constant_blob();
+    constants_internal_offset_.resize(
+        model->num_constants() - model->num_folded_constants());
+    secondary_cpu_constants_internal_offset_.resize(
+        model->num_constants() - model->num_folded_constants());
+    model->compute_constant_blob(
+        blob_size_,
+        constants_internal_offset_,
+        secondary_cpu_blob_size_,
+        secondary_cpu_constants_internal_offset_);
+    constant_folded_ = ConstantState::INITIALIZED;
+
+    for (auto& model : models_) {
+      model->update_constants_map(constants_map_);
+    }
+
+    in_spec_ = model->get_in_spec();
+    out_spec_ = model->get_out_spec();
+  }
+
+  void run(
+      AtenTensorHandle*
+          input_handles, // array of input AtenTensorHandle; handles
+                         // are stolen; the array itself is borrowed
+      AtenTensorHandle*
+          output_handles, // array for writing output AtenTensorHandle; handles
+                          // will be stolen by the caller; the array itself is
+                          // borrowed
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor) {
+    std::shared_lock model_lk(model_exec_mutex_);
+    auto* model = get_available_model();
+
+    ConstantState& const_folded =
+        use_secondary_ ? constant_folded_secondary_ : constant_folded_;
+    if (const_folded == ConstantState::INITIALIZED) {
+      // At this point, constant is not ready yet. We need to call constant
+      // folding before we execute the model. We obtain a unique lock at this
+      // point to make sure constant is ready for all.
+      model_lk.unlock();
+      std::unique_lock constants_folding_lk(model_exec_mutex_);
+      // Double locking to make sure constant folding is only ran once.
+      if (const_folded == ConstantState::INITIALIZED) {
+        auto folded_const_map = model->run_const_fold(
+            stream, proxy_executor, /* initialization = */ true);
+        update_constant_buffer(
+            std::move(folded_const_map),
+            /* use_inactive = */ false,
+            /* validate_full_update = */ false);
+        const_folded = ConstantState::FOLDED;
+      }
+      constants_folding_lk.unlock();
+      model_lk.lock();
+    } else if (const_folded != ConstantState::FOLDED) {
+      throw std::runtime_error(
+          "Unknown constant state: " + toStringConstantState(constant_folded_));
+    }
+
+    try {
+      model->run(input_handles, output_handles, stream, proxy_executor);
+    } catch (...) {
+      std::lock_guard lk(models_mutex_);
+      available_models_.push_back(model);
+      throw;
+    }
+
+    {
+      std::lock_guard lk(models_mutex_);
+      pending_models_.push_back(model);
+    }
+    pending_models_available_.notify_one();
+  }
+
+  // Non-thread-aware variant of run(). Obviously unsafe to use in a threaded
+  // environment :)
+  void run_single_threaded(
+      AtenTensorHandle*
+          input_handles, // array of input AtenTensorHandle; handles
+                         // are stolen; the array itself is borrowed
+      AtenTensorHandle*
+          output_handles, // array for writing output AtenTensorHandle; handles
+                          // will be stolen by the caller; the array itself is
+                          // borrowed
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor) {
+    auto* model = available_models_[0];
+
+    ConstantState& const_folded =
+        use_secondary_ ? constant_folded_secondary_ : constant_folded_;
+    if (const_folded == ConstantState::INITIALIZED) {
+      auto folded_const_map = model->run_const_fold(
+          stream, proxy_executor, /* initialization = */ true);
+      update_constant_buffer(
+          std::move(folded_const_map),
+          /* use_inactive = */ false,
+          /* validate_full_update = */ false);
+      const_folded = ConstantState::FOLDED;
+    } else if (const_folded != ConstantState::FOLDED) {
+      throw std::runtime_error(
+          "Unknown constant state: " + toStringConstantState(const_folded));
+    }
+
+    model->run_single_threaded(
+        input_handles, output_handles, stream, proxy_executor);
+  }
+
+  const std::unordered_map<std::string, AtenTensorHandle> extract_constants_map(
+      bool use_inactive) const {
+    size_t n_consts = this->num_constants();
+    std::unordered_map<std::string, AtenTensorHandle> ret;
+    ret.reserve(n_consts);
+
+    std::shared_ptr<ConstantMap> extract_map = constants_map_;
+    // Essentially a XOR
+    if (use_inactive != use_secondary_) {
+      extract_map = constants_map_secondary_;
+    }
+    for (size_t idx = 0; idx < n_consts; idx++) {
+      if (this->constant_from_folded(idx)) {
+        continue;
+      }
+
+      auto it = extract_map->find(this->constant_name(idx));
+      if (it != extract_map->end()) {
+        ret.emplace(this->constant_original_fqn(idx), it->second);
+        continue;
+      }
+    }
+
+    return ret;
+  }
+
+  size_t num_constants() const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->num_constants();
+  }
+
+  // retrieve the constant name of constants_info_[idx]
+  const char* constant_name(size_t idx) const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->constant_name(static_cast<int64_t>(idx));
+  }
+
+  // retrieve original FQN of constants_info_[idx]
+  const char* constant_original_fqn(size_t idx) const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->constant_original_fqn(static_cast<int64_t>(idx));
+  }
+
+  // retrieve whether constant is from folded of constants_info_[idx]
+  bool constant_from_folded(size_t idx) const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->constant_from_folded(static_cast<int64_t>(idx));
+  }
+
+  size_t constant_data_size(size_t idx) const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->constant_data_size(static_cast<int64_t>(idx));
+  }
+
+  // retrieve type of constants_info_[idx]
+  int32_t constant_type(size_t idx) const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->constant_type(static_cast<int64_t>(idx));
+  }
+
+  // retrieve dtype of constants_info_[idx]
+  int32_t constant_dtype(size_t idx) const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->constant_dtype(static_cast<int64_t>(idx));
+  }
+
+  uint64_t constant_blob_size() const {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->constant_blob_size();
+  }
+
+  void update_constants_from_blob(const uint8_t* weight_blob_ptr) {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No available models in container!");
+    }
+    return models_[0]->update_constants_from_blob(weight_blob_ptr);
+  }
+
+  void run_const_fold(
+      bool inactive_buffer,
+      DeviceStreamType stream,
+      AOTIProxyExecutorHandle proxy_executor) {
+    AOTInductorModel* model;
+    ConstantState& const_folded = inactive_buffer == use_secondary_
+        ? constant_folded_
+        : constant_folded_secondary_;
+    if (!inactive_buffer) {
+      // We would need to acquire a unique lock if we want to run constant
+      // folding on the active buffer.
+      std::unique_lock constants_folding_lk(model_exec_mutex_);
+      model = get_available_model();
+      try {
+        auto folded_const_map = model->run_const_fold(stream, proxy_executor);
+        update_constant_buffer(
+            std::move(folded_const_map),
+            /* use_inactive = */ false,
+            /* validate_full_update = */ false);
+        const_folded = ConstantState::FOLDED;
+      } catch (...) {
+        std::lock_guard lk(models_mutex_);
+        available_models_.push_back(model);
+        throw;
+      }
+    } else {
+      std::shared_lock model_lk(model_exec_mutex_);
+      model = get_available_model();
+
+      // We swap the constant mapping to the inactive buffer in the model to run
+      // const run.
+      auto constants_map = get_constants_map(/* get_inactive= */ true);
+      auto constants_array = get_constants_array(/* get_inactive= */ true);
+
+      try {
+        model->update_constants_map(
+            constants_map, /* remap_constants_array= */ false);
+        model->update_constants_array(constants_array);
+
+        auto folded_const_map = model->run_const_fold(stream, proxy_executor);
+        update_constant_buffer(
+            std::move(folded_const_map),
+            /* use_inactive = */ true,
+            /* validate_full_update = */ false);
+
+        // Swap back the model's constants mapping
+        constants_map = get_constants_map(/* get_inactive= */ false);
+        constants_array = get_constants_array(/* get_inactive= */ false);
+        model->update_constants_map(
+            constants_map, /* remap_constants_array= */ false);
+        model->update_constants_array(constants_array);
+        const_folded = ConstantState::FOLDED;
+      } catch (...) {
+        std::lock_guard lk(models_mutex_);
+        available_models_.push_back(model);
+        throw;
+      }
+    }
+
+    {
+      std::lock_guard lk(models_mutex_);
+      pending_models_.push_back(model);
+    }
+    pending_models_available_.notify_one();
+  }
+
+  bool _is_tensor_constant_type(const size_t idx) const {
+    auto constant_type = models_[0]->constant_type(static_cast<int64_t>(idx));
+    // We should skip constants
+    return constant_type == ConstantType::TensorConstant;
+  }
+
+  bool _is_buffer_type(const size_t idx) const {
+    auto constant_type = models_[0]->constant_type(static_cast<int64_t>(idx));
+    // Buffer can be optionally skipped, so if it not provided by upstream
+    // services, it is OK to relax the check.
+    return constant_type == ConstantType::Buffer;
+  }
+
+  bool _is_empty_parameter_type(const size_t idx) const {
+    auto constant_type = models_[0]->constant_type(static_cast<int64_t>(idx));
+    auto constant_data_size =
+        models_[0]->constant_data_size(static_cast<int64_t>(idx));
+    // Empty parameters are skipped and not provided by the upstream services,
+    // it is OK to skip.
+    return constant_type == ConstantType::Parameter && constant_data_size == 0;
+  }
+
+  bool _is_tensor_constant_or_buffer_type_or_empty_parameter(
+      const size_t idx) const {
+    return _is_tensor_constant_type(idx) || _is_buffer_type(idx) ||
+        _is_empty_parameter_type(idx);
+  }
+
+  void assert_all_constants(
+      const std::unordered_map<std::string, AtenTensorHandle>& constants_map) {
+    auto num_constants = models_[0]->num_constants();
+    for (size_t idx = 0; idx < num_constants; idx++) {
+      if (models_[0]->constant_from_folded(static_cast<int64_t>(idx))) {
+        continue;
+      }
+
+      auto constant_name =
+          std::string(models_[0]->constant_name(static_cast<int64_t>(idx)));
+      auto it = constants_map.find(constant_name);
+      if (it == constants_map.end()) {
+        if (_is_tensor_constant_or_buffer_type_or_empty_parameter(idx)) {
+          // tracing sometimes creates tensors that are non-existent in
+          // original graph. We could skip those and do a direct copy.
+          std::cerr << "[WARNING] Found constant or module state buffer or "
+                    << "empty module state parameter " << constant_name
+                    << " in model, but not provided by user!\n";
+          continue;
+        }
+        throw std::runtime_error(
+            std::string("Cannot find constants ") + constant_name +
+            std::string(" in constants_map!"));
+      }
+    }
+  }
+
+  // We directly take ownership from AtenTensorHandle if constants are moved.
+  void update_constant_buffer(
+      std::unordered_map<std::string, AtenTensorHandle>&& constants_map,
+      bool use_inactive,
+      bool validate_full_update) {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No model available in container!");
+    }
+    if (validate_full_update) {
+      assert_all_constants(constants_map);
+    }
+
+    ConstantState& const_folded = use_inactive == use_secondary_
+        ? constant_folded_
+        : constant_folded_secondary_;
+    const_folded = ConstantState::INITIALIZED;
+
+    auto original_constants_map = get_constants_map(!use_inactive);
+    auto constants_map_to_update = get_constants_map(use_inactive);
+
+    auto num_constants = models_[0]->num_constants();
+    for (size_t idx = 0; idx < num_constants; idx++) {
+      auto constant_name =
+          std::string(models_[0]->constant_name(static_cast<int64_t>(idx)));
+      auto it = constants_map.find(constant_name);
+      if (it == constants_map.end() &&
+          !(use_inactive && _is_tensor_constant_type(idx))) {
+        continue;
+      }
+
+      AtenTensorHandle tensor;
+      if (it == constants_map.end()) {
+        aoti_torch_clone(
+            original_constants_map->find(constant_name)->second.get(), &tensor);
+      } else {
+        tensor = it->second;
+      }
+
+      constants_map_to_update->insert_or_assign(
+          constant_name, RAIIAtenTensorHandle(tensor));
+    }
+    // Update the inactive constant array.
+    update_array_from_map(
+        get_constants_array(use_inactive), constants_map_to_update);
+  }
+
+  // This function updates the buffer for storing constants.
+  // It will update the buffer, the mapping and the array mapping.
+  void update_constant_buffer(
+      const std::unordered_map<std::string, AtenTensorHandle>& constants_map,
+      bool use_inactive,
+      bool validate_full_update,
+      bool user_managed = false) {
+    if (this->num_models() == 0) {
+      throw std::runtime_error("No model available in container!");
+    }
+    if (validate_full_update) {
+      assert_all_constants(constants_map);
+    }
+
+    // update_constant_buffer does not support mixed CPU/CUDA constants
+    int32_t model_device_type = models_[0]->get_device_type();
+    for (const auto& kv : constants_map) {
+      int32_t tensor_device_type = 0;
+      aoti_torch_get_device_type(kv.second, &tensor_device_type);
+      if (tensor_device_type != model_device_type) {
+        throw std::runtime_error(
+            "update_constant_buffer does not support mixed device constants. "
+            "Constant '" +
+            kv.first + "' has device type " +
+            std::to_string(tensor_device_type) +
+            " but model expects device type " +
+            std::to_string(model_device_type));
+      }
+    }
+
+    ConstantState& const_folded = use_inactive == use_secondary_
+        ? constant_folded_
+        : constant_folded_secondary_;
+    const_folded = ConstantState::INITIALIZED;
+
+    auto original_constants_map = get_constants_map(!use_inactive);
+    auto constants_map_to_update = get_constants_map(use_inactive);
+
+    auto num_constants = models_[0]->num_constants();
+    for (size_t idx = 0; idx < num_constants; idx++) {
+      auto constant_name =
+          std::string(models_[0]->constant_name(static_cast<int64_t>(idx)));
+      auto it = constants_map.find(constant_name);
+      if (it == constants_map.end() &&
+          !(use_inactive &&
+            _is_tensor_constant_or_buffer_type_or_empty_parameter(idx))) {
+        continue;
+      }
+
+      AtenTensorHandle tensor;
+      if (it == constants_map.end()) {
+        tensor = original_constants_map->find(constant_name)->second.get();
+      } else {
+        tensor = it->second;
+      }
+
+      if (user_managed) {
+        // If user managed, we pass in the pointer directly, and skip the
+        // copy.
+        constants_map_to_update->insert_or_assign(
+            constant_name,
+            MaybeOwningAtenTensorHandle(tensor, /* user_managed = */ true));
+        continue;
+      }
+
+      auto* constants_blob_ptr =
+          static_cast<uint8_t*>(get_constant_blob_ptr(use_inactive));
+
+      // Move the data to container handled blob.
+      uint8_t* internal_constants_ptr =
+          constants_blob_ptr + constants_internal_offset_[idx];
+      void* user_constant_ptr;
+      int64_t constant_size;
+      int64_t* stride;
+      int64_t offset;
+      aoti_torch_get_data_ptr(tensor, &user_constant_ptr);
+      aoti_torch_get_storage_size(tensor, &constant_size);
+      AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(tensor, &stride));
+      AOTI_TORCH_ERROR_CODE_CHECK(
+          aoti_torch_get_storage_offset(tensor, &offset));
+      auto dtype = models_[0]->constant_dtype(idx);
+
+#ifdef USE_XPU
+      sycl::queue* queue_ptr = nullptr;
+      aoti_torch_get_current_sycl_queue((void**)&queue_ptr);
+      queue_ptr
+          ->memcpy(internal_constants_ptr, user_constant_ptr, constant_size)
+          .wait();
+#elif USE_MPS
+      internal_constants_ptr = constants_blob_ptr;
+      aoti_torch_mps_copy_buffer(
+          user_constant_ptr,
+          constants_blob_ptr,
+          constant_size,
+          offset,
+          constants_internal_offset_[idx]);
+      // For mps tensors, all constants are stored in one buffer, with the
+      // offset being where the constant starts. So we want to change the
+      // constant tensor's offset to point to constants_internal_offset_[idx]
+      offset = constants_internal_offset_[idx] /
+          aoti_torch_dtype_element_size(dtype);
+#elif USE_CUDA
+      AOTI_RUNTIME_CUDA_CHECK(cudaMemcpy(
+          internal_constants_ptr,
+          user_constant_ptr,
+          constant_size,
+          cudaMemcpyDefault));
+#else
+      memcpy(internal_constants_ptr, user_constant_ptr, constant_size);
+#endif
+      // Generate Tensor from container handled blob.
+      // We extract stride and offset from provided Tensor since we do not
+      // guarantee that the tensor is contiguous.
+      AtenTensorHandle tensor_handle;
+      int device_type = models_[0]->get_device_type();
+      int device_idx = models_[0]->get_device_idx();
+      AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob(
+          internal_constants_ptr,
+          models_[0]->constant_ndim(idx),
+          models_[0]->constant_shape(idx),
+          stride,
+          offset,
+          dtype,
+          device_type,
+          device_idx,
+          &tensor_handle));
+
+      // Now place the tensor to constants_map. Note at this point the
+      // ownership of the tensor_handle will be taken over.
+      constants_map_to_update->insert_or_assign(
+          constant_name, RAIIAtenTensorHandle(tensor_handle));
+    }
+    // Update the inactive constant array.
+    update_array_from_map(
+        get_constants_array(use_inactive), constants_map_to_update);
+  }
+
+  void update_array_from_map(
+      const std::shared_ptr<std::vector<ConstantHandle>>& constants_array,
+      const std::shared_ptr<ConstantMap>& constants_map) {
+    auto num_constants = models_[0]->num_constants();
+    for (size_t idx = 0; idx < num_constants; idx++) {
+      if (constants_map->find(models_[0]->constant_name(
+              static_cast<int64_t>(idx))) != constants_map->end()) {
+        constants_array->at(idx) = ConstantHandle(
+            constants_map
+                ->find(models_[0]->constant_name(static_cast<int64_t>(idx)))
+                ->second);
+      }
+    }
+  }
+
+  void swap_constant_buffer() {
+    std::lock_guard unique_lk(model_exec_mutex_);
+
+    auto constants_map = get_constants_map(/* get_inactive= */ true);
+    auto constants_array = get_constants_array(/* get_inactive= */ true);
+
+    for (auto& model : models_) {
+      model->update_constants_map(
+          constants_map, /* remap_constants_array = */ false);
+      model->update_constants_array(constants_array);
+    }
+
+    use_secondary_ = !use_secondary_;
+  }
+
+  void free_inactive_constant_buffer() {
+    if (use_secondary_) {
+      constant_folded_ = ConstantState::NONE;
+      constant_blob_.reset();
+    } else {
+      constant_folded_secondary_ = ConstantState::NONE;
+      constant_blob_secondary_.reset();
+    }
+    // Free the internally held constants
+    int num_constants = static_cast<int>(models_[0]->num_constants());
+    std::shared_ptr<ConstantMap> to_free_map =
+        use_secondary_ ? constants_map_ : constants_map_secondary_;
+
+    for (int i = 0; i < num_constants; i++) {
+      if (models_[0]->constant_from_folded(i)) {
+        auto it = to_free_map->find(models_[0]->constant_name(i));
+        if (it != to_free_map->end()) {
+          it->second.reset();
+        }
+      }
+    }
+  }
+
+  size_t num_inputs() const {
+    return input_names_.size();
+  }
+
+  size_t num_outputs() const {
+    return output_names_.size();
+  }
+
+  const char* input_name(size_t idx) const {
+    return input_names_.at(idx).c_str();
+  }
+
+  const char* output_name(size_t idx) const {
+    return output_names_.at(idx).c_str();
+  }
+
+  size_t num_models() const {
+    return models_.size();
+  }
+
+  const char* get_in_spec() const {
+    return in_spec_;
+  }
+
+  const char* get_out_spec() const {
+    return out_spec_;
+  }
+
+ private:
+  std::vector<std::string> input_names_;
+  std::vector<std::string> output_names_;
+  const char* in_spec_;
+  const char* out_spec_;
+
+  // Holds the blob storage for constants' at::Tensor within the container.
+  // This blob of memory will be managed by the container.
+  RAIIDataPtr constant_blob_;
+  RAIIDataPtr constant_blob_secondary_;
+
+  size_t blob_size_;
+  std::vector<size_t> constants_internal_offset_;
+  size_t secondary_cpu_blob_size_;
+  std::vector<size_t> secondary_cpu_constants_internal_offset_;
+
+  // Determine which constants is being used for the model.
+  // If true,
+  // constants_map_secondary/constant_blob_secondary/constants_array_secondary
+  // is being used.
+  bool use_secondary_{false};
+
+  // Determine whether we have ran constant folding
+  ConstantState constant_folded_{ConstantState::NONE};
+  ConstantState constant_folded_secondary_{ConstantState::NONE};
+
+  // Holds the mapping of constants to at::Tensor.
+  // The underlying data of at::Tensor is in either constant_blob_ (for CUDA).
+  // or _binary_constants_bin_start (for CPU).
+  std::shared_ptr<ConstantMap> constants_map_;
+  std::shared_ptr<ConstantMap> constants_map_secondary_;
+
+  // Holds the indexed array of constant for faster lookup during runtime.
+  std::shared_ptr<std::vector<ConstantHandle>> constants_array_;
+  std::shared_ptr<std::vector<ConstantHandle>> constants_array_secondary_;
+
+  // Holds all the AOTInductorModel instances owned by this container.
+  std::vector<std::unique_ptr<AOTInductorModel>> models_;
+
+  // Holds the AOTInductorModel instances available for inference.
+  std::vector<AOTInductorModel*> available_models_;
+
+  // Holds the AOTInductorModel instances that have started running
+  // inference and can be placed onto available_models_ upon their
+  // completion.
+  std::deque<AOTInductorModel*> pending_models_;
+
+  // Protects available_models_ and pending_models_.
+  std::mutex models_mutex_;
+
+  // Notified whenever a model is placed onto pending_models_.
+  std::condition_variable pending_models_available_;
+
+  AOTInductorModel* get_available_model() {
+    std::unique_lock lk(models_mutex_);
+    if (available_models_.empty()) {
+      reclaim_finished_models(lk);
+    }
+    auto* result = available_models_.back();
+    available_models_.pop_back();
+    return result;
+  }
+
+  // This mutex is used to protect execution of model.
+  // We acquire the mutex in shared mode if we allow concurrent execution.
+  // We acquire the mutex in unique mode when we want exclusive access of the
+  // model. One such case is when we want to do a weight swapping. We want to
+  // make sure no one is executing the model.
+  std::shared_mutex model_exec_mutex_;
+
+  RAIIDataPtr allocate_constant_blob() {
+#if defined(USE_CUDA) || defined(USE_XPU) || defined(USE_MPS)
+    return RAII_gpuMalloc(blob_size_);
+#else
+    return RAII_cpuMalloc(blob_size_);
+#endif // USE_CUDA
+  }
+
+  void* get_constant_blob_ptr(bool get_inactive) {
+    if ((get_inactive && use_secondary_) ||
+        (!get_inactive && !use_secondary_)) {
+      if (!constant_blob_) {
+        constant_blob_ = allocate_constant_blob();
+      }
+      return constant_blob_.get();
+    } else {
+      if (!constant_blob_secondary_) {
+        constant_blob_secondary_ = allocate_constant_blob();
+      }
+      return constant_blob_secondary_.get();
+    }
+  }
+
+  std::shared_ptr<ConstantMap> get_constants_map(bool get_inactive) {
+    if ((get_inactive && use_secondary_) ||
+        (!get_inactive && !use_secondary_)) {
+      return constants_map_;
+    } else {
+      if (!constants_map_secondary_) {
+        constants_map_secondary_ = std::make_shared<ConstantMap>();
+      }
+      return constants_map_secondary_;
+    }
+  }
+
+  std::shared_ptr<std::vector<ConstantHandle>> get_constants_array(
+      bool get_inactive) {
+    if ((get_inactive && use_secondary_) ||
+        (!get_inactive && !use_secondary_)) {
+      return constants_array_;
+    } else {
+      if (!constants_array_secondary_) {
+        constants_array_secondary_ =
+            std::make_shared<std::vector<ConstantHandle>>(
+                models_[0]->num_constants());
+      }
+      return constants_array_secondary_;
+    }
+  }
+
+  void reclaim_finished_models(std::unique_lock<std::mutex>& lk) {
+#ifdef __aarch64__
+    // push finished model instances to the end of pending_models_
+    auto it = std::partition(
+        pending_models_.begin(),
+        pending_models_.end(),
+        [](AOTInductorModel* m) { return !m->is_finished(); });
+#else
+    // push finished model instances to the end of pending_models_
+    auto it = std::stable_partition(
+        pending_models_.begin(),
+        pending_models_.end(),
+        [](AOTInductorModel* m) { return !m->is_finished(); });
+#endif
+
+    if (it != pending_models_.end()) {
+      // We have finished model instances that can be pushed into
+      // available_models_ so that we don't have to be blocked on waiting
+      // the pending_models_available_ condition.
+      available_models_.insert(
+          available_models_.end(), it, pending_models_.end());
+      pending_models_.erase(it, pending_models_.end());
+      return;
+    }
+
+    pending_models_available_.wait(
+        lk, [this]() { return !pending_models_.empty(); });
+    // Let's make the schedule simple first. We always wait on the first
+    // pending_models_ to be complete.
+    auto* model = pending_models_.front();
+    pending_models_.pop_front();
+    lk.unlock();
+    try {
+      model->wait_for_completion();
+    } catch (...) {
+      lk.lock();
+      available_models_.push_back(model);
+      throw;
+    }
+    lk.lock();
+    available_models_.push_back(model);
+  }
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..3ba22f3a286f0cb7725f2d013ed0bc4c447a7ffb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h
@@ -0,0 +1,43 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/complex.h>
+#include <torch/csrc/inductor/aoti_runtime/utils.h>
+
+namespace torch::aot_inductor {
+
+template <typename T>
+inline RAIIAtenTensorHandle scalar_to_tensor_handle(T value) {
+  throw std::runtime_error("Unsupported scalar_to_tensor_handle");
+}
+
+// Specialize for supported C++ primitive types
+#define AOTI_RUNTIME_SCALAR_TO_TENSOR(dtype, ctype)                         \
+  template <>                                                               \
+  inline RAIIAtenTensorHandle scalar_to_tensor_handle<ctype>(ctype value) { \
+    AtenTensorHandle tensor_handle;                                         \
+    AOTI_TORCH_ERROR_CODE_CHECK(                                            \
+        aoti_torch_scalar_to_tensor_##dtype(value, &tensor_handle));        \
+    return RAIIAtenTensorHandle(tensor_handle);                             \
+  }
+
+AOTI_RUNTIME_SCALAR_TO_TENSOR(float32, float)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(float64, double)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint8, uint8_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint16, uint16_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint32, uint32_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(uint64, uint64_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(int8, int8_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(int16, int16_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(int32, int32_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(int64, int64_t)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(bool, bool)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(complex64, c10::complex<float>)
+AOTI_RUNTIME_SCALAR_TO_TENSOR(complex128, c10::complex<double>)
+#undef AOTI_RUNTIME_SCALAR_TO_TENSOR
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h
new file mode 100644
index 0000000000000000000000000000000000000000..e9f9bdab9c5270e6fc9aababcd5587ad54bfee5c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h
@@ -0,0 +1,183 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// NOLINT
+#pragma once
+#ifdef USE_XPU
+#include <c10/xpu/XPUFunctions.h>
+#include <level_zero/ze_api.h>
+#include <sycl/sycl.hpp>
+#include <fstream>
+#include <iostream>
+#include <string>
+
+#define ZE_CHECK(status)                                                  \
+  {                                                                       \
+    if (status != ZE_RESULT_SUCCESS) {                                    \
+      std::stringstream ss;                                               \
+      ss << "L0 runtime error: " << std::hex << std::uppercase << status; \
+      throw std::runtime_error(ss.str());                                 \
+    }                                                                     \
+  }
+
+static ze_module_handle_t _createModule(
+    const uint8_t* binaryPtr,
+    size_t binarySize,
+    bool isSpirv = false) {
+  sycl::device& syclDevice =
+      c10::xpu::get_raw_device(c10::xpu::current_device());
+  auto& syclContext = c10::xpu::get_device_context();
+  auto device =
+      sycl::get_native<sycl::backend::ext_oneapi_level_zero>(syclDevice);
+  auto context =
+      sycl::get_native<sycl::backend::ext_oneapi_level_zero>(syclContext);
+
+  const char* buildFlags = "";
+  const ze_module_format_t format =
+      isSpirv ? ZE_MODULE_FORMAT_IL_SPIRV : ZE_MODULE_FORMAT_NATIVE;
+  ze_module_desc_t moduleDescription = {};
+  moduleDescription.stype = ZE_STRUCTURE_TYPE_MODULE_DESC;
+  moduleDescription.format = format;
+  moduleDescription.inputSize = binarySize;
+  moduleDescription.pInputModule = (uint8_t*)binaryPtr;
+  moduleDescription.pBuildFlags = buildFlags;
+  ze_module_build_log_handle_t buildLog = nullptr;
+  ze_module_handle_t module = nullptr;
+  auto error_no = ZE_RESULT_SUCCESS;
+  error_no =
+      zeModuleCreate(context, device, &moduleDescription, &module, &buildLog);
+
+  if (error_no != ZE_RESULT_SUCCESS) {
+    size_t szLog = 0;
+    ZE_CHECK(zeModuleBuildLogGetString(buildLog, &szLog, nullptr));
+    char* strLog = (char*)malloc(szLog);
+    ZE_CHECK(zeModuleBuildLogGetString(buildLog, &szLog, strLog));
+    std::cerr << "L0 build module failed. Log: " << strLog << std::endl;
+    free(strLog);
+  }
+  if (buildLog) {
+    ZE_CHECK(zeModuleBuildLogDestroy(buildLog));
+  }
+  ZE_CHECK(error_no);
+  return module;
+}
+
+static std::unique_ptr<sycl::kernel> _createKernel(
+    ze_module_handle_t module,
+    const char* kernelName) {
+  assert(module);
+  assert(kernelName);
+  ze_kernel_handle_t kernel = nullptr;
+  ze_kernel_desc_t kernelDescription = {};
+  kernelDescription.stype = ZE_STRUCTURE_TYPE_KERNEL_DESC;
+  kernelDescription.pNext = nullptr;
+  kernelDescription.flags = ZE_KERNEL_FLAG_FORCE_RESIDENCY;
+  kernelDescription.pKernelName = kernelName;
+  ZE_CHECK(zeKernelCreate(module, &kernelDescription, &kernel));
+
+  auto& syclContext = c10::xpu::get_device_context();
+  auto mod = sycl::make_kernel_bundle<
+      sycl::backend::ext_oneapi_level_zero,
+      sycl::bundle_state::executable>(
+      {module, sycl::ext::oneapi::level_zero::ownership::transfer},
+      syclContext);
+  auto fun = sycl::make_kernel<sycl::backend::ext_oneapi_level_zero>(
+      {mod, kernel, sycl::ext::oneapi::level_zero::ownership::transfer},
+      syclContext);
+  return std::make_unique<sycl::kernel>(fun);
+}
+
+// GPU Cpp Wrapper API
+[[maybe_unused]] static std::unique_ptr<sycl::kernel> loadKernel(
+    std::string filePath,
+    const std::string& funcName,
+    uint32_t sharedMemBytes,
+    const std::optional<std::string>& binDir = std::nullopt) {
+  if (binDir) {
+    std::filesystem::path p1{*binDir};
+    std::filesystem::path p2{filePath};
+    filePath = (p1 / p2.filename()).string();
+  }
+
+  std::ifstream IFS(filePath.c_str(), std::ios::binary);
+  std::ostringstream OSS;
+  OSS << IFS.rdbuf();
+  std::string data(OSS.str());
+
+  bool isSpirv = filePath.size() >= 4 &&
+      filePath.compare(filePath.size() - 4, 4, ".spv") == 0;
+  auto mod = _createModule(
+      reinterpret_cast<const uint8_t*>(data.c_str()), data.size(), isSpirv);
+
+  return _createKernel(mod, funcName.c_str());
+}
+
+// GPU Cpp Wrapper API
+[[maybe_unused]] static std::unique_ptr<sycl::kernel> loadKernel(
+    const void* start,
+    const void* end,
+    const std::string& funcName,
+    uint32_t sharedMemBytes,
+    bool isSpirv) {
+  size_t size = reinterpret_cast<const uint8_t*>(end) -
+      reinterpret_cast<const uint8_t*>(start);
+
+  auto mod =
+      _createModule(reinterpret_cast<const uint8_t*>(start), size, isSpirv);
+
+  return _createKernel(mod, funcName.c_str());
+}
+
+// GPU Cpp Wrapper API
+[[maybe_unused]] static void launchKernel(
+    std::unique_ptr<sycl::kernel>& kernelPtr,
+    uint32_t gridX,
+    uint32_t gridY,
+    uint32_t gridZ,
+    uint32_t numWarps,
+    uint32_t sharedMemory,
+    void** params,
+    sycl::queue* queuePtr) {
+  uint32_t threadsPerWarp = kernelPtr->get_info<
+      sycl::info::kernel_device_specific::compile_sub_group_size>(
+      queuePtr->get_device());
+  if (threadsPerWarp == 0) {
+    threadsPerWarp = 32; // default to 32 if not set
+  }
+  std::string kernelName =
+      kernelPtr->get_info<sycl::info::kernel::function_name>();
+  uint32_t numParams = kernelPtr->get_info<sycl::info::kernel::num_args>();
+  size_t globalRangeX = gridX * threadsPerWarp * numWarps;
+  size_t globalRangeY = gridY;
+  size_t globalRangeZ = gridZ;
+  size_t localRangeX = numWarps * threadsPerWarp;
+  size_t localRangeY = 1;
+  size_t localRangeZ = 1;
+  sycl::range<3> globalRange(globalRangeZ, globalRangeY, globalRangeX);
+  sycl::range<3> localRange(localRangeZ, localRangeY, localRangeX);
+  sycl::nd_range<3> parallelWorkSize(globalRange, localRange);
+  if (sharedMemory) {
+    // numParams from sycl info  = user provided args + sharedMemoryBuffer
+    numParams -= 1;
+  }
+  // Submit the imported kernel.
+  auto cgf = [&](sycl::handler& cgh) {
+    for (uint32_t i = 0; i < numParams; ++i) {
+      cgh.set_arg(i, *(static_cast<void**>(params[i])));
+    }
+
+    if (sharedMemory > 0) {
+      constexpr int dimensions = 1;
+      using share_mem_t = sycl::local_accessor<int8_t, dimensions>;
+      share_mem_t localBuffer = share_mem_t(sharedMemory, cgh);
+      cgh.set_arg(numParams, localBuffer);
+      cgh.parallel_for(parallelWorkSize, *kernelPtr);
+    } else {
+      cgh.parallel_for(parallelWorkSize, *kernelPtr);
+    }
+  };
+  auto event = queuePtr->submit(cgf);
+}
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/thread_local.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/thread_local.h
new file mode 100644
index 0000000000000000000000000000000000000000..7aea78361dc5765eae78d748e3022cfbfa1ec63d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/thread_local.h
@@ -0,0 +1,165 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor.h>
+
+namespace torch::aot_inductor {
+
+template <typename T>
+struct ThreadLocalCachedOutputTensor;
+
+template <>
+struct ThreadLocalCachedOutputTensor<RAIIAtenTensorHandle> {
+  explicit ThreadLocalCachedOutputTensor(const RAIIAtenTensorHandle&) {}
+  void copy_data_from(const RAIIAtenTensorHandle& handle) {
+    throw std::runtime_error("can't happen");
+  }
+
+  AtenTensorHandle tensor() const {
+    throw std::runtime_error("can't happen");
+  }
+};
+
+template <>
+struct ThreadLocalCachedOutputTensor<AtenTensorHandle> {
+  explicit ThreadLocalCachedOutputTensor(const AtenTensorHandle&) {}
+  void copy_data_from(const AtenTensorHandle& handle) {
+    throw std::runtime_error("can't happen");
+  }
+
+  AtenTensorHandle tensor() const {
+    throw std::runtime_error("can't happen");
+  }
+};
+
+template <>
+struct ThreadLocalCachedOutputTensor<ConstantHandle> {
+  explicit ThreadLocalCachedOutputTensor(const ConstantHandle&) {}
+  void copy_data_from(const ConstantHandle& handle) {
+    throw std::runtime_error("can't happen");
+  }
+
+  AtenTensorHandle tensor() const {
+    throw std::runtime_error("can't happen");
+  }
+};
+
+template <typename T>
+struct ThreadLocalCachedOutputTensor<ArrayRefTensor<T>> {
+  explicit ThreadLocalCachedOutputTensor(const ArrayRefTensor<T>& t) {
+    realloc(t);
+  }
+
+  void copy_data_from(const ArrayRefTensor<T>& t) {
+    if (t.numel() > capacity_) {
+      realloc(t);
+    }
+    std::copy(t.data(), t.data() + t.numel(), storage_.get());
+  }
+
+  AtenTensorHandle tensor() const {
+    return tensor_.get();
+  }
+
+ private:
+  void realloc(const ArrayRefTensor<T>& t) {
+    capacity_ = t.numel();
+    // NOLINTNEXTLINE(*arrays*)
+    storage_ = std::make_unique<T[]>(t.numel());
+    AtenTensorHandle handle = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob(
+        storage_.get(),
+        t.sizes().size(),
+        t.sizes().data(),
+        t.strides().data(),
+        0,
+        aoti_torch_dtype<std::remove_const_t<T>>(),
+        t.device_type(),
+        t.device_idx(),
+        &handle));
+    tensor_ = handle;
+  }
+
+  // NOLINTNEXTLINE(*arrays*)
+  std::unique_ptr<T[]> storage_;
+  int64_t capacity_ = 0;
+  RAIIAtenTensorHandle tensor_;
+};
+
+template <typename T>
+struct ThreadLocalCachedOutputArray;
+
+// Just needs to compile, doesn't need to do anything.
+template <>
+struct ThreadLocalCachedOutputArray<RAIIAtenTensorHandle> {
+  explicit ThreadLocalCachedOutputArray(const RAIIAtenTensorHandle&) {
+    throw std::runtime_error("can't happen");
+  }
+
+  // Not supported yet! We would need to put contiguous() or
+  // expect_contiguous() into the ABI.
+  void copy_data_from(const RAIIAtenTensorHandle&) {
+    throw std::runtime_error("can't happen");
+  }
+
+  template <typename U>
+  ArrayRefTensor<U> arrayref_tensor() const {
+    throw std::runtime_error("can't happen");
+  }
+};
+
+// Just needs to compile, doesn't need to do anything.
+template <>
+struct ThreadLocalCachedOutputArray<ConstantHandle> {
+  explicit ThreadLocalCachedOutputArray(const ConstantHandle&) {
+    throw std::runtime_error("can't happen");
+  }
+
+  // Not supported yet! We would need to put contiguous() or
+  // expect_contiguous() into the ABI.
+  void copy_data_from(const ConstantHandle&) {
+    throw std::runtime_error("can't happen");
+  }
+
+  template <typename U>
+  ArrayRefTensor<U> arrayref_tensor() const {
+    throw std::runtime_error("can't happen");
+  }
+};
+
+template <typename T>
+struct ThreadLocalCachedOutputArray<ArrayRefTensor<T>> {
+  explicit ThreadLocalCachedOutputArray(const ArrayRefTensor<T>& t) {}
+
+  template <
+      typename U,
+      std::enable_if_t<
+          std::is_same_v<std::remove_const_t<T>, std::remove_const_t<U>>,
+          bool> = true>
+  ArrayRefTensor<T> arrayref_tensor() const {
+    return tensor_;
+  }
+
+  void copy_data_from(const ArrayRefTensor<T>& t) {
+    if (t.numel() > capacity_) {
+      capacity_ = t.numel();
+      // NOLINTNEXTLINE(*arrays*)
+      storage_ = std::make_unique<T[]>(capacity_);
+    }
+    std::copy(t.data(), t.data() + t.numel(), storage_.get());
+    tensor_ = t;
+    tensor_.set_arrayref(MiniArrayRef<T>(storage_.get(), t.numel()));
+  }
+
+ private:
+  // NOLINTNEXTLINE(*arrays*)
+  std::unique_ptr<T[]> storage_;
+  uint32_t capacity_ = 0;
+  ArrayRefTensor<T> tensor_;
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..1eb79da9f82d5bbb45f9e9452da1842cfacf7347
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils.h
@@ -0,0 +1,484 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <iostream>
+#include <memory>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/headeronly/util/shim_utils.h>
+
+#if defined(__GNUC__) || defined(__clang__)
+#define AOTI_NOINLINE __attribute__((noinline))
+#elif _MSC_VER
+#define AOTI_NOINLINE __declspec(noinline)
+#else
+#define AOTI_NOINLINE
+#endif
+
+#define AOTI_TORCH_ERROR_CODE_CHECK(call)                                  \
+  if ((call) != AOTI_TORCH_SUCCESS) {                                      \
+    torch::headeronly::detail::throw_exception(#call, __FILE__, __LINE__); \
+  }
+
+using AOTIRuntimeError = int32_t;
+#define AOTI_RUNTIME_SUCCESS 0
+#define AOTI_RUNTIME_FAILURE 1
+
+#define AOTI_RUNTIME_ERROR_CODE_CHECK(call)                                \
+  if ((call) != AOTI_RUNTIME_SUCCESS) {                                    \
+    torch::headeronly::detail::throw_exception(#call, __FILE__, __LINE__); \
+  }
+
+namespace torch::aot_inductor {
+
+using DeleterFnPtr = void (*)(void*);
+
+inline void noop_deleter(void* /*unused*/) {}
+
+inline void delete_record_function_object(void* ptr) {
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_record_function_end(
+      reinterpret_cast<AtenRecordFunctionHandle>(ptr)));
+}
+
+inline void delete_tensor_object(void* ptr) {
+  AOTI_TORCH_ERROR_CODE_CHECK(
+      aoti_torch_delete_tensor_object(reinterpret_cast<AtenTensorHandle>(ptr)));
+}
+
+inline void delete_c10_value_object(void* ptr) {
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_delete_c10_value_object(
+      reinterpret_cast<C10IValueHandle>(ptr)));
+}
+
+class RAIIAtenRecordFunctionHandle {
+ public:
+  RAIIAtenRecordFunctionHandle() : handle_(nullptr, noop_deleter) {}
+  RAIIAtenRecordFunctionHandle(const RAIIAtenRecordFunctionHandle& other) =
+      delete;
+  RAIIAtenRecordFunctionHandle& operator=(
+      const RAIIAtenRecordFunctionHandle& other) = delete;
+
+  // Initiate an RAII RecordFunction without Inputs
+  RAIIAtenRecordFunctionHandle(const char* name, IValueMapHandle kwargs)
+      : handle_(nullptr, delete_record_function_object) {
+    AtenRecordFunctionHandle tmp_handle = nullptr;
+    aoti_record_function_start(name, kwargs, nullptr, 0, &tmp_handle);
+    handle_.reset(tmp_handle);
+  }
+
+  // Initiate an RAII RecordFunction with Inputs
+  RAIIAtenRecordFunctionHandle(
+      const char* name,
+      IValueMapHandle kwargs,
+      std::vector<C10IValueHandle> inputs)
+      : handle_(nullptr, delete_record_function_object) {
+    AtenRecordFunctionHandle tmp_handle = nullptr;
+    aoti_record_function_start(
+        name, kwargs, inputs.data(), inputs.size(), &tmp_handle);
+    handle_.reset(tmp_handle);
+  }
+
+  // Steal the ownership from another RAIIAtenRecordFunctionHandle using
+  // std::move
+  RAIIAtenRecordFunctionHandle(RAIIAtenRecordFunctionHandle&& other) = default;
+  RAIIAtenRecordFunctionHandle& operator=(
+      RAIIAtenRecordFunctionHandle&& other) = default;
+
+  // Steal the ownership from raw AtenRecordFunctionHandle
+  RAIIAtenRecordFunctionHandle(AtenRecordFunctionHandle handle)
+      : handle_(handle, delete_record_function_object) {}
+
+  ~RAIIAtenRecordFunctionHandle() {
+    handle_.reset();
+  }
+
+  // Return a raw AtenRecordFunctionHandle to be used by aoti_torch functions
+  // Note: this function does NOT transfer the ownership of the handle
+  operator AtenRecordFunctionHandle() const {
+    return handle_.get();
+  }
+
+  AtenRecordFunctionHandle release() {
+    return handle_.release();
+  }
+
+  AtenRecordFunctionHandle get() const {
+    return handle_.get();
+  }
+
+  void reset() {
+    handle_.reset();
+  }
+
+ private:
+  std::unique_ptr<AtenRecordFunctionOpaque, DeleterFnPtr> handle_;
+};
+
+// RAIIAtenTensorHandle steals the tensor objects created by the libtorch C ABI
+class RAIIAtenTensorHandle {
+ public:
+  RAIIAtenTensorHandle() : handle_(nullptr, noop_deleter) {}
+  RAIIAtenTensorHandle(const RAIIAtenTensorHandle& other) = delete;
+  RAIIAtenTensorHandle& operator=(const RAIIAtenTensorHandle& other) = delete;
+
+  // Steal the ownership from another RAIIAtenTensorHandle using std::move
+  RAIIAtenTensorHandle(RAIIAtenTensorHandle&& other) = default;
+  RAIIAtenTensorHandle& operator=(RAIIAtenTensorHandle&& other) = default;
+
+  // Steal the ownership from raw AtenTensorHandle
+  RAIIAtenTensorHandle(AtenTensorHandle handle)
+      : handle_(handle, delete_tensor_object) {}
+
+  ~RAIIAtenTensorHandle() {
+    handle_.reset();
+  }
+
+  // Return a raw AtenTensorHandle to be used by aoti_torch functions
+  // Note: this function does NOT transfer the ownership of the handle
+  operator AtenTensorHandle() const {
+    return handle_.get();
+  }
+
+  AtenTensorHandle release() {
+    return handle_.release();
+  }
+
+  AtenTensorHandle get() const {
+    return handle_.get();
+  }
+
+  void reset() {
+    handle_.reset();
+  }
+
+  int64_t size(int64_t d) {
+    int64_t size = 0;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_size(handle_.get(), d, &size));
+    return size;
+  }
+
+  int64_t stride(int64_t d) {
+    int64_t stride = 0;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_stride(handle_.get(), d, &stride));
+    return stride;
+  }
+
+  int64_t storage_offset() {
+    int64_t storage_offset = 0;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_storage_offset(handle_.get(), &storage_offset));
+    return storage_offset;
+  }
+
+  void* data_ptr() const {
+    void* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_data_ptr(handle_.get(), &result));
+    return result;
+  }
+
+  int64_t* sizes() const {
+    int64_t* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle_.get(), &result));
+    return result;
+  }
+
+  int64_t* strides() const {
+    int64_t* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle_.get(), &result));
+    return result;
+  }
+
+ private:
+  std::unique_ptr<AtenTensorOpaque, DeleterFnPtr> handle_;
+};
+
+// RAIIC10IValueHandle steals the IValue objects created by the libtorch C ABI
+class RAIIC10IValueHandle {
+ public:
+  RAIIC10IValueHandle() : handle_(nullptr, noop_deleter) {}
+  RAIIC10IValueHandle(const RAIIC10IValueHandle& other) = delete;
+  RAIIC10IValueHandle& operator=(const RAIIC10IValueHandle& other) = delete;
+
+  // Steal the ownership from another RAIIC10IValueHandle using std::move
+  RAIIC10IValueHandle(RAIIC10IValueHandle&& other) = default;
+  RAIIC10IValueHandle& operator=(RAIIC10IValueHandle&& other) = default;
+
+  // Steal the ownership from raw C10IValueHandle
+  RAIIC10IValueHandle(C10IValueHandle handle)
+      : handle_(handle, delete_c10_value_object) {}
+
+  ~RAIIC10IValueHandle() {
+    handle_.reset();
+  }
+
+  // Return a raw C10IValueHandle to be used by aoti_torch functions
+  // Note: this function does NOT transfer the ownership of the handle
+  operator C10IValueHandle() const {
+    return handle_.get();
+  }
+
+  C10IValueHandle release() {
+    return handle_.release();
+  }
+
+  C10IValueHandle get() const {
+    return handle_.get();
+  }
+
+  void reset() {
+    handle_.reset();
+  }
+
+ private:
+  std::unique_ptr<C10IValueOpaque, DeleterFnPtr> handle_;
+};
+
+class MaybeOwningAtenTensorHandle {
+ public:
+  MaybeOwningAtenTensorHandle() : handle_(nullptr) {}
+  // We skip copy constructor as MaybeOwningAtenTensorHandle might be RAII which
+  // makes it undefined.
+  MaybeOwningAtenTensorHandle(const MaybeOwningAtenTensorHandle& other) =
+      delete;
+  MaybeOwningAtenTensorHandle& operator=(
+      const MaybeOwningAtenTensorHandle& other) = delete;
+
+  // Move constructor and move assignment operator
+  MaybeOwningAtenTensorHandle(MaybeOwningAtenTensorHandle&& other) = default;
+  MaybeOwningAtenTensorHandle& operator=(MaybeOwningAtenTensorHandle&& other) =
+      default;
+
+  // Steal the ownership from another RAIIAtenTensorHandle using std::move
+  MaybeOwningAtenTensorHandle(RAIIAtenTensorHandle&& other)
+      : raii_handle_(std::move(other)) {
+    handle_ = raii_handle_.get();
+  }
+  MaybeOwningAtenTensorHandle& operator=(RAIIAtenTensorHandle&& other) {
+    raii_handle_ = std::move(other);
+    handle_ = raii_handle_.get();
+    return *this;
+  }
+
+  // By default, steal the ownership from raw AtenTensorHandle
+  MaybeOwningAtenTensorHandle(AtenTensorHandle handle) : raii_handle_(handle) {
+    handle_ = raii_handle_.get();
+  }
+
+  // If user_managed is true, we do not steal the ownership.
+  MaybeOwningAtenTensorHandle(AtenTensorHandle handle, bool user_managed) {
+    if (user_managed) {
+      aoti_torch_new_tensor_handle(handle, &handle_);
+    } else {
+      raii_handle_ = RAIIAtenTensorHandle(handle);
+      handle_ = raii_handle_.get();
+    }
+  }
+
+  ~MaybeOwningAtenTensorHandle() {
+    // This is no-op if we don't hold raii_handle with the
+    // MaybeOwningAtenTensorHandle.
+    raii_handle_.reset();
+  }
+
+  // Return a raw AtenTensorHandle to be used by aoti_torch functions
+  // Note: this function does NOT transfer the ownership of the handle
+  operator AtenTensorHandle() const {
+    return handle_;
+  }
+
+  AtenTensorHandle release() {
+    if (raii_handle_) {
+      return raii_handle_.release();
+    } else {
+      AtenTensorHandle handle = handle_;
+      handle_ = nullptr;
+      return handle;
+    }
+  }
+
+  AtenTensorHandle get() const {
+    return handle_;
+  }
+
+  void reset() {
+    handle_ = nullptr;
+    raii_handle_.reset();
+  }
+
+  int64_t size(int64_t d) {
+    int64_t size = 0;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_size(handle_, d, &size));
+    return size;
+  }
+
+  int64_t stride(int64_t d) {
+    int64_t stride = 0;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_stride(handle_, d, &stride));
+    return stride;
+  }
+
+  int64_t storage_offset() {
+    int64_t storage_offset = 0;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_storage_offset(handle_, &storage_offset));
+    return storage_offset;
+  }
+
+  void* data_ptr() const {
+    void* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(handle_, &result));
+    return result;
+  }
+
+  int64_t* sizes() const {
+    int64_t* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle_, &result));
+    return result;
+  }
+
+  int64_t* strides() const {
+    int64_t* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle_, &result));
+    return result;
+  }
+
+ private:
+  // handle_ is the underlying AtenTensorHandle of raii_handle_ if raii_handle_
+  // exists. Otherwise it would just be the AtenTensorHandle passed in by users.
+  AtenTensorHandle handle_;
+  RAIIAtenTensorHandle raii_handle_;
+};
+
+// Steal the ownership from raw AtenTensorHandle to RAIIAtenTensorHandle
+inline std::vector<RAIIAtenTensorHandle> steal_from_raw_handles_to_raii_handles(
+    AtenTensorHandle* handles,
+    size_t size) {
+  std::vector<RAIIAtenTensorHandle> result;
+  result.reserve(size);
+  for (size_t i = 0; i < size; i++) {
+    result.emplace_back(handles[i]);
+    handles[i] = nullptr;
+  }
+  return result;
+}
+
+inline AtenTensorHandle reinterpret_tensor_wrapper(
+    AtenTensorHandle self,
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int64_t storage_offset) {
+  AtenTensorHandle result = nullptr;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch__reinterpret_tensor(
+      self, ndim, sizes_ptr, strides_ptr, storage_offset, &result));
+  return result;
+}
+
+inline void* get_data_ptr_wrapper(AtenTensorHandle tensor) {
+  void* result = nullptr;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(tensor, &result));
+  return result;
+}
+
+inline AtenTensorHandle unwrap_raii_handle_if_needed(
+    const RAIIAtenTensorHandle& handle) {
+  return handle.get();
+}
+
+inline RAIIAtenTensorHandle wrap_with_raii_handle_if_needed(
+    AtenTensorHandle handle) {
+  return RAIIAtenTensorHandle(handle);
+}
+
+class ConstantHandle {
+ public:
+  ConstantHandle() = default;
+
+  explicit ConstantHandle(AtenTensorHandle handle) : handle_(handle) {
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(handle_, &data_));
+  }
+
+  operator AtenTensorHandle() const {
+    return handle_;
+  }
+
+  AtenTensorHandle tensor() const {
+    return handle_;
+  }
+
+  AtenTensorHandle get() const {
+    return handle_;
+  }
+
+  void* data_ptr() const {
+    return data_;
+  }
+
+  int64_t* sizes() const {
+    int64_t* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle_, &result));
+    return result;
+  }
+
+  int64_t* strides() const {
+    int64_t* result = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle_, &result));
+    return result;
+  }
+
+ private:
+  AtenTensorHandle handle_{};
+  void* data_ = nullptr;
+};
+
+inline void* get_data_ptr_wrapper(const ConstantHandle& constant) {
+  return constant.data_ptr();
+}
+
+inline const ConstantHandle& unwrap_raii_handle_if_needed(
+    const ConstantHandle& handle) {
+  return handle;
+}
+
+// Shouldn't be called.
+inline AtenTensorHandle wrap_with_raii_handle_if_needed(
+    const ConstantHandle& handle) = delete;
+
+// DANGEROUS.  Do not call unless you explicitly intend to get a reference to a
+// temporary value, which will expire at the end of the current expression.
+// This should only be called in cases where the C-shim API expects an optional
+// input argument (passed by pointer), and a temporary needs to be passed to it.
+template <class T>
+T& temporary_reference(T&& t) {
+  return t;
+}
+
+#define CACHE_TORCH_DTYPE(typename) \
+  static auto cached_torch_dtype_##typename = aoti_torch_dtype_##typename()
+
+#define CACHE_TORCH_DEVICE(device)                \
+  static auto cached_torch_device_type_##device = \
+      aoti_torch_device_type_##device()
+
+#define CACHE_TORCH_LAYOUT(layout) \
+  static auto cached_torch_layout_##layout = aoti_torch_layout_##layout()
+
+#define CACHE_TORCH_MEMORY_FORMAT(format)           \
+  static auto cached_torch_memory_format_##format = \
+      aoti_torch_memory_format_##format()
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils_cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils_cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..3dce90bd4059e185b68ead3729c2bb729cf7cbf7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils_cuda.h
@@ -0,0 +1,68 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_CUDA
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+#include <torch/csrc/inductor/aoti_runtime/utils.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#ifndef USE_ROCM
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#endif
+
+namespace torch::aot_inductor {
+
+inline void delete_cuda_guard(void* ptr) {
+  AOTI_TORCH_ERROR_CODE_CHECK(
+      aoti_torch_delete_cuda_guard(reinterpret_cast<CUDAGuardHandle>(ptr)));
+}
+
+inline void delete_cuda_stream_guard(void* ptr) {
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_delete_cuda_stream_guard(
+      reinterpret_cast<CUDAStreamGuardHandle>(ptr)));
+}
+
+class AOTICudaGuard {
+ public:
+  AOTICudaGuard(int32_t device_index) : guard_(nullptr, delete_cuda_guard) {
+    CUDAGuardHandle ptr = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_create_cuda_guard(device_index, &ptr));
+    guard_.reset(ptr);
+  }
+
+  void set_index(int32_t device_index) {
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_cuda_guard_set_index(guard_.get(), device_index));
+  }
+
+ private:
+  std::unique_ptr<CUDAGuardOpaque, DeleterFnPtr> guard_;
+};
+
+class AOTICudaStreamGuard {
+ public:
+  AOTICudaStreamGuard(cudaStream_t stream, int32_t device_index)
+      : guard_(nullptr, delete_cuda_stream_guard) {
+    CUDAStreamGuardHandle ptr = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_create_cuda_stream_guard(stream, device_index, &ptr));
+    guard_.reset(ptr);
+  }
+
+ private:
+  std::unique_ptr<CUDAStreamGuardOpaque, DeleterFnPtr> guard_;
+};
+
+} // namespace torch::aot_inductor
+#endif // USE_CUDA
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils_xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils_xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..50149503b3117d912df41cabfe395af47bd7cbc5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_runtime/utils_xpu.h
@@ -0,0 +1,61 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_XPU
+// WARNING: Be careful when adding new includes here. This header will be used
+// in model.so, and should not refer to any aten/c10 headers except the stable
+// C ABI defined in torch/csrc/inductor/aoti_torch/c/shim.h. The same rule
+// applies to other files under torch/csrc/inductor/aoti_runtime/.
+#include <torch/csrc/inductor/aoti_runtime/utils.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim_xpu.h>
+
+namespace torch::aot_inductor {
+
+inline void delete_xpu_guard(void* ptr) {
+  AOTI_TORCH_ERROR_CODE_CHECK(
+      aoti_torch_delete_xpu_guard(reinterpret_cast<XPUGuardHandle>(ptr)));
+}
+
+inline void delete_xpu_stream_guard(void* ptr) {
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_delete_xpu_stream_guard(
+      reinterpret_cast<XPUStreamGuardHandle>(ptr)));
+}
+
+class AOTIXpuGuard {
+ public:
+  AOTIXpuGuard(int32_t device_index) : guard_(nullptr, delete_xpu_guard) {
+    XPUGuardHandle ptr = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_create_xpu_guard(device_index, &ptr));
+    guard_.reset(ptr);
+  }
+
+  void set_index(int32_t device_index) {
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_xpu_guard_set_index(guard_.get(), device_index));
+  }
+
+ private:
+  std::unique_ptr<XPUGuardOpaque, DeleterFnPtr> guard_;
+};
+
+class AOTIXpuStreamGuard {
+ public:
+  AOTIXpuStreamGuard(void* stream, int32_t device_index)
+      : guard_(nullptr, delete_xpu_stream_guard) {
+    XPUStreamGuardHandle ptr = nullptr;
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_create_xpu_stream_guard(stream, device_index, &ptr));
+    guard_.reset(ptr);
+  }
+
+ private:
+  std::unique_ptr<XPUStreamGuardOpaque, DeleterFnPtr> guard_;
+};
+
+} // namespace torch::aot_inductor
+#endif // USE_XPU
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/macros.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/macros.h
new file mode 100644
index 0000000000000000000000000000000000000000..e49cd39deac0c8f2f0a5939c1503f515abed0d04
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/macros.h
@@ -0,0 +1,66 @@
+#ifndef AOTI_TORCH_MACRO_H
+#define AOTI_TORCH_MACRO_H
+
+#include <stddef.h>
+#include <stdint.h>
+#ifdef __GNUC__
+#define AOTI_TORCH_EXPORT __attribute__((__visibility__("default")))
+#else // !__GNUC__
+#ifdef _WIN32
+// PyTorch2 doesn't currently work on Windows. Exporting these APIs can lead
+// to symbol clashes at link time if libtorch is included in a DLL and binary
+// that depends on the DLL. As a short term fix, we don't export the symbols.
+// In the long term, this will need to be addressed when Windows is supported.
+#ifdef OVRSOURCE
+// Do not export AOTI on Windows for internal builds
+#define AOTI_TORCH_EXPORT
+#else /* OVRSOURCE */
+#ifdef EXPORT_AOTI_FUNCTIONS
+#define AOTI_TORCH_EXPORT __declspec(dllexport)
+#else
+#define AOTI_TORCH_EXPORT __declspec(dllimport)
+#endif
+#endif /* OVRSOURCE */
+#else // !_WIN32
+#define AOTI_TORCH_EXPORT
+#endif // _WIN32
+#endif // __GNUC__
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+// AtenTensorHandle represents an abstract notion of Tensor that can be passed
+// between model.so and libtorch.so.  The contents of the structure itself
+// are private; model.so is not allowed to access any fields directly, it must
+// go through functions defined in this ABI.  Under the hood, this is
+// represented as at::Tensor*, but we reserve the right to change this (and in
+// fact, we probably should change it to at::TensorImpl* at least).
+//
+// An AtenTensorHandle can be owning (please check the API reference for exact
+// ownership/borrow semantics).  If you have an owning AtenTensorHandle
+// in model.so, you are obligated to aoti_torch_delete_tensor_object when you
+// are done.  You can use the helper C++ class RAIIAtenTensorHandle
+// (see aot_runtime/model.h) to ensure the deallocator is called in RAII style
+// (note that RAIIAtenTensorHandle is private to model.so, and never crosses
+// the ABI boundary.)
+struct AtenTensorOpaque;
+using AtenTensorHandle = AtenTensorOpaque*;
+
+struct AtenGeneratorOpaque;
+using AtenGeneratorHandle = AtenGeneratorOpaque*;
+
+struct AOTIProxyExecutorOpaque;
+using AOTIProxyExecutorHandle = AOTIProxyExecutorOpaque*;
+
+struct C10IValueOpaque;
+using C10IValueHandle = C10IValueOpaque*;
+
+using AOTITorchError = int32_t;
+#define AOTI_TORCH_SUCCESS 0
+#define AOTI_TORCH_FAILURE 1
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#endif // AOTI_TORCH_MACRO_H
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim.h
new file mode 100644
index 0000000000000000000000000000000000000000..f719828ba00d6c3e135267c6a42210483b9b1122
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim.h
@@ -0,0 +1,697 @@
+#ifndef AOTI_TORCH_SHIM
+#define AOTI_TORCH_SHIM
+
+#include <torch/csrc/inductor/aoti_torch/c/macros.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim_deprecated.h>
+#include <torch/headeronly/util/Exception.h>
+
+// This header defines a stable C API for certain ATen functionality in
+// libtorch. The AOTInductor compiled model.so will only refer to this header
+// instead of other headers from aten/c10, which means it will NOT be able to
+// directly use any data structures or call functions from libtorch.
+//
+// What problems are we trying to solve here?  Direct use of aten/c10 APIs
+// means use of C++ APIs on a library that doesn't have any ABI compatibility
+// guarantees.  However, we want model.so to remain usable across updates
+// to the PyTorch C++ libraries, which requires a stable ABI.  By introducing
+// a C shim layer, we can minimize the surface that will cause breakage. The
+// corresponding software stack can be illustrated as follows:
+//
+// |--------------------------------|
+// |     inference service code     |
+// |--------------------------------|
+// |           model.so             |
+// |--------------|-----------------|
+// |           <c shim>             |
+// |          libtorch.so           |
+// |--------------------------------|
+//
+// The general guidelines for the C API:
+//
+//  - No exceptions, return an explicit error code to be checked at call site
+//  - Only pointers (AtenTensorHandle counts), integers and floats in headers
+//
+// If you want to make changes to this header, you MUST MAINTAIN ABI
+// compatibility.  Typically, this means you will have to add a _v2 version
+// of a function that you, e.g., want to add a new function parameter to, and
+// maintain the old and new versions of the APIs until all old model.so
+// go out of use.
+
+// The following files are implemented in a header-only way and are guarded by
+// test/cpp/aoti_abi_check
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
+#include <torch/headeronly/util/complex.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Getter functions for retrieving various constants from the runtime, that
+// can subsequently be passed to other aoti_* functions.  By hiding these
+// behind functions, the precise value of device/dtype is NOT part of the
+// ABI contract.  (In practice, aten/c10 is pretty good about not renumbering
+// these, so we probably could later switch to having these in the ABI, if
+// desired for perf reasons.)
+AOTI_TORCH_EXPORT int32_t aoti_torch_device_type_cpu();
+AOTI_TORCH_EXPORT int32_t aoti_torch_device_type_cuda();
+AOTI_TORCH_EXPORT int32_t aoti_torch_device_type_meta();
+AOTI_TORCH_EXPORT int32_t aoti_torch_device_type_xpu();
+AOTI_TORCH_EXPORT int32_t aoti_torch_device_type_mps();
+AOTI_TORCH_EXPORT int32_t aoti_torch_device_type_privateuse1();
+
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e5m2();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e4m3fn();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e5m2fnuz();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e4m3fnuz();
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float8_e8m0fnu();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float4_e2m1fn_x2();
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_bfloat16();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float16();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float32();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_float64();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_uint8();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_uint16();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_uint32();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_uint64();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_int8();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_int16();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_int32();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_int64();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_bool();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_complex32();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_complex64();
+AOTI_TORCH_EXPORT int32_t aoti_torch_dtype_complex128();
+AOTI_TORCH_EXPORT size_t aoti_torch_dtype_element_size(int32_t dtype);
+
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout_strided();
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout_sparse_coo();
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout_sparse_csr();
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout_sparse_csc();
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout_sparse_bsr();
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout_sparse_bsc();
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout__mkldnn();
+AOTI_TORCH_EXPORT int32_t aoti_torch_layout_jagged();
+
+AOTI_TORCH_EXPORT int32_t aoti_torch_memory_format_contiguous_format();
+AOTI_TORCH_EXPORT int32_t aoti_torch_memory_format_channels_last();
+AOTI_TORCH_EXPORT int32_t aoti_torch_memory_format_channels_last_3d();
+AOTI_TORCH_EXPORT int32_t aoti_torch_memory_format_preserve_format();
+
+// Get TORCH_ABI_VERSION of the built libtorch.so
+AOTI_TORCH_EXPORT uint64_t aoti_torch_abi_version();
+
+// Functions for converting a single-element tensor to a scalar value
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_float16(AtenTensorHandle tensor, c10::Half* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_float32(AtenTensorHandle tensor, float* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_float64(AtenTensorHandle tensor, double* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_uint8(AtenTensorHandle tensor, uint8_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_uint16(AtenTensorHandle tensor, uint16_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_uint32(AtenTensorHandle tensor, uint32_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_uint64(AtenTensorHandle tensor, uint64_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_int8(AtenTensorHandle tensor, int8_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_int16(AtenTensorHandle tensor, int16_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_int32(AtenTensorHandle tensor, int32_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_int64(AtenTensorHandle tensor, int64_t* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_bool(AtenTensorHandle tensor, bool* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_item_bfloat16(AtenTensorHandle tensor, c10::BFloat16* ret_value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_item_complex64(
+    AtenTensorHandle tensor,
+    c10::complex<float>* ret_value);
+
+// Functions for wrapping a scalar value to a single-element tensor
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_float32(
+    float value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_float64(
+    double value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_uint8(
+    uint8_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_uint16(
+    uint16_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_uint32(
+    uint32_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_uint64(
+    uint64_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_int8(
+    int8_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_int16(
+    int16_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_int32(
+    int32_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_int64(
+    int64_t value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_scalar_to_tensor_bool(bool value, AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_complex64(
+    c10::complex<float> value,
+    AtenTensorHandle* ret_new_tensor);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scalar_to_tensor_complex128(
+    c10::complex<double> value,
+    AtenTensorHandle* ret_new_tensor);
+
+AOTI_TORCH_EXPORT bool aoti_torch_grad_mode_is_enabled();
+AOTI_TORCH_EXPORT void aoti_torch_grad_mode_set_enabled(bool enabled);
+
+// Free the tensor object
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_tensor_object(AtenTensorHandle tensor);
+
+// c10::IValue <int64_t> object conversion
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_int64_to_ivalue(int64_t val, C10IValueHandle* ivalue);
+
+// c10::IValue <const char** > object conversions
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_strlist_to_ivalue(
+    const char** val,
+    int64_t len,
+    C10IValueHandle* ivalue);
+
+// c10::IValue <const char* > object conversions
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_str_to_ivalue(const char* val, C10IValueHandle* ivalue);
+
+// c10::IValue <at::Tensor> object conversions
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_tensor_to_ivalue(AtenTensorHandle val, C10IValueHandle* ivalue);
+
+// Free the c10::IValue object
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_c10_value_object(C10IValueHandle handle);
+
+// Get a pointer to the underlying storage data
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_data_ptr(
+    AtenTensorHandle tensor,
+    void** ret_data_ptr // returns borrowed reference
+);
+
+// Get the nbytes of the underlying storage
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_storage_size(AtenTensorHandle tensor, int64_t* ret_size);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_dim(AtenTensorHandle tensor, int64_t* ret_dim);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_numel(AtenTensorHandle tensor, int64_t* ret_numel);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_storage_numel(AtenTensorHandle tensor, int64_t* ret_numel);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_sizes(
+    AtenTensorHandle tensor,
+    int64_t** ret_sizes // returns borrowed reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_size(AtenTensorHandle tensor, int64_t d, int64_t* ret_size);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_strides(
+    AtenTensorHandle tensor,
+    int64_t** ret_strides // returns borrowed reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_stride(AtenTensorHandle tensor, int64_t d, int64_t* ret_stride);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_dtype(AtenTensorHandle tensor, int32_t* ret_dtype);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_device_type(AtenTensorHandle tensor, int32_t* ret_device_type);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_device_index(AtenTensorHandle tensor, int32_t* ret_device_index);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_layout(AtenTensorHandle tensor, int32_t* ret_layout);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_storage_offset(
+    AtenTensorHandle tensor,
+    int64_t* ret_storage_offset);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_is_contiguous(AtenTensorHandle tensor, bool* ret_is_contiguous);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_is_defined(AtenTensorHandle tensor, bool* ret_is_defined);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_new_tensor_handle(
+    AtenTensorHandle orig_handle,
+    AtenTensorHandle* new_handle);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__alloc_from_pool(
+    AtenTensorHandle self,
+    int64_t offset_bytes,
+    int32_t dtype,
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    AtenTensorHandle* ret_new_tensor);
+
+// This function will create a new tensor object and its pointer is returned
+// through *out. The caller is responsible for wrapping the tensor pointer
+// with RAIIAtenTensorHandle which will call aoti_torch_delete_tensor_object
+// when going out of scope.
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__reinterpret_tensor(
+    AtenTensorHandle self,
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int64_t storage_offset,
+    AtenTensorHandle* ret_new_tensor // returns new reference
+);
+
+// This function will create a new tensor object and its pointer is returned
+// through *out. The caller is responsible for wrapping the tensor pointer
+// with RAIIAtenTensorHandle which will call aoti_torch_delete_tensor_object
+// when going out of scope.
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_empty_strided(
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int32_t dtype,
+    int32_t device_type,
+    int32_t device_index,
+    AtenTensorHandle* ret_new_tensor // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_empty_strided_pinned(
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int32_t dtype,
+    int32_t device_type,
+    int32_t device_index,
+    AtenTensorHandle* ret_new_tensor // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_as_strided(
+    AtenTensorHandle self,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    AtenTensorHandle* ret);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_tensor_from_blob(
+    void* data,
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int64_t storage_offset,
+    int32_t dtype,
+    int32_t device_type,
+    int32_t device_index,
+    AtenTensorHandle* ret // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_tensor_from_blob_v2(
+    void* data,
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int64_t storage_offset,
+    int32_t dtype,
+    int32_t device_type,
+    int32_t device_index,
+    AtenTensorHandle* ret, // returns new reference
+    int32_t layout,
+    const uint8_t* opaque_metadata,
+    int64_t opaque_metadata_size);
+
+// This function will create a new uninitialized tensor object
+// and its pointer is returned through *ret.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_new_uninitialized_tensor(AtenTensorHandle* ret);
+
+// WARNING: This will be deprecated. Use aoti_torch_copy_ instead.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_tensor_copy_(AtenTensorHandle src, AtenTensorHandle dst);
+
+// Make the tensor referred to by dst an alias for the tensor referred
+// to by src. The two tensors must still be deleted with
+// aoti_torch_delete_tensor separately (or not) as before the call.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_assign_tensors(AtenTensorHandle src, AtenTensorHandle dst);
+
+// Make a shallow copy of the tensor referred to by src and assign
+// it to the handle in the ret_dst. This is similar to the above
+// aoti_torch_assign_tensors function, but creates and sets the
+// ret_dst from within.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_assign_tensors_out(AtenTensorHandle src, AtenTensorHandle* ret_dst);
+
+// This function will create a new tensor object and its pointer is returned
+// through *ret. The caller is responsible for wrapping the tensor pointer
+// with RAIIAtenTensorHandle which will call aoti_torch_delete_tensor_object
+// when going out of scope.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_clone(AtenTensorHandle self, AtenTensorHandle* ret);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_clone_preserve_strides(AtenTensorHandle self, AtenTensorHandle* ret);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_copy_(
+    AtenTensorHandle self,
+    AtenTensorHandle src,
+    int32_t non_blocking);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__mm_plus_mm_out(
+    AtenTensorHandle out,
+    AtenTensorHandle a,
+    AtenTensorHandle b,
+    AtenTensorHandle c,
+    AtenTensorHandle d);
+
+// This will soon be deprecated after ao_quantization is complete.
+// Please refrain from using this or increasing callsites.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu_wrapped_fbgemm_pack_gemm_matrix_fp16(
+    AtenTensorHandle weight,
+    AtenTensorHandle* out);
+
+// This will soon be deprecated after ao_quantization is complete.
+// Please refrain from using this or increasing callsites.
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__wrapped_linear_prepack(
+    AtenTensorHandle weight,
+    AtenTensorHandle weight_scale,
+    AtenTensorHandle weight_zero_point,
+    AtenTensorHandle bias,
+    AtenTensorHandle* out);
+
+// This will soon be deprecated after ao_quantization is complete.
+// Please refrain from using this or increasing callsites.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu_wrapped_fbgemm_linear_fp16_weight(
+    AtenTensorHandle input,
+    AtenTensorHandle weight,
+    AtenTensorHandle bias, // optional argument
+    int64_t out_channel,
+    AtenTensorHandle* out);
+
+// This will soon be deprecated after ao_quantization is complete.
+// Please refrain from using this or increasing callsites.
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu__wrapped_quantized_linear_prepacked(
+    AtenTensorHandle input,
+    AtenTensorHandle input_scale,
+    AtenTensorHandle input_zero_point,
+    AtenTensorHandle weight,
+    AtenTensorHandle out_scale,
+    AtenTensorHandle out_zeropoint,
+    int64_t out_channel,
+    AtenTensorHandle* out);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_zero_(AtenTensorHandle self);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_check_inf_and_nan(const char* tensor_name, AtenTensorHandle tensor);
+
+struct AtenRecordFunctionOpaque;
+using AtenRecordFunctionHandle = AtenRecordFunctionOpaque*;
+
+struct IValueMapOpaque;
+using IValueMapHandle = IValueMapOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_record_function_start(
+    const char* name,
+    IValueMapHandle kwargs,
+    const C10IValueHandle* inputs,
+    const uint64_t n_inputs,
+    AtenRecordFunctionHandle* guard);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_record_function_end(AtenRecordFunctionHandle guard);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scatter_out(
+    AtenTensorHandle out,
+    AtenTensorHandle self,
+    int64_t dim,
+    AtenTensorHandle index,
+    AtenTensorHandle src);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_scatter_reduce_out(
+    AtenTensorHandle out,
+    AtenTensorHandle self,
+    int64_t dim,
+    AtenTensorHandle index,
+    AtenTensorHandle src,
+    const char* reduce,
+    int32_t include_self);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_index_put_out(
+    AtenTensorHandle out,
+    AtenTensorHandle self,
+    const AtenTensorHandle* indices,
+    const uint32_t num_indices,
+    const AtenTensorHandle values,
+    bool accumulate);
+
+AOTI_TORCH_EXPORT void aoti_torch_print_tensor_handle(
+    AtenTensorHandle self,
+    const char* msg);
+
+// When AOTI debug printer option is enabled, this function will be invoked to
+// torch pickle save the intermediate tensor for debugging purpose.
+AOTI_TORCH_EXPORT void aoti_torch_save_tensor_handle(
+    AtenTensorHandle self,
+    const char* tensor_name,
+    const char* launch_prefix,
+    const char* kernel_name);
+
+// helpers for converting between StableIValue and actual IValues
+using StableIValue = uint64_t;
+
+class TorchLibraryOpaque;
+using TorchLibraryHandle = TorchLibraryOpaque*;
+
+// stable corollary to torch::Library constructor with Kind::IMPL
+// will create a new torch::Library object on the heap
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_library_init_impl(
+    const char* ns,
+    const char* k,
+    const char* file,
+    uint32_t line,
+    TorchLibraryHandle* ret_new_torch_lib);
+
+// stable corollary to torch::Library constructor with Kind::DEF
+// will create a new torch::Library object on the heap
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_library_init_def(
+    const char* ns,
+    const char* file,
+    uint32_t line,
+    TorchLibraryHandle* ret_new_torch_lib);
+
+// stable corollary to torch::Library constructor with Kind::FRAGMENT
+// will create a new torch::Library object on the heap
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_library_init_fragment(
+    const char* ns,
+    const char* file,
+    uint32_t line,
+    TorchLibraryHandle* ret_new_torch_lib);
+
+// stable corollary to torch::Library method m.impl(), should be
+// called from StableLibrary
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_library_impl(
+    TorchLibraryHandle self,
+    const char* name,
+    void (*fn)(StableIValue*, uint64_t, uint64_t));
+
+// stable corollary to torch::Library method m.def(), should be
+// called from StableLibrary
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_library_def(TorchLibraryHandle self, const char* schema);
+
+// the above stable constructors for torch::Library add Library objects
+// to the heap. if you are calling those functions directly, please use
+// this function to free the Library's memory. The more user friendly
+// alternative is to use StableLibrary, which will free its handle upon
+// destruction
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_library_object(TorchLibraryHandle tlh);
+
+// calls the op overload defined by a given opName, overloadName, and a
+// stack of StableIValues. This call will populate any return values of the
+// op into the stack in their StableIValue form, with ret0 at index 0, ret1
+// at index 1, and so on.
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_call_dispatcher(
+    const char* opName,
+    const char* overloadName,
+    StableIValue* stack);
+
+// Device-generic guard for managing device context
+struct DeviceGuardOpaque;
+using DeviceGuardHandle = DeviceGuardOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_device_guard(
+    int32_t device_index,
+    DeviceGuardHandle* ret_guard // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_device_guard(DeviceGuardHandle guard);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_device_guard_set_index(
+    DeviceGuardHandle guard,
+    int32_t device_index);
+
+// Device-generic stream for managing stream objects
+struct StreamOpaque;
+using StreamHandle = StreamOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_delete_stream(StreamHandle stream);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_stream_id(StreamHandle stream, int64_t* ret_stream_id);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_current_stream(
+    int32_t device_index,
+    StreamHandle* ret_stream // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_current_device_index(int32_t* ret_device_index);
+
+#ifdef USE_CUDA
+
+struct CUDAGuardOpaque;
+using CUDAGuardHandle = CUDAGuardOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_cuda_guard(
+    int32_t device_index,
+    CUDAGuardHandle* ret_guard // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_cuda_guard(CUDAGuardHandle guard);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cuda_guard_set_index(CUDAGuardHandle guard, int32_t device_index);
+
+struct CUDAStreamGuardOpaque;
+using CUDAStreamGuardHandle = CUDAStreamGuardOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_cuda_stream_guard(
+    void* stream,
+    int32_t device_index,
+    CUDAStreamGuardHandle* ret_guard // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_cuda_stream_guard(CUDAStreamGuardHandle guard);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_current_cuda_stream(int32_t device_index, void** ret_stream);
+
+// CUDA memory allocation using CUDACachingAllocator
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_caching_allocator_raw_alloc(
+    uint64_t nbytes,
+    void** ret_ptr // returns raw GPU memory pointer
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cuda_caching_allocator_raw_delete(void* ptr);
+
+#endif // USE_CUDA
+
+// See `ProxyExecutor Design Note` in ir.py for more details
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_proxy_executor_call_function(
+    AOTIProxyExecutorHandle proxy_executor,
+    int extern_node_index,
+    int num_ints,
+    int64_t* flatten_int_args,
+    int num_tensors,
+    AtenTensorHandle* flatten_tensor_args);
+
+AOTI_TORCH_EXPORT void aoti_torch_check(
+    bool cond,
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* msg);
+
+#ifdef STRIP_ERROR_MESSAGES
+#define AOTI_TORCH_CHECK(cond, ...)                  \
+  if (!(cond)) {                                     \
+    aoti_torch_check(                                \
+        false,                                       \
+        __func__,                                    \
+        __FILE__,                                    \
+        static_cast<uint32_t>(__LINE__),             \
+        STD_TORCH_CHECK_MSG(cond, "", __VA_ARGS__)); \
+  }
+#else
+#define AOTI_TORCH_CHECK(cond, ...)                    \
+  if (!(cond)) {                                       \
+    aoti_torch_check(                                  \
+        false,                                         \
+        __func__,                                      \
+        __FILE__,                                      \
+        static_cast<uint32_t>(__LINE__),               \
+        STD_TORCH_CHECK_MSG(cond, "", ##__VA_ARGS__)); \
+  }
+#endif
+
+AOTI_TORCH_EXPORT void aoti_torch_warn(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* msg);
+
+#ifdef DISABLE_WARN
+#define AOTI_TORCH_WARN(...) ((void)0);
+#else
+#define AOTI_TORCH_WARN(...) \
+  aoti_torch_warn(           \
+      __func__, __FILE__, static_cast<uint32_t>(__LINE__), #__VA_ARGS__);
+#endif
+
+#ifdef __cplusplus
+} // extern "C"
+
+template <typename T>
+int32_t aoti_torch_dtype() = delete;
+
+#define DEFINE_DTYPE_SPECIALIZATION(ctype, typename) \
+  template <>                                        \
+  inline int32_t aoti_torch_dtype<ctype>() {         \
+    return aoti_torch_dtype_##typename();            \
+  }
+
+DEFINE_DTYPE_SPECIALIZATION(c10::BFloat16, bfloat16)
+DEFINE_DTYPE_SPECIALIZATION(c10::Half, float16)
+DEFINE_DTYPE_SPECIALIZATION(c10::complex<float>, complex64)
+DEFINE_DTYPE_SPECIALIZATION(float, float32)
+DEFINE_DTYPE_SPECIALIZATION(double, float64)
+DEFINE_DTYPE_SPECIALIZATION(uint8_t, uint8)
+DEFINE_DTYPE_SPECIALIZATION(int8_t, int8)
+DEFINE_DTYPE_SPECIALIZATION(int16_t, int16)
+DEFINE_DTYPE_SPECIALIZATION(int32_t, int32)
+DEFINE_DTYPE_SPECIALIZATION(int64_t, int64)
+DEFINE_DTYPE_SPECIALIZATION(bool, bool)
+
+#endif
+#endif // AOTI_TORCH_SHIM
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_cpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_cpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..5a10290decd1db529a924a89596e17c8e6829dcc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_cpu.h
@@ -0,0 +1,267 @@
+#ifndef AOTI_TORCH_SHIM_CPU
+#define AOTI_TORCH_SHIM_CPU
+
+#include <ATen/Config.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#if AT_MKLDNN_ENABLED()
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu_mkldnn__convolution_pointwise_binary(
+    AtenTensorHandle X,
+    AtenTensorHandle other,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    const char* binary_attr,
+    double* alpha,
+    const char** unary_attr,
+    const double** unary_scalars,
+    int64_t unary_scalars_len_,
+    const char** unary_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu_mkldnn__convolution_pointwise_binary_(
+    AtenTensorHandle other,
+    AtenTensorHandle X,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    const char* binary_attr,
+    double* alpha,
+    const char** unary_attr,
+    const double** unary_scalars,
+    int64_t unary_scalars_len_,
+    const char** unary_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_mkldnn__convolution_pointwise(
+    AtenTensorHandle X,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    const char* attr,
+    const double** scalars,
+    int64_t scalars_len_,
+    const char** algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu_mkldnn__convolution_transpose_pointwise(
+    AtenTensorHandle X,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* output_padding,
+    int64_t output_padding_len_,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    const char* attr,
+    const double** scalars,
+    int64_t scalars_len_,
+    const char** algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_mkldnn_rnn_layer(
+    AtenTensorHandle input,
+    AtenTensorHandle weight0,
+    AtenTensorHandle weight1,
+    AtenTensorHandle weight2,
+    AtenTensorHandle weight3,
+    AtenTensorHandle hx_,
+    AtenTensorHandle cx_,
+    int32_t reverse,
+    const int64_t* batch_sizes,
+    int64_t batch_sizes_len_,
+    int64_t mode,
+    int64_t hidden_size,
+    int64_t num_layers,
+    int32_t has_biases,
+    int32_t bidirectional,
+    int32_t batch_first,
+    int32_t train,
+    AtenTensorHandle* ret0,
+    AtenTensorHandle* ret1,
+    AtenTensorHandle* ret2,
+    AtenTensorHandle* ret3);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__linear_pointwise(
+    AtenTensorHandle X,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const char* attr,
+    const double** scalars,
+    int64_t scalars_len_,
+    const char** algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__linear_pointwise_binary(
+    AtenTensorHandle X,
+    AtenTensorHandle other,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const char* attr,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__qlinear_pointwise_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle* B,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    const char* post_op_name,
+    const double** post_op_args,
+    int64_t post_op_args_len_,
+    const char* post_op_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu__qlinear_pointwise_binary_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle* other,
+    AtenTensorHandle* B,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    double other_scale,
+    int64_t other_zero_point,
+    const char* binary_post_op,
+    double binary_alpha,
+    const char* unary_post_op,
+    const double** unary_post_op_args,
+    int64_t unary_post_op_args_len_,
+    const char* unary_post_op_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__qconv_pointwise_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle* B,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    const char* attr,
+    const double** post_op_args,
+    int64_t post_op_args_len_,
+    const char** algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_cpu__qconv2d_pointwise_binary_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle accum,
+    AtenTensorHandle* B,
+    const int64_t* stride_args,
+    int64_t stride_len_,
+    const int64_t* padding_args,
+    int64_t padding_len_,
+    const int64_t* dilation_args,
+    int64_t dilation_len_,
+    int64_t groups,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    double accum_scale,
+    int64_t accum_zero_point,
+    const char* binary_attr,
+    double* alpha,
+    const char** unary_attr,
+    const double** unary_scalars,
+    int64_t unary_scalars_len_,
+    const char** unary_algorithm,
+    AtenTensorHandle* ret0);
+
+#if AT_MKL_ENABLED()
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__mkl_linear(
+    AtenTensorHandle X,
+    AtenTensorHandle W,
+    AtenTensorHandle origin_W,
+    AtenTensorHandle* B,
+    int64_t prepack_batch_size,
+    AtenTensorHandle* ret0);
+
+#endif // AT_MKL_ENABLED
+
+#endif // AT_MKLDNN_ENABLED()
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__weight_int4pack_mm_cpu_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle w,
+    AtenTensorHandle qGroupSize,
+    AtenTensorHandle qScaleAndZeros,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__c10d_functional_all_reduce_(
+    AtenTensorHandle inp,
+    const char* reduce_op,
+    const char* group_name,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__c10d_functional_all_reduce(
+    AtenTensorHandle inp,
+    const char* reduce_op,
+    const char* group_name,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__c10d_functional_wait_tensor(
+    AtenTensorHandle inp,
+    AtenTensorHandle* ret0);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+#endif // AOTI_TORCH_SHIM_CPU
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_deprecated.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_deprecated.h
new file mode 100644
index 0000000000000000000000000000000000000000..964db6b0076c97bcbaac3bf5a831b6b87cec54c3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_deprecated.h
@@ -0,0 +1,199 @@
+#ifndef AOTI_TORCH_SHIM_DEPRECATED
+#define AOTI_TORCH_SHIM_DEPRECATED
+
+#include <torch/csrc/inductor/aoti_torch/c/macros.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+[[deprecated(
+    "aoti_torch__embedding_bag is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__embedding_bag(
+    AtenTensorHandle weight,
+    AtenTensorHandle indices,
+    AtenTensorHandle offsets,
+    int32_t scale_grad_by_freq,
+    int32_t mode,
+    int32_t sparse,
+    AtenTensorHandle per_sample_weights, // optional argument
+    int32_t include_last_offset,
+    int32_t padding_idx,
+    AtenTensorHandle* ret0, // returns new reference
+    AtenTensorHandle* ret1, // returns new reference
+    AtenTensorHandle* ret2, // returns new reference
+    AtenTensorHandle* ret3 // returns new reference
+);
+
+[[deprecated(
+    "aoti_torch__fft_c2c is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__fft_c2c(
+    AtenTensorHandle self,
+    const int64_t* dim_ptr,
+    int64_t dim_size,
+    int64_t normalization,
+    int32_t forward,
+    AtenTensorHandle* ret // returns new reference
+);
+
+[[deprecated(
+    "aoti_torch__scaled_mm is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__scaled_mm(
+    AtenTensorHandle self,
+    AtenTensorHandle mat2,
+    AtenTensorHandle bias,
+    int32_t* out_dtype,
+    AtenTensorHandle scale_a,
+    AtenTensorHandle scale_b,
+    AtenTensorHandle scale_result,
+    int8_t use_fast_accum,
+    AtenTensorHandle* ret0,
+    AtenTensorHandle* ret1);
+
+[[deprecated(
+    "aoti_torch__scaled_mm_v2 is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__scaled_mm_v2(
+    AtenTensorHandle self,
+    AtenTensorHandle mat2,
+    AtenTensorHandle scale_a,
+    AtenTensorHandle scale_b,
+    AtenTensorHandle bias,
+    AtenTensorHandle scale_result,
+    int32_t* out_dtype,
+    int8_t use_fast_accum,
+    AtenTensorHandle* ret0);
+
+[[deprecated(
+    "aoti_torch_addmm_out is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_addmm_out(
+    AtenTensorHandle out,
+    AtenTensorHandle self,
+    AtenTensorHandle mat1,
+    AtenTensorHandle mat2,
+    float beta,
+    float alpha);
+
+[[deprecated(
+    "aoti_torch_bmm is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_bmm_out(
+    AtenTensorHandle out,
+    AtenTensorHandle self,
+    AtenTensorHandle mat2);
+
+[[deprecated(
+    "aoti_torch_convolution is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_convolution(
+    AtenTensorHandle input,
+    AtenTensorHandle weight,
+    AtenTensorHandle bias, // optional argument
+    const int64_t* stride_ptr,
+    int64_t stride_size,
+    const int64_t* padding_ptr,
+    int64_t padding_size,
+    const int64_t* dilation_ptr,
+    int64_t dilation_size,
+    int transposed,
+    const int64_t* output_padding_ptr,
+    int64_t output_padding_size,
+    int64_t groups,
+    AtenTensorHandle* ret // returns new reference
+);
+
+[[deprecated(
+    "aoti_torch_mm_out is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mm_out(
+    AtenTensorHandle out,
+    AtenTensorHandle self,
+    AtenTensorHandle mat2);
+
+[[deprecated(
+    "aoti_torch_nonzero is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_nonzero(AtenTensorHandle self, AtenTensorHandle* out);
+
+[[deprecated(
+    "aoti_torch_repeat_interleave_Tensor is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_repeat_interleave_Tensor(
+    AtenTensorHandle repeats,
+    int64_t* output_size,
+    AtenTensorHandle* out);
+
+[[deprecated(
+    "aoti_torch_view_as_real is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_view_as_real(
+    AtenTensorHandle self,
+    AtenTensorHandle* ret // returns new reference
+);
+
+[[deprecated(
+    "aoti_torch_view_dtype is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_view_dtype(
+    AtenTensorHandle self,
+    int32_t dtype,
+    AtenTensorHandle* ret // returns new reference
+);
+
+[[deprecated(
+    "aoti_torch__scaled_dot_product_flash_attention is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch__scaled_dot_product_flash_attention(
+    AtenTensorHandle query,
+    AtenTensorHandle key,
+    AtenTensorHandle value,
+    double dropout_p,
+    bool is_causal,
+    bool return_debug_mask,
+    double scale,
+    AtenTensorHandle* ret0, // returns new reference
+    AtenTensorHandle* ret1, // returns new reference
+    AtenTensorHandle* ret2, // returns new reference
+    AtenTensorHandle* ret3, // returns new reference
+    int64_t* ret4,
+    int64_t* ret5,
+    AtenTensorHandle* ret6, // returns new reference
+    AtenTensorHandle* ret7, // returns new reference
+    AtenTensorHandle* ret8 // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch__scaled_dot_product_flash_attention_v2(
+    AtenTensorHandle query,
+    AtenTensorHandle key,
+    AtenTensorHandle value,
+    double dropout_p,
+    int is_causal,
+    int return_debug_mask,
+    double* scale, // optional argument
+    AtenTensorHandle* ret0, // returns new reference
+    AtenTensorHandle* ret1, // returns new reference
+    AtenTensorHandle* ret2, // returns new reference
+    AtenTensorHandle* ret3, // returns new reference
+    int64_t* ret4,
+    int64_t* ret5,
+    AtenTensorHandle* ret6, // returns new reference
+    AtenTensorHandle* ret7, // returns new reference
+    AtenTensorHandle* ret8 // returns new reference
+);
+
+[[deprecated(
+    "aoti_torch__scaled_dot_product_efficient_attention is deprecated and will be removed in future versions.")]]
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch__scaled_dot_product_efficient_attention(
+    AtenTensorHandle query,
+    AtenTensorHandle key,
+    AtenTensorHandle value,
+    AtenTensorHandle attn_bias, // optional argument
+    int compute_log_sumexp,
+    double dropout_p,
+    int is_causal,
+    double* scale, // optional argument
+    AtenTensorHandle* ret0, // returns new reference
+    AtenTensorHandle* ret1, // returns new reference
+    AtenTensorHandle* ret2, // returns new reference
+    AtenTensorHandle* ret3 // returns new reference
+);
+
+#ifdef __cplusplus
+} // extern "C"
+
+#endif
+#endif // AOTI_TORCH_SHIM_DEPRECATED
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_mps.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_mps.h
new file mode 100644
index 0000000000000000000000000000000000000000..2ab0057805121c902c23e18386cb6121073ebe92
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_mps.h
@@ -0,0 +1,103 @@
+#ifndef AOTI_TORCH_SHIM_MPS
+#define AOTI_TORCH_SHIM_MPS
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+struct AOTIMetalKernelFunctionOpaque;
+using AOTIMetalKernelFunctionHandle = AOTIMetalKernelFunctionOpaque*;
+
+struct AOTIMetalShaderLibraryOpaque;
+using AOTIMetalShaderLibraryHandle = AOTIMetalShaderLibraryOpaque*;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// MetalShaderLibrary functions
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_create_shader_library(
+    const char* metal_shader_source,
+    AOTIMetalShaderLibraryHandle* library_handle);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_delete_shader_library(
+    AOTIMetalShaderLibraryHandle library_handle);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_get_kernel_function(
+    AOTIMetalShaderLibraryHandle library_handle,
+    const char* kernel_name,
+    AOTIMetalKernelFunctionHandle* function_handle);
+
+// MetalKernelFunction functions
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_mps_start_encoding(AOTIMetalKernelFunctionHandle func);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_set_arg_tensor(
+    AOTIMetalKernelFunctionHandle func,
+    unsigned idx,
+    AtenTensorHandle tensor);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_set_arg_int(
+    AOTIMetalKernelFunctionHandle func,
+    unsigned idx,
+    int64_t val);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_dispatch_single(
+    AOTIMetalKernelFunctionHandle func,
+    uint64_t length);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_dispatch_single_with_group_size(
+    AOTIMetalKernelFunctionHandle func,
+    uint64_t length,
+    uint64_t group_size);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_dispatch_array(
+    AOTIMetalKernelFunctionHandle func,
+    const uint64_t* length,
+    size_t length_size);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_dispatch_array_with_group_size(
+    AOTIMetalKernelFunctionHandle func,
+    const uint64_t* length,
+    size_t length_size,
+    const uint64_t* group_size,
+    size_t group_size_size);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_mps_malloc(void** buffer, size_t num_bytes);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_free(void* ptr);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_memcpy(
+    void* buffer,
+    size_t constant_offset,
+    size_t bytes_read,
+    size_t data_size,
+    uint8_t* constants_start);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_copy_buffer(
+    void* src_buffer,
+    void* dst_buffer,
+    size_t data_size,
+    size_t src_offset,
+    size_t dst_offset);
+
+// C callback function type for command block execution
+typedef void (*aoti_torch_mps_command_block_callback_t)(
+    AOTIMetalKernelFunctionHandle func,
+    void* user_data);
+
+// Shared callback function for std::function trampoline
+AOTI_TORCH_EXPORT void aoti_torch_mps_shared_callback(
+    AOTIMetalKernelFunctionHandle func,
+    void* user_data);
+
+// Pure C version using function pointer and user data for trampoline pattern
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_run_command_block(
+    AOTIMetalKernelFunctionHandle func,
+    aoti_torch_mps_command_block_callback_t callback,
+    void* user_data);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#endif // AOTI_TORCH_SHIM_MPS
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..c25fe6443c9485c1c4b82760b3ccd4012fa03802
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/c/shim_xpu.h
@@ -0,0 +1,210 @@
+#ifndef AOTI_TORCH_SHIM_XPU
+#define AOTI_TORCH_SHIM_XPU
+
+#include <ATen/Config.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#ifdef USE_XPU
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+struct XPUGuardOpaque;
+using XPUGuardHandle = XPUGuardOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_xpu_guard(
+    int32_t device_index,
+    XPUGuardHandle* ret_guard // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_xpu_guard(XPUGuardHandle guard);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_xpu_guard_set_index(XPUGuardHandle guard, int32_t device_index);
+
+struct XPUStreamGuardOpaque;
+using XPUStreamGuardHandle = XPUStreamGuardOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_create_xpu_stream_guard(
+    void* stream,
+    int32_t device_index,
+    XPUStreamGuardHandle* ret_guard // returns new reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_delete_xpu_stream_guard(XPUStreamGuardHandle guard);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_current_xpu_stream(int32_t device_index, void** ret_stream);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_current_xpu_device(int32_t* device_index);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_set_current_xpu_device(const int32_t& device_index);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_current_sycl_queue(void** ret);
+
+#if AT_MKLDNN_ENABLED()
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_xpu_mkldnn__convolution_pointwise_binary(
+    AtenTensorHandle X,
+    AtenTensorHandle other,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    const char* binary_attr,
+    double* alpha,
+    const char** unary_attr,
+    const double** unary_scalars,
+    int64_t unary_scalars_len_,
+    const char** unary_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_mkldnn__convolution_pointwise(
+    AtenTensorHandle X,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    const char* attr,
+    const double** scalars,
+    int64_t scalars_len_,
+    const char** algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_xpu_mkldnn__convolution_pointwise_binary_(
+    AtenTensorHandle other,
+    AtenTensorHandle X,
+    AtenTensorHandle W,
+    AtenTensorHandle* B,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    const char* binary_attr,
+    double* alpha,
+    const char** unary_attr,
+    const double** unary_scalars,
+    int64_t unary_scalars_len_,
+    const char** unary_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__qlinear_pointwise_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle* B,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    const char* post_op_name,
+    const double** post_op_args,
+    int64_t post_op_args_len_,
+    const char* post_op_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_xpu__qlinear_pointwise_binary_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle* other,
+    AtenTensorHandle* B,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    double other_scale,
+    int64_t other_zero_point,
+    const char* binary_post_op,
+    double binary_alpha,
+    const char* unary_post_op,
+    const double** unary_post_op_args,
+    int64_t unary_post_op_args_len_,
+    const char* unary_post_op_algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__qconv_pointwise_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle* B,
+    const int64_t* stride,
+    int64_t stride_len_,
+    const int64_t* padding,
+    int64_t padding_len_,
+    const int64_t* dilation,
+    int64_t dilation_len_,
+    int64_t groups,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    const char* attr,
+    const double** post_op_args,
+    int64_t post_op_args_len_,
+    const char** algorithm,
+    AtenTensorHandle* ret0);
+
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_xpu__qconv2d_pointwise_binary_tensor(
+    AtenTensorHandle X,
+    AtenTensorHandle act_scale,
+    AtenTensorHandle act_zero_point,
+    AtenTensorHandle onednn_weight,
+    AtenTensorHandle weight_scales,
+    AtenTensorHandle weight_zero_points,
+    AtenTensorHandle accum,
+    AtenTensorHandle* B,
+    const int64_t* stride_args,
+    int64_t stride_len_,
+    const int64_t* padding_args,
+    int64_t padding_len_,
+    const int64_t* dilation_args,
+    int64_t dilation_len_,
+    int64_t groups,
+    double output_scale,
+    int64_t output_zero_point,
+    const int32_t* output_dtype,
+    double accum_scale,
+    int64_t accum_zero_point,
+    const char* binary_attr,
+    double* alpha,
+    const char** unary_attr,
+    const double** unary_scalars,
+    int64_t unary_scalars_len_,
+    const char** unary_algorithm,
+    AtenTensorHandle* ret0);
+
+#endif // AT_MKLDNN_ENABLED()
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#endif // USE_XPU
+#endif // AOTI_TORCH_SHIM_XPU
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_aten.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_aten.h
new file mode 100644
index 0000000000000000000000000000000000000000..96348d6a129be1c1be09b3e716edd6133315d840
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_aten.h
@@ -0,0 +1,28 @@
+
+
+// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND.
+// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details
+
+// This file corresponds to the aten_shimified_ops list in torchgen/aoti/fallback_ops.py
+
+
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_amax(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int32_t keepdim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_fill__Scalar(AtenTensorHandle self, double value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_full(const int64_t* size, int64_t size_len_, double fill_value, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_new_empty(AtenTensorHandle self, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_new_zeros(AtenTensorHandle self, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_aten_subtract_Tensor(AtenTensorHandle self, AtenTensorHandle other, double alpha, AtenTensorHandle* ret0);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..21ee7438165dc170d6c91de2f50080d82b33d62f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h
@@ -0,0 +1,168 @@
+
+
+// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND.
+// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details
+
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__adaptive_avg_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__adaptive_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__adaptive_avg_pool3d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__adaptive_avg_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__addmm_activation(AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha, int32_t use_gelu, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__cdist_backward(AtenTensorHandle grad, AtenTensorHandle x1, AtenTensorHandle x2, double p, AtenTensorHandle cdist, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__cdist_forward(AtenTensorHandle x1, AtenTensorHandle x2, double p, int64_t* compute_mode, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__dyn_quant_matmul_4bit(AtenTensorHandle inp, AtenTensorHandle packed_weights, int64_t block_size, int64_t in_features, int64_t out_features, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__dyn_quant_pack_4bit_weight(AtenTensorHandle weights, AtenTensorHandle scales_zeros, AtenTensorHandle* bias, int64_t block_size, int64_t in_features, int64_t out_features, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__efficientzerotensor(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__embedding_bag_dense_backward(AtenTensorHandle grad, AtenTensorHandle indices, AtenTensorHandle offset2bag, AtenTensorHandle bag_size, AtenTensorHandle maximum_indices, int64_t num_weights, int32_t scale_grad_by_freq, int64_t mode, AtenTensorHandle* per_sample_weights, int64_t padding_idx, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__embedding_bag_forward_only(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__embedding_bag_per_sample_weights_backward(AtenTensorHandle grad, AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, AtenTensorHandle offset2bag, int64_t mode, int64_t padding_idx, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fft_c2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t forward, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fft_r2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t onesided, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fused_moving_avg_obs_fq_helper(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__histogramdd_from_bin_cts(AtenTensorHandle self, const int64_t* bins, int64_t bins_len_, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__int_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__pdist_backward(AtenTensorHandle grad, AtenTensorHandle self, double p, AtenTensorHandle pdist, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__pdist_forward(AtenTensorHandle self, double p, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__scaled_dot_product_flash_attention_for_cpu(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, double dropout_p, int32_t is_causal, AtenTensorHandle* attn_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__scaled_dot_product_flash_attention_for_cpu_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle out, AtenTensorHandle logsumexp, double dropout_p, int32_t is_causal, AtenTensorHandle* attn_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__scaled_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scale_a, AtenTensorHandle scale_b, AtenTensorHandle* bias, AtenTensorHandle* scale_result, int32_t* out_dtype, int32_t use_fast_accum, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__scaled_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scale_a, AtenTensorHandle scale_b, AtenTensorHandle* bias, AtenTensorHandle* scale_result, int32_t* out_dtype, int32_t use_fast_accum);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__segment_reduce_backward(AtenTensorHandle grad, AtenTensorHandle output, AtenTensorHandle data, const char* reduce, AtenTensorHandle* lengths, AtenTensorHandle* offsets, int64_t axis, double* initial, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__trilinear(AtenTensorHandle i1, AtenTensorHandle i2, AtenTensorHandle i3, const int64_t* expand1, int64_t expand1_len_, const int64_t* expand2, int64_t expand2_len_, const int64_t* expand3, int64_t expand3_len_, const int64_t* sumdim, int64_t sumdim_len_, int64_t unroll_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu__weight_int8pack_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scales, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_abs(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_adaptive_max_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_adaptive_max_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_adaptive_max_pool3d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_adaptive_max_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_add_Scalar(AtenTensorHandle self, double other, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_add_Tensor(AtenTensorHandle self, AtenTensorHandle other, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_addbmm(AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_addmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_addmv(AtenTensorHandle self, AtenTensorHandle mat, AtenTensorHandle vec, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_angle(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_avg_pool2d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_avg_pool3d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_avg_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_baddbmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_bernoulli__Tensor(AtenTensorHandle self, AtenTensorHandle p, AtenGeneratorHandle* generator);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_bernoulli__float(AtenTensorHandle self, double p, AtenGeneratorHandle* generator);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_bmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_bucketize_Tensor(AtenTensorHandle self, AtenTensorHandle boundaries, int32_t out_int32, int32_t right, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_cat(const AtenTensorHandle* tensors, int64_t tensors_len_, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_cholesky_inverse(AtenTensorHandle self, int32_t upper, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_cholesky_solve(AtenTensorHandle self, AtenTensorHandle input2, int32_t upper, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_convolution(AtenTensorHandle input, AtenTensorHandle weight, AtenTensorHandle* bias, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_convolution_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle weight, const int64_t** bias_sizes, int64_t bias_sizes_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_cummax(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_cummin(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_cumprod(AtenTensorHandle self, int64_t dim, int32_t* dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_cumsum(AtenTensorHandle self, int64_t dim, int32_t* dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_exponential(AtenTensorHandle self, double lambd, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_fill__Scalar(AtenTensorHandle self, double value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_fractional_max_pool2d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle random_samples, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_fractional_max_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_fractional_max_pool3d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle random_samples, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_fractional_max_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_gcd(AtenTensorHandle self, AtenTensorHandle other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_geqrf(AtenTensorHandle self, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_grid_sampler_2d_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle grid, int64_t interpolation_mode, int64_t padding_mode, int32_t align_corners, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_hann_window(int64_t window_length, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_histc(AtenTensorHandle self, int64_t bins, double min, double max, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_histogram_bin_ct(AtenTensorHandle self, int64_t bins, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_index_Tensor(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_index_put(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle values, int32_t accumulate, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_index_reduce(AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle source, const char* reduce, int32_t include_self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_kthvalue(AtenTensorHandle self, int64_t k, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_logcumsumexp(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_lu_unpack(AtenTensorHandle LU_data, AtenTensorHandle LU_pivots, int32_t unpack_data, int32_t unpack_pivots, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_masked_scatter(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle source, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_masked_scatter_backward(AtenTensorHandle grad_output, AtenTensorHandle mask, const int64_t* sizes, int64_t sizes_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_masked_select(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_max_pool2d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_max_pool2d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_max_pool3d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_max_pool3d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_max_unpool2d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_max_unpool3d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_median(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_mode(AtenTensorHandle self, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_mul_Scalar(AtenTensorHandle self, double other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_mul_Tensor(AtenTensorHandle self, AtenTensorHandle other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_nanmedian(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_native_dropout(AtenTensorHandle input, double p, int32_t* train, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_nonzero(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_nonzero_static(AtenTensorHandle self, int64_t size, int64_t fill_value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_normal_functional(AtenTensorHandle self, double mean, double std, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_ormqr(AtenTensorHandle self, AtenTensorHandle input2, AtenTensorHandle input3, int32_t left, int32_t transpose, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_permute(AtenTensorHandle self, const int64_t* dims, int64_t dims_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_polar(AtenTensorHandle abs, AtenTensorHandle angle, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pow_Scalar(double self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pow_Tensor_Scalar(AtenTensorHandle self, double exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_replication_pad2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_reshape(AtenTensorHandle self, const int64_t* shape, int64_t shape_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_resize_(AtenTensorHandle self, const int64_t* size, int64_t size_len_, int32_t* memory_format);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_resize_as_(AtenTensorHandle self, AtenTensorHandle the_template, int32_t* memory_format);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_scatter_src_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle src);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_scatter_value_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, double value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_scatter_reduce_two_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle src, const char* reduce, int32_t include_self);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_searchsorted_Scalar(AtenTensorHandle sorted_sequence, double self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_searchsorted_Tensor(AtenTensorHandle sorted_sequence, AtenTensorHandle self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_segment_reduce(AtenTensorHandle data, const char* reduce, AtenTensorHandle* lengths, AtenTensorHandle* indices, AtenTensorHandle* offsets, int64_t axis, int32_t unsafe, double* initial, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_set__source_Tensor(AtenTensorHandle self, AtenTensorHandle source);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_slice_Tensor(AtenTensorHandle self, int64_t dim, int64_t* start, int64_t* end, int64_t step, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_soft_margin_loss_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle target, int64_t reduction, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_sort(AtenTensorHandle self, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_sort_stable(AtenTensorHandle self, int32_t* stable, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_squeeze_dim(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_topk(AtenTensorHandle self, int64_t k, int64_t dim, int32_t largest, int32_t sorted, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_triangular_solve(AtenTensorHandle self, AtenTensorHandle A, int32_t upper, int32_t transpose, int32_t unitriangular, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_uniform(AtenTensorHandle self, double from, double to, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_upsample_bicubic2d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_upsample_linear1d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_upsample_trilinear3d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_d, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_view_dtype(AtenTensorHandle self, int32_t dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_view_as_complex(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_view_as_real(AtenTensorHandle self, AtenTensorHandle* ret0);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..413a393f6fa7de6abbdca892666ee04f11c4e455
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h
@@ -0,0 +1,182 @@
+
+
+// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND.
+// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details
+
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__adaptive_avg_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__adaptive_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__adaptive_avg_pool3d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__adaptive_avg_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__addmm_activation(AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha, int32_t use_gelu, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__cdist_backward(AtenTensorHandle grad, AtenTensorHandle x1, AtenTensorHandle x2, double p, AtenTensorHandle cdist, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__cdist_forward(AtenTensorHandle x1, AtenTensorHandle x2, double p, int64_t* compute_mode, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__cudnn_rnn(AtenTensorHandle input, const AtenTensorHandle* weight, int64_t weight_len_, int64_t weight_stride0, AtenTensorHandle* weight_buf, AtenTensorHandle hx, AtenTensorHandle* cx, int64_t mode, int64_t hidden_size, int64_t proj_size, int64_t num_layers, int32_t batch_first, double dropout, int32_t train, int32_t bidirectional, const int64_t* batch_sizes, int64_t batch_sizes_len_, AtenTensorHandle* dropout_state, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__efficient_attention_backward(AtenTensorHandle grad_out_, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* bias, AtenTensorHandle out, AtenTensorHandle* cu_seqlens_q, AtenTensorHandle* cu_seqlens_k, int64_t max_seqlen_q, int64_t max_seqlen_k, AtenTensorHandle logsumexp, double dropout_p, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, int64_t custom_mask_type, int32_t bias_requires_grad, double* scale, int64_t* num_splits_key, int64_t* window_size, int32_t shared_storage_dqdkdv, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__efficient_attention_forward(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* bias, AtenTensorHandle* cu_seqlens_q, AtenTensorHandle* cu_seqlens_k, int64_t* max_seqlen_q, int64_t* max_seqlen_k, double dropout_p, int64_t custom_mask_type, int32_t compute_log_sumexp, double* scale, AtenTensorHandle* seqlen_k, int64_t* window_size, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__efficientzerotensor(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__embedding_bag_dense_backward(AtenTensorHandle grad, AtenTensorHandle indices, AtenTensorHandle offset2bag, AtenTensorHandle bag_size, AtenTensorHandle maximum_indices, int64_t num_weights, int32_t scale_grad_by_freq, int64_t mode, AtenTensorHandle* per_sample_weights, int64_t padding_idx, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__embedding_bag_forward_only(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__embedding_bag_per_sample_weights_backward(AtenTensorHandle grad, AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, AtenTensorHandle offset2bag, int64_t mode, int64_t padding_idx, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fft_c2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t forward, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fft_r2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t onesided, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle rng_state, AtenTensorHandle unused, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_v2(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* block_table, int64_t* num_splits, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_quantized(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, AtenTensorHandle* q_descale, AtenTensorHandle* k_descale, AtenTensorHandle* v_descale, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_no_dropout_inplace_v2(AtenTensorHandle out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* block_table, int64_t* num_splits, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__flash_attention_forward_no_dropout_inplace(AtenTensorHandle out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* cum_seq_q, AtenTensorHandle* cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, int64_t* window_size_left, int64_t* window_size_right, AtenTensorHandle* seqused_k, AtenTensorHandle* alibi_slopes, AtenTensorHandle* block_table, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fused_moving_avg_obs_fq_helper(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__int_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__pdist_backward(AtenTensorHandle grad, AtenTensorHandle self, double p, AtenTensorHandle pdist, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__pdist_forward(AtenTensorHandle self, double p, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_cudnn_attention(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, int32_t compute_log_sumexp, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_cudnn_attention_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, AtenTensorHandle attn_bias, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_efficient_attention(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, int32_t compute_log_sumexp, double dropout_p, int32_t is_causal, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_efficient_attention_backward(AtenTensorHandle grad_out_, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double dropout_p, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, int32_t is_causal, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_flash_attention(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_flash_attention_quantized(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* q_descale, AtenTensorHandle* k_descale, AtenTensorHandle* v_descale, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_flash_attention_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scale_a, AtenTensorHandle scale_b, AtenTensorHandle* offs, AtenTensorHandle* bias, AtenTensorHandle* scale_result, int32_t* out_dtype, int32_t use_fast_accum, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scale_a, AtenTensorHandle scale_b, AtenTensorHandle* bias, AtenTensorHandle* scale_result, int32_t* out_dtype, int32_t use_fast_accum, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__scaled_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scale_a, AtenTensorHandle scale_b, AtenTensorHandle* bias, AtenTensorHandle* scale_result, int32_t* out_dtype, int32_t use_fast_accum);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__segment_reduce_backward(AtenTensorHandle grad, AtenTensorHandle output, AtenTensorHandle data, const char* reduce, AtenTensorHandle* lengths, AtenTensorHandle* offsets, int64_t axis, double* initial, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__thnn_fused_lstm_cell(AtenTensorHandle input_gates, AtenTensorHandle hidden_gates, AtenTensorHandle cx, AtenTensorHandle* input_bias, AtenTensorHandle* hidden_bias, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__trilinear(AtenTensorHandle i1, AtenTensorHandle i2, AtenTensorHandle i3, const int64_t* expand1, int64_t expand1_len_, const int64_t* expand2, int64_t expand2_len_, const int64_t* expand3, int64_t expand3_len_, const int64_t* sumdim, int64_t sumdim_len_, int64_t unroll_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__weight_int4pack_mm(AtenTensorHandle self, AtenTensorHandle mat2, int64_t qGroupSize, AtenTensorHandle qScaleAndZeros, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda__weight_int8pack_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scales, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_abs(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_adaptive_max_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_adaptive_max_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_adaptive_max_pool3d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_adaptive_max_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_add_Scalar(AtenTensorHandle self, double other, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_add_Tensor(AtenTensorHandle self, AtenTensorHandle other, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_addbmm(AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_addmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_addmv(AtenTensorHandle self, AtenTensorHandle mat, AtenTensorHandle vec, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_angle(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_avg_pool2d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_avg_pool3d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_avg_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_baddbmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_bernoulli__Tensor(AtenTensorHandle self, AtenTensorHandle p, AtenGeneratorHandle* generator);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_bernoulli__float(AtenTensorHandle self, double p, AtenGeneratorHandle* generator);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_bmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_bucketize_Tensor(AtenTensorHandle self, AtenTensorHandle boundaries, int32_t out_int32, int32_t right, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_cat(const AtenTensorHandle* tensors, int64_t tensors_len_, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_cholesky_inverse(AtenTensorHandle self, int32_t upper, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_cholesky_solve(AtenTensorHandle self, AtenTensorHandle input2, int32_t upper, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_convolution(AtenTensorHandle input, AtenTensorHandle weight, AtenTensorHandle* bias, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_convolution_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle weight, const int64_t** bias_sizes, int64_t bias_sizes_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_cummax(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_cummin(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_cumprod(AtenTensorHandle self, int64_t dim, int32_t* dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_cumsum(AtenTensorHandle self, int64_t dim, int32_t* dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_exponential(AtenTensorHandle self, double lambd, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_fill__Scalar(AtenTensorHandle self, double value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_fractional_max_pool2d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle random_samples, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_fractional_max_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_fractional_max_pool3d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle random_samples, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_fractional_max_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_gcd(AtenTensorHandle self, AtenTensorHandle other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_geqrf(AtenTensorHandle self, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_grid_sampler_2d_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle grid, int64_t interpolation_mode, int64_t padding_mode, int32_t align_corners, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_hann_window(int64_t window_length, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_histc(AtenTensorHandle self, int64_t bins, double min, double max, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_index_put(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle values, int32_t accumulate, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_index_reduce(AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle source, const char* reduce, int32_t include_self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_kthvalue(AtenTensorHandle self, int64_t k, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_logcumsumexp(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_lu_unpack(AtenTensorHandle LU_data, AtenTensorHandle LU_pivots, int32_t unpack_data, int32_t unpack_pivots, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_masked_scatter(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle source, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_masked_scatter_backward(AtenTensorHandle grad_output, AtenTensorHandle mask, const int64_t* sizes, int64_t sizes_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_masked_select(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_max_pool2d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_max_pool2d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_max_pool3d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_max_pool3d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_max_unpool2d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_max_unpool3d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_median(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_mm_dtype_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2, int32_t out_dtype);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_mode(AtenTensorHandle self, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_mul_Scalar(AtenTensorHandle self, double other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_mul_Tensor(AtenTensorHandle self, AtenTensorHandle other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_nanmedian(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_native_dropout(AtenTensorHandle input, double p, int32_t* train, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_nonzero(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_nonzero_static(AtenTensorHandle self, int64_t size, int64_t fill_value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_normal_functional(AtenTensorHandle self, double mean, double std, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_ormqr(AtenTensorHandle self, AtenTensorHandle input2, AtenTensorHandle input3, int32_t left, int32_t transpose, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_permute(AtenTensorHandle self, const int64_t* dims, int64_t dims_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_polar(AtenTensorHandle abs, AtenTensorHandle angle, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pow_Scalar(double self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pow_Tensor_Scalar(AtenTensorHandle self, double exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_replication_pad2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_reshape(AtenTensorHandle self, const int64_t* shape, int64_t shape_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_resize_(AtenTensorHandle self, const int64_t* size, int64_t size_len_, int32_t* memory_format);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_resize_as_(AtenTensorHandle self, AtenTensorHandle the_template, int32_t* memory_format);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_scatter_src_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle src);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_scatter_value_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, double value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_scatter_reduce_two_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle src, const char* reduce, int32_t include_self);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_searchsorted_Scalar(AtenTensorHandle sorted_sequence, double self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_searchsorted_Tensor(AtenTensorHandle sorted_sequence, AtenTensorHandle self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_segment_reduce(AtenTensorHandle data, const char* reduce, AtenTensorHandle* lengths, AtenTensorHandle* indices, AtenTensorHandle* offsets, int64_t axis, int32_t unsafe, double* initial, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_set__source_Tensor(AtenTensorHandle self, AtenTensorHandle source);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_slice_Tensor(AtenTensorHandle self, int64_t dim, int64_t* start, int64_t* end, int64_t step, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_soft_margin_loss_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle target, int64_t reduction, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_sort(AtenTensorHandle self, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_sort_stable(AtenTensorHandle self, int32_t* stable, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_squeeze_dim(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_topk(AtenTensorHandle self, int64_t k, int64_t dim, int32_t largest, int32_t sorted, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_triangular_solve(AtenTensorHandle self, AtenTensorHandle A, int32_t upper, int32_t transpose, int32_t unitriangular, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_uniform(AtenTensorHandle self, double from, double to, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_upsample_bicubic2d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_upsample_linear1d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_upsample_trilinear3d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_d, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_view_dtype(AtenTensorHandle self, int32_t dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_view_as_complex(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_view_as_real(AtenTensorHandle self, AtenTensorHandle* ret0);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h
new file mode 100644
index 0000000000000000000000000000000000000000..f5fe2cb852d4f9ace1afd32ffa5b424c9854bf39
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h
@@ -0,0 +1,145 @@
+
+
+// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND.
+// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details
+
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__adaptive_avg_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__adaptive_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__cdist_forward(AtenTensorHandle x1, AtenTensorHandle x2, double p, int64_t* compute_mode, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__efficientzerotensor(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__embedding_bag(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__embedding_bag_dense_backward(AtenTensorHandle grad, AtenTensorHandle indices, AtenTensorHandle offset2bag, AtenTensorHandle bag_size, AtenTensorHandle maximum_indices, int64_t num_weights, int32_t scale_grad_by_freq, int64_t mode, AtenTensorHandle* per_sample_weights, int64_t padding_idx, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__embedding_bag_forward_only(AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, int32_t scale_grad_by_freq, int64_t mode, int32_t sparse, AtenTensorHandle* per_sample_weights, int32_t include_last_offset, int64_t padding_idx, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__embedding_bag_per_sample_weights_backward(AtenTensorHandle grad, AtenTensorHandle weight, AtenTensorHandle indices, AtenTensorHandle offsets, AtenTensorHandle offset2bag, int64_t mode, int64_t padding_idx, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fft_c2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t forward, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fft_r2c(AtenTensorHandle self, const int64_t* dim, int64_t dim_len_, int64_t normalization, int32_t onesided, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__histogramdd_from_bin_cts(AtenTensorHandle self, const int64_t* bins, int64_t bins_len_, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_attention_math_for_mps_v2(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_mask, double dropout_p, int32_t is_causal, AtenTensorHandle* dropout_mask, double* scale, int32_t enable_gqa, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_attention_math_for_mps(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_mask, double dropout_p, int32_t is_causal, AtenTensorHandle* dropout_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__trilinear(AtenTensorHandle i1, AtenTensorHandle i2, AtenTensorHandle i3, const int64_t* expand1, int64_t expand1_len_, const int64_t* expand2, int64_t expand2_len_, const int64_t* expand3, int64_t expand3_len_, const int64_t* sumdim, int64_t sumdim_len_, int64_t unroll_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__weight_int4pack_mm(AtenTensorHandle self, AtenTensorHandle mat2, int64_t qGroupSize, AtenTensorHandle qScaleAndZeros, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__weight_int8pack_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scales, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_abs(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_adaptive_max_pool2d(AtenTensorHandle self, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_adaptive_max_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_add_Scalar(AtenTensorHandle self, double other, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_add_Tensor(AtenTensorHandle self, AtenTensorHandle other, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_addbmm(AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_addmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_addmv(AtenTensorHandle self, AtenTensorHandle mat, AtenTensorHandle vec, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_angle(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_avg_pool2d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_avg_pool2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_avg_pool3d(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_avg_pool3d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, int32_t ceil_mode, int32_t count_include_pad, int64_t* divisor_override, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_baddbmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_bernoulli__Tensor(AtenTensorHandle self, AtenTensorHandle p, AtenGeneratorHandle* generator);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_bernoulli__float(AtenTensorHandle self, double p, AtenGeneratorHandle* generator);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_bmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_bucketize_Tensor(AtenTensorHandle self, AtenTensorHandle boundaries, int32_t out_int32, int32_t right, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_cat(const AtenTensorHandle* tensors, int64_t tensors_len_, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_cholesky_inverse(AtenTensorHandle self, int32_t upper, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_cholesky_solve(AtenTensorHandle self, AtenTensorHandle input2, int32_t upper, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_convolution(AtenTensorHandle input, AtenTensorHandle weight, AtenTensorHandle* bias, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_convolution_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle weight, const int64_t** bias_sizes, int64_t bias_sizes_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_cummax(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_cummin(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_cumprod(AtenTensorHandle self, int64_t dim, int32_t* dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_cumsum(AtenTensorHandle self, int64_t dim, int32_t* dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_exponential(AtenTensorHandle self, double lambd, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_fill__Scalar(AtenTensorHandle self, double value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_gcd(AtenTensorHandle self, AtenTensorHandle other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_hann_window(int64_t window_length, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_histc(AtenTensorHandle self, int64_t bins, double min, double max, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_histogram_bin_ct(AtenTensorHandle self, int64_t bins, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_index_Tensor(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_index_put(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle values, int32_t accumulate, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_index_reduce(AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle source, const char* reduce, int32_t include_self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_kthvalue(AtenTensorHandle self, int64_t k, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_logcumsumexp(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_lu_unpack(AtenTensorHandle LU_data, AtenTensorHandle LU_pivots, int32_t unpack_data, int32_t unpack_pivots, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_masked_scatter(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle source, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_masked_scatter_backward(AtenTensorHandle grad_output, AtenTensorHandle mask, const int64_t* sizes, int64_t sizes_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_masked_select(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_max_pool2d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_max_pool2d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_max_pool3d_with_indices(AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_max_pool3d_with_indices_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* kernel_size, int64_t kernel_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t ceil_mode, AtenTensorHandle indices, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_max_unpool2d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_max_unpool3d(AtenTensorHandle self, AtenTensorHandle indices, const int64_t* output_size, int64_t output_size_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_median(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_mul_Scalar(AtenTensorHandle self, double other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_mul_Tensor(AtenTensorHandle self, AtenTensorHandle other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nanmedian(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_native_dropout(AtenTensorHandle input, double p, int32_t* train, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nonzero(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nonzero_static(AtenTensorHandle self, int64_t size, int64_t fill_value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_normal_functional(AtenTensorHandle self, double mean, double std, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_permute(AtenTensorHandle self, const int64_t* dims, int64_t dims_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_polar(AtenTensorHandle abs, AtenTensorHandle angle, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pow_Scalar(double self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pow_Tensor_Scalar(AtenTensorHandle self, double exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_replication_pad2d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_reshape(AtenTensorHandle self, const int64_t* shape, int64_t shape_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_resize_(AtenTensorHandle self, const int64_t* size, int64_t size_len_, int32_t* memory_format);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_resize_as_(AtenTensorHandle self, AtenTensorHandle the_template, int32_t* memory_format);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_scatter_src_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle src);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_scatter_value_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, double value);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_scatter_reduce_two_out(AtenTensorHandle out, AtenTensorHandle self, int64_t dim, AtenTensorHandle index, AtenTensorHandle src, const char* reduce, int32_t include_self);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_searchsorted_Scalar(AtenTensorHandle sorted_sequence, double self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_searchsorted_Tensor(AtenTensorHandle sorted_sequence, AtenTensorHandle self, int32_t out_int32, int32_t right, const char** side, AtenTensorHandle* sorter, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_set__source_Tensor(AtenTensorHandle self, AtenTensorHandle source);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_slice_Tensor(AtenTensorHandle self, int64_t dim, int64_t* start, int64_t* end, int64_t step, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_soft_margin_loss_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle target, int64_t reduction, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_sort(AtenTensorHandle self, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_sort_stable(AtenTensorHandle self, int32_t* stable, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_squeeze_dim(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_topk(AtenTensorHandle self, int64_t k, int64_t dim, int32_t largest, int32_t sorted, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_triangular_solve(AtenTensorHandle self, AtenTensorHandle A, int32_t upper, int32_t transpose, int32_t unitriangular, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_uniform(AtenTensorHandle self, double from, double to, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_upsample_bicubic2d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_upsample_linear1d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_upsample_trilinear3d_backward(AtenTensorHandle grad_output, const int64_t* output_size, int64_t output_size_len_, const int64_t* input_size, int64_t input_size_len_, int32_t align_corners, double* scales_d, double* scales_h, double* scales_w, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_view_dtype(AtenTensorHandle self, int32_t dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_view_as_complex(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_view_as_real(AtenTensorHandle self, AtenTensorHandle* ret0);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..ea0cbfdc5659d4489db74954c46c772a9abff471
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h
@@ -0,0 +1,82 @@
+
+
+// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND.
+// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details
+
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__addmm_activation(AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha, int32_t use_gelu, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__fused_moving_avg_obs_fq_helper_functional(AtenTensorHandle self, AtenTensorHandle observer_on, AtenTensorHandle fake_quant_on, AtenTensorHandle running_min, AtenTensorHandle running_max, AtenTensorHandle scale, AtenTensorHandle zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, int32_t per_row_fake_quant, int32_t symmetric_quant, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, AtenTensorHandle* ret4, AtenTensorHandle* ret5);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__int_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_dot_product_flash_attention(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_dot_product_flash_attention_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scale_a, AtenTensorHandle scale_b, AtenTensorHandle* bias, AtenTensorHandle* scale_result, int32_t* out_dtype, int32_t use_fast_accum, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__scaled_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scale_a, AtenTensorHandle scale_b, AtenTensorHandle* bias, AtenTensorHandle* scale_result, int32_t* out_dtype, int32_t use_fast_accum);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__trilinear(AtenTensorHandle i1, AtenTensorHandle i2, AtenTensorHandle i3, const int64_t* expand1, int64_t expand1_len_, const int64_t* expand2, int64_t expand2_len_, const int64_t* expand3, int64_t expand3_len_, const int64_t* sumdim, int64_t sumdim_len_, int64_t unroll_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__weight_int4pack_mm_with_scales_and_zeros(AtenTensorHandle self, AtenTensorHandle mat2, int64_t qGroupSize, AtenTensorHandle qScale, AtenTensorHandle qZeros, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu__weight_int8pack_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle scales, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_abs(AtenTensorHandle self, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_add_Scalar(AtenTensorHandle self, double other, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_addbmm(AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_addmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat1, AtenTensorHandle mat2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_addmv(AtenTensorHandle self, AtenTensorHandle mat, AtenTensorHandle vec, double beta, double alpha, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_baddbmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle batch1, AtenTensorHandle batch2, double beta, double alpha);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_bmm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_cholesky_solve(AtenTensorHandle self, AtenTensorHandle input2, int32_t upper, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_convolution(AtenTensorHandle input, AtenTensorHandle weight, AtenTensorHandle* bias, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_convolution_backward(AtenTensorHandle grad_output, AtenTensorHandle input, AtenTensorHandle weight, const int64_t** bias_sizes, int64_t bias_sizes_len_, const int64_t* stride, int64_t stride_len_, const int64_t* padding, int64_t padding_len_, const int64_t* dilation, int64_t dilation_len_, int32_t transposed, const int64_t* output_padding, int64_t output_padding_len_, int64_t groups, const int32_t* output_mask, int64_t output_mask_len_, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_cummax(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_cummin(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_exponential(AtenTensorHandle self, double lambd, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_hann_window(int64_t window_length, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_index_put(AtenTensorHandle self, const AtenTensorHandle** indices, int64_t indices_len_, AtenTensorHandle values, int32_t accumulate, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_kthvalue(AtenTensorHandle self, int64_t k, int64_t dim, int32_t keepdim, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_logcumsumexp(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_masked_scatter(AtenTensorHandle self, AtenTensorHandle mask, AtenTensorHandle source, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_masked_scatter_backward(AtenTensorHandle grad_output, AtenTensorHandle mask, const int64_t* sizes, int64_t sizes_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_mm_dtype_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2, int32_t out_dtype);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_mm_out(AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_mul_Scalar(AtenTensorHandle self, double other, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_normal_functional(AtenTensorHandle self, double mean, double std, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_permute(AtenTensorHandle self, const int64_t* dims, int64_t dims_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_polar(AtenTensorHandle abs, AtenTensorHandle angle, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_reshape(AtenTensorHandle self, const int64_t* shape, int64_t shape_len_, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_resize_as_(AtenTensorHandle self, AtenTensorHandle the_template, int32_t* memory_format);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_slice_Tensor(AtenTensorHandle self, int64_t dim, int64_t* start, int64_t* end, int64_t step, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_soft_margin_loss_backward(AtenTensorHandle grad_output, AtenTensorHandle self, AtenTensorHandle target, int64_t reduction, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_sort(AtenTensorHandle self, int64_t dim, int32_t descending, AtenTensorHandle* ret0, AtenTensorHandle* ret1);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_squeeze_dim(AtenTensorHandle self, int64_t dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_uniform(AtenTensorHandle self, double from, double to, AtenGeneratorHandle* generator, AtenTensorHandle* ret0);
+AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_view_dtype(AtenTensorHandle self, int32_t dtype, AtenTensorHandle* ret0);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated_enum_converters.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated_enum_converters.h
new file mode 100644
index 0000000000000000000000000000000000000000..fa5e8a180ac3e3d00e2750f7db110c1e92f67cf8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/generated_enum_converters.h
@@ -0,0 +1,128 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// @generated by update_schema.py
+// clang-format off
+
+#pragma once
+
+#include <c10/core/Layout.h>
+#include <c10/core/MemoryFormat.h>
+#include <c10/core/ScalarType.h>
+#include <c10/util/Exception.h>
+
+// Converter functions from serialized enum values (torch._export.serde.schema)
+// to c10 enums. The serialized format has different enum values than c10.
+
+namespace torch::aot_inductor {
+
+inline c10::ScalarType convertSerializedScalarType(int serialized_value) {
+  constexpr int kInvalid = -1;
+  constexpr int kScalarTypeMap[] = {
+      kInvalid, // 0
+      static_cast<int>(c10::ScalarType::Byte), // 1
+      static_cast<int>(c10::ScalarType::Char), // 2
+      static_cast<int>(c10::ScalarType::Short), // 3
+      static_cast<int>(c10::ScalarType::Int), // 4
+      static_cast<int>(c10::ScalarType::Long), // 5
+      static_cast<int>(c10::ScalarType::Half), // 6
+      static_cast<int>(c10::ScalarType::Float), // 7
+      static_cast<int>(c10::ScalarType::Double), // 8
+      static_cast<int>(c10::ScalarType::ComplexHalf), // 9
+      static_cast<int>(c10::ScalarType::ComplexFloat), // 10
+      static_cast<int>(c10::ScalarType::ComplexDouble), // 11
+      static_cast<int>(c10::ScalarType::Bool), // 12
+      static_cast<int>(c10::ScalarType::BFloat16), // 13
+      kInvalid, // 14
+      kInvalid, // 15
+      kInvalid, // 16
+      kInvalid, // 17
+      kInvalid, // 18
+      kInvalid, // 19
+      kInvalid, // 20
+      kInvalid, // 21
+      kInvalid, // 22
+      kInvalid, // 23
+      kInvalid, // 24
+      kInvalid, // 25
+      kInvalid, // 26
+      kInvalid, // 27
+      static_cast<int>(c10::ScalarType::UInt16), // 28
+      static_cast<int>(c10::ScalarType::Float8_e4m3fn), // 29
+      static_cast<int>(c10::ScalarType::Float8_e5m2), // 30
+      static_cast<int>(c10::ScalarType::Float8_e4m3fnuz), // 31
+      static_cast<int>(c10::ScalarType::Float8_e5m2fnuz), // 32
+      static_cast<int>(c10::ScalarType::Float8_e8m0fnu), // 33
+      static_cast<int>(c10::ScalarType::UInt32), // 34
+      static_cast<int>(c10::ScalarType::UInt64), // 35
+  };
+  constexpr int kMapSize = sizeof(kScalarTypeMap) / sizeof(kScalarTypeMap[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized ScalarType value out of range: ",
+      serialized_value);
+  int result = kScalarTypeMap[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized ScalarType value: ",
+      serialized_value);
+  return static_cast<c10::ScalarType>(result);
+}
+
+
+inline c10::Layout convertSerializedLayout(int serialized_value) {
+  constexpr int kInvalid = -1;
+  constexpr int kLayoutMap[] = {
+      kInvalid, // 0
+      static_cast<int>(c10::Layout::Sparse), // 1
+      static_cast<int>(c10::Layout::SparseCsr), // 2
+      static_cast<int>(c10::Layout::SparseCsc), // 3
+      static_cast<int>(c10::Layout::SparseBsr), // 4
+      static_cast<int>(c10::Layout::SparseBsc), // 5
+      static_cast<int>(c10::Layout::Mkldnn), // 6
+      static_cast<int>(c10::Layout::Strided), // 7
+  };
+  constexpr int kMapSize = sizeof(kLayoutMap) / sizeof(kLayoutMap[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized Layout value out of range: ",
+      serialized_value);
+  int result = kLayoutMap[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized Layout value: ",
+      serialized_value);
+  return static_cast<c10::Layout>(result);
+}
+
+
+inline c10::MemoryFormat convertSerializedMemoryFormat(int serialized_value) {
+  constexpr int kInvalid = -1;
+  constexpr int kMemoryFormatMap[] = {
+      kInvalid, // 0
+      static_cast<int>(c10::MemoryFormat::Contiguous), // 1
+      static_cast<int>(c10::MemoryFormat::ChannelsLast), // 2
+      static_cast<int>(c10::MemoryFormat::ChannelsLast3d), // 3
+      static_cast<int>(c10::MemoryFormat::Preserve), // 4
+  };
+  constexpr int kMapSize = sizeof(kMemoryFormatMap) / sizeof(kMemoryFormatMap[0]);
+
+  TORCH_CHECK(
+      serialized_value >= 0 && serialized_value < kMapSize,
+      "Serialized MemoryFormat value out of range: ",
+      serialized_value);
+  int result = kMemoryFormatMap[serialized_value];
+  TORCH_CHECK(
+      result != kInvalid,
+      "Invalid serialized MemoryFormat value: ",
+      serialized_value);
+  return static_cast<c10::MemoryFormat>(result);
+}
+
+} // namespace torch::aot_inductor
+
+// clang-format on
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/mkldnn_tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/mkldnn_tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..57aa0645fb79e83c94f30a120fc6898062e89bce
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/mkldnn_tensor.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Tensor.h>
+
+namespace torch::aot_inductor {
+
+void* data_ptr_from_mkldnn(at::Tensor* mkldnn_tensor);
+
+at::Tensor mkldnn_tensor_from_data_ptr(
+    void* data_ptr,
+    at::IntArrayRef dims,
+    at::ScalarType dtype,
+    at::Device device,
+    const uint8_t* opaque_metadata,
+    int64_t opaque_metadata_size);
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/oss_proxy_executor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/oss_proxy_executor.h
new file mode 100644
index 0000000000000000000000000000000000000000..61428e2cef64ceb3c600d8a8dd03ac63e018c4b1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/oss_proxy_executor.h
@@ -0,0 +1,158 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/core/ivalue.h>
+#include <c10/macros/Export.h>
+#include <nlohmann/json.hpp>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/inductor/aoti_torch/proxy_executor.h>
+#include <torch/csrc/jit/api/function_impl.h> // @manual
+#include <iostream>
+#include <utility>
+
+namespace torch::aot_inductor {
+
+inline std::ostream& operator<<(std::ostream& os, DynamicArgType arg_type) {
+  os << static_cast<int>(arg_type);
+  return os;
+}
+
+struct OSSDynamicArg {
+  OSSDynamicArg(
+      int arg_index,
+      DynamicArgType arg_type,
+      int length,
+      std::optional<std::vector<std::string>> list_item_types = std::nullopt)
+      : arg_index(arg_index),
+        arg_type(arg_type),
+        length(length),
+        list_item_types(std::move(list_item_types)) {}
+  int arg_index;
+  DynamicArgType arg_type;
+  int length;
+  std::optional<std::vector<std::string>>
+      list_item_types; // only used for parsing list of optional tensors
+};
+
+struct OSSTorchBindArg {
+  OSSTorchBindArg(int arg_index, std::string arg_name)
+      : arg_index(arg_index), arg_name(std::move(arg_name)) {}
+  int arg_index;
+  // arg_name is used to find the corresponding IValue in customObjs_
+  std::string arg_name;
+};
+
+struct OSSOpKernel {
+  explicit OSSOpKernel(std::string target) : target_(std::move(target)) {}
+  // Explicitly declare copy and move constructors
+  OSSOpKernel(const OSSOpKernel&) = default;
+  OSSOpKernel(OSSOpKernel&&) = default;
+  // Explicitly declare copy and move assignment operators
+  OSSOpKernel& operator=(const OSSOpKernel&) = default;
+  OSSOpKernel& operator=(OSSOpKernel&&) = default;
+
+  std::string target_;
+  std::vector<OSSDynamicArg> dynamic_args_;
+  std::vector<OSSTorchBindArg> torchbind_args_;
+  std::vector<OSSDynamicArg> outputs_;
+  std::vector<c10::IValue> stack_;
+
+  int num_output_tensors() const {
+    int num_output_tensors = 0;
+    for (const auto& output : outputs_) {
+      if (isTensorType(output.arg_type)) {
+        num_output_tensors += output.length;
+      }
+    }
+    return num_output_tensors;
+  }
+
+  int num_output_ints() const {
+    int num_output_ints = 0;
+    for (const auto& output : outputs_) {
+      if (output.arg_type == DynamicArgType::IntType) {
+        num_output_ints += output.length;
+      }
+    }
+    return num_output_ints;
+  }
+
+  virtual void run(std::vector<c10::IValue>& stack) = 0;
+  virtual c10::FunctionSchema schema() const = 0;
+  virtual ~OSSOpKernel() = default;
+};
+
+struct OSSOpKernelOperator : public OSSOpKernel {
+  OSSOpKernelOperator(std::string target, c10::OperatorHandle op_handle)
+      : OSSOpKernel(std::move(target)), op_handle_(std::move(op_handle)) {}
+
+  c10::OperatorHandle op_handle_;
+  void run(std::vector<c10::IValue>& stack) override {
+    op_handle_.callBoxed(stack);
+  }
+
+  c10::FunctionSchema schema() const override {
+    return op_handle_.schema();
+  }
+};
+
+struct OSSCallTorchBindKernel : public OSSOpKernel {
+  OSSCallTorchBindKernel(std::string target, torch::jit::Function* method)
+      : OSSOpKernel(std::move(target)), method_(method) {}
+  torch::jit::Function* method_;
+  void run(std::vector<c10::IValue>& stack) override {
+    method_->run(stack);
+  }
+
+  c10::FunctionSchema schema() const override {
+    return method_->getSchema();
+  }
+};
+
+class OSSProxyExecutor : public ProxyExecutor {
+ public:
+  explicit OSSProxyExecutor(
+      const std::string& json_path,
+      bool is_cpu,
+      std::optional<std::unordered_map<std::string, c10::IValue>> custom_objs =
+          std::nullopt);
+
+  void call_function(
+      int extern_node_index,
+      int num_ints,
+      int64_t* flatten_int_args,
+      int num_tensors,
+      AtenTensorHandle* flatten_tensor_args) override;
+
+ private:
+  void prefill_stack_with_static_arguments(
+      size_t index,
+      const at::TypePtr& schema_arg_type,
+      const nlohmann::json& serialized_arg,
+      OSSOpKernel* op_kernel,
+      const std::string& torchbind_arg_name);
+
+  void get_input_info_from_serialized(
+      const std::vector<c10::Argument>& schema_args,
+      const nlohmann::json& serialized_node,
+      OSSOpKernel& op_kernel);
+
+  void get_output_info_from_serialized(
+      const std::vector<c10::Argument>& schema_returns,
+      const nlohmann::json& serialized_node,
+      OSSOpKernel& op_kernel);
+
+  std::unique_ptr<OSSCallTorchBindKernel> get_call_torch_bind_kernel(
+      const nlohmann::json& serialized_node);
+
+  std::vector<std::unique_ptr<OSSOpKernel>> op_kernels_;
+  std::unique_ptr<c10::Device> device_;
+  std::unordered_map<std::string, c10::IValue> custom_objs_;
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/proxy_executor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/proxy_executor.h
new file mode 100644
index 0000000000000000000000000000000000000000..b25a7e396c8c349810c56d59a20e800fa68d0882
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/proxy_executor.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <c10/macros/Export.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+namespace torch::aot_inductor {
+
+enum class DynamicArgType : int {
+  TensorType = 0,
+  ListTensorType = 1,
+  ListOptionalTensorType = 2,
+  IntType = 3,
+  ListIntType = 4,
+  NoneType = 5,
+};
+
+inline bool isTensorType(DynamicArgType arg_type) {
+  return arg_type == DynamicArgType::TensorType ||
+      arg_type == DynamicArgType::ListTensorType ||
+      arg_type == DynamicArgType::ListOptionalTensorType;
+}
+
+class ProxyExecutor {
+ public:
+  ProxyExecutor() = default;
+  virtual ~ProxyExecutor() = default;
+
+  virtual void call_function(
+      int extern_node_index,
+      int num_ints,
+      int64_t* flatten_int_args,
+      int num_tensors,
+      AtenTensorHandle* flatten_tensor_args) = 0;
+};
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/tensor_converter.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/tensor_converter.h
new file mode 100644
index 0000000000000000000000000000000000000000..4461d682ad3a5f70ae83161c6590cab63bf6a675
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/tensor_converter.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Tensor.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+namespace torch::aot_inductor {
+
+// Functions declared here are not meant to be called from the AOTInductor
+// generated model.so
+
+// unsafe_alloc_new_handles_from_tensors is used for allocating new aten
+// tensor objects and return them as a vector of AtenTensorHandle (raw
+// pointers), and those pointers will be stolen by model.so.
+TORCH_API std::vector<AtenTensorHandle> unsafe_alloc_new_handles_from_tensors(
+    const std::vector<at::Tensor>& tensors);
+
+// alloc_tensors_by_stealing_from_handles is used for creating a vector of aten
+// tensors by stealing from an array of handles. Only the handles are stolen,
+// and the array itself is borrowed.
+//
+// WARNING: Can NOT be called in model.so
+TORCH_API std::vector<at::Tensor> alloc_tensors_by_stealing_from_handles(
+    AtenTensorHandle* handles,
+    size_t length);
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..0e71766b3f0a24a6129eac1e08fec84fefbbabbf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/aoti_torch/utils.h
@@ -0,0 +1,240 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Generator.h>
+#include <ATen/Tensor.h>
+#include <ATen/core/List.h>
+#include <c10/core/DeviceType.h>
+#include <c10/core/SymIntArrayRef.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/Logging.h>
+#include <c10/util/OptionalArrayRef.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <optional>
+
+#define AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE(...)    \
+  try {                                                    \
+    __VA_ARGS__                                            \
+  } catch (const std::exception& e) {                      \
+    LOG(ERROR) << "Exception in aoti_torch: " << e.what(); \
+    return AOTI_TORCH_FAILURE;                             \
+  } catch (...) {                                          \
+    LOG(ERROR) << "Exception in aoti_torch: UNKNOWN";      \
+    return AOTI_TORCH_FAILURE;                             \
+  }                                                        \
+  return AOTI_TORCH_SUCCESS;
+
+namespace torch::aot_inductor {
+
+inline at::Tensor* tensor_handle_to_tensor_pointer(AtenTensorHandle handle) {
+  return reinterpret_cast<at::Tensor*>(handle);
+}
+
+inline AtenTensorHandle tensor_pointer_to_tensor_handle(at::Tensor* tensor) {
+  return reinterpret_cast<AtenTensorHandle>(tensor);
+}
+
+inline at::Tensor resolve_tensor_dispatch_flags(AtenTensorHandle handle) {
+  at::Tensor* tensor{tensor_handle_to_tensor_pointer(handle)};
+  if (tensor->is_conj() || tensor->is_neg()) {
+    // If the conjugation or negation dispatch flags are set, runtime dispatch
+    // handles them by cloning the tensor before passing them to the native ATen
+    // function.  Since the C-shim calls the native function directly, we have
+    // to handle the flags ourselves, or results will be silently incorrect.
+    return tensor->clone();
+  }
+  return *tensor;
+}
+
+inline std::optional<at::Tensor> resolve_tensor_dispatch_flags(
+    const AtenTensorHandle* handle) {
+  return handle ? std::make_optional(resolve_tensor_dispatch_flags(*handle))
+                : std::nullopt;
+}
+
+inline std::vector<at::Tensor> resolve_tensor_list_dispatch_flags(
+    const AtenTensorHandle* handle,
+    int64_t len) {
+  std::vector<at::Tensor> ret{};
+  ret.reserve(len);
+  for (int64_t i{0}; i < len; ++i) {
+    ret.emplace_back(resolve_tensor_dispatch_flags(handle[i]));
+  }
+  return ret;
+}
+
+inline std::vector<std::optional<at::Tensor>> resolve_tensor_list_dispatch_flags(
+    const AtenTensorHandle** handle,
+    int64_t len) {
+  std::vector<std::optional<at::Tensor>> ret{};
+  ret.reserve(len);
+  for (int64_t i{0}; i < len; ++i) {
+    ret.emplace_back(resolve_tensor_dispatch_flags(handle[i]));
+  }
+  return ret;
+}
+
+inline at::Generator* generator_handle_to_generator_pointer(
+    AtenGeneratorHandle handle) {
+  return reinterpret_cast<at::Generator*>(handle);
+}
+
+inline AtenGeneratorHandle generator_pointer_to_generator_handle(
+    at::Generator* generator) {
+  return reinterpret_cast<AtenGeneratorHandle>(generator);
+}
+
+inline AtenTensorHandle new_tensor_handle(at::Tensor&& tensor) {
+  at::Tensor* new_tensor = new at::Tensor(std::move(tensor));
+  return tensor_pointer_to_tensor_handle(new_tensor);
+}
+
+inline void assert_inf_and_nan(
+    const std::string& tensor_name,
+    at::Tensor& check_tensor) {
+  auto isnan_tensor = check_tensor.isnan();
+  if (isnan_tensor.any().item<bool>()) {
+    throw std::runtime_error("At least one NaN in " + tensor_name);
+  }
+  auto isinf_tensor = check_tensor.isinf();
+  if (isinf_tensor.any().item<bool>()) {
+    throw std::runtime_error("At least one INF in " + tensor_name);
+  }
+}
+
+// utility functions to convert a pointer to an optional value
+template <class T>
+inline std::optional<T> pointer_to_optional(T* ptr) {
+  return ptr ? std::make_optional(*ptr) : std::nullopt;
+}
+
+template <class T, class U, typename = std::enable_if_t<!std::is_same_v<T, U>>>
+inline std::optional<T> pointer_to_optional(U* ptr) {
+  return ptr ? std::make_optional<T>(T(*ptr)) : std::nullopt;
+}
+
+template <>
+inline std::optional<at::Tensor> pointer_to_optional(AtenTensorHandle* ptr) {
+  return ptr ? std::make_optional(*tensor_handle_to_tensor_pointer(*ptr))
+             : std::nullopt;
+}
+
+template <>
+inline std::optional<at::Tensor> pointer_to_optional(
+    const AtenTensorHandle* ptr) {
+  return ptr ? std::make_optional(*tensor_handle_to_tensor_pointer(*ptr))
+             : std::nullopt;
+}
+
+template <>
+inline std::optional<at::Generator> pointer_to_optional(
+    AtenGeneratorHandle* ptr) {
+  return ptr ? std::make_optional(*generator_handle_to_generator_pointer(*ptr))
+             : std::nullopt;
+}
+
+inline std::optional<c10::Device> pointer_to_optional_device(
+    int32_t* device_type,
+    int32_t device_index) {
+  return device_type ? std::make_optional(c10::Device(
+                           static_cast<c10::DeviceType>(*device_type),
+                           static_cast<c10::DeviceIndex>(device_index)))
+                     : std::nullopt;
+}
+
+// utility functions to convert a pointer to a list
+template <typename T>
+struct is_optional : std::false_type {};
+template <typename T>
+struct is_optional<std::optional<T>> : std::true_type {};
+
+template <class T>
+inline c10::ArrayRef<T> pointer_to_list(T* ptr, int64_t len) {
+  return c10::ArrayRef<T>(ptr, len);
+}
+
+template <
+    class T,
+    class U,
+    typename = std::enable_if_t<!std::is_same_v<T, U>>,
+    typename = std::enable_if_t<!is_optional<T>::value>>
+inline std::vector<T> pointer_to_list(U* ptr, int64_t len) {
+  // std::vector<T> will be implicitly converted to c10::ArrayRef<T> at the call
+  // site
+  std::vector<T> result;
+  result.reserve(len);
+  for (int64_t i = 0; i < len; i++) {
+    result.emplace_back(T(ptr[i]));
+  }
+  return result;
+}
+
+template <class T, class U, typename = std::enable_if_t<is_optional<T>::value>>
+inline std::vector<T> pointer_to_list(U** ptr, int64_t len) {
+  // Here U** denotes a list of optional arguments
+  // std::vector<T> will be implicitly converted to c10::ArrayRef<T> at the call
+  // site
+  std::vector<T> result;
+  result.reserve(len);
+  for (int64_t i = 0; i < len; i++) {
+    result.emplace_back(pointer_to_optional(ptr[i]));
+  }
+  return result;
+}
+
+template <>
+inline std::vector<at::Tensor> pointer_to_list(
+    const AtenTensorHandle* ptr,
+    int64_t len) {
+  std::vector<at::Tensor> result;
+  result.reserve(len);
+  for (int64_t i = 0; i < len; i++) {
+    result.emplace_back(*tensor_handle_to_tensor_pointer(ptr[i]));
+  }
+  return result;
+}
+
+template <>
+inline std::vector<std::optional<at::Tensor>> pointer_to_list(
+    const AtenTensorHandle** ptr,
+    int64_t len) {
+  std::vector<std::optional<at::Tensor>> result;
+  result.reserve(len);
+  for (int64_t i = 0; i < len; i++) {
+    result.emplace_back(pointer_to_optional<at::Tensor>(ptr[i]));
+  }
+  return result;
+}
+
+template <int N>
+inline std::array<bool, N> pointer_to_list(const int32_t* ptr) {
+  std::array<bool, N> result;
+  std::copy(ptr, ptr + N, result.begin());
+  return result;
+}
+
+// Utility function to convert a pointer to an optional list of values
+template <class T, class U>
+inline std::optional<c10::ArrayRef<T>> pointer_to_optional_list(
+    U** ptr,
+    int64_t len) {
+  return ptr
+      ? std::make_optional<c10::ArrayRef<T>>(pointer_to_list<T>(*ptr, len))
+      : std::nullopt;
+}
+
+template <typename T>
+static c10::List<T> convert_to_c10_List(const T* scalars, const int64_t len) {
+  c10::List<T> scalars_list;
+  scalars_list.reserve(len);
+  for (int64_t i = 0; i < len; i++) {
+    scalars_list.emplace_back(scalars[i]);
+  }
+  return scalars_list;
+}
+
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/array_ref_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/array_ref_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..6c732af17029ec3d90fff1b667b0609d3731848a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/array_ref_impl.h
@@ -0,0 +1,93 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h>
+#include <torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/thread_local.h>
+#include <torch/csrc/inductor/aoti_torch/utils.h>
+
+namespace torch::aot_inductor {
+template <typename T>
+void convert_output_to_handle(
+    const ArrayRefTensor<T>& output,
+    AtenTensorHandle& handle) {
+  handle = output.expensiveCopyToTensor();
+}
+
+template <typename... Ts, std::size_t... Is>
+void convert_outputs_to_handles_helper(
+    const std::tuple<ArrayRefTensor<Ts>...>& outputs,
+    AtenTensorHandle* output_handles,
+    std::index_sequence<Is...>) {
+  (convert_output_to_handle(std::get<Is>(outputs), output_handles[Is]), ...);
+}
+template <typename... Ts>
+void convert_outputs_to_handles(
+    const std::tuple<ArrayRefTensor<Ts>...>& outputs,
+    AtenTensorHandle* output_handles) {
+  convert_outputs_to_handles_helper(
+      outputs, output_handles, std::make_index_sequence<sizeof...(Ts)>());
+}
+
+template <typename T>
+void convert_handle_to_arrayref_tensor(
+    AtenTensorHandle handle,
+    ArrayRefTensor<T>& input) {
+  void* data_ptr;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(handle, &data_ptr));
+  int64_t dim;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_dim(handle, &dim));
+  int64_t numel;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_numel(handle, &numel));
+  int64_t* sizes;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(handle, &sizes));
+  int64_t* strides;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(handle, &strides));
+  int32_t dtype;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(handle, &dtype));
+  int32_t device_type;
+  AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_device_type(handle, &device_type));
+  int32_t device_index;
+  AOTI_TORCH_ERROR_CODE_CHECK(
+      aoti_torch_get_device_index(handle, &device_index));
+
+  input = ArrayRefTensor<T>(
+      MiniArrayRef<T>(reinterpret_cast<T*>(data_ptr), numel),
+      MiniArrayRef<const int64_t>(sizes, dim),
+      MiniArrayRef<const int64_t>(strides, dim),
+      device_type,
+      device_index);
+}
+
+template <typename... Ts, std::size_t... Is>
+void convert_handles_to_inputs_helper(
+    AtenTensorHandle* input_handles,
+    std::tuple<ArrayRefTensor<Ts>...>& inputs,
+    std::index_sequence<Is...>) {
+  (convert_handle_to_arrayref_tensor(input_handles[Is], std::get<Is>(inputs)),
+   ...);
+}
+
+template <typename... Ts>
+void convert_handles_to_inputs(
+    AtenTensorHandle* input_handles,
+    std::tuple<ArrayRefTensor<Ts>...>& inputs) {
+  convert_handles_to_inputs_helper(
+      input_handles, inputs, std::make_index_sequence<sizeof...(Ts)>());
+}
+
+template <typename T>
+void assert_numel(const ArrayRefTensor<T>& tensor, uint64_t numel) {
+  TORCH_CHECK(
+      tensor.numel() == numel,
+      "incorrect numel for input tensor. expected ",
+      numel,
+      ", got ",
+      tensor.numel());
+}
+} // namespace torch::aot_inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_prefix.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_prefix.h
new file mode 100644
index 0000000000000000000000000000000000000000..5d867caf8b49364f006675c01cc03ee89c21a851
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_prefix.h
@@ -0,0 +1,1328 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <omp.h>
+#include <algorithm>
+#include <atomic>
+#include <cmath>
+#include <cstdlib>
+#include <limits>
+#include <map>
+#include <memory>
+#include <optional>
+
+// WARNING: be extra careful when including more ATen/c10 header files here!
+// Because AOTInductor generated code will copy-paste this cpp_prefix.h for
+// the CPU backend, we have to make sure the used headers are implemented
+// in a header-only way, i.e. all the function and class definitions are
+// in .h files instead of .cpp files, to avoid ABI backward-compatibility
+// breakage.
+
+#include <ATen/NumericUtils.h>
+#include <ATen/core/PhiloxRNGEngine.h>
+
+#include <c10/util/BFloat16-math.h>
+#include <c10/util/BFloat16.h>
+#include <c10/util/Float8_e4m3fn.h>
+#include <c10/util/Float8_e4m3fnuz.h>
+#include <c10/util/Float8_e5m2.h>
+#include <c10/util/Float8_e5m2fnuz.h>
+#include <c10/util/Half.h>
+#include <c10/util/TypeCast.h>
+#include <c10/util/generic_math.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#if defined(CPU_CAPABILITY_AVX512) || defined(CPU_CAPABILITY_AVX2) ||  \
+    defined(CPU_CAPABILITY_ZVECTOR) || defined(CPU_CAPABILITY_NEON) || \
+    defined(CPU_CAPABILITY_VSX) || defined(CPU_CAPABILITY_SVE256)
+#define INDUCTOR_USE_VECTOR_TYPES() 1
+#else
+#define INDUCTOR_USE_VECTOR_TYPES() 0
+#endif
+
+#if INDUCTOR_USE_VECTOR_TYPES()
+#include <ATen/cpu/vec/functional.h>
+#include <ATen/cpu/vec/vec.h>
+#else
+// For calc_erfinv
+#include <ATen/native/Math.h>
+#endif
+
+template <typename T>
+struct Welford {
+  T mean = T(0);
+  T m2 = T(0);
+  // Use weight for tail cases since the index of each element in the vec may be
+  // different. A single index can not express masked welford reduction.
+  T weight = T(0);
+  uint64_t index = 0;
+};
+
+template <typename T>
+struct IsVecType : std::false_type {};
+
+template <typename T>
+struct IsVecMaskType : std::false_type {};
+
+#if INDUCTOR_USE_VECTOR_TYPES()
+template <typename T>
+struct IsVecType<at::vec::Vectorized<T>> : std::true_type {};
+template <typename T, int N>
+struct IsVecType<at::vec::VectorizedN<T, N>> : std::true_type {};
+
+template <typename T, int N>
+struct IsVecMaskType<at::vec::VecMask<T, N>> : std::true_type {};
+#endif
+
+template <typename T>
+struct GetScalarType {
+  using type = T;
+};
+
+#if INDUCTOR_USE_VECTOR_TYPES()
+template <typename T>
+struct GetScalarType<at::vec::Vectorized<T>> {
+  using type = T;
+};
+template <typename T, int N>
+struct GetScalarType<at::vec::VectorizedN<T, N>> {
+  using type = T;
+};
+#endif
+
+template <typename T, uint64_t kChunkSize>
+struct CascadeSumHelper {
+  // A data struct to help cascade summation:
+  std::vector<T> sum_stk{};
+  uint64_t depth{0}; // depth of sum_stk.
+  uint64_t num_chunks{0}; // number of chunks stored in sum_stk.
+  uint64_t index{0}; // index of the current data.
+  CascadeSumHelper() = default;
+  CascadeSumHelper(uint64_t N) {
+    uint64_t m = (N + kChunkSize - 1) / kChunkSize; // div up
+    depth = m > 0
+        ? static_cast<std::uint64_t>(ceil(log2(static_cast<double>(m))))
+        : 0;
+    if constexpr (IsVecType<T>::value) {
+      sum_stk.assign(
+          std::max(depth, static_cast<uint64_t>(1)),
+          T(typename T::value_type(0)));
+    } else {
+      sum_stk.assign(std::max(depth, static_cast<uint64_t>(1)), T(0));
+    }
+  }
+};
+
+template <typename T, uint64_t kChunkSize = 0>
+inline T cascade_sum_combine(T& data, CascadeSumHelper<T, kChunkSize>* c) {
+  // Note: In order to be consistent with other reductions in inductor,
+  // the returned value may be wrong and cascade_sum_final must be executed to
+  // get the final correct result. Inductor uses the reduction suffix to ensure
+  // that cascade_sum_final is called in the end.
+  c->sum_stk[0] = c->sum_stk[0] + data;
+  // Use cascade summation to improve numerical stability.
+  // https://en.wikipedia.org/wiki/Pairwise_summation
+  if (c->depth > 0) {
+    c->index++;
+    if (c->index == kChunkSize) {
+      c->num_chunks += 1;
+      c->index = 0;
+      uint64_t mask = c->num_chunks;
+      uint64_t j = 1;
+      for (; j < c->depth && (mask & 1) == 0; ++j) {
+        c->sum_stk[j] = c->sum_stk[j] + c->sum_stk[j - 1];
+        c->sum_stk[j - 1] = T(0);
+        mask >>= 1;
+      }
+      return c->sum_stk[j - 1];
+    }
+  }
+  return c->sum_stk[0];
+}
+
+template <typename T, uint64_t kChunkSize = 0>
+inline T cascade_sum_final(CascadeSumHelper<T, kChunkSize>* c) {
+  T result = c->sum_stk[0];
+  for (const auto i : c10::irange(1, c->depth)) {
+    result = result + c->sum_stk[i];
+  }
+  return result;
+}
+
+template <typename T, uint64_t kChunkSize>
+struct WelfordHelper {
+  // A data struct to help welford reduction:
+  // 1. Save the reciprocal of weights to avoid redundant divisions.
+  // 2. Save the welford stack, which is used to combine welford reduction
+  //    with cascade summation to improve numerical stability.
+  static std::vector<typename GetScalarType<T>::type> weight_recps;
+  std::vector<Welford<T>> welford_stk{};
+  uint64_t depth{0}; // depth of welford_stk.
+  uint64_t num_chunks{0}; // number of chunks stored in welford_stk.
+  WelfordHelper() = default;
+  WelfordHelper(uint64_t N) {
+    uint64_t m = (N + kChunkSize - 1) / kChunkSize; // div up
+    depth = m > 0
+        ? static_cast<std::uint64_t>(ceil(log2(static_cast<double>(m))))
+        : 0;
+    welford_stk.assign(depth, Welford<T>());
+  }
+};
+
+template <typename T, uint64_t kChunkSize>
+std::vector<typename GetScalarType<T>::type>
+    WelfordHelper<T, kChunkSize>::weight_recps = []() {
+      using scalar_t = typename GetScalarType<T>::type;
+      std::vector<scalar_t> temp(kChunkSize);
+      for (const auto i : c10::irange(kChunkSize)) {
+        temp[i] = scalar_t(static_cast<double>(1) / static_cast<double>(i + 1));
+      }
+      return temp;
+    }();
+
+template <typename T>
+Welford<T> welford_combine(
+    const Welford<T>& a,
+    const Welford<T>& b,
+    bool use_index = false) {
+  if (a.index == 0) {
+    return b;
+  }
+  if (b.index == 0) {
+    return a;
+  }
+  auto delta = b.mean - a.mean;
+  auto a_weight = use_index ? T(a.index) : a.weight;
+  auto b_weight = use_index ? T(b.index) : b.weight;
+  auto new_weight = a_weight + b_weight;
+  auto new_index = a.index + b.index;
+  auto wb_over_w = b_weight / new_weight;
+  if constexpr (IsVecType<T>::value) {
+    // Guard against division by zero
+    wb_over_w = T::blendv(wb_over_w, T(0), new_weight == T(0));
+  }
+  auto result = Welford<T>{
+      a.mean + delta * wb_over_w,
+      a.m2 + b.m2 + delta * delta * a_weight * wb_over_w,
+      new_weight,
+      new_index};
+  return result;
+}
+
+template <typename T, uint64_t kChunkSize = 0>
+Welford<T> welford_combine(
+    Welford<T>& acc,
+    T& data,
+    WelfordHelper<T, kChunkSize>* w = nullptr) {
+  // Combine welford reduction with cascade summation to improve numerical
+  // stability.
+  // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
+  // https://en.wikipedia.org/wiki/Pairwise_summation
+  if (w != nullptr && w->depth > 0 && acc.index == kChunkSize) {
+    w->welford_stk[0] = welford_combine(w->welford_stk[0], acc);
+    w->num_chunks += 1;
+    acc.mean = T(0);
+    acc.m2 = T(0);
+    acc.weight = T(0);
+    acc.index = 0;
+    uint64_t mask = w->num_chunks;
+    for (uint64_t j = 1; j < w->depth && (mask & 1) == 0; ++j) {
+      w->welford_stk[j] =
+          welford_combine(w->welford_stk[j], w->welford_stk[j - 1]);
+      w->welford_stk[j - 1] = Welford<T>();
+      mask >>= 1;
+    }
+  }
+  // Add a single data point
+  uint64_t new_index = acc.index + 1;
+  auto new_weight = acc.weight + T(1);
+  auto delta = data - acc.mean;
+  T new_mean;
+  // use new_index to fecth 1 / new_weight to avoid divisions
+  new_mean = acc.mean +
+      ((w == nullptr || acc.index >= w->weight_recps.size())
+           ? delta / new_weight
+           : delta * T(w->weight_recps[acc.index]));
+  auto new_delta = data - new_mean;
+  auto result =
+      Welford<T>{new_mean, acc.m2 + delta * new_delta, new_weight, new_index};
+  return result;
+}
+
+template <typename T, uint64_t kChunkSize>
+Welford<T> welford_combine(Welford<T>& acc, WelfordHelper<T, kChunkSize>* w) {
+  for (const auto i : c10::irange(w->depth)) {
+    acc = welford_combine(acc, w->welford_stk[i]);
+  }
+  return acc;
+}
+
+template <typename T>
+struct IndexValue {
+  int64_t index{};
+  T value;
+  IndexValue(int64_t idx, T val) : index(idx), value(val) {}
+  IndexValue() = default;
+};
+
+#if INDUCTOR_USE_VECTOR_TYPES()
+template <typename T, uint64_t kChunkSize = 0>
+Welford<T> welford_combine(
+    Welford<T>& acc,
+    T& data,
+    int64_t tail_size,
+    WelfordHelper<T, kChunkSize>* w = nullptr) {
+  auto out = welford_combine(acc, data, w);
+  return Welford<T>{
+      T::set(acc.mean, out.mean, tail_size),
+      T::set(acc.m2, out.m2, tail_size),
+      T::set(acc.weight, out.weight, tail_size),
+      out.index};
+}
+
+template <typename T, uint64_t kChunkSize = 0>
+inline T cascade_sum_combine(
+    T& data,
+    int64_t tail_size,
+    CascadeSumHelper<T, kChunkSize>* c) {
+  auto out = c->sum_stk[0] + data;
+  c->sum_stk[0] = T::set(c->sum_stk[0], out, tail_size);
+  if (c->depth > 0) {
+    c->index++;
+    if (c->index == kChunkSize) {
+      c->num_chunks += 1;
+      c->index = 0;
+      uint64_t mask = c->num_chunks;
+      uint64_t j = 1;
+      for (; j < c->depth && (mask & 1) == 0; ++j) {
+        c->sum_stk[j] = c->sum_stk[j] + c->sum_stk[j - 1];
+        c->sum_stk[j - 1] = T(0);
+        mask >>= 1;
+      }
+      return c->sum_stk[j - 1];
+    }
+  }
+  return c->sum_stk[0];
+}
+
+template <typename T>
+inline T max_masked_reduce(const T& a, const T& b, const int64_t tail_size) {
+  auto out = at::vec::maximum(a, b);
+  return T::set(a, out, tail_size);
+}
+
+template <>
+inline at::vec::VecMask<float, 1> max_masked_reduce(
+    const at::vec::VecMask<float, 1>& a,
+    const at::vec::VecMask<float, 1>& b,
+    const int64_t tail_size) {
+  auto out = a | b;
+  return at::vec::VecMask<float, 1>::set(a, out, tail_size);
+}
+
+template <typename T>
+inline T min_masked_reduce(const T& a, const T& b, const int64_t tail_size) {
+  auto out = at::vec::minimum(a, b);
+  return T::set(a, out, tail_size);
+}
+
+template <>
+inline at::vec::VecMask<float, 1> min_masked_reduce(
+    const at::vec::VecMask<float, 1>& a,
+    const at::vec::VecMask<float, 1>& b,
+    const int64_t tail_size) {
+  auto out = a & b;
+  return at::vec::VecMask<float, 1>::set(a, out, tail_size);
+}
+
+template <typename T>
+inline T sum_masked_reduce(const T& a, const T& b, const int64_t tail_size) {
+  auto out = a + b;
+  return T::set(a, out, tail_size);
+}
+
+template <>
+inline at::vec::VecMask<float, 1> sum_masked_reduce(
+    const at::vec::VecMask<float, 1>& a,
+    const at::vec::VecMask<float, 1>& b,
+    const int64_t tail_size) {
+  auto out = a | b;
+  return at::vec::VecMask<float, 1>::set(a, out, tail_size);
+}
+
+template <typename T>
+T prod_masked_reduce(const T& a, const T& b, const int64_t tail_size) {
+  auto out = a * b;
+  return T::set(a, out, tail_size);
+}
+
+template <typename T>
+T xor_sum_masked_reduce(const T& a, const T& b, const int64_t tail_size) {
+  auto out = a ^ b;
+  return T::set(a, out, tail_size);
+}
+
+template <typename T>
+T any_masked_reduce(const T& a, const T& b, const int64_t tail_size) {
+  auto out = a | b;
+  return T::set(a, out, tail_size);
+}
+#endif
+
+// Refer to
+// https://github.com/pytorch/pytorch/blob/b5b36cf0c4e1958f1ff25120f5d4beeef3288187/
+// aten/src/ATen/native/SharedReduceOps.h#L419-L445
+template <typename scalar_t>
+inline bool greater_or_nan(
+    scalar_t a,
+    scalar_t b,
+    int64_t idx_a,
+    int64_t idx_b) {
+  // If (a == b), then choose the one with lower idx, else max(a, b)
+  if (at::_isnan(a)) {
+    if (at::_isnan(b)) {
+      return idx_a < idx_b;
+    }
+    return true;
+  }
+  return (a == b) ? idx_a < idx_b : (a > b);
+}
+
+template <typename scalar_t>
+inline bool less_or_nan(scalar_t a, scalar_t b, int64_t idx_a, int64_t idx_b) {
+  // If (a == b), then choose the one with lower idx, else min(a, b)
+  if (at::_isnan(a)) {
+    if (at::_isnan(b)) {
+      return idx_a < idx_b;
+    }
+    return true;
+  }
+  return (a == b) ? idx_a < idx_b : (a < b);
+}
+
+template <typename T>
+inline IndexValue<T>& argmin_combine(
+    IndexValue<T>& a,
+    T next_value,
+    int64_t next_index) {
+  if (!(less_or_nan(a.value, next_value, a.index, next_index))) {
+    a.value = next_value;
+    a.index = next_index;
+  }
+  return a;
+}
+template <typename T>
+inline IndexValue<T>& argmax_combine(
+    IndexValue<T>& a,
+    T next_value,
+    int64_t next_index) {
+  if (!(greater_or_nan(a.value, next_value, a.index, next_index))) {
+    a.value = next_value;
+    a.index = next_index;
+  }
+  return a;
+}
+template <typename T>
+inline IndexValue<T>& argmin_combine(
+    IndexValue<T>& a,
+    const IndexValue<T>& next) {
+  return argmin_combine(a, next.value, next.index);
+}
+template <typename T>
+inline IndexValue<T>& argmax_combine(
+    IndexValue<T>& a,
+    const IndexValue<T>& next) {
+  return argmax_combine(a, next.value, next.index);
+}
+
+#if INDUCTOR_USE_VECTOR_TYPES()
+
+template <typename scalar_t>
+inline at::vec::Vectorized<scalar_t> div_floor_floating_vec(
+    const at::vec::Vectorized<scalar_t>& a,
+    const at::vec::Vectorized<scalar_t>& b) {
+  using vec_t = at::vec::Vectorized<scalar_t>;
+  const auto basic_div = a / b;
+  vec_t inf(std::numeric_limits<scalar_t>::infinity());
+  auto mod = a.fmod(b);
+  // Fixup for a case that isn't properly handled by Sleef_fmod
+  auto floor =
+      vec_t::blendv(a - mod, a, (basic_div.abs() == inf) & (a.abs() != inf));
+  auto div = floor / b;
+  const auto zero = vec_t(0);
+  auto mask = (mod != zero) & ((b < zero) ^ (mod < zero));
+  const auto one = vec_t(1);
+  div = vec_t::blendv(div, div - one, mask);
+  auto floordiv = div.floor();
+  mask = (div - floordiv) > vec_t(0.5);
+  floordiv = vec_t::blendv(floordiv, floordiv + one, mask);
+  floordiv = vec_t::blendv(floordiv, zero.copysign(basic_div), div == zero);
+  floordiv = vec_t::blendv(floordiv, basic_div, b == zero);
+  return floordiv;
+};
+
+template <typename scalar_t, int N>
+inline at::vec::VectorizedN<scalar_t, N> div_floor_floating_vec(
+    const at::vec::VectorizedN<scalar_t, N>& a,
+    const at::vec::VectorizedN<scalar_t, N>& b) {
+  at::vec::VectorizedN<scalar_t, N> result;
+#ifndef _MSC_VER
+#pragma unroll
+#endif
+  for (int i = 0; i < N; ++i) {
+    result[i] = div_floor_floating_vec(a[i], b[i]);
+  }
+  return result;
+}
+
+template <typename T, int NV, int NI>
+struct IndexValueVec {
+  at::vec::VectorizedN<T, NV> value;
+  at::vec::VectorizedN<int64_t, NI> index;
+
+  IndexValueVec(const T _value) {
+    value = at::vec::VectorizedN<T, NV>(_value);
+    index = at::vec::VectorizedN<int64_t, NI>(0);
+  };
+
+  IndexValueVec() {};
+};
+
+template <
+    typename T,
+    int NV,
+    int NI,
+    typename std::enable_if_t<at::vec::is_floating_point_v<T>, int> = 0>
+at::vec::VecMask<int64_t, NI> inline get_mask_for_argmin_argmax(
+    const at::vec::VecMask<T, NV>& vmask,
+    const IndexValueVec<T, NV, NI>& a,
+    const at::vec::VectorizedN<T, NV>& value,
+    const at::vec::VectorizedN<int64_t, NI>& index) {
+  /*
+  vec impl for less_or_nan and greater_or_nan
+  example for argmin:
+  a.value = [NaN, NaN, 0, 2, 1, 0]
+  value = [NaN, 0, 0, 1, 2, NaN]
+  vmask = [false, false, false, false, true, false]
+  all_nan_or_equal = [true, false, true, false, false, false]
+  imask = [a.index[0] < index[0], ..., a.index[-1] < index[-1]]
+  iv_mask = blendv (vmask, imask, all_nan_or_equal)
+          [a.index[0] < index[0], false, a.index[2] < index[2], false, true,
+  false] a_nan_b_not: [false, false, false, false, false, true] mask = iv_mask |
+  a_nan_b_not [a.index[0] < index[0], false, a.index[2] < index[2], false, true,
+  true]
+  */
+  using v_t = at::vec::VecMask<T, NV>;
+  using i_t = at::vec::VecMask<int64_t, NI>;
+  i_t vmask_itype = vmask.template cast<int64_t, NI>();
+  // use itype here since there is vec impl for operator~ for itype
+  // while there may not vec impl for vtype
+  v_t isnan_a = a.value.isnan();
+  i_t isnan_a_itype = isnan_a.template cast<int64_t, NI>();
+  v_t isnan_b = value.isnan();
+  i_t isnan_b_type = isnan_b.template cast<int64_t, NI>();
+  i_t all_nan_mask = isnan_a_itype & isnan_b_type;
+  v_t equal_mask = (a.value == value);
+  i_t equal_mask_itype = equal_mask.template cast<int64_t, NI>();
+  i_t all_nan_or_equal = all_nan_mask | equal_mask_itype;
+  i_t imask(a.index < index);
+  i_t iv_mask = i_t::blendv(vmask_itype, imask, all_nan_or_equal);
+  i_t isnan_a_notnan_b = isnan_a_itype & (~isnan_b_type);
+  return iv_mask | isnan_a_notnan_b;
+}
+
+template <
+    typename T,
+    int NV,
+    int NI,
+    typename std::enable_if_t<!at::vec::is_floating_point_v<T>, int> = 0>
+at::vec::VecMask<int64_t, NI> inline get_mask_for_argmin_argmax(
+    const at::vec::VecMask<T, NV>& vmask,
+    const IndexValueVec<T, NV, NI>& a,
+    const at::vec::VectorizedN<T, NV>& value,
+    const at::vec::VectorizedN<int64_t, NI>& index) {
+  using v_t = at::vec::VecMask<T, NV>;
+  using i_t = at::vec::VecMask<int64_t, NI>;
+  i_t vmask_itype = vmask.template cast<int64_t, NI>();
+  v_t equal_mask = (a.value == value);
+  i_t equal_mask_itype = equal_mask.template cast<int64_t, NI>();
+  i_t imask(a.index < index);
+  return i_t::blendv(vmask_itype, imask, equal_mask_itype);
+}
+
+template <typename T, int NV, int NI>
+inline IndexValueVec<T, NV, NI>& argmin_vec_impl(
+    IndexValueVec<T, NV, NI>& a,
+    at::vec::VectorizedN<T, NV> value,
+    at::vec::VectorizedN<int64_t, NI> index,
+    std::optional<int64_t> tail_size) {
+  at::vec::VecMask<T, NV> vmask(a.value < value);
+  at::vec::VecMask<int64_t, NI> final_mask =
+      get_mask_for_argmin_argmax<T, NV, NI>(vmask, a, value, index);
+  if (tail_size.has_value()) {
+    a.value = at::vec::VectorizedN<T, NV>::set(
+        a.value, at::vec::minimum(a.value, value), tail_size.value());
+    a.index = at::vec::VectorizedN<int64_t, NI>::set(
+        a.index,
+        at::vec::VecMask<int64_t, NI>::blendv(index, a.index, final_mask),
+        tail_size.value());
+  } else {
+    a.value = at::vec::minimum(a.value, value);
+    a.index = at::vec::VecMask<int64_t, NI>::blendv(index, a.index, final_mask);
+  }
+  return a;
+}
+
+template <typename T, int NV, int NI>
+inline IndexValueVec<T, NV, NI>& argmax_vec_impl(
+    IndexValueVec<T, NV, NI>& a,
+    at::vec::VectorizedN<T, NV> value,
+    at::vec::VectorizedN<int64_t, NI> index,
+    std::optional<int64_t> tail_size) {
+  at::vec::VecMask<T, NV> vmask(a.value > value);
+  at::vec::VecMask<int64_t, NI> final_mask =
+      get_mask_for_argmin_argmax<T, NV, NI>(vmask, a, value, index);
+  if (tail_size.has_value()) {
+    a.value = at::vec::VectorizedN<T, NV>::set(
+        a.value, at::vec::maximum(a.value, value), tail_size.value());
+    a.index = at::vec::VectorizedN<int64_t, NI>::set(
+        a.index,
+        at::vec::VecMask<int64_t, NI>::blendv(index, a.index, final_mask),
+        tail_size.value());
+  } else {
+    a.value = at::vec::maximum(a.value, value);
+    a.index = at::vec::VecMask<int64_t, NI>::blendv(index, a.index, final_mask);
+  }
+  return a;
+}
+
+template <typename T, int NI, bool horizontal>
+inline at::vec::VectorizedN<int64_t, NI> create_index(int64_t next_index) {
+  at::vec::VectorizedN<int64_t, NI> next_idx;
+  if constexpr (horizontal) {
+    next_idx = at::vec::VectorizedN<int64_t, NI>::arange(next_index, 1);
+  } else {
+    next_idx = at::vec::VectorizedN<int64_t, NI>(next_index);
+  }
+  return next_idx;
+}
+
+template <typename T, int NV, int NI, bool horizontal>
+inline IndexValueVec<T, NV, NI>& argmin_combine_vec(
+    IndexValueVec<T, NV, NI>& a,
+    at::vec::VectorizedN<T, NV> next_value,
+    int64_t next_index,
+    std::optional<int64_t> tail_size = std::nullopt) {
+  auto next_idx = create_index<T, NI, horizontal>(next_index);
+  return argmin_vec_impl(a, next_value, next_idx, tail_size);
+}
+
+template <typename T, int NV, int NI, bool horizontal>
+inline IndexValueVec<T, NV, NI>& argmax_combine_vec(
+    IndexValueVec<T, NV, NI>& a,
+    at::vec::VectorizedN<T, NV> next_value,
+    int64_t next_index,
+    std::optional<int64_t> tail_size = std::nullopt) {
+  auto next_idx = create_index<T, NI, horizontal>(next_index);
+  return argmax_vec_impl(a, next_value, next_idx, tail_size);
+}
+
+template <typename T, int NV, int NI>
+inline IndexValue<T> argmin_vec_reduce_all(
+    const IndexValueVec<T, NV, NI>& vec) {
+  constexpr int len = at::vec::VectorizedN<T, NV>::size();
+  __at_align__ T tmpval[len];
+  __at_align__ int64_t tmpidx[len];
+  vec.value.store(tmpval);
+  vec.index.store(tmpidx);
+  IndexValue res = IndexValue<T>(tmpidx[0], tmpval[0]);
+  for (int i = 1; i < len; i++) {
+    res = argmin_combine(res, tmpval[i], tmpidx[i]);
+  }
+  return res;
+}
+
+template <typename T, int NV, int NI>
+inline IndexValue<T> argmax_vec_reduce_all(
+    const IndexValueVec<T, NV, NI>& vec) {
+  constexpr int len = at::vec::VectorizedN<T, NV>::size();
+  __at_align__ T tmpval[len];
+  __at_align__ int64_t tmpidx[len];
+  vec.value.store(tmpval);
+  vec.index.store(tmpidx);
+  IndexValue res = IndexValue<T>(tmpidx[0], tmpval[0]);
+  for (int i = 1; i < len; i++) {
+    res = argmax_combine(res, tmpval[i], tmpidx[i]);
+  }
+  return res;
+}
+
+template <typename T, int NV, int NI>
+inline IndexValueVec<T, NV, NI>& argmin_combine_vec(
+    IndexValueVec<T, NV, NI>& vec_a,
+    const IndexValueVec<T, NV, NI>& vec_b,
+    std::optional<int64_t> tail_size = std::nullopt) {
+  return argmin_vec_impl(vec_a, vec_b.value, vec_b.index, tail_size);
+}
+
+template <typename T, int NV, int NI>
+inline IndexValueVec<T, NV, NI>& argmax_combine_vec(
+    IndexValueVec<T, NV, NI>& vec_a,
+    const IndexValueVec<T, NV, NI>& vec_b,
+    std::optional<int64_t> tail_size = std::nullopt) {
+  return argmax_vec_impl(vec_a, vec_b.value, vec_b.index, tail_size);
+}
+
+template <typename scalar_t>
+inline at::vec::Vectorized<scalar_t> vec_shuffle_down(
+    at::vec::Vectorized<scalar_t> x,
+    size_t n) {
+  using Vec = at::vec::Vectorized<scalar_t>;
+  alignas(alignof(Vec)) scalar_t array[Vec::size()];
+  x.store(array);
+  for (size_t i = 0; i + n < Vec::size(); i += 2 * n) {
+    array[i] = array[i + n];
+  }
+  return Vec::loadu(array);
+}
+
+#ifdef CPU_CAPABILITY_AVX2
+inline at::vec::Vectorized<float> vec_shuffle_down(
+    at::vec::Vectorized<float> x,
+    size_t n) {
+  using vec_t = at::vec::Vectorized<float>;
+#define SHUFFLE_MASK(z, y, x, w) ((z << 6) | (y << 4) | (x << 2) | w)
+  switch (n) {
+    case 1:
+      return vec_t(_mm256_permute_ps(x, SHUFFLE_MASK(1, 1, 3, 3)));
+    case 2:
+      return vec_t(_mm256_permute_ps(x, SHUFFLE_MASK(2, 2, 2, 2)));
+    case 4:
+      return vec_t(_mm256_permute2f128_ps(x, x, SHUFFLE_MASK(1, 1, 1, 1)));
+  }
+
+  TORCH_CHECK(false, "Unhandled vec_shuffle_down value ", n);
+}
+#endif
+
+#ifdef CPU_CAPABILITY_AVX512
+inline at::vec::Vectorized<float> vec_shuffle_down(
+    at::vec::Vectorized<float> x,
+    size_t n) {
+  using vec_t = at::vec::Vectorized<float>;
+#define SHUFFLE_MASK(z, y, x, w) ((z << 6) | (y << 4) | (x << 2) | w)
+  switch (n) {
+    case 1:
+      return vec_t(_mm512_permute_ps(x, SHUFFLE_MASK(1, 1, 3, 3)));
+    case 2:
+      return vec_t(_mm512_permute_ps(x, SHUFFLE_MASK(2, 2, 2, 2)));
+    case 4:
+      return vec_t(_mm512_permutexvar_ps(
+          _mm512_set_epi32(
+              12, 12, 12, 12, 12, 12, 12, 12, 4, 4, 4, 4, 4, 4, 4, 4),
+          x));
+    case 8:
+      return vec_t(_mm512_permutexvar_ps(
+          _mm512_set_epi32(8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8), x));
+  }
+
+  TORCH_CHECK(false, "Unhandled vec_shuffle_down value ", n);
+}
+#endif
+
+template <typename scalar_t>
+Welford<scalar_t> welford_vec_reduce_all(
+    Welford<at::vec::Vectorized<scalar_t>> acc) {
+  using Vec = at::vec::Vectorized<scalar_t>;
+  Welford<scalar_t> result;
+  if (acc.index == 0) {
+    return result;
+  }
+  // if all values of acc.weight are same as index,
+  // use index to reduce to save the overhead of vec_shuffle_down for acc.weight
+  bool use_index = (acc.weight - Vec(acc.index)).zero_mask() ==
+      static_cast<int>((1 << Vec::size()) - 1);
+  for (size_t n = 1; n < Vec::size(); n *= 2) {
+    auto shuffled = Welford<Vec>{
+        vec_shuffle_down(acc.mean, n),
+        vec_shuffle_down(acc.m2, n),
+        use_index ? Vec(0) : vec_shuffle_down(acc.weight, n),
+        acc.index};
+    acc = welford_combine(acc, shuffled, use_index);
+  }
+
+  alignas(alignof(Vec)) scalar_t array[Vec::size()];
+  acc.mean.store(array);
+  result.mean = array[0];
+
+  acc.m2.store(array);
+  result.m2 = array[0];
+
+  acc.weight.store(array);
+  result.weight = array[0];
+  result.index = result.weight;
+
+  return result;
+}
+
+template <typename scalar_t>
+Welford<scalar_t> welford_vec_reduce_all(
+    Welford<at::vec::VectorizedN<scalar_t, 2>> acc) {
+  auto Welford0 = Welford<at::vec::Vectorized<scalar_t>>{
+      acc.mean[0], acc.m2[0], acc.weight[0], acc.index};
+  auto Welford1 = Welford<at::vec::Vectorized<scalar_t>>{
+      acc.mean[1], acc.m2[1], acc.weight[1], acc.index};
+  return welford_vec_reduce_all(welford_combine(Welford0, Welford1));
+}
+#endif
+
+template <typename T, typename U>
+inline typename std::common_type_t<T, U> mod(T a, U b) {
+  return a % b;
+}
+template <>
+inline float mod(float a, float b) {
+  return std::fmod(a, b);
+}
+template <>
+inline double mod(double a, double b) {
+  return std::fmod(a, b);
+}
+
+template <typename scalar_t>
+inline scalar_t max_propagate_nan(scalar_t a, scalar_t b) {
+  if (at::_isnan(a)) {
+    return a;
+  }
+  return a > b ? a : b;
+}
+
+template <typename scalar_t>
+inline scalar_t min_propagate_nan(scalar_t a, scalar_t b) {
+  if (at::_isnan(a)) {
+    return a;
+  }
+  return a < b ? a : b;
+}
+
+constexpr float uint32_to_uniform_float(uint32_t value) {
+  // maximum value such that `MAX_INT * scale < 1.0` (with float rounding)
+  constexpr float scale = 4.6566127342e-10;
+  return static_cast<float>(value & 0x7FFFFFFF) * scale;
+}
+
+inline float normalized_rand_cpu(uint32_t seed, uint32_t offset) {
+  return uint32_to_uniform_float(at::Philox4_32(seed, 0, offset)());
+}
+
+inline float randn_cpu(uint32_t seed, uint32_t offset) {
+  at::Philox4_32 engine(seed, 0, offset);
+  return engine.randn(10);
+}
+
+inline int64_t randint64_cpu(
+    uint32_t seed,
+    uint32_t offset,
+    int64_t low,
+    int64_t high) {
+  auto gen = at::Philox4_32(seed, 0, offset);
+  uint64_t r0 = gen();
+  uint64_t r1 = gen();
+  uint64_t result = r0 | (r1 << 32);
+  return static_cast<int64_t>(result % (high - low)) + low;
+}
+
+template <typename T>
+struct AsIntegerType {
+  typedef T type;
+};
+template <>
+struct AsIntegerType<float> {
+  typedef uint32_t type;
+};
+template <>
+struct AsIntegerType<double> {
+  typedef uint64_t type;
+};
+template <>
+struct AsIntegerType<at::BFloat16> {
+  typedef uint16_t type;
+};
+
+template <typename T>
+typename std::enable_if_t<
+    !c10::is_reduced_floating_point_v<T>,
+    T> inline fetch_value(volatile T* addr) {
+  return *addr;
+}
+
+template <typename T>
+typename std::enable_if_t<
+    c10::is_reduced_floating_point_v<T>,
+    T> inline fetch_value(volatile T* addr) {
+  return T(addr->x, T::from_bits());
+}
+
+template <typename T>
+typename std::enable_if_t<!std::is_integral_v<T>> atomic_add(
+    volatile T* addr,
+    T offset) {
+  typedef typename AsIntegerType<T>::type alt_type;
+
+  static_assert(
+      sizeof(std::atomic<alt_type>) == sizeof(T), "std::atomic issue");
+
+  alt_type expected;
+
+  alt_type desired;
+
+  std::atomic<alt_type>* atomic_addr = (std::atomic<alt_type>*)addr;
+  do {
+    T val = fetch_value(addr);
+    reinterpret_cast<T*>(&expected)[0] = val;
+    reinterpret_cast<T*>(&desired)[0] = val + offset;
+  } while (!atomic_addr->compare_exchange_weak(
+      expected, desired, std::memory_order_relaxed));
+}
+
+// Since C++20 float is supported by fetch_add, but the performance may not
+// better than compare_exchange_weak, which can be checked by microbenchmark
+// inductor_cpu_atomic.py
+template <typename T>
+typename std::enable_if_t<std::is_integral_v<T>> atomic_add(
+    volatile T* addr,
+    T offset) {
+  static_assert(sizeof(std::atomic<T>) == sizeof(T), "std::atomic issue");
+  std::atomic<T>* atomic_addr = (std::atomic<T>*)addr;
+  atomic_addr->fetch_add(offset, std::memory_order_relaxed);
+}
+
+#if INDUCTOR_USE_VECTOR_TYPES()
+template <typename T, int NI, int NV>
+void atomic_add_vec(
+    T* addr,
+    at::vec::VectorizedN<int64_t, NI> index,
+    at::vec::VectorizedN<T, NV> offset,
+    std::optional<int64_t> tail_size = std::nullopt) {
+  constexpr int len = at::vec::VectorizedN<int64_t, NI>::size();
+  static_assert(len <= at::vec::VectorizedN<T, NV>::size());
+  __at_align__ std::array<T, len> tmpbuf;
+  __at_align__ std::array<int64_t, len> tmpidx;
+  offset.store(tmpbuf.data(), len);
+  index.store(tmpidx.data(), len);
+  int size = tail_size.has_value() ? tail_size.value() : len;
+  for (int i = 0; i < size; i++) {
+    atomic_add(addr + tmpidx[i], tmpbuf[i]);
+  }
+}
+
+template <typename T, bool atomic_add>
+struct transpose_mxn_helper;
+
+template <typename T>
+struct transpose_mxn_helper<T, true> {
+  static void call(
+      const T* src,
+      int64_t ld_src,
+      T* dst,
+      int64_t ld_dst,
+      int M,
+      int N) {
+    for (int i = 0; i < M; i++) {
+      for (int j = 0; j < N; j++) {
+        atomic_add(&dst[j * ld_dst + i], src[i * ld_src + j]);
+      }
+    }
+  }
+};
+
+template <typename T>
+struct transpose_mxn_helper<T, false> {
+  static void call(
+      const T* src,
+      int64_t ld_src,
+      T* dst,
+      int64_t ld_dst,
+      int M,
+      int N) {
+    at::vec::transpose_mxn<T>(src, ld_src, dst, ld_dst, M, N);
+  }
+};
+
+template <typename T, bool atomic_add>
+inline void transpose_mxn(
+    const T* src,
+    int64_t ld_src,
+    T* dst,
+    int64_t ld_dst,
+    int M,
+    int N) {
+  transpose_mxn_helper<T, atomic_add>::call(src, ld_src, dst, ld_dst, M, N);
+}
+
+template <typename T, int M, int N, bool atomic_add>
+inline void transpose_mxn(
+    const T* src,
+    int64_t ld_src,
+    T* dst,
+    int64_t ld_dst) {
+  transpose_mxn<T, atomic_add>(src, ld_src, dst, ld_dst, M, N);
+}
+#endif
+
+// NOLINTBEGIN(*-avoid-c-arrays)
+inline std::tuple<std::shared_ptr<int64_t[]>, int> _get_factors(
+    int64_t number) {
+  int count = 0;
+  for (auto i = static_cast<int64_t>(std::sqrt(number)); i > 0; --i) {
+    if (number % i == 0) {
+      count += 2;
+    }
+  }
+  auto factors = std::shared_ptr<int64_t[]>(new int64_t[count]);
+  int index = 0;
+  for (auto i = static_cast<int64_t>(std::sqrt(number)); i > 0; --i) {
+    if (number % i == 0) {
+      factors[index++] = number / i;
+      factors[index++] = i;
+    }
+  }
+  return std::make_tuple(factors, count);
+}
+
+inline std::tuple<std::shared_ptr<int64_t[]>, int> get_factors(int64_t number) {
+  thread_local std::map<int64_t, std::tuple<std::shared_ptr<int64_t[]>, int>>
+      cache;
+  auto it = cache.find(number);
+  if (it != cache.end()) {
+    return it->second;
+  } else {
+    auto factors = _get_factors(number);
+    cache[number] = factors;
+    return factors;
+  }
+}
+// NOLINTEND(*-avoid-c-arrays)
+
+inline void _mm_get_thread_blocking(
+    int num_threads,
+    int max_k_slices,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t Mr,
+    int64_t Nr,
+    int64_t Kr,
+    int64_t& Mt,
+    int64_t& Nt,
+    int64_t& Kt) {
+  // see NOTE [Thread blocking in Cpp GEMM] for heuristics
+  Mt = Nt = Kt = 0;
+
+  auto get_blocking = [](int64_t m_factor,
+                         int64_t n_factor,
+                         int64_t k_factor,
+                         int64_t m_blocks,
+                         int64_t n_blocks,
+                         int64_t k_blocks) {
+    int64_t thread_block_k = (k_blocks + k_factor - 1) / k_factor;
+    int64_t thread_block_n = (n_blocks + n_factor - 1) / n_factor;
+    int64_t thread_block_m = (m_blocks + m_factor - 1) / m_factor;
+    return std::make_tuple(thread_block_m, thread_block_n, thread_block_k);
+  };
+
+  auto is_better_blocking = [=](int64_t Mt_,
+                                int64_t Nt_,
+                                int64_t Kt_,
+                                int64_t Mt,
+                                int64_t Nt,
+                                int64_t Kt) {
+    return Mt == 0 || Kt_ < Kt || Mt_ * Mr + Nt_ * Nr < Mt * Mr + Nt * Nr;
+  };
+
+  int64_t m_blocks = (M + Mr - 1) / Mr;
+  int64_t n_blocks = (N + Nr - 1) / Nr;
+  int64_t k_blocks = (K + Kr - 1) / Kr;
+
+  auto [factors, count] = get_factors(num_threads);
+  assert(count > 0);
+
+  for (int i = 0; i < count; ++i) {
+    int64_t n_factor = factors[i];
+    int64_t m_factor = num_threads / n_factor;
+    if (n_blocks >= n_factor && m_blocks >= m_factor) {
+      auto [Mt_, Nt_, Kt_] =
+          get_blocking(m_factor, n_factor, 1, m_blocks, n_blocks, k_blocks);
+      if (is_better_blocking(Mt_, Nt_, Kt_, Mt, Nt, Kt)) {
+        std::tie(Mt, Nt, Kt) = std::make_tuple(Mt_, Nt_, Kt_);
+      }
+    }
+  }
+
+  if (Mt != 0) {
+    return;
+  }
+
+  for (int i = 0; i < count; ++i) {
+    int64_t k_factor = factors[i];
+    if (k_blocks >= k_factor &&
+        (max_k_slices == 0 || k_factor <= max_k_slices)) {
+      auto [mxn_factors, mxn_count] = get_factors(num_threads / k_factor);
+      for (int j = 0; j < mxn_count; ++j) {
+        int64_t n_factor = mxn_factors[j];
+        int64_t m_factor = num_threads / (k_factor * n_factor);
+        if (n_blocks >= n_factor && m_blocks >= m_factor) {
+          auto [Mt_, Nt_, Kt_] = get_blocking(
+              m_factor, n_factor, k_factor, m_blocks, n_blocks, k_blocks);
+          if (is_better_blocking(Mt_, Nt_, Kt_, Mt, Nt, Kt)) {
+            std::tie(Mt, Nt, Kt) = std::make_tuple(Mt_, Nt_, Kt_);
+          }
+        }
+      }
+    }
+  }
+
+  if (Mt != 0) {
+    return;
+  }
+
+  for (int i = 0; i < count; ++i) {
+    int64_t n_factor = factors[i];
+    int64_t m_factor = num_threads / n_factor;
+    if (n_blocks >= n_factor || m_blocks >= m_factor) {
+      auto [Mt_, Nt_, Kt_] =
+          get_blocking(m_factor, n_factor, 1, m_blocks, n_blocks, k_blocks);
+      if (is_better_blocking(Mt_, Nt_, Kt_, Mt, Nt, Kt)) {
+        std::tie(Mt, Nt, Kt) = std::make_tuple(Mt_, Nt_, Kt_);
+      }
+    }
+  }
+
+  assert(Mt != 0);
+}
+
+inline void mm_get_thread_blocking(
+    int num_threads,
+    int max_k_slices,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t Mr,
+    int64_t Nr,
+    int64_t Kr,
+    int64_t& Mt,
+    int64_t& Nt,
+    int64_t& Kt) {
+  thread_local std::map<
+      std::
+          tuple<int, int, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t>,
+      std::tuple<int64_t, int64_t, int64_t>>
+      cache;
+  auto key = std::make_tuple(num_threads, max_k_slices, M, N, K, Mr, Nr, Kr);
+  auto it = cache.find(key);
+  if (it != cache.end()) {
+    std::tie(Mt, Nt, Kt) = it->second;
+    return;
+  } else {
+    _mm_get_thread_blocking(
+        num_threads, max_k_slices, M, N, K, Mr, Nr, Kr, Mt, Nt, Kt);
+    cache[key] = std::make_tuple(Mt, Nt, Kt);
+  }
+}
+
+// NOLINTBEGIN(*-narrowing-conversions)
+template <typename X_t, typename W_t>
+void _mm_get_cache_blocking(
+    int num_threads,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t Mr,
+    int64_t Nr,
+    int64_t Kr,
+    int64_t Mt_blocks,
+    int64_t Nt_blocks,
+    int64_t Kt_blocks,
+    int64_t& Mc_blocks,
+    int64_t& Nc_blocks,
+    int64_t& Kc_blocks,
+    uint32_t L1_cache_size,
+    uint32_t L2_cache_size) {
+  // See NOTE [CPP GEMM Cache Blocking Algorithm] for the cache blocking
+  // algorithm.
+  // TODO(jgong5): cache cache blocking results
+  // TODO: tune the factor here
+  float L1_limit_factor = 0.8;
+  float L2_limit_factor = 0.5;
+
+  auto L1 = L1_cache_size * L1_limit_factor;
+  auto L2 = L2_cache_size * L2_limit_factor;
+
+  constexpr size_t num_byte_A = sizeof(X_t);
+  constexpr size_t num_byte_B = sizeof(W_t);
+
+  int64_t size_cache_B = Kr * Kt_blocks * Nr * num_byte_B;
+  Kc_blocks = Kt_blocks;
+  if (size_cache_B > L1) {
+    Kc_blocks = (int64_t)std::floor(L1 / (Kr * Nr * num_byte_B));
+  }
+
+  float min_Mc_ratio = 2;
+  int64_t min_Mc_blocks = std::ceil(min_Mc_ratio * Mr / Nr);
+  auto Kt_bytes = Kt_blocks * Kr * num_byte_A;
+  if (min_Mc_blocks * Mr * Kt_bytes < L2) {
+    Mc_blocks = std::min(Mt_blocks, (int64_t)std::floor(L2 / (Mr * Kt_bytes)));
+    Nc_blocks = 1;
+  } else {
+    Mc_blocks = Mt_blocks;
+    Nc_blocks =
+        std::min((int64_t)std::ceil((float)Mc_blocks * Mr / Nr), Nt_blocks);
+    auto Nc_bytes = Nc_blocks * Nr * 4;
+    auto Kc_bytes = Kc_blocks * Kr * num_byte_A;
+    if (Mc_blocks * Mr * (Kc_bytes + Nc_bytes) > L2) {
+      auto M_max = (std::sqrt(Kc_bytes * Kc_bytes + 16 * L2) - Kc_bytes) / 8;
+      if (M_max < Mc_blocks * Mr) {
+        Mc_blocks = (int64_t)std::floor(M_max / Mr);
+        Nc_blocks =
+            std::min((int64_t)std::ceil((float)Mc_blocks * Mr / Nr), Nt_blocks);
+      }
+    }
+  }
+}
+// NOLINTEND(*-narrowing-conversions)
+
+template <typename X_t, typename W_t>
+void mm_get_cache_blocking(
+    int num_threads,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t Mr,
+    int64_t Nr,
+    int64_t Kr,
+    int64_t Mt_blocks,
+    int64_t Nt_blocks,
+    int64_t Kt_blocks,
+    int64_t& Mc_blocks,
+    int64_t& Nc_blocks,
+    int64_t& Kc_blocks,
+    uint32_t L1_cache_size,
+    uint32_t L2_cache_size) {
+  thread_local std::map<
+      std::tuple<
+          int,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t,
+          int64_t>,
+      std::tuple<int64_t, int64_t, int64_t>>
+      cache;
+  auto key = std::make_tuple(
+      num_threads,
+      M,
+      N,
+      K,
+      Mr,
+      Nr,
+      Kr,
+      Mt_blocks,
+      Nt_blocks,
+      Kt_blocks,
+      L1_cache_size,
+      L2_cache_size);
+  auto it = cache.find(key);
+  if (it != cache.end()) {
+    std::tie(Mc_blocks, Nc_blocks, Kc_blocks) = it->second;
+    return;
+  } else {
+    _mm_get_cache_blocking<X_t, W_t>(
+        num_threads,
+        M,
+        N,
+        K,
+        Mr,
+        Nr,
+        Kr,
+        Mt_blocks,
+        Nt_blocks,
+        Kt_blocks,
+        Mc_blocks,
+        Nc_blocks,
+        Kc_blocks,
+        L1_cache_size,
+        L2_cache_size);
+    cache[key] = std::make_tuple(Mc_blocks, Nc_blocks, Kc_blocks);
+  }
+}
+
+struct amx_tilecfg {
+  uint8_t palette_id{0};
+  uint8_t start_row{0};
+  std::array<uint8_t, 14> reserved_0{};
+  std::array<uint16_t, 16> colsb{};
+  std::array<uint8_t, 16> rows{};
+};
+
+class AMXState {
+ private:
+  amx_tilecfg tilecfg_{};
+  uint8_t rows_{0};
+  uint16_t colsb_{0};
+  uint8_t num_tile_rows_{0};
+  uint8_t num_tile_columns_{0};
+
+ public:
+  AMXState() = default;
+
+  inline void configure(
+      uint8_t rows,
+      uint16_t colsb,
+      uint8_t num_tile_rows,
+      uint8_t num_tile_columns,
+      void (*loadconfig)(const amx_tilecfg&)) {
+    if (tilecfg_.palette_id == 1 && rows_ == rows && colsb_ == colsb &&
+        num_tile_rows_ == num_tile_rows &&
+        num_tile_columns_ == num_tile_columns) {
+      return;
+    }
+    tilecfg_.palette_id = 1;
+    rows_ = rows;
+    colsb_ = colsb;
+    num_tile_rows_ = num_tile_rows;
+    num_tile_columns_ = num_tile_columns;
+    const auto num_c_tiles = num_tile_rows * num_tile_columns;
+    // For C
+    for (int i = 0; i < num_c_tiles; i++) {
+      tilecfg_.rows[i] = rows;
+      tilecfg_.colsb[i] = 64;
+    }
+    // For A
+    for (int i = 0; i < num_tile_rows; i++) {
+      tilecfg_.rows[i + num_c_tiles] = rows;
+      tilecfg_.colsb[i + num_c_tiles] = colsb;
+    }
+    // For B
+    for (int i = 0; i < num_tile_columns; i++) {
+      tilecfg_.rows[i + num_c_tiles + num_tile_rows] = colsb / 4;
+      tilecfg_.colsb[i + num_c_tiles + num_tile_rows] = 64;
+    }
+    loadconfig(tilecfg_);
+  }
+
+  inline void release(void (*tile_release)()) {
+    tilecfg_.palette_id = 0;
+    tile_release();
+  }
+};
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/array_ref.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/array_ref.h
new file mode 100644
index 0000000000000000000000000000000000000000..3ff757120d49fa01155e8e20626bff05d9a0c9fe
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/array_ref.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runtime/arrayref_tensor.h>
+#include <torch/csrc/inductor/aoti_runtime/thread_local.h>
+#include <torch/csrc/inductor/array_ref_impl.h>
+#include <torch/csrc/inductor/cpp_wrapper/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/common.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/common.h
new file mode 100644
index 0000000000000000000000000000000000000000..d8c5697757770bde1bf1cce3c0d166c6f3a0f17a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/common.h
@@ -0,0 +1,95 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <filesystem>
+#include <optional>
+#include <utility>
+
+#include <Python.h>
+#if __has_include(<pybind11/gil_simple.h>)
+#include <pybind11/gil_simple.h>
+#else
+// pybind11 < 3.0: gil_simple.h does not exist yet.
+#define PYBIND11_SIMPLE_GIL_MANAGEMENT
+#include <pybind11/gil.h>
+// Provide the _simple aliases so generated code works with either version.
+namespace pybind11 {
+using gil_scoped_acquire_simple = gil_scoped_acquire;
+using gil_scoped_release_simple = gil_scoped_release;
+} // namespace pybind11
+#endif
+
+// Required for custom op dispatch via the stable ABI
+#include <torch/csrc/stable/library.h>
+
+#ifdef TORCH_INDUCTOR_PRECOMPILE_HEADERS
+// Include some often-used cpp_wrapper headers, for precompiling.
+#include <c10/util/BFloat16.h>
+#include <torch/csrc/Device.h>
+#include <torch/csrc/DynamicTypes.h>
+#include <torch/csrc/utils/pythoncapi_compat.h>
+#include <torch/csrc/utils/tensor_memoryformats.h>
+#endif
+
+namespace py = pybind11; // NOLINT(misc-unused-alias-decls)
+
+class RAIIPyObject {
+ public:
+  RAIIPyObject() = default;
+  // steals a reference to a PyObject
+  RAIIPyObject(PyObject* obj) : obj_{obj} {}
+  RAIIPyObject(const RAIIPyObject& other) : obj_{other.obj_} {
+    Py_XINCREF(obj_);
+  }
+  RAIIPyObject(RAIIPyObject&& other) noexcept {
+    // refcount doesn't change, and obj_ is currently nullptr
+    std::swap(obj_, other.obj_);
+  }
+  ~RAIIPyObject() {
+    Py_XDECREF(obj_);
+  }
+  RAIIPyObject& operator=(const RAIIPyObject& other) {
+    if (this != &other) {
+      Py_XDECREF(obj_);
+      obj_ = other.obj_;
+      Py_XINCREF(obj_);
+    }
+    return *this;
+  }
+  RAIIPyObject& operator=(RAIIPyObject&& other) noexcept {
+    // refcount to the current object decreases, but refcount to other.obj_ is
+    // the same
+    Py_XDECREF(obj_);
+    obj_ = std::exchange(other.obj_, nullptr);
+    return *this;
+  }
+  operator bool() const noexcept {
+    return obj_;
+  }
+  operator PyObject*() {
+    return obj_;
+  }
+  PyObject* get() {
+    return obj_;
+  }
+
+ private:
+  PyObject* obj_{nullptr};
+};
+
+#include <torch/csrc/inductor/aoti_runtime/device_utils.h>
+#include <torch/csrc/inductor/aoti_runtime/utils.h>
+using namespace torch::aot_inductor;
+
+#include <c10/util/generic_math.h>
+#include <torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h>
+
+// Round up to the nearest multiple of 64
+[[maybe_unused]] inline int64_t align(int64_t nbytes) {
+  return (nbytes + 64 - 1) & -64;
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/cpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/cpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..e98e987de358efd0ad1c810d837cae8205bbb33c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/cpu.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/cpp_wrapper/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..209d9b49c4f1dbcd60764858ba2a1d4be7f394be
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/cuda.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/inductor/cpp_wrapper/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h>
+#include <torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h>
+
+#ifdef TORCH_INDUCTOR_PRECOMPILE_HEADERS
+#include <ATen/cuda/EmptyTensor.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..98fd14fb45fc6138446fcc80b9148b3471d4ebde
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h
@@ -0,0 +1,8 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..2f9e7104c1b2a16f8d4a0dbd920e6a2816f059a2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runtime/utils_cuda.h>
+#include <torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/mps.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/mps.h
new file mode 100644
index 0000000000000000000000000000000000000000..7711a5e019966c970818c5e7082a5b99f22856e1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/mps.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim_mps.h>
+#include <torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..eda98fd2052767ac360e31245fa94a9e350cbdaf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h
@@ -0,0 +1,10 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h>
+#include <torch/csrc/inductor/aoti_runtime/utils_xpu.h>
+#include <torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h
new file mode 100644
index 0000000000000000000000000000000000000000..a3bb04b421d2b02d5533fedf9b11f5bb95a96453
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h
@@ -0,0 +1,206 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+
+#include <torch/csrc/inductor/aoti_torch/utils.h>
+#include <torch/csrc/inductor/cpp_wrapper/common.h>
+#if defined(USE_XPU)
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h>
+#else
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h>
+#endif
+
+struct LazyKernelCompileResult {
+  std::string cubin_path;
+  std::string mangled_name;
+  int num_warps;
+  int shared_mem;
+  int xblock;
+  int yblock;
+  int zblock;
+  int r0block;
+  int rsplit;
+  int rsplit_size;
+  int config_index;
+  int global_scratch;
+  int profile_scratch;
+};
+
+static PyObject* (*_THPVariable_Wrap)(const at::TensorBase&) = nullptr;
+static int32_t (*_THPUtils_unpackInt)(PyObject*) = nullptr;
+
+// Cached module and function references
+static PyObject* triton_lazy_compile_module = nullptr;
+static PyObject* start_kernel_compile = nullptr;
+static PyObject* run_triton_kernel_with_autotune = nullptr;
+
+// Per-module dict for pending kernel compile results (avoids global state
+// collisions when multiple compiled modules produce kernels with the same
+// name).
+static PyObject* _module_pending_kernels = nullptr;
+
+static inline void loadLazyCompileFuncs() {
+  if (triton_lazy_compile_module == nullptr) {
+    triton_lazy_compile_module =
+        PyImport_ImportModule("torch._inductor.runtime.triton_lazy_compile");
+    AOTI_TORCH_CHECK(
+        triton_lazy_compile_module, "Failed to import triton_lazy_compile");
+
+    start_kernel_compile = PyObject_GetAttrString(
+        triton_lazy_compile_module, "start_kernel_compile");
+    AOTI_TORCH_CHECK(
+        start_kernel_compile, "Failed to get start_kernel_compile");
+
+    run_triton_kernel_with_autotune = PyObject_GetAttrString(
+        triton_lazy_compile_module, "run_triton_kernel_with_autotune");
+    AOTI_TORCH_CHECK(
+        run_triton_kernel_with_autotune,
+        "Failed to get run_triton_kernel_with_autotune");
+
+    RAIIPyObject guards_mod = PyImport_ImportModule("torch._C._dynamo.guards");
+    AOTI_TORCH_CHECK(guards_mod, "Failed to import torch._C._dynamo.guards");
+
+    RAIIPyObject wrap_addr =
+        PyObject_GetAttrString(guards_mod, "_torchinductor_thp_variable_wrap");
+    AOTI_TORCH_CHECK(
+        wrap_addr, "Failed to get _torchinductor_thp_variable_wrap");
+    _THPVariable_Wrap = reinterpret_cast<decltype(_THPVariable_Wrap)>(
+        PyLong_AsVoidPtr(wrap_addr));
+    AOTI_TORCH_CHECK(_THPVariable_Wrap, "THPVariable_Wrap not resolved");
+
+    RAIIPyObject unpack_addr = PyObject_GetAttrString(
+        guards_mod, "_torchinductor_thputils_unpack_int");
+    AOTI_TORCH_CHECK(
+        unpack_addr, "Failed to get _torchinductor_thputils_unpack_int");
+    _THPUtils_unpackInt = reinterpret_cast<decltype(_THPUtils_unpackInt)>(
+        PyLong_AsVoidPtr(unpack_addr));
+    AOTI_TORCH_CHECK(_THPUtils_unpackInt, "THPUtils_unpackInt not resolved");
+  }
+}
+
+static inline std::string getStringAttr(PyObject* obj, const char* attr) {
+  RAIIPyObject val = PyObject_GetAttrString(obj, attr);
+  AOTI_TORCH_CHECK(val, "Failed to get attribute");
+  return PyUnicode_AsUTF8(val);
+}
+
+static inline int getIntAttr(PyObject* obj, const char* attr) {
+  RAIIPyObject val = PyObject_GetAttrString(obj, attr);
+  AOTI_TORCH_CHECK(val, "Failed to get attribute");
+  return _THPUtils_unpackInt(val);
+}
+
+static inline int getOptionalIntAttr(
+    PyObject* obj,
+    const char* attr,
+    int sentinel = -1) {
+  RAIIPyObject val = PyObject_GetAttrString(obj, attr);
+  AOTI_TORCH_CHECK(val, "Failed to get attribute");
+  return (val.get() != Py_None) ? _THPUtils_unpackInt(val) : sentinel;
+}
+
+static inline LazyKernelCompileResult extractCompileResult(PyObject* result) {
+  LazyKernelCompileResult compile_result;
+  compile_result.cubin_path = getStringAttr(result, "cubin_path");
+  compile_result.mangled_name = getStringAttr(result, "mangled_name");
+  compile_result.num_warps = getIntAttr(result, "num_warps");
+  compile_result.shared_mem = getIntAttr(result, "shared_mem");
+  compile_result.xblock = getIntAttr(result, "xblock");
+  compile_result.yblock = getIntAttr(result, "yblock");
+  compile_result.zblock = getIntAttr(result, "zblock");
+  compile_result.r0block = getIntAttr(result, "r0block");
+  compile_result.rsplit = getIntAttr(result, "rsplit");
+  compile_result.rsplit_size = getIntAttr(result, "rsplit_size");
+  compile_result.config_index = getOptionalIntAttr(result, "config_index");
+  compile_result.global_scratch = getOptionalIntAttr(result, "global_scratch");
+  compile_result.profile_scratch =
+      getOptionalIntAttr(result, "profile_scratch");
+  return compile_result;
+}
+
+template <typename T>
+static inline PyObject* convertArgToPython(const T& arg) {
+  using DecayedT = std::decay_t<T>;
+  if constexpr (std::is_same_v<DecayedT, AtenTensorHandle>) {
+    at::Tensor* tensor_ptr =
+        torch::aot_inductor::tensor_handle_to_tensor_pointer(arg);
+    return _THPVariable_Wrap(*tensor_ptr);
+  } else if constexpr (std::is_same_v<
+                           DecayedT,
+                           torch::aot_inductor::RAIIAtenTensorHandle>) {
+    at::Tensor* tensor_ptr =
+        torch::aot_inductor::tensor_handle_to_tensor_pointer(arg.get());
+    return _THPVariable_Wrap(*tensor_ptr);
+  } else if constexpr (std::is_same_v<DecayedT, bool>) {
+    PyObject* py_arg = arg ? Py_True : Py_False;
+    Py_INCREF(py_arg);
+    return py_arg;
+  } else if constexpr (std::is_integral_v<DecayedT>) {
+    return PyLong_FromLongLong(static_cast<long long>(arg));
+  } else if constexpr (std::is_floating_point_v<DecayedT>) {
+    return PyFloat_FromDouble(static_cast<double>(arg));
+  } else {
+    AOTI_TORCH_CHECK(false, "Invalid input type to convertArgToPython");
+  }
+}
+
+template <typename... Args>
+static inline LazyKernelCompileResult runTritonKernelWithAutotune(
+    PyObject* pending_kernels,
+    const std::string& kernel_name,
+    void* stream,
+    const Args&... kernel_args) {
+  py::gil_scoped_acquire_simple acquire;
+
+  constexpr size_t num_args = sizeof...(Args);
+  RAIIPyObject py_args_list = PyList_New(num_args);
+  AOTI_TORCH_CHECK(py_args_list, "Failed to create args list");
+
+  size_t idx = 0;
+  auto add_arg = [&py_args_list, &idx](PyObject* py_arg) {
+    AOTI_TORCH_CHECK(py_arg, "Failed to convert argument");
+    PyList_SetItem(py_args_list, idx++, py_arg);
+  };
+  // Use array pack-expansion instead of a fold expression to avoid
+  // hitting the compiler's expression-nesting limit when there are
+  // hundreds of kernel arguments (e.g. combo kernels).
+  int dummy[] = {0, (add_arg(convertArgToPython(kernel_args)), 0)...};
+  (void)dummy;
+
+  RAIIPyObject call_args = PyTuple_Pack(
+      4,
+      pending_kernels,
+      PyUnicode_FromString(kernel_name.c_str()),
+      PyLong_FromVoidPtr(stream),
+      py_args_list.get());
+  AOTI_TORCH_CHECK(call_args, "Failed to create call args");
+
+  RAIIPyObject result =
+      PyObject_CallObject(run_triton_kernel_with_autotune, call_args);
+  AOTI_TORCH_CHECK(result, "Failed to run kernel with autotuning");
+
+  return extractCompileResult(result);
+}
+
+static inline void startKernelCompile(
+    PyObject* pending_kernels,
+    const std::string& kernel_name,
+    const std::string& kernel_source) {
+  py::gil_scoped_acquire_simple acquire;
+
+  RAIIPyObject py_name = PyUnicode_FromString(kernel_name.c_str());
+  RAIIPyObject py_source = PyUnicode_FromString(kernel_source.c_str());
+  AOTI_TORCH_CHECK(py_name && py_source, "Failed to create Python args");
+
+  RAIIPyObject call_args =
+      PyTuple_Pack(3, pending_kernels, py_name.get(), py_source.get());
+  AOTI_TORCH_CHECK(call_args, "Failed to create call args");
+
+  RAIIPyObject result = PyObject_CallObject(start_kernel_compile, call_args);
+  AOTI_TORCH_CHECK(result, "Failed to start kernel compilation");
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/mps.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/mps.h
new file mode 100644
index 0000000000000000000000000000000000000000..00d4464071f64f8c46df042813e6bd56222a304e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/mps.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/cpp_wrapper/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/mps.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..705efd22032156ac5125638d286741568d0c25bd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/cpp_wrapper/xpu.h
@@ -0,0 +1,10 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/inductor/cpp_wrapper/common.h>
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h>
+#include <torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/inductor_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/inductor_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..77eef6c69f5eb4725fcf08489bf57e27dd97ae30
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/inductor_ops.h
@@ -0,0 +1,44 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Tensor.h>
+
+namespace torch::inductor {
+
+TORCH_API at::Tensor _mm_plus_mm_out(
+    at::Tensor& out,
+    const at::Tensor& a,
+    const at::Tensor& b,
+    const at::Tensor& c,
+    const at::Tensor& d);
+
+// After adding _mm_plus_mm_out, this should not be exposed and called by model
+// code. Keeping it around for backward compatibility. Will be deprecated later.
+TORCH_API at::Tensor _mm_plus_mm(
+    const at::Tensor& a,
+    const at::Tensor& b,
+    const at::Tensor& c,
+    const at::Tensor& d,
+    at::Tensor& out);
+
+TORCH_API at::Tensor _alloc_from_pool(
+    const at::Tensor& self,
+    int64_t offset_bytes,
+    at::ScalarType dtype,
+    at::IntArrayRef size,
+    at::IntArrayRef stride);
+
+// Similar to as_strided with the following differences
+// - offset is added to the existing offset (rather than replacing it)
+// - view tracking is disabled similar to unsafe_view
+TORCH_API at::Tensor _reinterpret_tensor(
+    const at::Tensor& self,
+    at::IntArrayRef size,
+    at::IntArrayRef stride,
+    int64_t offset_increment = 0);
+
+} // namespace torch::inductor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/static_launcher/cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/static_launcher/cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..4015bd2d535659c1e6b2cad92a71cd04af2117ac
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/static_launcher/cuda.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#if defined(USE_CUDA)
+#include <torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h>
+#include <torch/csrc/python_headers.h>
+
+bool StaticCudaLauncher_init(PyObject* module);
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/static_launcher/xpu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/static_launcher/xpu.h
new file mode 100644
index 0000000000000000000000000000000000000000..91fc6de742003b33c79b3ebdb2083e37533c1b6a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/inductor/static_launcher/xpu.h
@@ -0,0 +1,11 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#if defined(USE_XPU)
+#include <torch/csrc/python_headers.h>
+
+bool StaticXpuLauncher_init(PyObject* module);
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/instruction_counter/Module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/instruction_counter/Module.h
new file mode 100644
index 0000000000000000000000000000000000000000..b7fbac85c979734a709aef89819ae609e54b6686
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/instruction_counter/Module.h
@@ -0,0 +1,13 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/python_headers.h>
+
+namespace torch::instruction_counter {
+
+void initModule(PyObject* module);
+
+} // namespace torch::instruction_counter
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/itt.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/itt.h
new file mode 100644
index 0000000000000000000000000000000000000000..64d9824c3b551da4ee761a2ab213bb8f7ce143b6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/itt.h
@@ -0,0 +1,13 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef ITT_H
+#define ITT_H
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::profiler {
+void initIttBindings(PyObject* module); // namespace torch::profiler
+}
+#endif // ITT_H
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/itt_wrapper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/itt_wrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..3a556190443fc84b9eb72fb5f6f251149425b318
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/itt_wrapper.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef PROFILER_ITT_H
+#define PROFILER_ITT_H
+#include <c10/macros/Export.h>
+
+namespace torch::profiler {
+TORCH_API bool itt_is_available();
+TORCH_API void itt_range_push(const char* msg);
+TORCH_API void itt_range_pop();
+TORCH_API void itt_mark(const char* msg);
+} // namespace torch::profiler
+
+#endif // PROFILER_ITT_H
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/compilation_unit.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/compilation_unit.h
new file mode 100644
index 0000000000000000000000000000000000000000..e6264f6f992a61123df6647a90176397261bc47f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/compilation_unit.h
@@ -0,0 +1,356 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/function.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/jit/api/function_impl.h>
+#include <torch/csrc/jit/frontend/name_mangler.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/runtime/graph_executor.h>
+
+#include <torch/csrc/Export.h>
+
+#include <ATen/core/function_schema.h>
+#include <ATen/core/qualified_name.h>
+#include <c10/util/ArrayRef.h>
+#include <optional>
+
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <ostream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace torch::jit {
+
+struct Def;
+struct Property;
+struct ClassDef;
+struct SugaredValue;
+struct Resolver;
+
+using ResolverPtr = std::shared_ptr<Resolver>;
+struct Self {
+  virtual ~Self() = default;
+  virtual std::shared_ptr<SugaredValue> makeSugared(Value* v) const = 0;
+  virtual ClassTypePtr getClassType() const = 0;
+};
+
+// A CompilationUnit is a list of named Functions
+// with helper methods to iterate the list or invoke the function.
+// Classes have a CompilationUnit holding the class methods,
+// and Modules have a CompilationUnit holding the Functions that
+// are used to implement their Methods
+
+struct TORCH_API CompilationUnit {
+  enum class FunctionType { Method, Hook, PreHook };
+  // constructor that takes a set of functions to compile using the native
+  // resolver
+  explicit CompilationUnit(const std::string& source);
+  CompilationUnit() = default;
+
+  CompilationUnit& operator=(CompilationUnit&&) = default;
+  CompilationUnit(CompilationUnit&&) = default;
+  CompilationUnit& operator=(const CompilationUnit&) = delete;
+  CompilationUnit(const CompilationUnit&) = delete;
+
+  Function* find_function(const c10::QualifiedName& name) const {
+    auto it = dict_.find(name);
+    if (it == dict_.end()) {
+      return nullptr;
+    }
+    return functions_[it->second].get();
+  }
+
+  Function& get_function(const c10::QualifiedName& name) const {
+    if (auto r = find_function(name)) {
+      return *r;
+    }
+    TORCH_CHECK(false, "attempted to get undefined function ", name.name());
+  }
+
+  void set_optimized(bool o) {
+    TORCH_WARN(
+        "CompilationUnit::set_optimized() is deprecated and has no effect. "
+        "Please use setGraphExecutorOptimize()");
+  }
+
+  bool is_optimized() const {
+    TORCH_WARN(
+        "CompilationUnit::is_optimized() is deprecated and always returns true. "
+        "Please use getGraphExecutorOptimize()");
+    return true;
+  }
+
+  // for historic reasons, these are defined in ir_emitter.cpp
+  // Returns the list of Functions just defined.
+  std::vector<Function*> define(
+      const std::optional<c10::QualifiedName>& prefix,
+      const std::vector<Property>& properties,
+      const std::vector<ResolverPtr>& propResolvers,
+      const std::vector<Def>& definitions,
+      const std::vector<ResolverPtr>&
+          defResolvers, /* determines how we handle free
+                     variables in each definition*/
+      // if non-null, the first argument to each def, is bound to this value
+      const Self* self,
+      // see [name mangling]
+      bool shouldMangle = false,
+      std::optional<size_t> operator_set_version = std::nullopt);
+
+  void define_hooks(
+      const std::optional<c10::QualifiedName>& prefix,
+      const std::vector<Def>& hookDefs,
+      const std::vector<ResolverPtr>& hookResolvers,
+      const std::vector<Def>& preHookDefs,
+      const std::vector<ResolverPtr>& preHookResolvers,
+      const Self* self,
+      bool shouldMangle = false);
+
+  // same as above but parse the definitions from source
+  // Returns the list of Functions just defined.
+  std::vector<Function*> define(
+      // prefix namespace to put all the defined functions into
+      const std::optional<c10::QualifiedName>& prefix,
+      const std::string& source,
+      const ResolverPtr& resolver,
+      const Self* self);
+
+  void define_interface(
+      const c10::QualifiedName& qualifiedName,
+      const ClassDef& classDef,
+      ResolverPtr rcb,
+      bool is_module = false);
+
+  Function* create_function(
+      c10::QualifiedName name,
+      std::shared_ptr<Graph> graph,
+      bool shouldMangle = false) {
+    if (shouldMangle) {
+      name = mangle(name);
+    }
+    auto fn = std::make_unique<GraphFunction>(
+        std::move(name), std::move(graph), nullptr);
+    auto ret = fn.get();
+    register_function(std::move(fn));
+    return ret;
+  }
+
+  std::vector<Function*> get_functions() const {
+    return fmap(functions_, [](const std::unique_ptr<Function>& fn) {
+      return fn.get();
+    });
+  }
+
+  /// Run a method from this compilation.
+  ///
+  /// For example:
+  /// @code
+  ///   IValue output = module->run("relu_script", a, b);
+  /// @endcode
+  ///
+  /// To get a compile a module from a source string, see torch::jit::compile
+  ///
+  /// @param method_name The name of the method to run
+  /// @param args Arguments to be passed to the method
+  /// @return An IValue containing the return value (or values if it is a tuple)
+  /// from the method
+  template <typename... Types>
+  IValue run_method(const c10::QualifiedName& method_name, Types&&... args) {
+    return get_function(method_name)({IValue(std::forward<Types>(args))...});
+  }
+
+  void drop_all_functions() {
+    dict_.clear();
+    functions_.clear();
+  }
+
+  /**
+   * Register a class as being owned by this compilation unit.
+   */
+  void register_type(c10::NamedTypePtr namedType) {
+    // TODO: class types cannot be redefined because we have no way right now
+    // of invalidating their methods. NamedTuples are fine though, since they
+    // don't have methods.
+    TORCH_CHECK(
+        0 == classDict_.count(*namedType->name()),
+        "class '",
+        namedType->name()->qualifiedName(),
+        "' already defined.");
+    classes_.push_back(std::move(namedType));
+    classDict_[*classes_.back()->name()] = classes_.size() - 1;
+  }
+
+  c10::ClassTypePtr get_class(const c10::QualifiedName& name) const {
+    auto type = get_type(name);
+    if (!type) {
+      return nullptr;
+    }
+    return type->cast<c10::ClassType>();
+  }
+
+  c10::InterfaceTypePtr get_interface(const c10::QualifiedName& name) const {
+    auto type = get_type(name);
+    if (!type) {
+      return nullptr;
+    }
+    return type->cast<c10::InterfaceType>();
+  }
+
+  c10::TupleTypePtr get_named_tuple(const c10::QualifiedName& name) const {
+    for (const auto& cls : classes_) {
+      if (cls->name()->qualifiedName() == name.qualifiedName()) {
+        return cls->expect<TupleType>();
+      }
+    }
+    return nullptr;
+  }
+
+  c10::NamedTypePtr get_type(const c10::QualifiedName& name) const {
+    auto it = classDict_.find(name);
+    if (it == classDict_.end()) {
+      return nullptr;
+    }
+    return classes_[it->second];
+  }
+
+  // For testing: clear all Python-defined classes to ensure that unit tests
+  // have isolation.
+  void _clear_python_cu() {
+    // Delete all the associated class methods
+    for (const auto& type : classes_) {
+      if (auto cls = type->cast<ClassType>()) {
+        for (auto method : cls->methods()) {
+          // Tombstone the method in the compilation unit.
+          // Don't erase because the dict_
+          auto it = dict_.find(method->qualname());
+          if (it != dict_.end()) {
+            functions_[it->second] = nullptr;
+            // Erase in our big lookup table
+            dict_.erase(it);
+          }
+        }
+        // Classes can have multiple pointers to the same hook,
+        // need to make sure to not delete it twice
+        std::unordered_set<Function*> hooks_to_delete;
+        for (const auto& hook : cls->getForwardHooks()) {
+          hooks_to_delete.insert(hook);
+        }
+        for (const auto& pre_hook : cls->getForwardPreHooks()) {
+          hooks_to_delete.insert(pre_hook);
+        }
+        for (const auto& hook : hooks_to_delete) {
+          // Tombstone the hook in the compilation unit.
+          auto it = dict_.find(hook->qualname());
+          if (it != dict_.end()) {
+            functions_[it->second] = nullptr;
+            // Erase in our big lookup table
+            dict_.erase(it);
+          }
+        }
+      }
+    }
+    classes_.clear();
+    classDict_.clear();
+  }
+
+  // [Internal Only] Remove method.
+  // Note Used for freezing.
+  void unsafeRemoveMethod(const c10::QualifiedName& method_name) {
+    auto it = dict_.find(method_name);
+    TORCH_CHECK(
+        it != dict_.end(),
+        "method '",
+        method_name.qualifiedName(),
+        "' does not exist.");
+    functions_[it->second] = nullptr;
+    dict_.erase(it);
+  }
+
+  // [name mangling] All code objects must have a unique qualified name in a
+  // CompilationUnit. In Python, sometimes functions won't have unique qualified
+  // name (for example, nested functions). So we mangle Python functions to
+  // ensure that they are uniquely named.
+  //
+  // We also use mangling to distinguish different Module instances. Since each
+  // Module is a singleton class instance, different instances of the same
+  // Python Module will have different types but the same qualified name.
+  c10::QualifiedName mangle(const c10::QualifiedName& name) const {
+    auto mangled = name;
+    while (get_type(mangled) || find_function(mangled)) {
+      mangled = mangler_.mangle(mangled);
+    }
+    return mangled;
+  }
+
+ private:
+  std::unique_ptr<Function> define(
+      const std::optional<c10::QualifiedName>& prefix,
+      const Def& def,
+      const ResolverPtr& resolver,
+      const Self* self,
+      const std::unordered_map<std::string, Function*>& function_table,
+      bool shouldMangle = false,
+      FunctionType type = FunctionType::Method,
+      std::optional<size_t> version = std::nullopt) const;
+
+  // Define a property on \p self.
+  struct PropertyPair;
+  PropertyPair define_property(
+      const std::optional<c10::QualifiedName>& prefix,
+      const Property& prop,
+      const ResolverPtr& resolver,
+      const Self* self,
+      const std::unordered_map<std::string, Function*>& function_table,
+      bool shouldMangle = false) const;
+
+  Function& register_function(std::unique_ptr<Function> fn) {
+    TORCH_CHECK(
+        0 == dict_.count(fn->qualname().qualifiedName()),
+        "method '",
+        fn->qualname().qualifiedName(),
+        "' already defined.");
+    functions_.emplace_back(std::move(fn));
+    dict_[functions_.back()->qualname()] = functions_.size() - 1;
+    return *functions_.back();
+  }
+  std::vector<std::unique_ptr<Function>> functions_;
+  // for fast lookup
+  std::unordered_map<c10::QualifiedName, size_t> dict_;
+  std::unordered_map<c10::QualifiedName, size_t> classDict_;
+
+  // [class ownership] Right now there are two relationships between classes
+  // and compilation units:
+  // 1. Classes have compilation units internally that hold their methods.
+  // 2. On load, the TypePtrs of any imported classes are owned by the main
+  // module's compilation unit.
+  std::vector<c10::NamedTypePtr> classes_;
+
+  mutable NameMangler mangler_;
+};
+
+// An owning pointer to a Function. Just a pair of a raw Function ptr and it's
+// owning CU. We need this because pybind requires a ref-counted way to refer to
+// Functions.
+struct StrongFunctionPtr {
+  StrongFunctionPtr(std::shared_ptr<CompilationUnit> cu, Function* function)
+      : cu_(std::move(cu)), function_(function) {
+    TORCH_INTERNAL_ASSERT(cu_);
+    TORCH_INTERNAL_ASSERT(function_);
+  }
+  std::shared_ptr<CompilationUnit> cu_;
+  Function* function_;
+};
+
+namespace script {
+// We once had a `script::` namespace that was deleted. This is for backcompat
+// of the public API; new code should not use this type alias.
+using CompilationUnit = ::torch::jit::CompilationUnit;
+} // namespace script
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/function_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/function_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..e311563890e13e39c5382b582953433131ab62cd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/function_impl.h
@@ -0,0 +1,185 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/function.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/runtime/graph_executor.h>
+
+namespace torch::jit {
+
+struct TORCH_API GraphFunction : public Function {
+  GraphFunction(
+      c10::QualifiedName name,
+      std::shared_ptr<Graph> graph,
+      std::function<void(GraphFunction&)> function_creator,
+      std::optional<ExecutorExecutionMode> executor_execution_mode =
+          std::nullopt)
+      : name_(std::move(name)),
+        graph_(std::move(graph)),
+        executor_execution_mode_(executor_execution_mode),
+        function_creator_(std::move(function_creator)) {}
+
+  bool isGraphFunction() const override {
+    return true;
+  }
+
+  void run(Stack& stack) override;
+
+  std::function<void(GraphFunction&)> function_creator() const {
+    return function_creator_;
+  }
+
+  c10::intrusive_ptr<c10::ivalue::Future> runAsync(
+      Stack& stack,
+      TaskLauncher taskLauncher = at::launch) override;
+
+  std::shared_ptr<Graph> graph() const {
+    return graph_;
+  }
+
+  std::shared_ptr<Graph> optimized_graph() const;
+
+  const c10::QualifiedName& qualname() const override {
+    return name_;
+  }
+
+  // private/unstable api. sets the initial execution mode
+  // will not affect executor if there is an existing executor
+  // created for this function
+  void _set_initial_executor_execution_mode(ExecutorExecutionMode mode) {
+    executor_execution_mode_ = mode;
+  }
+  // private/unstable api. sets flag of whether or not to ignore amp.
+  // will not affect executor if there is an existing executor
+  // created for this function
+  void _set_ignore_amp(bool ignore_amp) {
+    force_no_amp_ = ignore_amp;
+  }
+
+  // if this isn't yet defined, run its method_creator function
+  void ensure_defined() override;
+
+  size_t num_inputs() const override {
+    return graph()->inputs().size();
+  }
+
+  Function& setSchema(FunctionSchema schema) override {
+    schema_ = std::make_unique<FunctionSchema>(std::move(schema));
+    return *this;
+  }
+
+  const FunctionSchema& getSchema() const override;
+
+  GraphExecutorState getDebugState() {
+    return get_executor().getDebugState();
+  }
+
+  bool is_optimized() const {
+    TORCH_WARN(
+        "GraphFunction::is_optimized() is deprecated and always returns true. "
+        "Please use getGraphExecutorOptimize()");
+    return true;
+  }
+
+  void check_single_output() {
+    TORCH_CHECK(
+        graph()->outputs().size() == 1,
+        "Method (but not graphs in general) require a single output. Use None/Tuple for 0 or 2+ outputs");
+  }
+
+  GraphExecutor& get_executor() {
+    ensure_defined();
+    std::lock_guard<std::recursive_mutex> lock(compile_mutex);
+    auto& executor = executors_[currentSpecialization()];
+    if (executor) {
+      return *executor;
+    }
+    check_single_output();
+    const std::string& name = name_.name();
+    std::shared_ptr<Graph> opt_graph = optimized_graph();
+    if (!executor_execution_mode_) {
+      executor = GraphExecutor(opt_graph, name);
+    } else {
+      executor = GraphExecutor(opt_graph, name, *executor_execution_mode_);
+    }
+    return *executor;
+  }
+
+  using Function::call;
+  bool call(
+      Stack& stack,
+      std::optional<size_t> bailOut,
+      c10::function_ref<void(const Code&)> f) override {
+    f(get_executor().getPlanFor(stack, bailOut).code);
+    return true;
+  }
+
+  void clear_optimized_graphs() {
+    optimized_graphs_.fill(nullptr);
+  }
+
+ private:
+  enum SpecializationKey {
+    AutocastOff,
+    CpuAutocastOn,
+    GpuAutocastOn,
+    CpuGpuAutocastOn,
+
+    // This provides the number of specializations
+    // (Must be last entry)
+    TotalCount
+  };
+
+  SpecializationKey currentSpecialization() const;
+
+ private:
+  c10::QualifiedName name_;
+  // The original, non-optimized graph
+  std::shared_ptr<Graph> graph_; // for debugging and for inlining
+
+  // allows users to specify Simple/Profiling Executor for function
+  // TODO: add more executors
+  mutable std::optional<ExecutorExecutionMode> executor_execution_mode_;
+
+  // if invoked on a graph that has already traced through amp
+  // don't invoke amp pass
+  mutable bool force_no_amp_ = false;
+  // Optimized graph, computed lazily. Used for inlining.
+  mutable std::array<std::shared_ptr<Graph>, SpecializationKey::TotalCount>
+      optimized_graphs_;
+
+  // GraphFunctions are invocable from multiple threads, so this lock needs to
+  // be held when we're initializing graph executor for the first time or
+  // computing the optimized graph. We're using reentrant mutex so that we don't
+  // need to worry about causing a deadlock by calling one method from another
+  // (e.g. optimized_graph() from get_executor()).
+  mutable std::recursive_mutex compile_mutex;
+
+  // executor_[0] - autocast off
+  // executor_[1] - autocast cpu on
+  // executor_[2] - autocast gpu on
+  // executor_[3] - autocast cpu & gpu on
+  std::array<std::optional<GraphExecutor>, SpecializationKey::TotalCount>
+      executors_;
+
+  // an optional function that actually creates the method when
+  // ensure_defined() is called. This is used by the compiler so
+  // that it can construct methods out of order
+  std::function<void(GraphFunction&)> function_creator_;
+
+  // if absent, then we generate a default schema based on the graph
+  // mutable because getSchema caches the default schema if one is requested
+  // before a call to setSchema
+  mutable std::unique_ptr<FunctionSchema> schema_;
+};
+
+// Short hands for dynamic_cast<GraphFunction*>.
+TORCH_API GraphFunction* tryToGraphFunction(Function& /*function*/) noexcept;
+TORCH_API GraphFunction& toGraphFunction(Function& /*function*/);
+TORCH_API const GraphFunction& toGraphFunction(const Function& /*function*/);
+} // namespace torch::jit
+C10_DECLARE_bool(torch_jit_do_not_store_optimized_graph);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/method.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/method.h
new file mode 100644
index 0000000000000000000000000000000000000000..d138f8f847d2d0074f0b1669022fa0f4b3e811f6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/method.h
@@ -0,0 +1,91 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/function.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/api/include/torch/imethod.h>
+#include <torch/csrc/jit/api/function_impl.h>
+
+namespace torch::jit {
+
+using ObjectPtr = c10::intrusive_ptr<c10::ivalue::Object>;
+
+// A method in a module, e.g. f in:
+//
+// class M(ScriptModule):
+//   @script_method
+//   def f(self, x):
+//     ...
+// Note: because Method/Module are exposed to python these
+// classes use python method naming conventions
+struct TORCH_API Method : public torch::IMethod {
+  Method(ObjectPtr owner, Function* function);
+
+  // the module that contains this method.
+  Module owner() const;
+  // the raw objectptr that owns this method, for when the method is owned by a
+  // torchbind object.
+  ObjectPtr raw_owner() const;
+  void run(Stack& stack);
+  void run(Stack&& stack) {
+    run(stack);
+  }
+
+  c10::IValue operator()(
+      std::vector<c10::IValue> stack,
+      const Kwargs& kwargs = Kwargs()) const override;
+
+  // Run method async. Invocation on this function would invokes a JIT
+  // interpreter that executes ops inline, one by one, on caller's thread. A
+  // model can utilize async op, i.e. `fork`, to launch an asynchronous task
+  // which will be launched on provided `taskLauncher`.
+  c10::intrusive_ptr<c10::ivalue::Future> run_async(
+      std::vector<c10::IValue> stack,
+      const Kwargs& kwargs = Kwargs(),
+      TaskLauncher taskLauncher = at::launch);
+
+  std::shared_ptr<Graph> graph() const {
+    return toGraphFunction(*function_).graph();
+  }
+
+  const std::string& name() const override {
+    return function_->name();
+  }
+
+  size_t num_inputs() const {
+    return function_->num_inputs();
+  }
+
+  GraphExecutor& get_executor() {
+    return toGraphFunction(*function_).get_executor();
+  }
+
+  Function& function() const {
+    return *function_;
+  }
+
+ private:
+  void setArgumentNames(
+      std::vector<std::string>& /*argumentNames*/ /*argumentNamesOut*/)
+      const override;
+
+  // Methods are uniqued owned by a single module. This raw pointer allows
+  // looking up the module.
+  ObjectPtr owner_;
+
+  // Underlying unbound function
+  Function* function_;
+};
+
+namespace script {
+// We once had a `script::` namespace that was deleted. This is for backcompat
+// of the public API; new code should not use this type alias.
+using Method = ::torch::jit::Method;
+} // namespace script
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/module.h
new file mode 100644
index 0000000000000000000000000000000000000000..385c1ec489fc9d43caee073b604b4d5c777c286e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/module.h
@@ -0,0 +1,690 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Exception.h>
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/jit/api/object.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/named_value.h>
+#include <torch/csrc/jit/runtime/argument_spec.h>
+#include <torch/csrc/jit/runtime/graph_executor.h>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/api/include/torch/ordered_dict.h>
+#include <torch/csrc/jit/api/compilation_unit.h>
+
+#include <ATen/core/function_schema.h>
+#include <ATen/core/qualified_name.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/irange.h>
+#include <optional>
+
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <ostream>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+// This file contains classes which assist in desugaring Python style
+// modules and their methods into flattened graphs which don't have any
+// function calls.
+
+namespace torch::jit {
+
+using ::c10::Argument;
+using ::c10::FunctionSchema;
+using ::c10::QualifiedName;
+// Map which stores filename to content.
+using ExtraFilesMap = std::unordered_map<std::string, std::string>;
+
+using ModulePtr = c10::intrusive_ptr<c10::ivalue::Object>;
+
+struct Module;
+
+template <typename T>
+struct slot_list_impl;
+
+template <typename T>
+struct Named {
+  std::string name;
+  T value;
+};
+
+using NameModule = Named<Module>;
+using NameValue = Named<IValue>;
+using NameTensor = Named<at::Tensor>;
+
+namespace detail {
+struct TORCH_API ModulePolicy;
+struct TORCH_API ParameterPolicy;
+struct TORCH_API AttributePolicy;
+struct TORCH_API BufferPolicy;
+template <typename P>
+struct NamedPolicy;
+} // namespace detail
+
+using module_list = slot_list_impl<detail::ModulePolicy>;
+using named_module_list =
+    slot_list_impl<detail::NamedPolicy<detail::ModulePolicy>>;
+
+using parameter_list = slot_list_impl<detail::ParameterPolicy>;
+using named_parameter_list =
+    slot_list_impl<detail::NamedPolicy<detail::ParameterPolicy>>;
+
+using attribute_list = slot_list_impl<detail::AttributePolicy>;
+using named_attribute_list =
+    slot_list_impl<detail::NamedPolicy<detail::AttributePolicy>>;
+
+using buffer_list = slot_list_impl<detail::BufferPolicy>;
+using named_buffer_list =
+    slot_list_impl<detail::NamedPolicy<detail::BufferPolicy>>;
+
+using ModuleLookup = std::function<Module(const std::vector<std::string>&)>;
+
+struct TORCH_API Module : public Object {
+  explicit Module(c10::QualifiedName class_name);
+  Module(std::shared_ptr<CompilationUnit> cu, const c10::ClassTypePtr& type);
+  Module() = default;
+  Module(const Module&) = default;
+  Module& operator=(const Module&) = default;
+  Module(Module&&) noexcept = default;
+  Module& operator=(Module&&) noexcept = default;
+  Module(
+      c10::QualifiedName /*class_name*/,
+      std::shared_ptr<CompilationUnit> cu,
+      bool shouldMangle = false);
+  Module(ModulePtr module_value) : Object(std::move(module_value)) {}
+  ~Module() = default;
+
+  void set_optimized(bool o) {
+    TORCH_WARN(
+        "Module::set_optimized() is deprecated and has no effect. "
+        "Please use setGraphExecutorOptimize()");
+  }
+
+  bool is_optimized() const {
+    TORCH_WARN(
+        "Module::is_optimized() is deprecated and always returns true. "
+        "Please use getGraphExecutorOptimize()");
+    return true;
+  }
+
+  IValue forward(std::vector<IValue> inputs, const Kwargs& kwargs = Kwargs()) {
+    return get_method("forward")(std::move(inputs), kwargs);
+  }
+
+  // In script modules, buffers are Tensors attribute that are _not_ registered
+  // as parameters. This is different than in nn.Module where there is a special
+  // register_buffer method. With this simplification, we only need to track
+  // whether a slot is a parameter to be able to classify it.
+  void register_buffer(const std::string& name, at::Tensor v) {
+    bool is_param = false;
+    bool is_buffer = true;
+    std::lock_guard<std::mutex> lock(*register_mutex_);
+    type()->addOrCheckAttribute(name, TensorType::get(), is_param, is_buffer);
+    _ivalue()->setAttr(name, std::move(v));
+  }
+
+  void register_parameter(
+      const std::string& name,
+      at::Tensor v,
+      bool is_buffer) {
+    std::lock_guard<std::mutex> lock(*register_mutex_);
+    type()->addOrCheckAttribute(name, TensorType::get(), !is_buffer, is_buffer);
+    _ivalue()->setAttr(name, std::move(v));
+  }
+
+  void register_attribute(
+      const std::string& name,
+      const TypePtr& t,
+      IValue v,
+      bool is_param = false,
+      bool is_buffer = false) {
+    type()->addOrCheckAttribute(name, t, is_param, is_buffer);
+    _ivalue()->setAttr(name, std::move(v));
+  }
+
+  void register_module(const std::string& name, const Module& module) {
+    type()->addOrCheckAttribute(name, module.type());
+    _ivalue()->setAttr(name, module._ivalue());
+  }
+
+  void apply(const std::function<void(Module&)>& fn);
+
+  buffer_list buffers(bool recurse = true) const;
+  named_buffer_list named_buffers(bool recurse = true) const;
+
+  module_list children() const; // direct modules
+  named_module_list named_children() const;
+  module_list modules() const; // all modules, including this one, recursively
+  named_module_list named_modules() const;
+
+  // all tensors involved in gradient optimization
+  parameter_list parameters(bool recurse = true) const;
+  named_parameter_list named_parameters(bool recurse = true) const;
+
+  // all members of the object, similar to iterating over dir(obj) in python
+  attribute_list attributes(bool recurse = true) const;
+  named_attribute_list named_attributes(bool recurse = true) const;
+
+  void dump(
+      bool print_method_bodies,
+      bool print_attr_values,
+      bool print_param_values) const;
+
+  std::string dump_to_str(
+      bool print_method_bodies,
+      bool print_attr_values,
+      bool print_param_values) const;
+
+  /// Enables "training" mode.
+  void train(bool on = true);
+  /// Calls train(false) to enable "eval" mode.
+  /// Do not override this method, override `train()` instead.
+  void eval() {
+    train(/*on=*/false);
+  }
+  /// True if the module is in training mode.
+  bool is_training() const {
+    return attr("training", true).toBool();
+  }
+
+  /// Recursively casts all parameters to the given `dtype` and `device`.
+  ///
+  /// If `non_blocking` is true and the source is in pinned memory and
+  /// destination is on the GPU or vice versa, the copy is performed
+  /// asynchronously with respect to the host. Otherwise, the argument has no
+  /// effect.
+  void to(at::Device device, at::ScalarType dtype, bool non_blocking = false);
+
+  /// Recursively casts all parameters to the given dtype.
+  ///
+  /// If `non_blocking` is true and the source is in pinned memory and
+  /// destination is on the GPU or vice versa, the copy is performed
+  /// asynchronously with respect to the host. Otherwise, the argument has no
+  /// effect.
+  void to(at::ScalarType dtype, bool non_blocking = false);
+
+  /// Recursively moves all parameters to the given device.
+  ///
+  /// If `non_blocking` is true and the source is in pinned memory and
+  /// destination is on the GPU or vice versa, the copy is performed
+  /// asynchronously with respect to the host. Otherwise, the argument has no
+  /// effect.
+  void to(at::Device device, bool non_blocking = false);
+
+  void save(
+      std::ostream& out,
+      const ExtraFilesMap& extra_files = ExtraFilesMap()) const;
+
+  void save(
+      const std::string& filename,
+      const ExtraFilesMap& extra_files = ExtraFilesMap()) const;
+
+  void _save_for_mobile(
+      std::ostream& out,
+      const ExtraFilesMap& extra_files = ExtraFilesMap(),
+      bool save_mobile_debug_info = false,
+      bool use_flatbuffer = false) const;
+
+  void _save_for_mobile(
+      const std::string& filename,
+      const ExtraFilesMap& extra_files = ExtraFilesMap(),
+      bool save_mobile_debug_info = false,
+      bool use_flatbuffer = false) const;
+
+  Module copy() const;
+
+  Module deepcopy(std::optional<at::Device> device = std::nullopt) const;
+
+  // Clones both the underlying `ClassType` and the module instance(data), this
+  // function creates a new `ClassType` and returns a new instance that has the
+  // same data as the current instance but with the new type, shared ClassType
+  // will be preserved as well
+  Module clone(bool inplace = false) const;
+
+  // Clones both the underlying `ClassType` and the module instance(data), this
+  // function creates a new `ClassType` and returns a new instance that has the
+  // same data as the current instance but with the new type, shared ClassType
+  // will be preserved as well. Also allows the caller to specify a set of
+  // method and attribute names to not clone.
+  Module clone(
+      bool inplace,
+      const std::unordered_set<std::string>& ignored_method,
+      const std::unordered_set<std::string>& ignored_attributes) const;
+
+  void clone_method(const Module& orig, const std::string& name);
+
+  IValue operator()(std::vector<IValue> inputs);
+
+  template <typename... Types>
+  IValue create_class(const c10::QualifiedName& name, Types&&... args) const {
+    return create_class(name, {IValue(std::forward<Types>(args))...});
+  }
+
+  IValue create_class(const c10::QualifiedName& name, Stack stack) const;
+
+  inline bool operator==(const Module& y) const noexcept {
+    return _ivalue() == y._ivalue();
+  }
+
+  void set_delete_memory(std::shared_ptr<char> delete_mem) {
+    mem_to_delete_ = std::move(delete_mem);
+  }
+
+  // A set of functions to maintain input shapes through torch.jit.save and
+  // torch.jit.load. It only works on tensors and lists/dicts of tensors
+  // because tracing is only supported by these types.
+  void store_traced_inputs(
+      const std::string& func_name,
+      std::vector<IValue> inputs) {
+    if (inputs.empty()) {
+      return;
+    }
+    auto c10_inputs = c10::impl::GenericList(AnyType::get());
+    for (IValue& value : inputs) {
+      // Not checking whether this is traceable type as that is already checked
+      // higher up in the stack and changing that would require a larger
+      // restructuring.
+      c10_inputs.emplace_back(std::move(value));
+    }
+    traced_inputs_.insert_or_assign(func_name, c10_inputs);
+  }
+
+  c10::Dict<std::string, c10::impl::GenericList> retrieve_traced_inputs()
+      const {
+    return traced_inputs_;
+  }
+
+ private:
+  Module clone_impl(
+      std::unordered_map<TypePtr, TypePtr>& type_remap,
+      bool inplace,
+      IValue::HashIdentityIValueMap memo,
+      const std::unordered_set<std::string>& ignored_methods,
+      const std::unordered_set<std::string>& ignored_attributes) const;
+
+  void clone_method(
+      const Module& orig,
+      const Function& method,
+      const std::unordered_map<TypePtr, TypePtr>& type_remap);
+
+  c10::QualifiedName getNameForMethod(std::string basename) const {
+    return QualifiedName(*type()->name(), std::move(basename));
+  }
+
+  void to_impl(
+      const std::optional<at::Device>& device,
+      const std::optional<at::ScalarType>& dtype,
+      bool non_blocking);
+
+  // Extra handle for the module to delete when itself is deleted
+  std::shared_ptr<char> mem_to_delete_;
+
+  // Map of function names to the traced inputs that they have been traced with
+  c10::Dict<std::string, c10::impl::GenericList> traced_inputs_;
+
+  // Mutex to keep registering buffer or parameter thread safe.
+  std::shared_ptr<std::mutex> register_mutex_ = std::make_shared<std::mutex>();
+};
+
+// C++ equivalent api of `torch.jit.freeze`. See documentation there for
+// details.
+TORCH_API Module freeze(
+    const Module& module,
+    const std::optional<std::vector<std::string>>& preserved_attrs =
+        std::nullopt,
+    bool optimize_numerics = true);
+
+// C++ equivalent api of `torch.jit.optimize_for_inference`. See documentation
+// there for details.
+TORCH_API Module optimize_for_inference(
+    Module& module,
+    const std::vector<std::string>& other_methods = {});
+
+enum class FusionBehavior { STATIC, DYNAMIC };
+
+using FusionStrategy = std::vector<std::pair<FusionBehavior, size_t>>;
+// clang-format off
+/*
+Sets the type and number of specializations that can occur during fusion.
+
+Usage: provide a list of pairs (type, depth) where type is one of STATIC or DYNAMIC
+and depth is an integer.
+
+Behavior - static vs dynamic:
+    In STATIC fusion, fused ops are compiled to have fixed input shapes. The shape is determined
+    based on some initial profiling runs.
+    In DYNAMIC fusion, fused ops are compiled to have variable input shapes, so that multiple
+    shapes are possible.
+
+In both cases, we also recompile on new striding behavior, device, or dtype.
+
+Behavior - fallback functions & depth:
+    When an input doesn't match the format required by the specialized compiled op, it will run
+    a fallback function. Fallback functions are recursively be compiled and specialized based
+    on the observed tensor shapes. Since compilation can be slow, the "depth" parameter is provided to
+    limit the number of specializations that can be compiled, before giving up on recompiling and
+    falling back to a completely un-fused, un-specialized implementation.
+
+The list of (type, depth) pairs controls the type of specializations and the number of
+specializations. For example: [(STATIC, 2), (DYNAMIC, 2)] indicates that the first
+two specializations will use static fusions, the following two specializations will use
+dynamic fusion, and any inputs that satisfy none of the 4 options will run an
+unfused implementation.
+
+NB: in the future, if more as more fusion backends are added there may be more granular
+apis for specific fusers.
+*/
+// clang-format on
+TORCH_API FusionStrategy getFusionStrategy();
+// returns previous strategy
+TORCH_API FusionStrategy setFusionStrategy(FusionStrategy& fusion_strategy);
+
+namespace detail {
+
+struct TORCH_API SlotCursor {
+  Module module_;
+  int64_t i_; // slot offset, -1 indicates the module itself
+};
+
+} // namespace detail
+
+// This iterator allows the (optionally recursive) enumeration of
+// the  members of a Module. It performs a depth-first pre-order
+// traversal of the module. The Policy template parameter determines
+// which slots of the object should be included. For instance,
+// when iterating parameters, we return the parameter tensors,
+// but skip modules, buffers, and other attributes.
+// See ModulePolicy for comments about Policy object's API.
+template <typename Policy>
+struct slot_iterator_impl {
+  using SlotCursor = detail::SlotCursor;
+  using value_type = typename Policy::value_type;
+  slot_iterator_impl(
+      Module root,
+      bool recurse, // if true, do a depth-first search, otherwise, just look at
+                    // slots of root
+      bool return_module) // if true include root itself as the first thing
+                          // visited (used in modules())
+      : cursors_({SlotCursor{std::move(root), return_module ? -1 : 0}}),
+        recurse_(recurse) {
+    // advance iterator to first valid element (or the end, if empty)
+    while_not_valid_next();
+  }
+  // empty cursors_, represents end of iteration
+  slot_iterator_impl() : recurse_(false) {}
+  value_type operator*() const {
+    return Policy::create(cursors_, cur());
+  }
+  value_type operator->() const {
+    return **this;
+  }
+  slot_iterator_impl& operator++() {
+    next_valid();
+    return *this;
+  }
+  slot_iterator_impl operator++(int) {
+    // this is really expensive, should we delete it so people don't use it
+    // instead of prefix?
+    slot_iterator_impl old = *this;
+    ++(*this);
+    return old;
+  }
+
+ private:
+  // return_module() is a corner case where instead of returning a submodule
+  // of root, we are returning root itself, because we are iterating modules(),
+  // which contains the root module itself.
+  // It is represented with a single SlotCursor whose index is -1.
+  bool return_module() const {
+    return top().i_ == -1;
+  }
+  const SlotCursor& top() const {
+    return cursors_.back();
+  }
+  SlotCursor& top() {
+    return cursors_.back();
+  }
+  IValue cur() const {
+    return return_module() ? top().module_._ivalue()
+                           : top().module_._ivalue()->getSlot(top().i_);
+  }
+
+  // advance to the next slot in a depth first pre-order traversal of the
+  // modules slots. This function does not guarantee the next slot is a
+  // valid element of the iteration. That is done by valid().
+  // invariant: !cursors_.empty()
+  void next() {
+    // we just returned the module itself, advance i_ to 0 so we are now
+    // at the first slot of the module.
+    if (return_module()) {
+      ++top().i_;
+      return;
+    }
+    // the last traversal action advanced beyond the number of slots in the
+    // module so continue the iteration in the parent.
+    if (top().i_ >= int64_t(top().module_._ivalue()->type()->numAttributes())) {
+      cursors_.pop_back();
+      if (!cursors_.empty()) {
+        ++top().i_;
+      }
+      return;
+    }
+    // if the current thing is a module, we have to scan it for recursive
+    // traversals. We do this by adding a new SlotCursor to track the traversal.
+    if (recurse_ &&
+        top().module_._ivalue()->type()->getAttribute(top().i_)->is_module()) {
+      cursors_.emplace_back(SlotCursor{cur().toModule(), 0});
+      return;
+    }
+    // common case: advance to the next slot.
+    ++top().i_;
+  }
+  // is the current position of the iterator a valid one?
+  // otherwise, we have to continue advancing.
+  bool valid() const {
+    return top().i_ <
+        int64_t(top().module_._ivalue()->type()->numAttributes()) &&
+        Policy::valid(
+               top().module_._ivalue()->type(),
+               top().i_,
+               top().module_._ivalue()->getSlot(top().i_));
+  }
+  void while_not_valid_next() {
+    // advance iteration until we are either at the end (cursors_.empty())
+    // or in a valid state. return_module() is a special case,
+    // and is always considered valid, regardless of Policy, because it is
+    // it is only true when we are iterating modules.
+    while (!cursors_.empty() && !return_module() && !valid()) {
+      next();
+    }
+  }
+  void next_valid() {
+    // avoid crashing if this is empty
+    if (cursors_.empty()) {
+      return;
+    }
+    // advance to next element, which is maybe not valid
+    next();
+    while_not_valid_next();
+  }
+
+  std::vector<SlotCursor> cursors_;
+  bool recurse_;
+
+  friend inline bool operator!=(
+      const slot_iterator_impl<Policy>& a,
+      const slot_iterator_impl<Policy>& b) {
+    // we are finished iteration when we have no more iteration SlotCursors.
+    // end is always an empty iterator with no cursors.
+    return (a.cursors_.empty() != b.cursors_.empty());
+  }
+};
+
+// This type represents lists of parameters, attributes, and
+// submodules contained in the module. It is abstract because
+// they are not stored directly in std::vectors but inside the
+// module's IValue object itself.
+template <typename Policy>
+struct slot_list_impl {
+  using iterator = slot_iterator_impl<Policy>;
+  using const_iterator = slot_iterator_impl<Policy>;
+  using value_type = typename iterator::value_type;
+  slot_iterator_impl<Policy> begin() const {
+    return slot_iterator_impl<Policy>(module_, recurse_, return_module_);
+  }
+  slot_iterator_impl<Policy> end() const {
+    return slot_iterator_impl<Policy>();
+  }
+  size_t size() const {
+    if (!size_) {
+      size_ = size_t(0);
+      for ([[maybe_unused]] const value_type& _ : *(this)) {
+        ++*size_;
+      }
+    }
+    return *size_;
+  }
+
+  slot_list_impl(Module module, bool recurse, bool return_module)
+      : module_(std::move(module)),
+        recurse_(recurse),
+        return_module_(return_module),
+        size_(std::nullopt) {
+    if (!recurse && !return_module && Policy::all_slots) {
+      size_ = module_.num_slots();
+    }
+  }
+
+ private:
+  Module module_;
+  bool recurse_;
+  bool return_module_;
+  // size of this list, cached on first request
+  // when we need to filter the slot list
+  mutable std::optional<size_t> size_;
+  friend struct Module;
+};
+
+namespace detail {
+
+// slot_iterator_impl always iterate over all the slots in a module,
+// the Policy template argument determines slots should be returned and their
+// types
+struct TORCH_API ModulePolicy {
+  // the type of the value being returned
+  using value_type = Module;
+
+  // the logic for creating the type being returned, given the raw IValue
+  // of that object.
+  static value_type create(
+      const std::vector<detail::SlotCursor>& cursors,
+      IValue v) {
+    return Module(std::move(v).toObject());
+  }
+  // is slot i in typ something that this iterator should return, otherwise,
+  // we skip it.
+  static bool valid(const ClassTypePtr& typ, size_t i, const IValue& v) {
+    return typ->getAttribute(i)->is_module();
+  }
+  // are we going to return everything? If so, we can optimize the calculate
+  // of the size of the list.
+  static constexpr bool all_slots = false;
+};
+
+struct TORCH_API ParameterPolicy {
+  using value_type = at::Tensor;
+  static value_type create(
+      const std::vector<detail::SlotCursor>& cursors,
+      IValue v) {
+    return std::move(v).toTensor();
+  }
+  static bool valid(const ClassTypePtr& typ, size_t i, const IValue& v) {
+    return typ->is_parameter(i) && v.isTensor();
+  }
+  static constexpr bool all_slots = false;
+};
+
+struct TORCH_API BufferPolicy {
+  using value_type = at::Tensor;
+  static value_type create(
+      const std::vector<detail::SlotCursor>& cursors,
+      IValue v) {
+    return std::move(v).toTensor();
+  }
+  static bool valid(const ClassTypePtr& typ, size_t i, const IValue& v) {
+    return typ->getAttribute(i)->isSubtypeOf(*TensorType::get()) &&
+        typ->is_buffer(i);
+  }
+  static constexpr bool all_slots = false;
+};
+
+struct TORCH_API AttributePolicy {
+  using value_type = IValue;
+  static value_type create(
+      const std::vector<detail::SlotCursor>& cursors,
+      IValue v) {
+    return v;
+  }
+  static bool valid(const ClassTypePtr& typ, size_t i, const IValue& v) {
+    return true;
+  }
+  static constexpr bool all_slots = true;
+};
+
+// take a Policy object, and make a version of it that returns the slot.
+// along with the fully qualified name of that slot. This is used for the named_
+// variants like named_parameters().
+template <typename Policy>
+struct NamedPolicy {
+  using value_type = Named<typename Policy::value_type>;
+  static value_type create(
+      const std::vector<detail::SlotCursor>& cursors,
+      IValue v) {
+    std::string name;
+    if (cursors.size() == 1) {
+      name = (cursors.back().i_ == -1) ? "" : nameFragment(cursors.back());
+    } else {
+      std::ostringstream ss;
+      for (const auto i : c10::irange(cursors.size())) {
+        if (i > 0) {
+          ss << '.';
+        }
+        ss << nameFragment(cursors[i]);
+      }
+      name = ss.str();
+    }
+    return value_type{std::move(name), Policy::create(cursors, std::move(v))};
+  }
+  static bool valid(const ClassTypePtr& t, size_t i, const IValue& v) {
+    return Policy::valid(t, i, v);
+  }
+  static constexpr bool all_slots = Policy::all_slots;
+
+ private:
+  static std::string nameFragment(const detail::SlotCursor& f) {
+    return f.module_.type()->getAttributeName(f.i_);
+  }
+};
+
+} // namespace detail
+
+TORCH_API bool& getInlineEverythingMode();
+
+namespace script {
+// We once had a `script::` namespace that was deleted. This is for backcompat
+// of the public API; new code should not use this type alias.
+using Module = ::torch::jit::Module;
+using ExtraFilesMap = ::torch::jit::ExtraFilesMap;
+} // namespace script
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/object.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/object.h
new file mode 100644
index 0000000000000000000000000000000000000000..f25e599974b138172c593fe2c3f9f9fac2e26397
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/api/object.h
@@ -0,0 +1,205 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/functional.h>
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/api/method.h>
+#include <optional>
+
+#include <utility>
+
+namespace torch::jit {
+
+struct Resolver;
+using ResolverPtr = std::shared_ptr<Resolver>;
+
+using ObjectPtr = c10::intrusive_ptr<c10::ivalue::Object>;
+
+// Throw this in C++ land if `attr` fails. This will be converted to a Python
+// AttributeError by the Python binding code
+class ObjectAttributeError : public std::runtime_error {
+ public:
+  ObjectAttributeError(const std::string& what) : std::runtime_error(what) {}
+};
+
+struct TORCH_API Object {
+  Object() = default;
+  Object(const Object&) = default;
+  Object& operator=(const Object&) = default;
+  Object(Object&&) noexcept = default;
+  Object& operator=(Object&&) noexcept = default;
+  Object(ObjectPtr _ivalue) : _ivalue_(std::move(_ivalue)) {}
+  Object(std::shared_ptr<CompilationUnit> cu, const c10::ClassTypePtr& type);
+  Object(
+      c10::QualifiedName,
+      std::shared_ptr<CompilationUnit> cu,
+      bool shouldMangle = false);
+
+  ObjectPtr _ivalue() const {
+    TORCH_INTERNAL_ASSERT(_ivalue_);
+    return _ivalue_;
+  }
+
+  c10::ClassTypePtr type() const {
+    return _ivalue()->type();
+  }
+
+  struct Property {
+    std::string name;
+    Method getter_func;
+    std::optional<Method> setter_func;
+  };
+
+  void setattr(const std::string& name, c10::IValue v) {
+    if (_ivalue()->type()->hasConstant(name)) {
+      TORCH_CHECK(
+          false,
+          "Can't set constant '",
+          name,
+          "' which has value:",
+          _ivalue()->type()->getConstant(name));
+    } else if (auto slot = _ivalue()->type()->findAttributeSlot(name)) {
+      const c10::TypePtr& expected = _ivalue()->type()->getAttribute(*slot);
+      TORCH_CHECK(
+          v.type()->isSubtypeOf(*expected),
+          "Expected a value of type '",
+          expected->repr_str(),
+          "' for field '",
+          name,
+          "', but found '",
+          v.type()->repr_str(),
+          "'");
+      _ivalue()->setSlot(*slot, std::move(v));
+    } else {
+      TORCH_CHECK(false, "Module has no attribute '", name, "'");
+    }
+  }
+
+  c10::IValue attr(const std::string& name) const {
+    if (auto r = _ivalue()->type()->findAttributeSlot(name)) {
+      return _ivalue()->getSlot(*r);
+    }
+    if (auto r = _ivalue()->type()->findConstantSlot(name)) {
+      return _ivalue()->type()->getConstant(*r);
+    }
+    std::stringstream err;
+    err << _ivalue()->type()->repr_str() << " does not have a field with name '"
+        << name.c_str() << "'";
+    throw ObjectAttributeError(err.str());
+  }
+
+  c10::IValue attr(const std::string& name, c10::IValue or_else) const {
+    if (auto r = _ivalue()->type()->findAttributeSlot(name)) {
+      return _ivalue()->getSlot(*r);
+    }
+    if (auto r = _ivalue()->type()->findConstantSlot(name)) {
+      return _ivalue()->type()->getConstant(*r);
+    }
+    return or_else;
+  }
+
+  bool hasattr(const std::string& name) const {
+    return _ivalue()->type()->hasAttribute(name) ||
+        _ivalue()->type()->hasConstant(name);
+  }
+
+  // each object owns its methods. The reference returned here
+  // is guaranteed to stay valid until this module has been destroyed
+  Method get_method(const std::string& name) const {
+    if (auto method = find_method(name)) {
+      return *method;
+    }
+    TORCH_CHECK(false, "Method '", name, "' is not defined.");
+  }
+
+  const std::vector<Method> get_methods() const {
+    return c10::fmap(type()->methods(), [&](Function* func) {
+      return Method(_ivalue(), func);
+    });
+  }
+
+  bool has_property(const std::string& name) const {
+    for (const auto& prop : type()->properties()) {
+      if (prop.name == name) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  const Property get_property(const std::string& name) const {
+    for (const auto& prop : type()->properties()) {
+      if (prop.name == name) {
+        std::optional<Method> setter = std::nullopt;
+        if (prop.setter) {
+          setter = Method(_ivalue(), prop.setter);
+        }
+        return Property{
+            prop.name, Method(_ivalue(), prop.getter), std::move(setter)};
+      }
+    }
+    TORCH_CHECK(false, "Property '", name, "' is not defined.");
+  }
+
+  const std::vector<Property> get_properties() const {
+    return c10::fmap(type()->properties(), [&](ClassType::Property prop) {
+      std::optional<Method> setter = std::nullopt;
+      if (prop.setter) {
+        setter = Method(_ivalue(), prop.setter);
+      }
+      return Property{
+          std::move(prop.name),
+          Method(_ivalue(), prop.getter),
+          std::move(setter)};
+    });
+  }
+
+  std::optional<Method> find_method(const std::string& basename) const;
+
+  /// Run a method from this module.
+  ///
+  /// For example:
+  /// @code
+  ///   IValue output = module->run("relu_script", a, b);
+  /// @endcode
+  ///
+  /// To get a compile a module from a source string, see torch::jit::compile
+  ///
+  /// @param method_name The name of the method to run
+  /// @param args Arguments to be passed to the method
+  /// @return An IValue containing the return value (or values if it is a tuple)
+  /// from the method
+  template <typename... Types>
+  IValue run_method(const std::string& method_name, Types&&... args) {
+    return get_method(method_name)({IValue(std::forward<Types>(args))...});
+  }
+
+  // so that C++ users can easily add methods
+  void define(const std::string& src, const ResolverPtr& resolver = nullptr);
+
+  size_t num_slots() const {
+    return _ivalue()->slots().size();
+  }
+
+  // shallow copy the object
+  Object copy() const;
+
+  // Copies all the attributes of the object recursively without creating new
+  // `ClassType`, including deepcopy of Tensors
+  Object deepcopy() const;
+
+ private:
+  // mutable be we lazily initialize in module_object.
+  mutable ObjectPtr _ivalue_;
+};
+
+namespace script {
+// We once had a `script::` namespace that was deleted. This is for backcompat
+// of the public API; new code should not use this type alias.
+using Object = ::torch::jit::Object;
+} // namespace script
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend.h
new file mode 100644
index 0000000000000000000000000000000000000000..cea04920023b6876ac4c8123c4b887b88d457fbd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend.h
@@ -0,0 +1,119 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/builtin_function.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/jit/backends/backend_interface.h>
+#include <torch/custom_class.h>
+
+namespace torch::jit {
+namespace {
+inline c10::FunctionSchema getIsAvailableSchema() {
+  c10::Argument self("self", c10::AnyType::get());
+  c10::Argument available("available", c10::BoolType::get());
+  c10::FunctionSchema preprocessor_schema(
+      "is_available",
+      /*overload_name=*/"",
+      /*arguments=*/{self},
+      /*returns=*/{available});
+  return preprocessor_schema;
+}
+
+constexpr static auto kBackendsNamespace = "__backends__";
+
+inline c10::FunctionSchema getCompileSchema() {
+  c10::Argument self("self", c10::AnyType::get());
+  c10::Argument mod("processed", c10::AnyType::get());
+  auto any_dict_ty =
+      c10::DictType::create(c10::StringType::get(), c10::AnyType::get());
+  c10::Argument method_compile_spec("method_compile_spec", any_dict_ty);
+  c10::Argument handles("handles", any_dict_ty);
+
+  c10::FunctionSchema compile_schema(
+      "compile",
+      /*overload_name=*/"",
+      /*arguments=*/{self, mod, method_compile_spec},
+      /*returns=*/{handles});
+  return compile_schema;
+}
+
+inline c10::FunctionSchema getExecuteSchema() {
+  auto any_list_ty = c10::ListType::create(c10::AnyType::get());
+  c10::Argument self("self", c10::AnyType::get());
+  c10::Argument handle("handle", c10::AnyType::get());
+  c10::Argument input("input", any_list_ty);
+  c10::Argument output("output", any_list_ty);
+  return c10::FunctionSchema(
+      "execute",
+      /*overload_name=*/"",
+      /*arguments=*/{self, handle, input},
+      /*returns=*/{output});
+}
+
+template <typename TBackendInterface>
+std::function<void(Stack&)> getIsAvailableFunc() {
+  return [](Stack& stack) {
+    auto self = pop(stack).toCustomClass<TBackendInterface>();
+    auto ret = self->is_available();
+    push(stack, ret);
+  };
+}
+
+template <typename TBackendInterface>
+std::function<void(Stack&)> getCompileFunc() {
+  return [](Stack& stack) {
+    auto method_compile_spec = pop(stack).toGenericDict();
+    auto processed = pop(stack);
+    auto self = pop(stack).toCustomClass<TBackendInterface>();
+    auto ret = self->compile(processed, method_compile_spec);
+    push(stack, ret);
+  };
+}
+
+template <typename TBackendInterface>
+std::function<void(Stack&)> getExecuteFunc() {
+  return [](Stack& stack) {
+    auto args = pop(stack);
+    auto handle = pop(stack);
+    auto self = pop(stack);
+    auto backend = self.toCustomClass<TBackendInterface>();
+    auto res = backend->execute(handle, args.toList());
+    push(stack, res);
+  };
+}
+} // namespace
+
+// Static registration API for backends.
+template <class TBackendInterface>
+class backend {
+  static_assert(
+      std::is_base_of_v<PyTorchBackendInterface, TBackendInterface>,
+      "torch::jit::backend<T> requires T to inherit from PyTorchBackendInterface");
+  std::string backend_name_;
+
+ public:
+  // Registers a new backend with /p name, and the given /p preprocess
+  // function.
+  backend(const std::string& name) : backend_name_(name) {
+    static auto cls = torch::class_<TBackendInterface>(kBackendsNamespace, name)
+                          .def(torch::init<>())
+                          ._def_unboxed(
+                              "is_available",
+                              getIsAvailableFunc<TBackendInterface>(),
+                              getIsAvailableSchema())
+                          ._def_unboxed(
+                              "compile",
+                              getCompileFunc<TBackendInterface>(),
+                              getCompileSchema())
+                          ._def_unboxed(
+                              "execute",
+                              getExecuteFunc<TBackendInterface>(),
+                              getExecuteSchema());
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_debug_handler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_debug_handler.h
new file mode 100644
index 0000000000000000000000000000000000000000..ec124d0cf8ae0cbee9a38a575c49c22e2712164d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_debug_handler.h
@@ -0,0 +1,143 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/ivalue.h>
+
+#include <torch/csrc/jit/backends/backend_detail.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/scope.h>
+
+#include <atomic>
+
+namespace torch::jit {
+
+/*
+ *  BackendDebugHandleManager is responsible for issuing debug handles to
+ *  backends. Debug handles are associated with nodes of a graph.
+ *  BackendDebugHandleManager also maintains a map
+ *  [debug-handle, DebugInfoTuple = {source range, inlined callstack ptr]} that
+ *  will help generate a callstack for exception raised using debug handles.
+ *  Effectively debug handles are something that is given to backend and later
+ *  when an exception occurs in the backend, backend can tell, using debug
+ *  handle, that an exception occurred here. Then the runtime can generate
+ *  callstack corresponding to the exception.
+ *  There are two parts to BackendDebugHandleManager:
+ *  1. static std::atomic debug_handle
+ *  2. Map of [debug-handle, DebugInfoTuple]
+ *
+ *  About 1:
+ *  Why do they have to be unique. The reason is that by ensuring
+ *  uniqueness of debug handles, we remove the burden of another layer of
+ *  mapping where we need to say this set of debug handles were generated for
+ *  this lowered module or this bytecode function. This simplifies the API for
+ *  serialization since debug handles can uniquely identify DebugInfoTuple.
+ *  Thus simplifies the runtime API for throwing exception. Exception throwing
+ *  only needs to know debug_handle and not which module or method threw it.
+ *  There are 2 issues to keep in mind, though,for static std::atomic
+ *  debug_handle: A. Performance implications of using atomic variable. However
+ *  this is only used for compilation so we assume to absorb some of that
+ *  penalty. Plus if there is no contention then we should have less to worry
+ *  about. B. If repeated compilation is part of a long running process then we
+ *  may overflow int64_t. We may detect and fail on this. For now this is not
+ *  done.
+ *
+ *  Now about 2:
+ *  There are two usecases for [debug-handle, DebugInfoTuple]
+ *  A. During bytecode generation the DebugInfoTuple corresponding to the nodes
+ *  of the inlined graph being serialized, are stored in this object and a
+ *  unique debug handle is returned. This unique debug handle is stored in
+ *  mobile_debug info for pytorch lite models. It will be used for raising
+ *  exceptions as well as profiling. B. During backend lowering, each backend's
+ *  preprocess/compile method can compile method's graph and serialize those
+ *  methods. Once the method is lowered to backend, graph is essentially lost.
+ *  Without access to graph it is hard to generate model level debug info. Thus
+ *  the debug handles provide a way to map nodes of the graph to the model level
+ *  debug info.
+ *
+ *  During byte-code model serialization, [debug-handle, DebugInfoTuple] is
+ *  serialized. Now we know a. debug handles and b. how to map debug handles to
+ *  model source code. Thus we can either do eager symbolication by converting
+ *  debug handles to corresponding source code at runtime, or do lazy
+ *  symbolicattion offline.
+ *
+ *  Note that it is not necessary to serialize [debug-handle, DebugInfoTuple]
+ *  corresponding to lowered backend if the lowering process, that is
+ *  preprocess/compile, and execution happens in the same session, then eager
+ *  symbolication can be employed.
+ *
+ *  Now how does BackendDebugHandleManager capture all of the above?
+ *  By providing two API.
+ *  1. getNextDebugHandle which given a Node* returns a unique debug handle,
+ *     that will uniquely identify DebugInfoTuple.
+ *     and
+ *  2. getCallStackPtrMap which returns the map
+ *     [debug-handle, DebugInfoTuple]
+ *
+ *  1 provides debug handles to backends and 2 provides runtime a way to map
+ *  debug handles to source level debug info.
+ *
+ *  So why does debug handle map to DebugInfoTuple = {source range and inlined
+ *  cs}? {debug_handle, source_range_tag, serialized_callstack} Take this
+ *  example: class L(nn.Module): def __init__(self) -> None:
+ *      ...
+ *    def forward(self, x):
+ *      return x * 5
+ *  class M(nn.Module):
+ *    def __init__(self) -> None:
+ *      ...
+ *    def forward(self, x):
+ *      return x - 2
+ *  class N(nn.Module):
+ *    def __init__(self) -> None:
+ *      self.m = M()
+ *    def forward(self, x):
+ *      return self.m(x) + 3
+ *  m = torch.jit.script(N())
+ *  Once you inline m's forward method, m.forward.graph will look something
+ *  like this
+ *  graph(%self...):
+ *   %x = aten::mul(..)
+ *   %x = aten::sub(x, ..)
+ *   %y = aten::add(x, ..)
+ *   ..
+ *  Inlined callstack ptr for these two nodes will look like:
+ *  aten::mul's inlined CS (callstack): [N.forward, source range] -> [M.forward,
+ *  source range] aten::sub's inlined CS (callstack): [N.forward, source range]
+ *  aten::add's inlined CS: null
+ *  mul node's inlined CS contains only information about the callsites' source
+ *  range The information about mul node's source range ('return x * 5') is not
+ *  available in its inlined CS. It is rather part of node's source range
+ *  instead of inlined CS. Thus to get full stack: [N.forward, source range] ->
+ *  [M.forward, source range] -> [aten::mul's source range] We need to track
+ *  mul's source range and inlined CS both.
+ */
+
+using BackendDebugInfoMapType =
+    std::unordered_map<torch::jit::DebugHandleType, DebugInfoTuple>;
+
+/*
+ * This class is used to generate debug info map.
+ * backend's preprocess will call generate_debug_handles (see
+ * backend_detail.cpp), which uses debug_handle_manager to generate debug
+ * handles. When lowering process finishes, calling stopRecording will
+ * return debug info map from debug_handle_manager
+ */
+class TORCH_API BackendDebugInfoRecorder {
+ public:
+  BackendDebugInfoRecorder() = default;
+  int64_t getNextDebugHandle(const Node* node);
+  // Reason this is not done as RAII is that work done in stopRecording
+  // can throw, and throwing with dtor will call terminate and thus voids any
+  // exception catching at a higher level.
+  BackendDebugInfoMapType stopRecording();
+  NodeToDebugHandle generate_debug_handles(const std::shared_ptr<Graph>& graph);
+
+ private:
+  static std::atomic<DebugHandleType> unique_debug_handle_;
+  BackendDebugInfoMapType handles_to_inlined_callstack_ptrs_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_debug_info.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_debug_info.h
new file mode 100644
index 0000000000000000000000000000000000000000..b2ff9a3fe801206fba4bf40538e5770a3ae493e4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_debug_info.h
@@ -0,0 +1,68 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifndef BUILD_LITE_INTERPRETER
+#include <torch/csrc/jit/backends/backend_debug_handler.h>
+#endif
+#include <torch/custom_class.h>
+
+namespace torch::jit {
+
+constexpr static auto kBackendUtilsNamespace = "backendutils";
+constexpr static auto kBackendDebugInfoClass = "BackendDebugInfo";
+
+#ifndef BUILD_LITE_INTERPRETER
+/*
+ * Custom class for holding debug information in lowered modules, intended
+ * purely for keeping this information to be later serialized outside of the
+ * lowered module itself.
+ * Its usage pattern is:
+ * 1. LoweredModule declares an instance of this class in __backend_debug_info
+ * 2. During serialization, __backend_debug_info is used to obtain the debug
+ *    information.
+ * 3. The contents of LoweredModule.__backend_debug_info are not serialized
+ *    within the LoweredModule itself.
+ */
+class TORCH_API PyTorchBackendDebugInfo : public torch::CustomClassHolder {
+ public:
+  PyTorchBackendDebugInfo() = default;
+
+  std::optional<BackendDebugInfoMapType>& getDebugInfoMap() {
+    return debug_info_map_;
+  }
+
+  void setDebugInfoMap(BackendDebugInfoMapType&& debug_info_map) {
+    debug_info_map_ = std::move(debug_info_map);
+  }
+
+ private:
+  std::optional<BackendDebugInfoMapType> debug_info_map_;
+};
+
+#else
+
+/*
+ * Dummy instance exists for the following reason:
+ * __backend_debug_info is of type BackendDebugInfo which is a torchbind'
+ * class backed by cpp class PyTorchBackendDebugInfo.
+ * PyTorchBackendDebugInfo, depends on ir.h., scope.h, source_range etc.
+ * We dont include this on lite interpreter side. Thus on lite interpreter side
+ * we cannot have valid definition of PyTorchBackendDebugInfo. However we do not
+ * need valid instance of __backend_debug_info in lite interpreter anyway as we
+ * dont serialize this info as part of LowerdModule as mentioned ealrier.
+ * However since LoweredModule has registered attribute of __backend_debug_info
+ * we still need to make sure that BackendDebugInfo is registered with
+ * TorchScript. However in this instance it does not have to be backed by
+ * PyTorchBackendDebugInfo, so we create a dummy PyTorchBackendDebugInfoDummy
+ * just for this purpose.
+ */
+class PyTorchBackendDebugInfoDummy : public torch::CustomClassHolder {
+ public:
+  PyTorchBackendDebugInfoDummy() = default;
+};
+#endif
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_detail.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_detail.h
new file mode 100644
index 0000000000000000000000000000000000000000..cca52f2866881927fa9db1b8f35cb20be87a5183
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_detail.h
@@ -0,0 +1,44 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+
+#include <ATen/core/jit_type.h>
+
+#include <functional>
+
+namespace torch::jit {
+
+using DebugHandleType = int64_t;
+
+using NodeToDebugHandle = std::unordered_map<Node*, DebugHandleType>;
+
+using BackendDebugHandleGenerator =
+    std::function<NodeToDebugHandle(const std::shared_ptr<Graph>&)>;
+
+namespace detail {
+
+using BackendPreprocessFunction = std::function<c10::IValue(
+    const Module&,
+    const c10::Dict<IValue, IValue>&,
+    const BackendDebugHandleGenerator& generate_debug_handles)>;
+
+TORCH_API void registerBackendPreprocessFunction(
+    const std::string& name,
+    const BackendPreprocessFunction& preprocess);
+
+bool hasBackendPreprocessFunction(const std::string& name);
+
+BackendPreprocessFunction getBackendPreprocessFunction(const std::string& name);
+
+TORCH_API Module codegen_backend_module(
+    const std::string& backend_name,
+    const Module& orig_module,
+    const c10::Dict<IValue, IValue>& method_compile_spec,
+    const c10::DictTypePtr& any_dict_ty);
+} // namespace detail
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_exception.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_exception.h
new file mode 100644
index 0000000000000000000000000000000000000000..caec46ee4666f0c8867a0326ad9ed9ccb1567599
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_exception.h
@@ -0,0 +1,62 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// @allow-raw-throw
+#pragma once
+#include <c10/util/Exception.h>
+
+#include <utility>
+
+namespace c10 {
+class TORCH_API BackendRuntimeException : public c10::Error {
+ public:
+  // Use debug_handle to throw exception
+  BackendRuntimeException(
+      SourceLocation loc,
+      std::string msg,
+      int64_t debug_handle)
+      : c10::Error(loc, std::move(msg)) {
+    debug_handles.push_back(debug_handle);
+  }
+  // If rethrowing, can push another debug_handle
+  // This is useful in couple of scenarios.
+  // 1. A submodule is lowered and lite interpreter has CallMethod
+  //    to lowered module's method. In this case lowered module will throw with
+  //    a handle, plus there will be another debug handle corresponding
+  //    to the CallMethod node in lite interpreter. Both together give complete
+  //    trace. This function allows lite interpreter to rethrow with debug
+  //    handle it has for CallMethod.
+  // 2. Another scenarios is when lite interpreter can make function calls or
+  //    the lowered backend also has function call ability. Thus we have
+  //    multiple function frames. Now we need a stack of handles to symbolicate
+  //    entire stack trace.
+  void pushDebugHandle(int64_t debug_handle) {
+    debug_handles.push_back(debug_handle);
+  }
+  const std::vector<int64_t>& getDebugHandles() {
+    return debug_handles;
+  }
+
+ private:
+  // Stores stack of debug handles.
+  std::vector<int64_t> debug_handles;
+};
+
+} // namespace c10
+#define TORCH_DELEGATED_BACKEND_THROW(cond, msg, debug_handle) \
+  if (C10_UNLIKELY_OR_CONST(!(cond))) {                        \
+    throw ::c10::BackendRuntimeException(                      \
+        {__func__, __FILE__, static_cast<uint32_t>(__LINE__)}, \
+        msg,                                                   \
+        debug_handle);                                         \
+  }
+
+#define TORCH_DELEGATED_BACKEND_RETHROW(e, debug_handle) \
+  do {                                                   \
+    e.pushDebugHandle(debug_handle);                     \
+    throw;                                               \
+  } while (false)
+
+#define DEBUG_HANDLE_UNKNOWN -1
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_init.h
new file mode 100644
index 0000000000000000000000000000000000000000..bc490802ff882f7c10b304af7803b80d1c511b9e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_init.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/python/pybind.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::jit {
+// Initialize Python bindings for JIT to_<backend> functions.
+void initJitBackendBindings(PyObject* module);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_interface.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_interface.h
new file mode 100644
index 0000000000000000000000000000000000000000..5f7056a86d0628c1a861465c1cc30d46fc2d1db7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_interface.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/custom_class.h>
+
+namespace torch::jit {
+
+// Interface for a JIT backend.
+class TORCH_API PyTorchBackendInterface : public torch::CustomClassHolder {
+ public:
+  PyTorchBackendInterface() noexcept;
+  ~PyTorchBackendInterface() override;
+
+  // Returns true if the backend is available to process delegation calls.
+  virtual bool is_available() = 0;
+
+  // Compile the module contained in \p processed using the details provided in
+  // \p method_compile_spec for each module method that should be compiled for
+  // the backend. \p method_compile_spec should be of type Dict<string, Any>.
+  // \returns a dictionary of type Dict<string, Any> that contains a backend
+  // handle each method that can run on the backend (i.e. each key in \p
+  // method_compile_spec).
+  virtual c10::impl::GenericDict compile(
+      c10::IValue processed,
+      c10::impl::GenericDict method_compile_spec) = 0;
+
+  // Execute the method specified by \p handle using \p inputs. \returns the
+  // outputs as a tuple.
+  virtual c10::impl::GenericList execute(
+      c10::IValue handle,
+      c10::impl::GenericList inputs) = 0;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_preprocess.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_preprocess.h
new file mode 100644
index 0000000000000000000000000000000000000000..f0241ec96ef63b32e94b6116f8fcc31df5da9f51
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_preprocess.h
@@ -0,0 +1,21 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/backends/backend_detail.h>
+namespace torch::jit {
+class backend_preprocess_register {
+  std::string backend_name_;
+
+ public:
+  backend_preprocess_register(
+      const std::string& name,
+      const detail::BackendPreprocessFunction& preprocess)
+      : backend_name_(name) {
+    detail::registerBackendPreprocessFunction(name, preprocess);
+  }
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_resolver.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_resolver.h
new file mode 100644
index 0000000000000000000000000000000000000000..aee7fac6ddfb3fe942a82ad33074e672a3422821
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/backend_resolver.h
@@ -0,0 +1,13 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/frontend/resolver.h>
+
+namespace torch::jit {
+// Create a Resolver for use in generating LoweredModules for specific backends.
+TORCH_API std::shared_ptr<Resolver> loweredModuleResolver();
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/cpp/context.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/cpp/context.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ac2655639b3b0dc4a087056e90a7f75a593f226
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/cpp/context.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef PTM_COREML_Context_h
+#define PTM_COREML_Context_h
+
+#include <string>
+
+namespace torch::jit::mobile::coreml {
+
+struct ContextInterface {
+  virtual ~ContextInterface() = default;
+  virtual void setModelCacheDirectory(std::string path) = 0;
+};
+
+class BackendRegistrar {
+ public:
+  explicit BackendRegistrar(ContextInterface* ctx);
+};
+
+void setModelCacheDirectory(std::string path);
+
+} // namespace torch::jit::mobile::coreml
+
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLCompiler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLCompiler.h
new file mode 100644
index 0000000000000000000000000000000000000000..1b040c52c64f22364741a4c926686ec3579edd3b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLCompiler.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#import <CoreML/CoreML.h>
+
+#include <string>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface PTMCoreMLCompiler : NSObject
+
++ (void)setCacheDirectory:(const std::string&)dir;
+
++ (NSString*)cacheDirectory;
+
++ (BOOL)compileModel:(const std::string&)modelSpecs modelID:(const std::string&)modelID;
+
++ (nullable MLModel*)loadModel:(const std::string)modelID
+                       backend:(const std::string)backend
+             allowLowPrecision:(BOOL)allowLowPrecision
+                         error:(NSError**)error;
+
+@end
+
+NS_ASSUME_NONNULL_END
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h
new file mode 100644
index 0000000000000000000000000000000000000000..5a79337260dcc3099d395a1639f39420796ab6b0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#import <torch/csrc/jit/backends/coreml/objc/PTMCoreMLFeatureProvider.h>
+
+#import <CoreML/CoreML.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface PTMCoreMLExecutor : NSObject
+
+@property(atomic, strong) MLModel* model;
+
+- (instancetype)initWithFeatureNames:(NSArray<NSString*>*)featureNames;
+
+- (void)setInputs:(c10::impl::GenericList)inputs;
+
+- (id<MLFeatureProvider>)forward:(NSError**)error;
+
+@end
+
+NS_ASSUME_NONNULL_END
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLFeatureProvider.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLFeatureProvider.h
new file mode 100644
index 0000000000000000000000000000000000000000..c0e536370b6ee5209dc88ae869539bad7a423260
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLFeatureProvider.h
@@ -0,0 +1,21 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#import <ATen/ATen.h>
+#import <CoreML/CoreML.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface PTMCoreMLFeatureProvider : NSObject<MLFeatureProvider>
+
+- (instancetype)initWithFeatureNames:(NSSet<NSString*>*)featureNames;
+
+- (void)clearInputTensors;
+
+- (void)setInputTensor:(const at::Tensor&)tensor forFeatureName:(NSString*)name;
+
+@end
+
+NS_ASSUME_NONNULL_END
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLModelWrapper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLModelWrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..5ee77404da4e52f35be363da00dbf4779b75af89
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLModelWrapper.h
@@ -0,0 +1,46 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h>
+#include <torch/csrc/jit/backends/coreml/objc/PTMCoreMLTensorSpec.h>
+
+namespace torch {
+namespace jit {
+namespace mobile {
+namespace coreml {
+
+class MLModelWrapper : public CustomClassHolder {
+ public:
+  PTMCoreMLExecutor* executor;
+  std::vector<TensorSpec> outputs;
+
+  MLModelWrapper() = delete;
+
+  MLModelWrapper(PTMCoreMLExecutor* executor) : executor(executor) {
+    [executor retain];
+  }
+
+  MLModelWrapper(const MLModelWrapper& oldObject) {
+    executor = oldObject.executor;
+    outputs = oldObject.outputs;
+    [executor retain];
+  }
+
+  MLModelWrapper(MLModelWrapper&& oldObject) {
+    executor = oldObject.executor;
+    outputs = oldObject.outputs;
+    [executor retain];
+  }
+
+  ~MLModelWrapper() {
+    [executor release];
+  }
+};
+
+} // namespace coreml
+} // namespace mobile
+} // namespace jit
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLTensorSpec.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLTensorSpec.h
new file mode 100644
index 0000000000000000000000000000000000000000..7537f743d938199c4d13c8f8aca01c6e5b0c231b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLTensorSpec.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <c10/core/ScalarType.h>
+#import <nlohmann/json.hpp>
+
+#include <string>
+
+namespace torch::jit::mobile::coreml {
+
+struct TensorSpec {
+  std::string name;
+  c10::ScalarType dtype = c10::ScalarType::Float;
+};
+
+static inline c10::ScalarType scalar_type(const std::string& type_string) {
+  if (type_string == "0") {
+    return c10::ScalarType::Float;
+  } else if (type_string == "1") {
+    return c10::ScalarType::Double;
+  } else if (type_string == "2") {
+    return c10::ScalarType::Int;
+  } else if (type_string == "3") {
+    return c10::ScalarType::Long;
+  }
+  return c10::ScalarType::Undefined;
+}
+
+} // namespace torch::jit::mobile::coreml
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/compiler/xnn_compiler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/compiler/xnn_compiler.h
new file mode 100644
index 0000000000000000000000000000000000000000..61bd88c05345fc626ef96149efa97d5235afa52b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/compiler/xnn_compiler.h
@@ -0,0 +1,29 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <caffe2/torch/csrc/jit/backends/xnnpack/executor/xnn_executor.h>
+#include <xnnpack.h>
+#include <memory>
+#include <vector>
+
+namespace torch::jit::xnnpack::delegate {
+
+class XNNCompiler {
+ public:
+  // Takes Flatbuffer Serialized XNNPack Model and rebuilds the xnn-subgraph
+  // returns an executor object that holds the xnn runtime object which we
+  // can then use to set inputs and run inference using the xnn graph.
+  static void compileModel(
+      const void* buffer_pointer,
+      size_t num_bytes,
+      XNNExecutor* executor);
+};
+
+} // namespace torch::jit::xnnpack::delegate
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/executor/xnn_executor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/executor/xnn_executor.h
new file mode 100644
index 0000000000000000000000000000000000000000..376de821a60acabd138d32b375e1c833bd077886
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/executor/xnn_executor.h
@@ -0,0 +1,73 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+#include <xnnpack.h>
+#include <memory>
+#include <vector>
+
+namespace torch::jit::xnnpack::delegate {
+
+class XNNExecutor {
+ private:
+  std::unique_ptr<xnn_runtime, decltype(&xnn_delete_runtime)> runtime_{
+      nullptr,
+      &xnn_delete_runtime};
+  std::vector<uint32_t> input_ids_;
+  std::vector<uint32_t> output_ids_;
+  std::vector<xnn_external_value> externals_;
+
+ public:
+  XNNExecutor() = default;
+
+  template <typename T>
+  bool set_inputs(std::vector<T*>& inputs, std::vector<T*>& outputs) {
+    externals_.clear();
+
+    if (inputs.size() != input_ids_.size()) {
+      return false;
+    }
+
+    for (int i = 0; i < inputs.size(); i++) {
+      externals_.emplace_back(xnn_external_value{input_ids_[i], inputs[i]});
+    }
+
+    if (outputs.size() != output_ids_.size()) {
+      return false;
+    }
+
+    for (int i = 0; i < outputs.size(); i++) {
+      externals_.emplace_back(xnn_external_value{output_ids_[i], outputs[i]});
+    }
+
+    return true;
+  }
+
+  bool forward() {
+    xnn_status status =
+        xnn_setup_runtime(runtime_.get(), externals_.size(), externals_.data());
+
+    if (status != xnn_status_success) {
+      return false;
+    }
+
+    status = xnn_invoke_runtime(runtime_.get());
+
+    if (status != xnn_status_success) {
+      return false;
+    }
+
+    return true;
+  }
+
+  friend class XNNCompiler;
+};
+
+} // namespace torch::jit::xnnpack::delegate
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/serialization/serializer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/serialization/serializer.h
new file mode 100644
index 0000000000000000000000000000000000000000..1ca44842bad03c04d60d83a20a67ae8f845a1213
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/serialization/serializer.h
@@ -0,0 +1,94 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <torch/csrc/jit/backends/xnnpack/serialization/schema_generated.h>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace torch {
+namespace jit {
+namespace xnnpack {
+namespace delegate {
+
+using namespace fb_xnnpack; // Specified in the schema
+
+class XNNSerializer {
+ public:
+  // Constructors
+  // initial buffersize of 1024 which will grow
+  // automatically, constant buffer and buffer sizes initialized with dummy
+  // values as 0 index is reserved for non-constant tensors
+  XNNSerializer() : XNNSerializer(1024) {}
+
+  explicit XNNSerializer(size_t bufferSize)
+      : _builder(bufferSize),
+        _nodes(),
+        _values(),
+        _constantBuffer({CreateBuffer(
+            _builder,
+            {})}), // index 0 is reserved for non-const data
+        _bufferSizes({0}) {}
+
+  // Serializing Nodes
+
+  // Serialize add node, we are serializing the argument needed to call
+  // xnn_define_add2. Serializing these values, and at run time we build
+  // the graph by re running xnn_define_add2
+  void serializeAddNode(
+      uint32_t input1_id,
+      uint32_t input2_id,
+      uint32_t output_id,
+      uint32_t flags);
+
+  // Serializing Values
+  void serializeTensorValue(
+      uint32_t xnn_datatype,
+      size_t num_dims,
+      std::vector<size_t> dims,
+      size_t buffer_data_idx,
+      uint32_t external_id,
+      uint32_t flags,
+      uint32_t id_out);
+
+  // finish and serialize xnngraph returning serialized data
+  std::string finishAndSerialize(
+      std::vector<uint32_t> input_ids,
+      std::vector<uint32_t> output_ids,
+      size_t num_extern_ids);
+
+  // decoupled data serialization with tensor values. This way constant tensor
+  // data can be referenced by multiple intermediate tensors. This call
+  // serializes the num_bytes of the data_ptr and returns the index it was
+  // placed in.
+  size_t serializeData(const uint8_t* data_ptr, size_t num_bytes);
+
+ private:
+  // xnnpack version we are serializing
+  const char* _version_sha1 = "ae108ef49aa5623b896fc93d4298c49d1750d9ba";
+
+  // flatbuffer objects we will create and serialize together to create xnngraph
+  flatbuffers_fbsource::FlatBufferBuilder _builder;
+
+  // Vector of the serialized xnnpack nodes
+  std::vector<flatbuffers_fbsource::Offset<XNode>> _nodes;
+
+  // Vector of the serialized xnnpack values
+  std::vector<flatbuffers_fbsource::Offset<XValue>> _values;
+
+  std::vector<flatbuffers_fbsource::Offset<Buffer>> _constantBuffer;
+  std::vector<uint32_t> _bufferSizes;
+};
+
+} // namespace delegate
+} // namespace xnnpack
+} // namespace jit
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/xnnpack_graph_builder.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/xnnpack_graph_builder.h
new file mode 100644
index 0000000000000000000000000000000000000000..369d56f8d9a33971d3047c5c499fe014a4dea7c9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/backends/xnnpack/xnnpack_graph_builder.h
@@ -0,0 +1,102 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <ATen/Functions.h>
+#include <ATen/Utils.h>
+#include <torch/torch.h>
+#include <xnnpack.h>
+#include <unordered_set>
+#include <vector>
+
+#include <torch/csrc/jit/backends/xnnpack/serialization/serializer.h>
+
+namespace torch {
+namespace jit {
+namespace xnnpack {
+namespace delegate {
+
+class XNNGraph {
+ private:
+  const float output_min = -std::numeric_limits<float>::infinity();
+  const float output_max = std::numeric_limits<float>::infinity();
+
+  // serializer class
+  XNNSerializer _serializer;
+  // xnn subgraph
+  xnn_subgraph_t _subgraph_ptr;
+  // Set of all the tensor values throughout the jit graph
+  std::unordered_set<torch::jit::Value*> _intermediate_tensors;
+  // Set of all the tensor values mapped to the xnnpack ids
+  std::unordered_map<torch::jit::Value*, uint32_t> _val_to_ids;
+  // Vector containing the torch valued inputs/outputs,
+  // must be ordered to preserve the order of input/outputs
+  std::vector<torch::jit::Value*> _inputs;
+  std::vector<torch::jit::Value*> _outputs;
+
+  // Graph passes for optimizing and tracing torchscript graph
+  // Essentially massaging the graph into a digestiable format for
+  // xnnpack graph lowering.
+  std::shared_ptr<torch::jit::Graph> optimizeAndTraceGraph(
+      std::shared_ptr<torch::jit::Graph> graph,
+      std::vector<c10::IValue>& example_inputs);
+
+  // Gather all the intermediate tensor values within a graph. This
+  // skips through all prim constants. The purpose of this is for defining
+  // the tensor values beforehand for the xnnpack subgraph.
+  void gatherTensorValues(std::shared_ptr<torch::jit::Graph>& graph);
+
+  // Gathers the tensor values in a give node
+  void gatherNodeInputs(torch::jit::Node& node);
+
+  // Helper function to determine if a jit value is a graph input
+  bool isGraphInput(torch::jit::Value* val);
+
+  // Helper function to determine if a jit value is a graph output
+  bool isGraphOutput(torch::jit::Value* val);
+
+  // Defines all xnnpack nodes for the nodes in the graph
+  void defineAllNodes(std::shared_ptr<torch::jit::Graph>& graph);
+
+  // Defines all xnn tensor values used throughout the graph
+  void defineAllTensorValues();
+
+  // Makes a pass through the graph and throws if any ops are unsupported
+  void checkOpsToDelegate(std::shared_ptr<torch::jit::Graph>& graph);
+
+ public:
+  XNNGraph() : _serializer(), _subgraph_ptr(nullptr) {
+    xnn_status status = xnn_initialize(/*allocator =*/nullptr);
+    TORCH_CHECK(xnn_status_success == status, "Failed to initialize xnnpack");
+  }
+
+  ~XNNGraph() {
+    xnn_deinitialize();
+    if (_subgraph_ptr != nullptr) {
+      xnn_delete_subgraph(_subgraph_ptr);
+    }
+  }
+
+  void buildXNNGraph(
+      std::shared_ptr<torch::jit::Graph>& graph,
+      std::vector<c10::IValue> example_inputs);
+
+  void runGraphOnInputs(
+      std::vector<at::Tensor> tensor_inputs,
+      std::vector<at::Tensor> tensor_outputs);
+
+  std::string serializedXNNGraph();
+
+  std::vector<std::vector<long>> getGraphOutputShapes();
+};
+
+} // namespace delegate
+} // namespace xnnpack
+} // namespace jit
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/cuda/interface.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/cuda/interface.h
new file mode 100644
index 0000000000000000000000000000000000000000..211a2fe3a749f934e4c9347b46fb0c6eb111e7f6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/cuda/interface.h
@@ -0,0 +1,57 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/pass_manager.h>
+#include <torch/csrc/jit/runtime/profiling_record.h>
+
+/*
+ * This file contains APIs for cuda fuser;
+ *
+ * We use an empty static struct to hold the function pointers, which are
+ * registered separately. This is to support cpu-only compilation.
+ * Registration is done in torch/csrc/jit/codegen/cuda/register_interface.cpp
+ */
+
+namespace torch::jit::fuser::cuda {
+
+TORCH_API std::atomic<bool>& getCudaFusionGuardMode();
+
+TORCH_API bool getSingletonFusion();
+TORCH_API bool setSingletonFusion(bool value);
+TORCH_API bool getHorizontalFusion();
+TORCH_API bool setHorizontalFusion(bool value);
+
+// dummy struct to allow API registration
+struct CudaFuserInterface {
+  void (*fn_compile_n)(Node*) = nullptr;
+  void (*fn_run_n_s)(const Node*, Stack&) = nullptr;
+  void (*fn_fuse_graph)(std::shared_ptr<Graph>&) = nullptr;
+  bool (*fn_can_fuse_n)(const Node*) = nullptr;
+  void (*fn_insert_profile_inodes)(ProfilingRecord* pr) = nullptr;
+  bool (*fn_profile_n)(const Node*) = nullptr;
+  bool (*fn_skip_n)(const std::string&, bool flip) = nullptr;
+};
+
+// Get interface, this is used by registration and user facing API internally
+TORCH_API CudaFuserInterface* getFuserInterface();
+
+TORCH_API void compileFusionGroup(Node* fusion_node);
+TORCH_API void runFusionGroup(const Node* fusion_node, Stack& stack);
+TORCH_API void fuseGraph(std::shared_ptr<Graph>& /*graph*/);
+TORCH_API bool canFuseNode(const Node* node);
+TORCH_API void InsertProfileNodesForCUDAFuser(ProfilingRecord* pr);
+TORCH_API bool profileNode(const Node* node);
+
+TORCH_API bool skipNode(const std::string& symbol_str, bool flip = true);
+
+TORCH_API bool isEnabled();
+TORCH_API bool setEnabled(bool is_enabled);
+TORCH_API bool canBeEnabled();
+
+} // namespace torch::jit::fuser::cuda
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/arg_spec.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/arg_spec.h
new file mode 100644
index 0000000000000000000000000000000000000000..3621030aed0ee08ccf973d0d64ebe2592118a851
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/arg_spec.h
@@ -0,0 +1,60 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/ATen.h>
+#include <ATen/core/functional.h> // fmap
+#include <c10/util/hash.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/tensor_desc.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace torch::jit::fuser {
+
+// Describes the (runtime) arguments to a kernel.
+// ArgSpecs are also used as keys to lookup instantiated kernels, so
+//  they are hashable.
+// Note: the device to run on is included in the arg spec because kernels
+//  are compiled per-device.
+struct TORCH_API ArgSpec {
+  ArgSpec(at::TensorList inputs, const int _device)
+      : descs_{c10::fmap<TensorDesc>(inputs)},
+        hash_code_{c10::get_hash(_device, inputs.size(), descs_)},
+        device_{_device} {}
+
+  // (Common) hash function
+  static size_t hash(const ArgSpec& spec) {
+    return spec.hash_code_;
+  }
+
+  // Comparators
+  bool operator==(const ArgSpec& other) const {
+    return (descs_ == other.descs_ && device_ == other.device_);
+  }
+
+  bool operator!=(const ArgSpec& spec) const {
+    return !(*this == spec);
+  }
+
+  // Getters
+  size_t hashCode() const {
+    return hash_code_;
+  }
+  const std::vector<TensorDesc>& descs() const {
+    return descs_;
+  }
+  int device() const {
+    return device_;
+  }
+
+ private:
+  std::vector<TensorDesc> descs_;
+  size_t hash_code_;
+  int device_;
+};
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/codegen.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/codegen.h
new file mode 100644
index 0000000000000000000000000000000000000000..1cc359481bf7cb9eb5b1bf72b1f7043bf0612309
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/codegen.h
@@ -0,0 +1,29 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/arg_spec.h>
+#include <torch/csrc/jit/codegen/fuser/partition_desc.h>
+#include <torch/csrc/jit/codegen/fuser/tensor_desc.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <string>
+#include <vector>
+
+namespace torch::jit::fuser {
+
+// Creates a CPU or CUDA kernel for the given graph.
+// Returns the C++ or CUDA string implementing the kernel.
+TORCH_API std::string generateKernel(
+    const std::string& name,
+    const Graph& graph,
+    const std::vector<std::pair<const Value*, const std::optional<TensorDesc>>>&
+        inputs,
+    const std::vector<std::pair<const Value*, const TensorDesc>>& outputs,
+    const bool use_cuda);
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/compiler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/compiler.h
new file mode 100644
index 0000000000000000000000000000000000000000..e76959805a5cdee9d94582b09f76d3750e0ecdc4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/compiler.h
@@ -0,0 +1,61 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/arg_spec.h>
+#include <torch/csrc/jit/codegen/fuser/fused_kernel.h>
+#include <torch/csrc/jit/codegen/fuser/interface.h>
+#include <torch/csrc/jit/codegen/fuser/kernel_spec.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace torch::jit::fuser {
+
+// Performs device-independent "upfront" compilation of the given fusion_group,
+// if it has not been registered already.
+// Returns a key that can be used to run the fusion later
+TORCH_API int64_t registerFusion(const Node* fusion_group);
+
+// Performs device-specific "runtime" compilation of the given kernel
+//  with the runtime arguments specified in ArgSpec.
+//  Outputs are allocated using map_size on the specified device.
+TORCH_API std::shared_ptr<FusedKernel> compileKernel(
+    const KernelSpec& spec,
+    const ArgSpec& arg_spec,
+    const std::vector<int64_t>& map_size,
+    const at::Device& device);
+
+TORCH_API size_t nCompiledKernels();
+
+TORCH_API int debugFuser();
+
+using FusedKernelConstructor = std::function<std::shared_ptr<FusedKernel>(
+    int16_t device,
+    std::string name,
+    std::string code,
+    std::vector<TensorDesc> input_desc,
+    std::vector<TensorDesc> output_desc,
+    std::vector<PartitionDesc> chunk_desc,
+    std::vector<PartitionDesc> concat_desc,
+    bool has_random)>;
+
+TORCH_API void registerFusionBackend(
+    at::Device::Type backend_type,
+    FusedKernelConstructor ctor);
+TORCH_API bool hasFusionBackend(at::Device::Type backend_type);
+struct TORCH_API RegisterFusionBackend{RegisterFusionBackend(
+    at::Device::Type backend_type,
+    FusedKernelConstructor ctor){
+    registerFusionBackend(backend_type, std::move(ctor));
+} // namespace torch::jit::fuser
+}
+;
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.h
new file mode 100644
index 0000000000000000000000000000000000000000..13e37fe47a442853f902b18967222dfd930cec2c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.h
@@ -0,0 +1,44 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/DynamicLibrary.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/fused_kernel.h>
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+namespace torch::jit::fuser::cpu {
+
+// Represents a compiled CPU kernel and the metadata necessary to run it
+struct TORCH_API FusedKernelCPU : public FusedKernel {
+  FusedKernelCPU(
+      std::string name,
+      std::string code,
+      std::vector<TensorDesc> input_desc,
+      std::vector<TensorDesc> output_desc,
+      std::vector<PartitionDesc> chunk_desc,
+      std::vector<PartitionDesc> concat_desc,
+      bool has_random);
+
+  at::Backend backend() const override {
+    return at::Backend::CPU;
+  }
+
+  void launch_raw(const uint32_t numel, std::vector<void*>& arguments)
+      const override {
+    kernel(numel, arguments.data());
+  }
+
+ private:
+  std::unique_ptr<at::DynamicLibrary> so_lib;
+  void (*kernel)(uint32_t, void**) = nullptr;
+};
+
+} // namespace torch::jit::fuser::cpu
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/resource_strings.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/resource_strings.h
new file mode 100644
index 0000000000000000000000000000000000000000..62c3008b31ff5ecacfca5541068ced287dcb84cc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/resource_strings.h
@@ -0,0 +1,106 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/code_template.h>
+
+namespace torch::jit::fuser::cpu {
+
+/*with type_as not checking type of its input, a fusion group can have non-fp32
+tensor as input. Correct code for this case is generated, however, nvrtc does
+not know how to handle int*_t integer types, so typedefs help it handle those
+cases*/
+
+static auto type_declarations_template = at::jit::CodeTemplate(R"(
+
+#define POS_INFINITY INFINITY
+#define NEG_INFINITY -INFINITY
+
+typedef ${IndexType} IndexType;
+template<typename T, size_t N>
+struct TensorInfo {
+  T* data;
+  IndexType sizes[N];
+  IndexType strides[N];
+};
+template<typename T>
+struct TensorInfo<T, 0> {
+  T * data;
+};
+)");
+
+static auto cpu_compilation_unit_template = at::jit::CodeTemplate(R"(
+#include <math.h>
+#include <cstddef>
+#include <cstdint>
+
+double rsqrt(double x) {
+  return 1.0/sqrt(x);
+}
+
+float rsqrtf(float x) {
+  return 1.0f/sqrtf(x);
+}
+
+double frac(double x) {
+  return x - trunc(x);
+}
+
+float fracf(float x) {
+  return x - truncf(x);
+}
+
+${type_declarations}
+
+#ifdef _MSC_VER
+template<size_t n> struct int_of_size;
+
+#define DEFINE_INT_OF_SIZE(int_t) \
+template<> struct int_of_size<sizeof(int_t)> { using type = int_t; }
+
+DEFINE_INT_OF_SIZE(int64_t);
+DEFINE_INT_OF_SIZE(int32_t);
+DEFINE_INT_OF_SIZE(int16_t);
+DEFINE_INT_OF_SIZE(int8_t);
+
+#undef DEFINE_INT_OF_SIZE
+
+template <typename T>
+using int_same_size_t = typename int_of_size<sizeof(T)>::type;
+
+#define IndexTypeLoop int_same_size_t<IndexType>
+#define ToIndexTypeLoop(x) static_cast<IndexTypeLoop>(x)
+#else
+#define IndexTypeLoop IndexType
+#define ToIndexTypeLoop(x) x
+#endif
+
+#define OMP_THRESHOLD 100000
+static void ${kernelName}_kernel(IndexType totalElements, ${formals}) {
+  #pragma omp parallel for if(totalElements > OMP_THRESHOLD)
+  for (IndexTypeLoop linearIndex = 0;
+        linearIndex < ToIndexTypeLoop(totalElements);
+        linearIndex += 1) {
+      // Convert `linearIndex` into an offset of tensor:
+      ${tensorOffsets}
+      // calculate the results
+      ${kernelBody}
+    }
+}
+
+#ifdef _WIN32
+#define JIT_API __declspec(dllexport)
+#else
+#define JIT_API
+#endif
+
+extern "C"
+JIT_API void ${kernelName}(IndexType totalElements, void ** args) {
+  ${kernelName}_kernel(totalElements ${,argument_loads});
+}
+)");
+
+} // namespace torch::jit::fuser::cpu
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/temp_file.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/temp_file.h
new file mode 100644
index 0000000000000000000000000000000000000000..726ca1e7cc63e1fa32a04b9ca9995aad365ff778
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cpu/temp_file.h
@@ -0,0 +1,140 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Utils.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/Export.h>
+
+#ifdef _WIN32
+#include <WinError.h>
+#include <c10/util/Unicode.h>
+#include <c10/util/win32-headers.h>
+#include <fcntl.h>
+#include <io.h>
+#include <process.h>
+#include <stdio.h>
+#include <sys/stat.h>
+#include <random>
+#else
+#include <unistd.h>
+#endif
+
+#include <string>
+#include <vector>
+
+namespace torch::jit::fuser::cpu {
+
+#ifdef _MSC_VER
+inline int wmkstemps(wchar_t* tmpl, int suffix_len) {
+  int len;
+  wchar_t* name;
+  int fd = -1;
+  int save_errno = errno;
+
+  len = wcslen(tmpl);
+  if (len < 6 + suffix_len ||
+      wcsncmp(&tmpl[len - 6 - suffix_len], L"XXXXXX", 6)) {
+    return -1;
+  }
+
+  name = &tmpl[len - 6 - suffix_len];
+
+  std::random_device rd;
+  do {
+    for (unsigned i = 0; i < 6; ++i) {
+      name[i] = "abcdefghijklmnopqrstuvwxyz0123456789"[rd() % 36];
+    }
+
+    fd = _wopen(tmpl, _O_RDWR | _O_CREAT | _O_EXCL, _S_IWRITE | _S_IREAD);
+  } while (errno == EEXIST);
+
+  if (fd >= 0) {
+    errno = save_errno;
+    return fd;
+  } else {
+    return -1;
+  }
+}
+#endif
+
+struct TempFile {
+  AT_DISALLOW_COPY_AND_ASSIGN(TempFile);
+
+  TempFile(const std::string& t, int suffix) {
+#ifdef _MSC_VER
+    auto wt = c10::u8u16(t);
+    std::vector<wchar_t> tt(wt.c_str(), wt.c_str() + wt.size() + 1);
+    int fd = wmkstemps(tt.data(), suffix);
+    AT_ASSERT(fd != -1);
+    file_ = _wfdopen(fd, L"r+");
+    auto wname = std::wstring(tt.begin(), tt.end() - 1);
+    name_ = c10::u16u8(wname);
+#else
+    // mkstemps edits its first argument in places
+    // so we make a copy of the string here, including null terminator
+    std::vector<char> tt(t.c_str(), t.c_str() + t.size() + 1);
+    int fd = mkstemps(tt.data(), suffix);
+    AT_ASSERT(fd != -1);
+    file_ = fdopen(fd, "r+");
+    // - 1 because tt.size() includes the null terminator,
+    // but std::string does not expect one
+    name_ = std::string(tt.begin(), tt.end() - 1);
+#endif
+  }
+
+  const std::string& name() const {
+    return name_;
+  }
+
+  void sync() {
+    fflush(file_);
+  }
+
+  void write(const std::string& str) {
+    size_t result = fwrite(str.c_str(), 1, str.size(), file_);
+    AT_ASSERT(str.size() == result);
+  }
+
+#ifdef _MSC_VER
+  void close() {
+    if (file_ != nullptr) {
+      fclose(file_);
+    }
+    file_ = nullptr;
+  }
+#endif
+
+  FILE* file() {
+    return file_;
+  }
+
+  ~TempFile() {
+#ifdef _MSC_VER
+    if (file_ != nullptr) {
+      fclose(file_);
+    }
+    auto wname = c10::u8u16(name_);
+    if (!wname.empty() && _waccess(wname.c_str(), 0) != -1) {
+      _wunlink(wname.c_str());
+    }
+#else
+    if (file_ != nullptr) {
+      // unlink first to ensure another mkstemps doesn't
+      // race between close and unlink
+      unlink(name_.c_str());
+      fclose(file_);
+    }
+#endif
+  }
+
+ private:
+  FILE* file_ = nullptr;
+  std::string name_;
+};
+
+} // namespace torch::jit::fuser::cpu
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.h
new file mode 100644
index 0000000000000000000000000000000000000000..85c23b394ab73e4cce98eb632a4979059f2429d3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.h
@@ -0,0 +1,64 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/fused_kernel.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <nvrtc.h>
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace torch::jit::fuser::cuda {
+
+// query codegen output arch and target
+TORCH_CUDA_CU_API void codegenOutputQuery(
+    const cudaDeviceProp* const prop,
+    int& major,
+    int& minor,
+    bool& compile_to_sass);
+
+// A class holding metadata for an actual CUDA function.
+// Note: CUDA functions are per device.
+struct TORCH_CUDA_CU_API FusedKernelCUDA
+    : public ::torch::jit::fuser::FusedKernel {
+  FusedKernelCUDA(
+      at::DeviceIndex device,
+      std::string name,
+      std::string code,
+      std::vector<TensorDesc> input_desc,
+      std::vector<TensorDesc> output_desc,
+      std::vector<PartitionDesc> chunk_desc,
+      std::vector<PartitionDesc> concat_desc,
+      bool has_random);
+
+  ~FusedKernelCUDA() override;
+
+  void launch_raw(const uint32_t numel, std::vector<void*>& arguments)
+      const override;
+
+  at::Backend backend() const override {
+    return at::Backend::CUDA;
+  }
+
+ private:
+  static constexpr auto kBlockSize = 128;
+
+  // Note: per device to store device properties and compute launch heuristics
+  //  Acquiring these values at launch time would be too slow
+  at::DeviceIndex device_;
+  int maxBlocks_{};
+  cudaDeviceProp* prop_{};
+  std::vector<char> ptx_;
+  CUmodule module_{};
+  CUfunction function_{};
+};
+
+} // namespace torch::jit::fuser::cuda
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h
new file mode 100644
index 0000000000000000000000000000000000000000..b495251ad391d9d9a86ab1c8867e84dfd1f07f07
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h
@@ -0,0 +1,417 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/code_template.h>
+#include <torch/csrc/Export.h>
+
+namespace torch::jit::fuser::cuda {
+
+/*with type_as not checking type of its input, a fusion group can have non-fp32
+tensor as input. Correct code for this case is generated, however, nvrtc does
+not know how to handle int*_t integer types, so typedefs help it handle those
+cases*/
+
+static constexpr auto bfloat16_type_string = "__nv_bfloat16";
+
+#if defined(USE_ROCM) && ROCM_VERSION < 70000
+static auto type_declarations_template = at::jit::CodeTemplate(R"(
+${HalfHeader}
+${BFloat16Header}
+${RandHeader}
+
+#define NAN __int_as_float(0x7fffffff)
+#define POS_INFINITY __int_as_float(0x7f800000)
+#define NEG_INFINITY __int_as_float(0xff800000)
+
+typedef ${IndexType} IndexType;
+template<typename T, size_t N>
+struct TensorInfo {
+  T* data;
+  IndexType sizes[N];
+  IndexType strides[N];
+};
+template<typename T>
+struct TensorInfo<T, 0> {
+  T * data;
+};
+)");
+#else
+static auto type_declarations_template = at::jit::CodeTemplate(R"(
+typedef unsigned char uint8_t;
+typedef signed char int8_t;
+typedef short int  int16_t;
+typedef long long int int64_t;
+typedef unsigned long long int uint64_t;
+${HalfHeader}
+${BFloat16Header}
+${RandHeader}
+
+#define NAN __int_as_float(0x7fffffff)
+#define POS_INFINITY __int_as_float(0x7f800000)
+#define NEG_INFINITY __int_as_float(0xff800000)
+
+typedef ${IndexType} IndexType;
+template<typename T, size_t N>
+struct TensorInfo {
+  T* data;
+  IndexType sizes[N];
+  IndexType strides[N];
+};
+template<typename T>
+struct TensorInfo<T, 0> {
+  T * data;
+};
+)");
+#endif
+
+// We rewrite the code for philox RNG from curand as nvrtc couldn't resolve the
+// curand header correctly.
+constexpr auto rand_support_literal = R"(
+
+  class Philox {
+  public:
+    __device__ inline Philox(unsigned long long seed,
+                             unsigned long long subsequence,
+                             unsigned long long offset) {
+      key.x = (unsigned int)seed;
+      key.y = (unsigned int)(seed >> 32);
+      counter = make_uint4(0, 0, 0, 0);
+      counter.z = (unsigned int)(subsequence);
+      counter.w = (unsigned int)(subsequence >> 32);
+      STATE = 0;
+      incr_n(offset / 4);
+    }
+
+    __device__ inline unsigned long operator()() {
+      if(STATE == 0) {
+        uint4 counter_ = counter;
+        uint2 key_ = key;
+        for(int i = 0; i < 9; i++) {
+          counter_ = single_round(counter_, key_);
+          key_.x += (kPhilox10A); key_.y += (kPhilox10B);
+        }
+        output = single_round(counter_, key_);
+        incr();
+      }
+      unsigned long ret;
+      switch(STATE) {
+        case 0: ret = output.x; break;
+        case 1: ret = output.y; break;
+        case 2: ret = output.z; break;
+        case 3: ret = output.w; break;
+      }
+      STATE = (STATE + 1) % 4;
+      return ret;
+    }
+
+  private:
+    uint4 counter;
+    uint4 output;
+    uint2 key;
+    unsigned int STATE;
+    __device__ inline void incr_n(unsigned long long n) {
+      unsigned int nlo = (unsigned int)(n);
+      unsigned int nhi = (unsigned int)(n >> 32);
+      counter.x += nlo;
+      if (counter.x < nlo)
+        nhi++;
+      counter.y += nhi;
+      if (nhi <= counter.y)
+        return;
+      if (++counter.z)
+        return;
+      ++counter.w;
+    }
+    __device__ inline void incr() {
+      if (++counter.x)
+        return;
+      if (++counter.y)
+        return;
+      if (++counter.z)
+        return;
+      ++counter.w;
+    }
+    __device__ unsigned int mulhilo32(unsigned int a, unsigned int b,
+                                      unsigned int *result_high) {
+      *result_high = __umulhi(a, b);
+      return a*b;
+    }
+
+    __device__ inline uint4 single_round(uint4 ctr, uint2 key) {
+      unsigned int hi0;
+      unsigned int hi1;
+      unsigned int lo0 = mulhilo32(kPhiloxSA, ctr.x, &hi0);
+      unsigned int lo1 = mulhilo32(kPhiloxSB, ctr.z, &hi1);
+
+      uint4 ret = {hi1 ^ ctr.y ^ key.x, lo1, hi0 ^ ctr.w ^ key.y, lo0};
+      return ret;
+    }
+
+    static const unsigned long kPhilox10A = 0x9E3779B9;
+    static const unsigned long kPhilox10B = 0xBB67AE85;
+    static const unsigned long kPhiloxSA = 0xD2511F53;
+    static const unsigned long kPhiloxSB = 0xCD9E8D57;
+  };
+
+  // Inverse of 2^32.
+  #define M_RAN_INVM32 2.3283064e-10f
+  __device__  __inline__ float uniform(unsigned int x) {
+    return x * M_RAN_INVM32;
+  }
+)";
+
+constexpr auto rand_param =
+    ",unsigned long long seed, unsigned long long offset";
+
+constexpr auto rand_init = R"(
+  int idx = blockIdx.x*blockDim.x + threadIdx.x;
+  Philox rnd(seed, idx, offset);
+)";
+
+static auto cuda_compilation_unit_template = at::jit::CodeTemplate(R"(
+${type_declarations}
+
+extern "C" __global__
+void ${kernelName}(IndexType totalElements, ${formals} ${RandParam}) {
+  ${RandInit}
+  // check whether do vectorized load/store and allocate buffer
+  bool flag_vec4 = true;
+  ${tensorChecks}
+  if (flag_vec4) {
+    for (IndexType linearIndex = 4 * (blockIdx.x * blockDim.x + threadIdx.x);
+         linearIndex < totalElements;
+         linearIndex += 4 * gridDim.x * blockDim.x) {
+      // Convert `linearIndex` into an offset of tensor as it is:
+      ${tensorOffsets}
+      // load 4 at a time
+      ${kernelLoad}
+      #pragma unroll 4
+      for (int i=0; i<4; i++) {
+        // calculate the results
+        ${kernelBody_vec4}
+      }
+      // store 4 at a time
+      ${kernelStore}
+    }
+  } else {
+    for (IndexType linearIndex = blockIdx.x * blockDim.x + threadIdx.x;
+         linearIndex < totalElements;
+         linearIndex += gridDim.x * blockDim.x) {
+      // Convert `linearIndex` into an offset of tensor:
+      ${tensorOffsets}
+      // calculate the results
+      ${kernelBody}
+    }
+  }
+}
+)");
+
+// This snippet enables half support in the jit. Following the pattern for
+// reductions, fp16 input data is immediately upconverted to float
+// with __half2float(). All mathematical operations are done on float
+// values, and if needed the intermediate float representation is
+// converted to half with __float2half() when writing to a half tensor.
+#if defined(USE_ROCM)
+constexpr auto half_support_literal =
+    R"(
+typedef __half half;
+)";
+#else
+constexpr auto half_support_literal =
+    R"(
+#define __HALF_TO_US(var) *(reinterpret_cast<unsigned short *>(&(var)))
+#define __HALF_TO_CUS(var) *(reinterpret_cast<const unsigned short *>(&(var)))
+#if defined(__cplusplus)
+  struct __align__(2) __half {
+    __host__ __device__ __half() { }
+
+  protected:
+    unsigned short __x;
+  };
+
+  /* All intrinsic functions are only available to nvcc compilers */
+  #if defined(__CUDACC__)
+    /* Definitions of intrinsics */
+    __device__ __half __float2half(const float f) {
+      __half val;
+      asm("{  cvt.rn.f16.f32 %0, %1;}\n" : "=h"(__HALF_TO_US(val)) : "f"(f));
+      return val;
+    }
+
+    __device__ float __half2float(const __half h) {
+      float val;
+      asm("{  cvt.f32.f16 %0, %1;}\n" : "=f"(val) : "h"(__HALF_TO_CUS(h)));
+      return val;
+    }
+)"
+    // MSVC's preprocessor (but not the standard compiler) has a bug
+    // where it incorrectly tokenizes raw string literals, ending when it sees a
+    // " this causes the #endif in this string literal to be treated as a
+    // preprocessor token which, in turn, cause sccache on windows CI to fail.
+    // See https://godbolt.org/z/eVTIJq as an example.
+    // This workaround uses string-pasting to separate the " and the #endif into
+    // different strings
+    R"(
+  #endif /* defined(__CUDACC__) */
+#endif /* defined(__cplusplus) */
+#undef __HALF_TO_US
+#undef __HALF_TO_CUS
+
+typedef __half half;
+)";
+#endif
+
+#if defined(USE_ROCM)
+
+#if ROCM_VERSION >= 70000
+#define BF16_UINT32_DEF "typedef unsigned int uint32_t;\n"
+#else
+#define BF16_UINT32_DEF ""
+#endif
+
+constexpr auto bfloat16_support_literal =
+    R"(
+#ifndef __align__
+#define __align__(x) __attribute__((aligned(x)))
+#endif
+)" BF16_UINT32_DEF R"(
+typedef struct __align__(2) {
+  unsigned short x;
+}
+__nv_bfloat16_raw;
+
+#if defined(__cplusplus)
+struct __align__(2) __nv_bfloat16 {
+  __host__ __device__ __nv_bfloat16() {}
+
+  __host__ __device__ __nv_bfloat16& operator=(const __nv_bfloat16_raw& hr) {
+    __x = hr.x;
+    return *this;
+  }
+
+  unsigned short __x;
+};
+
+__device__ unsigned short __internal_float2bfloat16(
+    const float f,
+    unsigned int& sign,
+    unsigned int& remainder) {
+  unsigned int x;
+
+  x = __float_as_uint(f);
+
+  if ((x & 0x7fffffffU) > 0x7f800000U) {
+    sign = 0U;
+    remainder = 0U;
+    return static_cast<unsigned short>(0x7fffU);
+  }
+  sign = x >> 31;
+  remainder = x << 16;
+  return static_cast<unsigned short>(x >> 16);
+}
+
+/* Definitions of intrinsics */
+__device__ __nv_bfloat16 __float2bfloat16(const float a) {
+  __nv_bfloat16 val;
+  __nv_bfloat16_raw r;
+  unsigned int sign;
+  unsigned int remainder;
+  r.x = __internal_float2bfloat16(a, sign, remainder);
+  if ((remainder > 0x80000000U) ||
+      ((remainder == 0x80000000U) && ((r.x & 0x1U) != 0U))) {
+    r.x++;
+  }
+  val = r;
+  return val;
+}
+
+__device__ float __bfloat162float(const __nv_bfloat16 a) {
+  union
+  {
+      uint32_t int32;
+      float    fp32;
+  } u = {uint32_t(a.__x) << 16};
+  return u.fp32;
+}
+#endif /* defined(__cplusplus) */
+)";
+#else
+constexpr auto bfloat16_support_literal =
+    R"(
+#define __BFLOAT16_TO_US(var) *(reinterpret_cast<unsigned short*>(&(var)))
+#define __BFLOAT16_TO_CUS(var) \
+  *(reinterpret_cast<const unsigned short*>(&(var)))
+
+typedef struct __align__(2) {
+  unsigned short x;
+}
+__nv_bfloat16_raw;
+
+#if defined(__cplusplus)
+struct __align__(2) __nv_bfloat16 {
+  __host__ __device__ __nv_bfloat16() {}
+
+  __host__ __device__ __nv_bfloat16& operator=(const __nv_bfloat16_raw& hr) {
+    __x = hr.x;
+    return *this;
+  }
+
+ protected:
+  unsigned short __x;
+};
+
+#if defined(__CUDACC__)
+__device__ unsigned short __internal_float2bfloat16(
+    const float f,
+    unsigned int& sign,
+    unsigned int& remainder) {
+  unsigned int x;
+
+  x = __float_as_uint(f);
+
+  if ((x & 0x7fffffffU) > 0x7f800000U) {
+    sign = 0U;
+    remainder = 0U;
+    return static_cast<unsigned short>(0x7fffU);
+  }
+  sign = x >> 31;
+  remainder = x << 16;
+  return static_cast<unsigned short>(x >> 16);
+}
+
+/* Definitions of intrinsics */
+__device__ __nv_bfloat16 __float2bfloat16(const float a) {
+  __nv_bfloat16 val;
+#if __CUDA_ARCH__ >= 800
+  asm("{  cvt.rn.bf16.f32 %0, %1;}\n" : "=h"(__BFLOAT16_TO_US(val)) : "f"(a));
+#else
+  __nv_bfloat16_raw r;
+  unsigned int sign;
+  unsigned int remainder;
+  r.x = __internal_float2bfloat16(a, sign, remainder);
+  if ((remainder > 0x80000000U) ||
+      ((remainder == 0x80000000U) && ((r.x & 0x1U) != 0U))) {
+    r.x++;
+  }
+  val = r;
+#endif
+  return val;
+}
+
+__device__ float __bfloat162float(const __nv_bfloat16 a) {
+  float val;
+  asm("{ mov.b32 %0, {0,%1};}\n" : "=f"(val) : "h"(__BFLOAT16_TO_CUS(a)));
+  return val;
+}
+#endif /* defined(__CUDACC__) */
+#endif /* defined(__cplusplus) */
+#undef __BFLOAT16_TO_US
+#undef __BFLOAT16_TO_CUS
+)";
+#endif
+
+} // namespace torch::jit::fuser::cuda
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/executor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/executor.h
new file mode 100644
index 0000000000000000000000000000000000000000..aa6c691a0b807558cb4f5ce030dd156afb128a18
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/executor.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/fused_kernel.h>
+#include <torch/csrc/jit/codegen/fuser/kernel_spec.h>
+
+#include <cstdint>
+
+namespace torch::jit::fuser {
+
+// Runs the fusion associated with the key (see registerFusion() in interface.h)
+// on the inputs taken from the given Stack.
+TORCH_API bool runFusion(
+    const int64_t key,
+    Stack& stack,
+    std::string* code_out = nullptr);
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/fallback.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/fallback.h
new file mode 100644
index 0000000000000000000000000000000000000000..393127a6e99dc29fa41aecbb36d7e437e3a5bd51
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/fallback.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/stack.h>
+
+#include <cstdlib>
+
+namespace torch::jit::fuser {
+
+void runFallback(int64_t key, Stack& stack);
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/fused_kernel.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/fused_kernel.h
new file mode 100644
index 0000000000000000000000000000000000000000..ca2adedd7b196d1db9dddca4b3319c51e7060226
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/fused_kernel.h
@@ -0,0 +1,103 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Utils.h>
+#include <torch/csrc/jit/codegen/fuser/partition_desc.h>
+#include <torch/csrc/jit/codegen/fuser/tensor_desc.h>
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace torch::jit::fuser {
+
+struct FusedKernel {
+  AT_DISALLOW_COPY_AND_ASSIGN(FusedKernel);
+
+  FusedKernel(
+      std::string name,
+      std::string code,
+      std::vector<TensorDesc> input_desc,
+      std::vector<TensorDesc> output_desc,
+      std::vector<PartitionDesc> chunk_desc,
+      std::vector<PartitionDesc> concat_desc,
+      bool has_random)
+      : name_(std::move(name)),
+        code_(std::move(code)),
+        input_desc_(std::move(input_desc)),
+        output_desc_(std::move(output_desc)),
+        chunk_desc_(std::move(chunk_desc)),
+        concat_desc_(std::move(concat_desc)),
+        has_random_(has_random) {}
+
+  virtual ~FusedKernel() = default;
+
+  // arguments is a list of pointers to the arguments for the compiled CUDA/CPU
+  // code.
+  // The format of arguments is suitable for directly passing to a call to
+  // cuLaunchKernel as the kernel arguments.
+  // Currently the first argument is a pointer to numel (for passing to
+  // CUDA code), and the remainder are pointers to the TensorInfo<T> structs
+  // that compiled code uses to load Tensor data.
+  // launch_with_tensors handles packing at::Tensors into this arguments array.
+  // CPU code uses the same convention so that launch_with_tensors can be
+  // shared.
+  virtual void launch_raw(const uint32_t numel, std::vector<void*>& arguments)
+      const = 0;
+  virtual at::Backend backend() const = 0;
+
+  // Getters
+  const std::string& name() const {
+    return name_;
+  }
+  const std::string& code() const {
+    return code_;
+  }
+  const std::vector<TensorDesc>& inputDesc() const {
+    return input_desc_;
+  }
+  const std::vector<TensorDesc>& outputDesc() const {
+    return output_desc_;
+  }
+  const std::vector<PartitionDesc>& chunkDesc() const {
+    return chunk_desc_;
+  }
+  const std::vector<PartitionDesc>& concatDesc() const {
+    return concat_desc_;
+  }
+  bool hasRandom() const {
+    return has_random_;
+  }
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::string name_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::string code_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::vector<TensorDesc> input_desc_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::vector<TensorDesc> output_desc_;
+
+  // same size as input_desc, describes whether an
+  // input should be broken into subtensors (chunks)
+  // to be consumed by the fusion group
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::vector<PartitionDesc> chunk_desc_;
+
+  // same size as output_desc, describes whether
+  // an output is actually a concatenation of
+  // many subtensors that the fusion group produces
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const std::vector<PartitionDesc> concat_desc_;
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const bool has_random_;
+};
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/interface.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/interface.h
new file mode 100644
index 0000000000000000000000000000000000000000..516e192a0fb38a1c4af1cd56a9bff94509335347
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/interface.h
@@ -0,0 +1,59 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+namespace torch::jit {
+
+constexpr int kCPUDevice = -1;
+
+// Assigns a "key" to the given fusion_group that it can use to run its
+// fusion later (via runFusion() below).
+TORCH_API int64_t registerFusion(const Node* fusion_group);
+
+// Runs the fusion corresponding to the given key on the inputs
+// found on the stack. Outputs are placed on the same stack.
+// In some cases a fusion cannot be run and a fallback path where
+// PyTorch's interpreter runs the graph instead is attempted.
+TORCH_API void runFusion(const int64_t key, Stack& stack);
+
+// True if the respective devices can fuse, false otherwise
+TORCH_API bool canFuseOnCPU();
+TORCH_API bool canFuseOnGPU();
+
+// Sets whether fusion on the CPU is allowed (disabled by default due to
+// flakiness)
+TORCH_API void overrideCanFuseOnCPU(bool value);
+
+// Sets whether fusion on CPU must use LLVM Codegen and not SimplieIREval
+TORCH_API void overrideMustUseLLVMOnCPU(bool value);
+
+// Sets whether fusion on the GPU is allowed (enabled by default)
+TORCH_API void overrideCanFuseOnGPU(bool value);
+
+// Treats the given graph as a fusion group and launches it on the
+// specified device with the given inputs.
+// Returns the outputs.
+TORCH_API std::vector<at::Tensor> debugLaunchGraph(
+    Graph& graph,
+    at::ArrayRef<at::Tensor> inputs);
+
+// Treats the given graph as a fusion group and returns the generated code.
+TORCH_API std::string debugGetFusedKernelCode(
+    Graph& graph,
+    at::ArrayRef<at::Tensor> inputs);
+
+TORCH_API size_t nCompiledKernels();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/kernel_cache.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/kernel_cache.h
new file mode 100644
index 0000000000000000000000000000000000000000..d2446f6aa8af57c66c8eeef1c5198cf5199e966f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/kernel_cache.h
@@ -0,0 +1,38 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/kernel_spec.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <cstdint>
+#include <functional>
+#include <optional>
+
+namespace torch::jit::fuser {
+
+// A thread-safe cache interface.
+
+// Normalizes the graph by canonicalizing and erasing shape information
+TORCH_API std::shared_ptr<Graph> normalizeGraphForCache(
+    const std::shared_ptr<Graph>& graph);
+
+// Stores the given graph, returning the key used to access it
+TORCH_API int64_t store(std::shared_ptr<Graph> graph);
+
+// Given a graph, find a KernelSpec based on it
+TORCH_API std::optional<KernelSpec*> lookupGraph(
+    const std::shared_ptr<Graph>& graph);
+
+// Returns the graph corresponding to the given key (if it exists)
+TORCH_API std::optional<KernelSpec*> retrieve(const int64_t key);
+
+// Returns the size of the fusion key -> KernelSpec cache.
+// Only used for testing.
+TORCH_API int64_t debugNumCachedKernelSpecs();
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/kernel_spec.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/kernel_spec.h
new file mode 100644
index 0000000000000000000000000000000000000000..a84bcc7b3b7c18de8bda1dfc1905a44f77cead26
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/kernel_spec.h
@@ -0,0 +1,149 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/arg_spec.h>
+#include <torch/csrc/jit/codegen/fuser/fused_kernel.h>
+#include <torch/csrc/jit/codegen/fuser/interface.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <optional>
+
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+
+namespace torch::jit::fuser {
+
+// Helper struct containing partition information: the number of tensors
+// created and the dimension the partitioning is performed on.
+// Note: created during upfront compilation, once the tensors are known
+// at runtime the partition info is logically combined with the tensor
+// descriptions to create PartitionDesc objects.
+struct TORCH_API PartitionInfo {
+  PartitionInfo(const int64_t _nSubTensors, const int64_t _dim)
+      : nSubTensors_{_nSubTensors}, dim_{_dim} {}
+
+  int64_t nSubTensors() const {
+    return nSubTensors_;
+  }
+  int64_t dim() const {
+    return dim_;
+  }
+
+ private:
+  int64_t nSubTensors_;
+  int64_t dim_;
+};
+
+// "Kernel Specification." - Contains device-independent fusion information.
+// Each kernel specification contains a map of instantiated generated functions
+// that implement some or most of its functionality. Multiple generated
+// functions are needed by each abstract specification because of different
+// devices (cpu vs gpu, different gpus) and different inputs (int vs float,
+// contiguous vs discontiguous).
+// Note: uses a mutex to control access to its kernel store
+// Note: unordered containers do not invalidate references/pointers on
+//   rehashing, which is critical for thread-safety.
+// TODO: allow abstract kernels to use multiple generated kernels
+// TODO: allow abstract kernels to reuse generated kernels from common pool
+struct TORCH_API KernelSpec {
+  // Note: assumes the spec is a single block
+  // Note: This is the appropriate place to generalize if you want to add other
+  //  passes to upfront compilation that walk the graph.
+  KernelSpec(const int64_t _key, const std::shared_ptr<Graph>& _graph)
+      : key_{_key},
+        graph_{_graph},
+        code_{_graph, "<fused code>"},
+        nInputs_{_graph->inputs().size()}
+
+  {
+    // No need to iterate over reference since n is pointer
+    for (const auto n : graph_->nodes()) {
+      static_assert(std::is_pointer_v<decltype(n)>, "n must be a pointer");
+      if (n->kind() == aten::rand_like) {
+        has_random_ = true;
+        break;
+      }
+    }
+    nTensorInputs_ = std::count_if(
+        graph_->inputs().begin(), graph_->inputs().end(), [](const Value* v) {
+          return v->type()->isSubtypeOf(*TensorType::get());
+        });
+  }
+
+  // Getters
+  int64_t key() const {
+    return key_;
+  }
+  std::shared_ptr<Graph> graph() const {
+    return graph_;
+  }
+  const Code& code() const {
+    return code_;
+  }
+  int64_t nInputs() const {
+    return nInputs_;
+  }
+  int64_t nTensorInputs() const {
+    return nTensorInputs_;
+  }
+
+  std::vector<std::vector<int64_t>>& inputBroadcastGroups() {
+    return inputBroadcastGroups_;
+  }
+  const std::vector<std::vector<int64_t>>& inputBroadcastGroups() const {
+    return inputBroadcastGroups_;
+  }
+
+  std::vector<PartitionInfo>& inputChunks() {
+    return inputChunks_;
+  }
+  const std::vector<PartitionInfo>& inputChunks() const {
+    return inputChunks_;
+  }
+
+  bool hasRandom() const {
+    return has_random_;
+  }
+
+  // Cache functions
+  std::optional<std::shared_ptr<FusedKernel>> findKernel(
+      const ArgSpec& arg_spec) const {
+    std::lock_guard<std::mutex> guard{mutex_};
+    const auto it = kernels_.find(arg_spec);
+    if (it == kernels_.end())
+      return std::nullopt;
+    return it->second;
+  }
+  void cacheKernel(
+      const ArgSpec& arg_spec,
+      const std::shared_ptr<FusedKernel>& kernel) const {
+    std::lock_guard<std::mutex> guard{mutex_};
+    kernels_.emplace(arg_spec, kernel);
+  }
+
+ private:
+  int64_t key_;
+  std::shared_ptr<Graph> graph_;
+  Code code_;
+  uint64_t nInputs_;
+  uint64_t nTensorInputs_{};
+  std::vector<std::vector<int64_t>> inputBroadcastGroups_;
+  std::vector<PartitionInfo> inputChunks_;
+  bool has_random_{false};
+  mutable std::mutex mutex_;
+  mutable std::
+      unordered_map<ArgSpec, std::shared_ptr<FusedKernel>, c10::hash<ArgSpec>>
+          kernels_;
+};
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/partition_desc.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/partition_desc.h
new file mode 100644
index 0000000000000000000000000000000000000000..6eee194e80d162044d5065be72d9e2797df3db2f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/partition_desc.h
@@ -0,0 +1,63 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/Exception.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/codegen/fuser/tensor_desc.h>
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+namespace torch::jit::fuser {
+
+// Descriptor for chunk-ing an input tensor into subtensors
+// OR concat-ing an output tensor from subtensors
+// Note: default constructed used for tensors that do not participate in
+// chunk or cat operations.
+struct TORCH_API PartitionDesc {
+  PartitionDesc() : nSubTensors_{1}, dim_{0} {}
+
+  PartitionDesc(const TensorDesc& _desc, size_t _nSubTensors, size_t _dim)
+      : nSubTensors_{_nSubTensors}, dim_{_dim} {
+    AT_ASSERT(nSubTensors_ > 1);
+    std::vector<bool> cont = _desc.contiguity;
+    if (dim_ > 0) {
+      // when we narrow the concatenated output/chunked input
+      // we make the size[dim] smaller while keeping the stride[dim] the same,
+      // meaning: stride[dim - 1] != stride[dim]*size[dim]
+      // so dim - 1 is no longer contiguous
+      cont[dim_ - 1] = false;
+    }
+    subTensorDesc_ = std::make_shared<TensorDesc>(_desc.scalar_type, cont);
+  }
+
+  bool isNoop() const {
+    return (nSubTensors_ == 1);
+  }
+  size_t nSubTensors() const {
+    return nSubTensors_;
+  }
+  size_t dim() const {
+    return dim_;
+  }
+  std::shared_ptr<TensorDesc> subTensorDesc() {
+    return subTensorDesc_;
+  }
+  const std::shared_ptr<TensorDesc> subTensorDesc() const {
+    return subTensorDesc_;
+  }
+
+ private:
+  size_t nSubTensors_; // == 1 for tensors that should not be operated on via
+                       // chunk/cat
+  size_t dim_; // dimension along which the chunk/concat occurs
+  std::shared_ptr<TensorDesc>
+      subTensorDesc_; // descriptor for the subtensor, if it exists
+};
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/tensor_desc.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/tensor_desc.h
new file mode 100644
index 0000000000000000000000000000000000000000..0376875925a04b26615051b34d72ff4bf481898f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/tensor_desc.h
@@ -0,0 +1,103 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/jit_type.h>
+#include <c10/util/Exception.h>
+#include <c10/util/hash.h>
+#include <torch/csrc/Export.h>
+
+#include <algorithm>
+#include <ostream>
+#include <vector>
+
+namespace torch::jit::fuser {
+
+// type information needed by the compiler for input/outputs
+// contiguity[i] is true if the dim i is contiguous with dim i + 1.
+// contiguity.back() == true means strides.back() == 1.
+struct TORCH_API TensorDesc {
+  at::ScalarType scalar_type;
+  std::vector<bool> contiguity;
+
+  TensorDesc(const at::ScalarType& type, const std::vector<bool>& contiguity)
+      : scalar_type{type}, contiguity{contiguity} {
+    if (contiguity.empty()) {
+      nDim_ = 0;
+    } else {
+      nDim_ = std::count(contiguity.begin(), contiguity.end(), false) +
+          (lastIsContiguous() ? 1 : 0);
+    }
+  }
+
+  // Delegating constructors
+  TensorDesc(
+      const at::ScalarType& type,
+      const at::IntArrayRef& sizes,
+      const at::IntArrayRef& strides)
+      : TensorDesc(type, TensorDesc::findContiguous(sizes, strides)) {}
+
+  TensorDesc(const at::Tensor& t)
+      : TensorDesc(t.scalar_type(), t.sizes(), t.strides()) {}
+
+  TensorDesc(const c10::TensorTypePtr& type)
+      : TensorDesc(
+            type->scalarType().value(),
+            type->sizes().concrete_sizes().value(),
+            type->strides().concrete_sizes().value()) {}
+
+  // number of dimensions after contiguity compression
+  size_t nDim() const {
+    return nDim_;
+  }
+
+  // True iff innermost stride is 1
+  bool lastIsContiguous() const {
+    return (contiguity.empty() || contiguity.back());
+  }
+
+  static std::vector<bool> findContiguous(
+      const at::IntArrayRef& sizes,
+      const at::IntArrayRef& strides) {
+    AT_ASSERT(sizes.size() == strides.size());
+    std::vector<bool> cont(sizes.size());
+    for (size_t i = 0; i < sizes.size(); ++i) {
+      const auto expected_stride =
+          (i + 1 < sizes.size()) ? sizes[i + 1] * strides[i + 1] : 1;
+      cont[i] = (strides[i] == expected_stride);
+    }
+    return cont;
+  }
+
+  bool operator==(const TensorDesc& desc) const {
+    return scalar_type == desc.scalar_type && contiguity == desc.contiguity;
+  }
+
+  bool operator!=(const TensorDesc& desc) const {
+    return !(*this == desc);
+  }
+
+  static size_t hash(const TensorDesc& spec) {
+    return c10::get_hash(
+        spec.scalar_type,
+        spec.nDim_,
+        std::hash<std::vector<bool>>{}(spec.contiguity));
+  }
+
+ private:
+  size_t nDim_;
+};
+
+inline std::ostream& operator<<(std::ostream& out, const TensorDesc& d) {
+  out << d.scalar_type << '[';
+  for (const auto b : d.contiguity)
+    out << b << ';';
+  out << ']';
+  return out;
+}
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/tensor_info.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/tensor_info.h
new file mode 100644
index 0000000000000000000000000000000000000000..df2c1e12963bdaa58b56e688a3d6b950ee4e2f2d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/fuser/tensor_info.h
@@ -0,0 +1,29 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace torch::jit::fuser {
+
+// Host-side view of TensorInfo
+// Note dims[0] - we need to dynamically allocate the dims.
+struct TORCH_API TensorInfo {
+  uint32_t* sizes(size_t nDim) {
+    return &sizes_strides[0];
+  }
+  uint32_t* strides(size_t nDim) {
+    return &sizes_strides[nDim];
+  }
+
+  void* data;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+  uint32_t sizes_strides[0];
+};
+
+} // namespace torch::jit::fuser
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/LlgaTensorImpl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/LlgaTensorImpl.h
new file mode 100644
index 0000000000000000000000000000000000000000..72f267a4cffaffb73f8d10d2b1b129efe5cb159f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/LlgaTensorImpl.h
@@ -0,0 +1,277 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Config.h>
+
+#include <oneapi/dnnl/dnnl_graph.hpp>
+#include <torch/csrc/jit/ir/ir.h>
+#include <utility>
+
+namespace torch::jit::fuser::onednn {
+
+// Engine represents a device and its context. From the device kind, the engine
+// knows how to generate code for the target device and what kind of device
+// object to be expected. The device id ensures that there is a unique engine
+// being created for each device. The device handle passed from PyTorch allows
+// oneDNN Graph implementation to work on the device specified by PyTorch, which
+// is currently CPU, so we only have one engine.
+// Ref:
+// https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onednn/source/graph/programming_model#engine
+struct Engine {
+  // CPU engine singleton
+  static dnnl::engine& getEngine();
+  Engine(const Engine&) = delete;
+  void operator=(const Engine&) = delete;
+};
+
+// Stream is the logical abstraction for execution units. It is created on top
+// of oneDNN Graph engine. A compiled oneDNN Graph partition is submitted to a
+// stream for execution.
+struct Stream {
+  // CPU stream singleton
+  static dnnl::stream& getStream();
+  Stream(const Stream&) = delete;
+  void operator=(const Stream&) = delete;
+};
+
+struct LlgaTensorDesc {
+  using desc = dnnl::graph::logical_tensor;
+
+  LlgaTensorDesc(
+      size_t tid,
+      std::vector<int64_t> sizes,
+      std::vector<int64_t> strides,
+      desc::data_type dtype,
+      desc::property_type property_type)
+      : tid_(tid),
+        sizes_(std::move(sizes)),
+        strides_(std::move(strides)),
+        dtype_(dtype),
+        property_type_(property_type),
+        layout_type_(desc::layout_type::strided),
+        layout_id_(-1) {}
+
+  LlgaTensorDesc(const desc& t)
+      : tid_(t.get_id()),
+        sizes_(t.get_dims()),
+        strides_({-1}),
+        dtype_(t.get_data_type()),
+        property_type_(t.get_property_type()),
+        layout_type_(t.get_layout_type()),
+        layout_id_(-1) {
+    if (is_opaque()) {
+      layout_id_ = t.get_layout_id();
+    }
+    if (is_strided()) {
+      strides_ = t.get_strides();
+    }
+  }
+
+  LlgaTensorDesc(const torch::jit::Value* v)
+      : LlgaTensorDesc(
+            v->unique(),
+            {},
+            {},
+            desc::data_type::f32,
+            get_property_type(v)) {
+    if (v->type()->isSubtypeOf(TensorType::get())) {
+      auto tt = v->type()->cast<TensorType>();
+
+      if (tt->scalarType()) {
+        dtype_ = getLlgaDataType(tt->scalarType().value());
+      }
+
+      auto sizes = tt->sizes();
+      if (sizes.sizes()) {
+        for (auto d : *sizes.sizes()) {
+          sizes_.push_back(d.value_or(DNNL_GRAPH_UNKNOWN_DIM));
+        }
+      }
+
+      auto strides = tt->strides();
+      if (strides.sizes()) {
+        for (auto d : *strides.sizes()) {
+          strides_.push_back(d.value_or(DNNL_GRAPH_UNKNOWN_DIM));
+        }
+      }
+    }
+  }
+
+  LlgaTensorDesc supplementTensorInfo(const at::Tensor& t) const;
+
+  desc::data_type getLlgaDataType(at::ScalarType dt) const;
+
+  at::ScalarType aten_scalar_type() const;
+
+  const std::vector<int64_t>& sizes() const {
+    return sizes_;
+  }
+
+  const std::vector<int64_t>& strides() const {
+    TORCH_CHECK(!is_opaque(), "Cannot get strides on opaque layout");
+    return strides_;
+  }
+
+  size_t tid() const {
+    return tid_;
+  }
+
+  LlgaTensorDesc tid(uint64_t new_id) const {
+    auto ret = *this;
+    ret.tid_ = new_id;
+    return ret;
+  }
+
+  desc::data_type dtype() const {
+    return dtype_;
+  }
+
+  LlgaTensorDesc dtype(desc::data_type new_dtype) const {
+    return LlgaTensorDesc(tid_, sizes_, strides_, new_dtype, property_type_);
+  }
+
+  desc::layout_type layout_type() const {
+    return layout_type_;
+  }
+
+  LlgaTensorDesc layout_type(desc::layout_type new_layout_type) {
+    auto ret = *this;
+    ret.layout_type_ = new_layout_type;
+    return ret;
+  }
+
+  desc::property_type get_property_type(const torch::jit::Value* v) {
+    switch (v->node()->kind()) {
+      case prim::Constant:
+        return desc::property_type::constant;
+      default:
+        return desc::property_type::variable;
+    }
+  }
+
+  LlgaTensorDesc any() {
+    return layout_type(desc::layout_type::any);
+  }
+
+  size_t storage_size() const {
+    return logical_tensor().get_mem_size();
+  }
+
+  desc logical_tensor() const {
+    if (is_dimensionality_unknown()) {
+      return desc(
+          tid_, dtype_, DNNL_GRAPH_UNKNOWN_NDIMS, layout_type_, property_type_);
+    } else if (is_opaque()) {
+      return desc(tid_, dtype_, sizes_, layout_id_, property_type_);
+    } else if (is_any()) {
+      return desc(tid_, dtype_, sizes_, layout_type_, property_type_);
+    } else {
+      return desc(tid_, dtype_, sizes_, strides_, property_type_);
+    }
+  }
+
+  bool is_strided() const {
+    return layout_type_ == desc::layout_type::strided;
+  }
+
+  bool is_any() const {
+    return layout_type_ == desc::layout_type::any;
+  }
+
+  bool is_opaque() const {
+    return layout_type_ == desc::layout_type::opaque;
+  }
+
+  bool operator==(const LlgaTensorDesc& desc) const {
+    return tid_ == desc.tid_ && sizes_ == desc.sizes_ &&
+        dtype_ == desc.dtype_ && layout_type_ == desc.layout_type_ &&
+        ((is_opaque() && layout_id_ == desc.layout_id_) ||
+         strides_ == desc.strides_);
+  }
+
+  bool operator!=(const LlgaTensorDesc& desc) const {
+    return (tid_ != desc.tid_) || (sizes_ != desc.sizes_) ||
+        (dtype_ != desc.dtype_) || (layout_type_ != desc.layout_type_) ||
+        !((is_opaque() && (layout_id_ == desc.layout_id_)) ||
+          (strides_ == desc.strides_));
+  }
+
+  static size_t hash(const LlgaTensorDesc& desc) {
+    return c10::get_hash(
+        desc.tid_,
+        desc.sizes_,
+        desc.dtype_,
+        desc.layout_type_,
+        desc.layout_id_);
+  }
+
+  void set_compute_inplace() {
+    compute_inplace_ = true;
+  }
+
+  void set_input_tensor_index(size_t index) {
+    input_tensor_index_ = index;
+  }
+
+  bool reuses_input_tensor() {
+    return compute_inplace_;
+  }
+
+  size_t get_input_tensor_index() {
+    return input_tensor_index_;
+  }
+
+ private:
+  bool is_dimensionality_unknown() const {
+    return sizes_.empty();
+  }
+
+  size_t tid_;
+  std::vector<int64_t> sizes_;
+  std::vector<int64_t> strides_;
+  desc::data_type dtype_;
+  desc::property_type property_type_;
+  desc::layout_type layout_type_;
+  size_t layout_id_;
+  // If this is an output tensor, and querying the compiled partition would
+  // determine that this tensor would reuse its input tensor, then
+  // compute_inplace would be true, and input_tensor_index would be the index of
+  // the corresponding input tensor in inputSpecs_ of the LlgaKernel object.
+  bool compute_inplace_ = false;
+  size_t input_tensor_index_{};
+};
+
+// Initially, oneDNN Graph also used to have blocked layout for tensors between
+// partitions, and the LlgaTensorImpl wrapper helped us bypass guard checks.
+// oneDNN Graph has switched over to using strided tensors between partitions,
+// but this wrapper still helps us bypass guard checks because the strides of
+// tensors between partitions would be different from the ones the guard is
+// otherwise expecting.
+struct TORCH_API LlgaTensorImpl : public c10::TensorImpl {
+  LlgaTensorImpl(
+      at::Storage&& storage,
+      const caffe2::TypeMeta& data_type,
+      const LlgaTensorDesc& desc);
+
+  const LlgaTensorDesc& desc() const {
+    return desc_;
+  }
+
+  static at::Tensor llga_to_aten_tensor(LlgaTensorImpl* llgaImpl);
+
+ private:
+  LlgaTensorDesc desc_;
+};
+
+at::Tensor empty_llga(
+    const LlgaTensorDesc& desc,
+    const c10::TensorOptions& options);
+
+dnnl::graph::tensor llga_from_aten_tensor(const at::Tensor& tensor);
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/decompose_silu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/decompose_silu.h
new file mode 100644
index 0000000000000000000000000000000000000000..24d20864e42cd77fa0540f8eb6378962af11f52c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/decompose_silu.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+void DecomposeSiluForLLGA(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/defer_size_check.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/defer_size_check.h
new file mode 100644
index 0000000000000000000000000000000000000000..0bb55003e88bb6c9d7b484ed761540a12a48e99a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/defer_size_check.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+void DeferSizeCheck(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/graph_fuser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/graph_fuser.h
new file mode 100644
index 0000000000000000000000000000000000000000..8f14c5e33a9b3ac58392caaddb6b620c3444fc33
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/graph_fuser.h
@@ -0,0 +1,52 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/codegen/onednn/graph_helper.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+struct WorkBlock : public std::pair<Node*, Node*> {
+  using pair::pair;
+
+  Node* begin() {
+    return this->first;
+  }
+  Node* end() {
+    return this->second;
+  }
+};
+
+class GraphRewriter {
+ public:
+  GraphRewriter(Block* block, std::shared_ptr<Graph> graph, AliasDb& aliasDb)
+      : block_(block),
+        graph_(std::move(graph)),
+        aliasDb_(aliasDb),
+        llgaHelper_(graph_) {}
+
+  void cleanupSubgraphs();
+  void buildupSubgraphs();
+
+ private:
+  Block* block_;
+  std::shared_ptr<Graph> graph_;
+  AliasDb& aliasDb_;
+  LlgaGraphHelper llgaHelper_;
+  std::vector<WorkBlock> buildWorkBlocks();
+  std::pair<graph_node_list::iterator, bool> scanNode(
+      Node* consumer,
+      graph_node_list::iterator workblock_begin);
+  std::optional<Node*> tryMerge(Node* consumer, Node* producer);
+};
+
+// This pass creates the subgraphs for oneDNN Graph Fusion Nodes.
+// Its code-structure has been vastly inspired from
+// torch/csrc/jit/passes/create_autodiff_subgraphs.cpp
+void CreateLlgaSubgraphs(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/graph_helper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/graph_helper.h
new file mode 100644
index 0000000000000000000000000000000000000000..582db7e71cf4856f2710c8e5aefe19b4383f957f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/graph_helper.h
@@ -0,0 +1,103 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <oneapi/dnnl/dnnl_graph.hpp>
+#include <torch/csrc/jit/codegen/onednn/operator.h>
+#include <torch/csrc/jit/ir/alias_analysis.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+#define STRIDED_LAYOUT 0
+#define OPAQUE_LAYOUT 1
+
+struct OpPartitionMap {
+  void add(uint64_t opId, uint64_t partitionId) {
+    opmap_[opId] = partitionId;
+  }
+  void add(Node* n, uint64_t partitionId) {
+    add(Operator::getId(n), partitionId);
+  }
+  bool has(uint64_t opId) {
+    return opmap_.count(opId) > 0;
+  }
+  bool has(Node* n) {
+    return has(Operator::getId(n));
+  }
+  uint64_t get(uint64_t opId) {
+    return opmap_[opId];
+  }
+  uint64_t get(Node* n) {
+    auto opId = Operator::getId(n);
+    TORCH_CHECK(
+        has(opId),
+        "Node ",
+        n->kind().toQualString(),
+        " does not belong to any LLGA partition");
+    return get(opId);
+  }
+
+ private:
+  std::unordered_map<uint64_t, uint64_t> opmap_;
+};
+
+class LlgaGraphHelper {
+ public:
+  LlgaGraphHelper(
+      const std::shared_ptr<Graph>& graph,
+      dnnl::graph::partition::policy policy =
+          dnnl::graph::partition::policy::fusion);
+
+  bool shouldMerge(Node* toMerge, Node* subgraph);
+
+  bool shouldConsiderForMerge(Node* node);
+
+  bool checkForSingleOpPartition(Node* node);
+
+  Node* createSingletonSubgraph(Node* n, AliasDb& db);
+
+  void mergeNodeIntoSubgraph(Node* toMerge, Node* subgraphNode, AliasDb& db);
+
+  void unmergeIfAnyNodeIsMissing(Node* subgraphNode);
+
+  static bool isLlgaSubgraph(const Node* node);
+
+  Operator makeEltwiseOp(Node* node, dnnl::graph::op::kind kind);
+
+  Operator makeBinaryOp(Node* node, dnnl::graph::op::kind kind);
+
+  std::vector<dnnl::graph::partition> getPartitions() const;
+
+  std::map<size_t, Value*> getTensorIdToValue() const;
+
+  Operator createOperator(Node* node);
+
+ private:
+  size_t countSupportedOps(const std::shared_ptr<Graph>& graph) const;
+  std::unique_ptr<dnnl::graph::graph> dnnl_graph_ = nullptr;
+  std::unique_ptr<torch::jit::AliasDb> aliasDb_ = nullptr;
+  OpPartitionMap opToOwningPartition_;
+  std::vector<dnnl::graph::partition> partitions_;
+  std::map<size_t, Value*>
+      tensorIdToValue_; // map from tensorId to torch::jit::Value
+};
+
+class LlgaNodeWrapper {
+ public:
+  LlgaNodeWrapper(const Node* node);
+
+  void setOpaqueLayout(size_t offset);
+
+  bool useOpaqueLayout(size_t offset) const;
+
+  friend class LlgaGraphHelper;
+
+ private:
+  Node* n;
+};
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/guard_shape.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/guard_shape.h
new file mode 100644
index 0000000000000000000000000000000000000000..73ca360ff573d42805ce3d65f350d9f5f8c9433f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/guard_shape.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+void prepareFusionGroupAndGuardOutputs(Block* block);
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/interface.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/interface.h
new file mode 100644
index 0000000000000000000000000000000000000000..68cc22c7d582f3589ab429c621424d86d73ae30d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/interface.h
@@ -0,0 +1,63 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/Config.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/pass_manager.h>
+
+namespace torch::jit {
+namespace fuser::onednn {
+
+static std::atomic<bool> onednn_enabled{false};
+
+static std::atomic<bool>& getLlgaEnabled() {
+  return onednn_enabled;
+}
+
+C10_EXPORT void fuseGraph(std::shared_ptr<Graph>& g);
+
+} // namespace fuser::onednn
+
+struct C10_EXPORT RegisterLlgaFuseGraph
+    : public PassManager<RegisterLlgaFuseGraph> {
+  static bool setEnabled(bool enabled) {
+    TORCH_CHECK(
+        AT_MKLDNN_ENABLED(),
+        "Running oneDNN Graph fuser is only supported with MKLDNN builds.");
+    bool oldState = fuser::onednn::getLlgaEnabled();
+    fuser::onednn::getLlgaEnabled() = enabled;
+    if (enabled) {
+      registerPass(fuser::onednn::fuseGraph);
+    } else {
+      clearPass();
+    }
+    return oldState;
+  }
+
+  static bool isEnabled() {
+    return fuser::onednn::getLlgaEnabled();
+  }
+
+  // override PassManager::registerPass to register pre-pass
+  static bool registerPass(GraphPass p) {
+    if (!isRegistered()) {
+      passID(registerPrePass(std::move(p)), true);
+      isRegistered(true);
+      return false;
+    }
+    return true;
+  }
+
+  // override PassManager::clearPass to clear pre-pass
+  static void clearPass() {
+    if (isRegistered()) {
+      clearPrePass(passID());
+      isRegistered(true);
+    }
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/kernel.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/kernel.h
new file mode 100644
index 0000000000000000000000000000000000000000..f2b358fed94ea1697fb2c5be46f1fb30d42eca75
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/kernel.h
@@ -0,0 +1,94 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <unordered_map>
+
+#include <oneapi/dnnl/dnnl_graph.hpp>
+#include <torch/csrc/jit/codegen/onednn/LlgaTensorImpl.h>
+#include <torch/csrc/jit/codegen/onednn/graph_helper.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/runtime/interpreter.h>
+
+#include <c10/util/CallOnce.h>
+
+namespace torch::jit::fuser::onednn {
+
+using ArgSpec = LlgaTensorDesc;
+using ArgSpecs = std::vector<ArgSpec>;
+using RunArg = dnnl::graph::tensor;
+using RunArgs = std::vector<RunArg>;
+using TensorArgs = std::vector<at::Tensor>;
+
+class LlgaKernel {
+ public:
+  explicit LlgaKernel(const Node* fusionNode);
+
+  void run(Stack& stack);
+
+  void initialize(const TensorArgs& inputs);
+
+  const std::string& debugName() const {
+    return debugName_;
+  }
+
+ private:
+  bool useOpaqueLayout(size_t offset) const;
+
+  // PyTorch copy constants inside the subgraph instead of referencing them.
+  // Constants inputs to the partition are no longer in the graph->inputs().
+  // Need use the tid retrieved from the partition to find the missing
+  // constant inputs.
+  void initializeConstantInputs();
+
+  ArgSpecs initializeInputSpecs(const TensorArgs& inputs);
+
+  ArgSpecs initializeOutputSpecs() const;
+
+  dnnl::graph::compiled_partition compile(
+      const dnnl::graph::partition& partition);
+
+  std::map<size_t, int64_t> initializeTensorIdToOccurrence() const;
+
+  std::tuple<RunArgs, RunArgs> prepareRunArgs(
+      const TensorArgs& inputs,
+      TensorArgs& outputs) const;
+
+  static std::string genDebugName() {
+    static size_t debugId = 0;
+    return "LlgaPartition_" + std::to_string(debugId++);
+  }
+
+  static dnnl::graph::logical_tensor toLogicalTensor(const ArgSpec& s) {
+    return s.logical_tensor();
+  }
+
+  at::Device device_ = at::kCPU;
+  const Node* fusionNode_;
+  std::shared_ptr<Graph> graph_;
+  int64_t nGraphInputs_ = 0; // number of inputs to graph_ on the IR
+  int64_t nOutputs_ = 0;
+  std::map<size_t, Value*> tensorIdToValue_;
+  std::vector<int64_t> runArgsIdx_;
+  dnnl::graph::partition partition_;
+  // nPartitionInputs_ is the actual number of inputs to partition_ of graph_
+  // needed by the backend.
+  // nPartitionInputs_ = nGraphInputs_ + constantInputs_.size() since Constant
+  // inputs are copied to the inside of the subgraph
+  int64_t nPartitionInputs_;
+  dnnl::graph::compiled_partition compilation_;
+  std::set<size_t> initializedInputIds_;
+  std::vector<Value*> constantValues_;
+  TensorArgs constantInputs_;
+  ArgSpecs inputSpecs_;
+  ArgSpecs outputSpecs_;
+  std::vector<dnnl::graph::logical_tensor> constantLogicalTensors_;
+  std::string debugName_;
+  c10::once_flag initialized_flag;
+  bool is_initialized_ = false;
+};
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/layout_propagation.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/layout_propagation.h
new file mode 100644
index 0000000000000000000000000000000000000000..a654d8e7d15afb46e8c142d1f918ade6a1d20770
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/layout_propagation.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+void PropagateLayout(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/operator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/operator.h
new file mode 100644
index 0000000000000000000000000000000000000000..ab289941e48a7087b30ae65efb9e1da3be51baab
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/operator.h
@@ -0,0 +1,151 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <oneapi/dnnl/dnnl_graph.hpp>
+#include <torch/csrc/jit/codegen/onednn/LlgaTensorImpl.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+class Operator {
+ public:
+  Operator(const Node* node, dnnl::graph::op::kind kind)
+      : n(node), o(getId(node), kind, node->kind().toQualString()), k(kind) {}
+
+  // Returns output index if the Value is a graph output.
+  // Otherwise returns -1
+  int32_t graphOutputIdx(Value* v) {
+    int32_t i = 0;
+    for (const Value* output : v->owningGraph()->outputs()) {
+      if (v == output) {
+        return i;
+      }
+      i++;
+    }
+    return -1;
+  }
+
+  Operator& setInputValue(Value* v) {
+    if (v->mustNotBeNone()) {
+      if (v->type()->kind() == c10::TensorType::Kind) {
+        o.add_input(createLogicalTensor(v));
+      }
+    }
+    return *this;
+  }
+
+  Operator& setInput(size_t offset) {
+    return setInputValue(n->input(offset));
+  }
+
+  template <typename... Ts>
+  Operator& setInput(size_t offset, Ts... other) {
+    setInput(offset);
+    return setInput(other...);
+  }
+
+  Operator& setOutputValue(Value* v) {
+    if (v->mustNotBeNone()) {
+      o.add_output(createLogicalTensor(v));
+    }
+    return *this;
+  }
+
+  // setOutputValue & setOutput require a pointer to the LLGA graph, as output
+  // logical tensors that are graph outputs should be connected to an End LLGA
+  // op. A value of NULL can be provided for the graph pointer in order to
+  // maintain the legacy functionality of this function.
+  Operator& setOutputValue(Value* v, std::unique_ptr<dnnl::graph::graph>& g) {
+    if (v->mustNotBeNone()) {
+      auto output_tensor = createLogicalTensor(v);
+      o.add_output(output_tensor);
+      if (g) {
+        int32_t outputIndex = graphOutputIdx(v);
+        if (outputIndex != -1) {
+          dnnl::graph::op newEndNode(
+              LONG_MAX - outputIndex,
+              dnnl::graph::op::kind::End,
+              "EndNodeForGraphOutput");
+          newEndNode.add_input(output_tensor);
+          g->add_op(newEndNode);
+        }
+      }
+    }
+    return *this;
+  }
+
+  Operator& setOutput(std::unique_ptr<dnnl::graph::graph>& g, size_t offset) {
+    return setOutputValue(n->output(offset), g);
+  }
+
+  Operator& setOutput(size_t offset) {
+    return setOutputValue(n->output(offset));
+  }
+
+  template <typename... Ts>
+  Operator& setOutput(
+      std::unique_ptr<dnnl::graph::graph>& g,
+      size_t offset,
+      Ts... other) {
+    setOutput(g, offset);
+    return setOutput(g, other...);
+  }
+
+  template <typename Attr>
+  Operator& setAttr(dnnl::graph::op::attr name, Attr&& attr) {
+    o.set_attr(name, std::forward<Attr>(attr));
+    return *this;
+  }
+
+  template <typename F>
+  Operator& setAttr(dnnl::graph::op::attr name, const F& fn, size_t offset) {
+    return setAttr(name, fn(n, offset));
+  }
+
+  static float ScalarToFloat(const Node* node, size_t offset) {
+    return toIValue(node->input(offset))->toScalar().to<float>();
+  }
+
+  static std::vector<int64_t> Ints(const Node* node, size_t offset) {
+    return toIValue(node->input(offset))->toIntVector();
+  }
+
+  static int64_t Int(const Node* node, size_t offset) {
+    return toIValue(node->input(offset))->toInt();
+  }
+
+  static float Float(const Node* node, size_t offset) {
+    return static_cast<float>(toIValue(node->input(offset))->toDouble());
+  }
+
+  static bool Bool(const Node* node, size_t offset) {
+    return toIValue(node->input(offset))->toBool();
+  }
+
+  static uint64_t getId(const Node* node) {
+    return reinterpret_cast<uint64_t>(node); // cast node address as op id
+  }
+
+  dnnl::graph::op::kind kind() const {
+    return k;
+  }
+
+  dnnl::graph::op llgaOp() const {
+    return o;
+  }
+
+ private:
+  dnnl::graph::logical_tensor createLogicalTensor(Value* value) const {
+    return LlgaTensorDesc(value).logical_tensor();
+  }
+
+  const Node* n;
+  dnnl::graph::op o;
+  dnnl::graph::op::kind k;
+};
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/prepare_binary.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/prepare_binary.h
new file mode 100644
index 0000000000000000000000000000000000000000..7e46d4d447b4922e96b66dbbf32b8d011af7cbae
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/codegen/onednn/prepare_binary.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::fuser::onednn {
+
+// Prepare binary ops for LLGA
+//
+// The pass does the following:
+//
+// - Convert scalar input of aten::add and aten::mul into Float tensor with
+//   dimension [1]
+//
+// - Decompose fused add into aten::mul + aten::add when alpha != 1.0
+//
+// - Eliminate identity add/mul, i.e., tensor + 0, tensor * 1
+//
+void PrepareBinaryForLLGA(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit::fuser::onednn
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/cuda/cuda.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/cuda/cuda.h
new file mode 100644
index 0000000000000000000000000000000000000000..ff30c27553f68809ee4008c5c94c87d80c139042
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/cuda/cuda.h
@@ -0,0 +1,184 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <ATen/cuda/CUDAEvent.h>
+#include <c10/core/Device.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/custom_class.h>
+
+namespace torch::jit {
+
+class CUDAEvent;
+// This class is a wrapper around c10::cuda::CUDAStream.
+// It is needed because TorchBind does not support all of the argument types
+// for c10::cuda::CUDAStream. For more details, please refer to
+// c10/cuda/CUDAStream.h.
+class CUDAStream final : public CustomClassHolder {
+ public:
+  CUDAStream(
+      std::optional<c10::Device> device = std::nullopt,
+      int64_t priority = 0) {
+    c10::DeviceIndex device_index =
+        device.has_value() ? device->index() : c10::cuda::current_device();
+    stream_ = std::make_unique<c10::cuda::CUDAStream>(
+        c10::cuda::getStreamFromPool(static_cast<int>(priority), device_index));
+  }
+
+  CUDAStream(c10::cuda::CUDAStream s) {
+    stream_ = std::make_unique<c10::cuda::CUDAStream>(s);
+  }
+
+  bool query() {
+    return stream_->query();
+  }
+
+  c10::intrusive_ptr<CUDAEvent> recordEvent(
+      c10::intrusive_ptr<CUDAEvent> event);
+
+  void synchronize() {
+    stream_->synchronize();
+  }
+
+  void waitEvent(const c10::intrusive_ptr<CUDAEvent>& event);
+
+  void waitStream(const c10::intrusive_ptr<CUDAStream>& stream);
+
+  /// Get the CUDA device index that this stream is associated with.
+  int64_t device_index() const {
+    return stream_->device_index();
+  }
+
+  /// Get the full Device that this stream is associated with.  The Device
+  /// is guaranteed to be a CUDA device.
+  c10::Device device() const {
+    return stream_->device();
+  }
+
+  /// Return the stream ID corresponding to this particular stream.
+  int64_t id() const {
+    return stream_->id();
+  }
+
+ private:
+  std::unique_ptr<c10::cuda::CUDAStream> stream_;
+  friend class CUDAEvent;
+};
+
+// This class is a wrapper around at::cuda::CUDAStream.
+// It is needed because TorchBind does not support all of the argument types
+// for at::cuda::CUDAEvent. For more details, please refer to
+// aten/src/ATen/cuda/CUDAEvent.h.
+class CUDAEvent final : public CustomClassHolder {
+ public:
+  CUDAEvent(
+      bool enable_timing = false,
+      bool blocking = false,
+      bool interprocess = false) {
+    int flags = cudaEventDisableTiming;
+    if (enable_timing) {
+      flags = cudaEventDefault;
+    }
+    if (blocking) {
+      flags |= cudaEventBlockingSync;
+    }
+    if (interprocess) {
+      TORCH_CHECK(!enable_timing);
+      flags |= cudaEventInterprocess;
+    }
+
+    event_ = std::make_unique<at::cuda::CUDAEvent>(flags);
+  }
+
+  double elapsedTime(const c10::intrusive_ptr<CUDAEvent>& end) {
+    return event_->elapsed_time(*end->event_);
+  }
+
+  std::string ipcHandle() {
+    cudaIpcEventHandle_t handle{};
+    event_->ipc_handle(&handle);
+    std::string str_handle((const char*)&handle, sizeof(handle));
+    return str_handle;
+  }
+
+  bool query() {
+    return event_->query();
+  }
+
+  void record(const c10::intrusive_ptr<CUDAStream>& stream);
+
+  void synchronize() {
+    event_->synchronize();
+  }
+  void wait(const c10::intrusive_ptr<CUDAStream>& stream);
+
+ private:
+  void recordInternal(CUDAStream* stream);
+  std::unique_ptr<at::cuda::CUDAEvent> event_;
+
+  friend class CUDAStream;
+};
+
+inline c10::intrusive_ptr<CUDAEvent> CUDAStream::recordEvent(
+    c10::intrusive_ptr<CUDAEvent> event) {
+  if (!event) {
+    event = c10::make_intrusive<CUDAEvent>();
+  }
+
+  event->recordInternal(this);
+  return event;
+}
+
+inline void CUDAStream::waitEvent(const c10::intrusive_ptr<CUDAEvent>& event) {
+  event->event_->block(*stream_);
+}
+
+inline void CUDAStream::waitStream(
+    const c10::intrusive_ptr<CUDAStream>& stream) {
+  auto ev = c10::make_intrusive<CUDAEvent>();
+  stream->recordEvent(ev);
+  waitEvent(ev);
+}
+
+inline void CUDAEvent::record(const c10::intrusive_ptr<CUDAStream>& stream) {
+  event_->record(*stream->stream_);
+}
+
+inline void CUDAEvent::recordInternal(CUDAStream* stream) {
+  event_->record(*stream->stream_);
+}
+
+inline void CUDAEvent::wait(const c10::intrusive_ptr<CUDAStream>& stream) {
+  event_->block(*stream->stream_);
+}
+
+TORCH_LIBRARY(cuda, m) {
+  auto stream_class = m.class_<torch::jit::CUDAStream>("Stream").def(
+      torch::init<std::optional<c10::Device>, int64_t>(),
+      "",
+      {torch::arg("device") = std::nullopt, torch::arg("priority") = 0});
+  auto event_class = m.class_<torch::jit::CUDAEvent>("Event").def(
+      torch::init<bool, bool, bool>(),
+      "",
+      {torch::arg("enable_timing") = false,
+       torch::arg("blocking") = false,
+       torch::arg("interprocess") = false});
+
+  stream_class.def("query", &CUDAStream::query)
+      .def("record_event", &CUDAStream::recordEvent)
+      .def("synchronize", &CUDAStream::synchronize)
+      .def("wait_event", &CUDAStream::waitEvent)
+      .def("wait_stream", &CUDAStream::waitStream)
+      .def("device_index", &CUDAStream::device_index)
+      .def_property("device", &CUDAStream::device)
+      .def("id", &CUDAStream::id);
+
+  event_class.def("elapsed_time", &CUDAEvent::elapsedTime)
+      .def("query", &CUDAEvent::query)
+      .def("record", &CUDAEvent::record)
+      .def("synchronize", &CUDAEvent::synchronize)
+      .def("wait", &CUDAEvent::wait);
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/builtin_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/builtin_functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..ed412b4acd8d06928e751733ec40f109863fb149
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/builtin_functions.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/api/module.h>
+
+namespace torch::jit {
+
+TORCH_API const std::vector<Function*>& getAllBuiltinFunctionsFor(Symbol name);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/canonicalize_modified_loop.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/canonicalize_modified_loop.h
new file mode 100644
index 0000000000000000000000000000000000000000..3dc39e392a8954600332b8e79ddf0a0e6a7b6a2f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/canonicalize_modified_loop.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <memory>
+
+#include <torch/csrc/Export.h>
+
+namespace torch::jit {
+
+struct Graph;
+
+// Transforms loops so that they can be represented as python
+// for or while loops
+TORCH_API void CanonicalizeModifiedLoops(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/concrete_module_type.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/concrete_module_type.h
new file mode 100644
index 0000000000000000000000000000000000000000..65e0aff09acc68c5c93d6078d9ccc58bef79ad00
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/concrete_module_type.h
@@ -0,0 +1,244 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace torch::jit {
+
+enum class IterableModuleKind { NONE, LIST, DICT, PARAMLIST, PARAMDICT };
+class ConcreteModuleType;
+
+// You can think of an nn.Module as a template that corresponds to a family of
+// JIT types. The template "arguments" are things like the constant values.
+// e.g.
+//   class M(nn.Module):
+//        __constants__ = ["const"]
+//        ...
+//
+// Is similar to writing the following in C++:
+//
+//    template<TConst>
+//    class M {
+//       ...
+//    }
+//
+// We need to consider each different member of the type family a different JIT
+// type because, e.g. different constant values lead to different versions of
+// the same method.
+//
+// ConcreteModuleType corresponds to a single member of the type family, with
+// all template arguments fully specified. Two Modules that share a
+// ConcreteModuleType can share a JIT type, and vice versa.
+//
+// Why not just use a JIT type to represent concrete types? Because constants,
+// function attributes, etc. are currently not representable in the type system,
+// so this acts a non-first-class way of tracking concrete types.
+//
+// ConcreteModuleType is also the source of truth for servicing all
+// ModuleValue::attr calls. This is so we can guarantee that if two Module's
+// share a JIT type (and thus a ConcreteModuleType), then they behave the same
+// way when you access attributes on them.
+
+// ConcreteModuleType has two phases.
+// 1. Creation: First we build it up, during the ScriptModule conversion
+// process. This is represented by ConcreteModuleTypeBuilder.
+//    ...then the converter calls ConcreteModuleTypeBuilder::build(), producing
+//    a
+//       ConcreteModuleType ready for querying.
+// 2. Querying: We use ConcreteModuleType as a source of truth for
+// ModuleValue::attr calls during method compilation.
+
+// Represents a concrete type during in the process for construction. We use
+// this to decide whether we can share types between modules.
+class VISIBILITY_HIDDEN ConcreteModuleTypeBuilder {
+ public:
+  explicit ConcreteModuleTypeBuilder(py::object pyClass) {
+    TORCH_INTERNAL_ASSERT(pyClass);
+    pyClass_ = std::move(pyClass);
+  }
+
+  void addConstant(std::string name, py::object value);
+  void addConstant(std::string name, IValue value);
+  void addAttribute(
+      std::string name,
+      const TypePtr& type,
+      bool isParameter,
+      bool isBuffer);
+  void addFunctionAttribute(
+      std::string name,
+      const TypePtr& type,
+      py::object pyFunction);
+
+  void addModule(std::string name, std::shared_ptr<ConcreteModuleType> meta);
+
+  void addForwardHook(py::object hook);
+  void addForwardPreHook(py::object pre_hook);
+
+  void addOverload(
+      std::string methodName,
+      std::vector<std::string> overloadedMethodNames);
+  void addBuiltinFunction(std::string name, const std::string& symbol_name);
+  void addFailedAttribute(std::string name, std::string failureReason);
+  void addIgnoredAttribute(std::string name);
+  void setIterableModuleKind(IterableModuleKind kind);
+
+  // If a ConcreteModuleType is poisoned, it will never compare equal to any
+  // other concrete type
+  void setPoisoned();
+
+  std::shared_ptr<ConcreteModuleType> build() const {
+    return std::make_shared<ConcreteModuleType>(*this);
+  }
+
+  // This determines whether two modules can share a type. The container structs
+  // used by ConcreteModuleType have been defined such that operator==
+  // implements a meaningful comparison in that context.
+  bool equals(const ConcreteModuleTypeBuilder& other) const;
+
+  struct FunctionAttribute {
+    FunctionTypePtr function_;
+    py::object pyFunction_;
+
+    friend bool operator==(
+        const FunctionAttribute& lhs,
+        const FunctionAttribute& rhs) {
+      // Functions are not first class, so we can't do type comparison like a
+      // regular attribute. So we do a pointer equality check on the actual
+      // Python function object.
+      return lhs.pyFunction_.is(rhs.pyFunction_);
+    }
+  };
+
+  struct Attribute {
+    Attribute(TypePtr type, bool isParam, bool isBuffer)
+        : type_(std::move(type)), isParam_(isParam), isBuffer_(isBuffer) {}
+
+    friend bool operator==(const Attribute& lhs, const Attribute& rhs) {
+      return *(lhs.type_) == *(rhs.type_) && lhs.isParam_ == rhs.isParam_;
+    }
+    TypePtr type_;
+    bool isParam_;
+    bool isBuffer_;
+  };
+
+  struct ModuleInfo {
+    ModuleInfo(std::string name, std::shared_ptr<ConcreteModuleType> meta)
+        : name_(std::move(name)), meta_(std::move(meta)) {}
+
+    friend bool operator==(const ModuleInfo& lhs, const ModuleInfo& rhs);
+
+    std::string name_;
+    std::shared_ptr<ConcreteModuleType> meta_;
+  };
+
+ private:
+  ConcreteModuleTypeBuilder() = default;
+  ClassTypePtr createTypeFromThis() const;
+
+  // If true, this type will never compare equally to anything else. This is
+  // used if we want to ensure that this type is not shared (for example, if it
+  // came from a traced module)
+  bool isPoisoned_ = false;
+
+  // The value of any constants defined by the module.
+  std::unordered_map<std::string, IValue> constants_;
+  // The types of any attributes
+  OrderedDict<std::string, Attribute> attributes_;
+  // Overloads, in the same format as `__overloads__` in Python
+  std::unordered_map<std::string, std::vector<std::string>> overloads_;
+  // Any attributes we failed to convert to TorchScript, along with a hint as to
+  // why
+  std::unordered_map<std::string, std::string> failedAttributes_;
+  // Any attributes that were marked as ignored. They cannot be used in
+  // TorchScript but can still be used in ignored function in Python.
+  std::unordered_set<std::string> ignoredAttributes_;
+  // Any function attributes. These are special right now because functions are
+  // not first-class in the type system.
+  std::unordered_map<std::string, FunctionAttribute> functionAttributes_;
+  // Function attributes that are calls to builtin functions. These get
+  // de-sugared directly into the corresponding aten:: call. The map is
+  // attribute name -> aten symbol name
+  std::unordered_map<std::string, c10::Symbol> builtinFunctions_;
+  // The concrete types of any submodules
+  std::vector<ModuleInfo> modules_;
+  // Hooks to be called before/after forward when the module
+  // is called directly. Used to ensure modules have different types
+  // when they have different python hooks
+  // Actual hooks are added to ClassType directly during compilation
+  std::vector<py::object> forwardHooks_;
+  std::vector<py::object> forwardPreHooks_;
+
+  // If something is a ModuleDict/ModuleList, it means:
+  //   1. The order of the submodules matters for comparing the type
+  //   2. The compiler is allowed to treat it like a dict/tuple
+  IterableModuleKind iterableModuleKind_ = IterableModuleKind::NONE;
+
+  // The original `nn.Module` class that we derived this ScriptModule from.
+  py::object pyClass_;
+
+  // NOTE: If you ever add any more state to this struct, you need to make sure
+  // operator== still makes sense!
+  friend ConcreteModuleType;
+};
+
+// Represents a finalized concrete type, used to service ModuleValue::attr calls
+// during method compilation.
+class VISIBILITY_HIDDEN ConcreteModuleType {
+ public:
+  explicit ConcreteModuleType(ConcreteModuleTypeBuilder data);
+
+  static std::shared_ptr<ConcreteModuleType> fromJitType(TypePtr type);
+
+  TypePtr getJitType() const;
+  std::optional<py::object> getPyClass() const;
+  IterableModuleKind getIterableModuleKind() const;
+  std::optional<std::vector<std::string>> findOverloads(
+      const std::string& name) const;
+  std::optional<Function*> findFunctionAttribute(const std::string& name) const;
+  std::optional<c10::Symbol> findBuiltinFunction(const std::string& name) const;
+  std::shared_ptr<ConcreteModuleType> findSubmoduleConcreteType(
+      const std::string& name) const;
+  std::optional<std::string> findFailedAttribute(const std::string& name) const;
+  bool isIgnoredAttribute(const std::string& name) const;
+
+  // These getters are only here to return things as types that can be
+  // automatically converted by pybind.
+  std::unordered_map<std::string, py::object> getConstantsPy() const;
+  std::unordered_map<std::string, std::pair<TypePtr, bool>> getAttributesPy()
+      const;
+  std::vector<std::pair<std::string, std::shared_ptr<ConcreteModuleType>>>
+  getModulesPy() const;
+
+  bool equals(const ConcreteModuleType& other) const {
+    if (jitType_ == other.jitType_) {
+      // If the computed types are the same, these modules can (obviously) share
+      // a type.
+      return true;
+    }
+
+    return data_.equals(other.data_);
+  }
+  bool equals(const ConcreteModuleTypeBuilder& other) const {
+    return data_.equals(other);
+  }
+
+  void dump() const;
+
+ private:
+  ConcreteModuleType() = default;
+
+  // The JIT type derived from this ConcreteModuleType.
+  ConcreteModuleTypeBuilder data_;
+  TypePtr jitType_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/convert_to_ssa.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/convert_to_ssa.h
new file mode 100644
index 0000000000000000000000000000000000000000..d9a3677aa2ef21c2e8a2ee0f35778832f9eb49e8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/convert_to_ssa.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <functional>
+#include <memory>
+#include <string>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Convert a graph with Loads & Stores into SSA form
+TORCH_API void ConvertToSSA(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/edit_distance.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/edit_distance.h
new file mode 100644
index 0000000000000000000000000000000000000000..711ddee3edd41e1c3e6c0ccccf5ab08c9bb8568b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/edit_distance.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <cstddef>
+
+namespace torch::jit {
+
+TORCH_API size_t ComputeEditDistance(
+    const char* word1,
+    const char* word2,
+    size_t maxEditDistance);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/error_report.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/error_report.h
new file mode 100644
index 0000000000000000000000000000000000000000..041831252736664f16f0547d34e27e2fa6175942
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/error_report.h
@@ -0,0 +1,92 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/frontend/tree.h>
+#include <mutex>
+
+namespace torch::jit {
+
+struct Call {
+  std::string fn_name;
+  SourceRange caller_range;
+};
+
+struct TORCH_API ErrorReport : public std::exception {
+  ErrorReport(const ErrorReport& e);
+
+  explicit ErrorReport(const SourceRange& r);
+  explicit ErrorReport(const TreeRef& tree) : ErrorReport(tree->range()) {}
+  explicit ErrorReport(const Token& tok) : ErrorReport(tok.range) {}
+
+  const char* what() const noexcept override;
+
+  class TORCH_API Calls {
+   private:
+    std::vector<Call> calls_;
+    mutable std::mutex mutex_;
+
+   public:
+    void push_back(Call call) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      calls_.push_back(std::move(call));
+    }
+
+    void pop_back() {
+      std::lock_guard<std::mutex> lock(mutex_);
+      calls_.pop_back();
+    }
+
+    bool empty() const {
+      std::lock_guard<std::mutex> lock(mutex_);
+      return calls_.empty();
+    }
+
+    void update_pending_range(const SourceRange& range) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      calls_.back().caller_range = range;
+    }
+
+    std::vector<Call> get_stack() const {
+      std::lock_guard<std::mutex> lock(mutex_);
+      return calls_;
+    }
+  };
+
+  struct TORCH_API CallStack {
+    // These functions are used to report why a function was being compiled
+    // (i.e. what was the call stack of user functions at compilation time that
+    // led to this error)
+    CallStack(const std::string& name, const SourceRange& range);
+    ~CallStack();
+
+    // Change the range that is relevant for the current function (i.e. after
+    // each successful expression compilation, change it to the next expression)
+    static void update_pending_range(const SourceRange& range);
+
+   private:
+    std::shared_ptr<Calls> source_callstack_;
+  };
+
+  static std::string current_call_stack();
+
+ private:
+  template <typename T>
+  friend const ErrorReport& operator<<(const ErrorReport& e, const T& t);
+
+  mutable std::stringstream ss;
+  OwnedSourceRange context;
+  mutable std::string the_message;
+  std::vector<Call> error_stack;
+};
+
+template <typename T>
+const ErrorReport& operator<<(const ErrorReport& e, const T& t) {
+  e.ss << t;
+  return e;
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/exit_transforms.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/exit_transforms.h
new file mode 100644
index 0000000000000000000000000000000000000000..c33e7cb04d7c6cbc3e330f61b6af9ee90b72d8b2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/exit_transforms.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void TransformExits(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/function_schema_parser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/function_schema_parser.h
new file mode 100644
index 0000000000000000000000000000000000000000..2626cfb7f96042a67194589f4f5a47a214bd2113
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/function_schema_parser.h
@@ -0,0 +1,28 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/function_schema.h>
+#include <c10/macros/Macros.h>
+#include <string>
+#include <variant>
+
+namespace torch::jit {
+
+// allow_typevars: If true, we assume that lowercase types that we don't
+// understand are type variables. This is only needed for TorchScript (and not
+// not needed for custom ops).
+// If false, we disallow typevars, except in certain cases for BC reason (i.e.
+// your op is in the aten or prim namespace).
+TORCH_API std::variant<c10::OperatorName, c10::FunctionSchema> parseSchemaOrName(
+    const std::string& schemaOrName,
+    bool allow_typevars = true);
+TORCH_API c10::FunctionSchema parseSchema(
+    const std::string& schema,
+    bool allow_typevars = true);
+TORCH_API c10::OperatorName parseName(const std::string& name);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/inline_loop_condition.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/inline_loop_condition.h
new file mode 100644
index 0000000000000000000000000000000000000000..6dba54dc8a69cc489da51be53469538908f4b849
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/inline_loop_condition.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <functional>
+#include <memory>
+#include <string>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void InlineLoopCondition(std::shared_ptr<Graph>& graph);
+TORCH_API void InlineBlockBeforeNode(Node* before_node, Block* block);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/ir_emitter.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/ir_emitter.h
new file mode 100644
index 0000000000000000000000000000000000000000..4e723618bab2d7f129dcd3a40f11d00237f66d2d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/ir_emitter.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <functional>
+#include <memory>
+#include <string>
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/frontend/error_report.h>
+#include <torch/csrc/jit/frontend/resolver.h>
+#include <torch/csrc/jit/frontend/sugared_value.h>
+#include <torch/csrc/jit/frontend/tree_views.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void runCleanupPasses(std::shared_ptr<Graph>& to_clean);
+
+TORCH_API bool meaningfulName(const std::string& name);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/lexer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/lexer.h
new file mode 100644
index 0000000000000000000000000000000000000000..d99be3ca74a0c21f9b0b284c19f6a462b194c937
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/lexer.h
@@ -0,0 +1,568 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/macros/Macros.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/parser_constants.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/frontend/strtod.h>
+#include <algorithm>
+#include <clocale>
+#include <cstdlib>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <vector>
+
+namespace torch::jit {
+
+// single character tokens are just the character itself '+'
+// multi-character tokens need an entry here
+// if the third entry is not the empty string, it is used
+// in the lexer to match this token.
+
+// These kinds are also used in Tree.h as the kind of the AST node.
+// Some kinds TK_APPLY, TK_LIST are only used in the AST and are not seen in the
+// lexer.
+
+#define TC_FORALL_TOKEN_KINDS(_)                 \
+  _(TK_EOF, "eof", "")                           \
+  _(TK_WHITESPACE, "whitespace", "")             \
+  _(TK_WHITESPACE_EOF, "whitespace_eof", "")     \
+  _(TK_NUMBER, "number", "")                     \
+  _(TK_NEWLINE, "newline", "")                   \
+  _(TK_INDENT, "indent", "")                     \
+  _(TK_DEDENT, "dedent", "")                     \
+  _(TK_DEF, "def", "def")                        \
+  _(TK_EQUIVALENT, "equivalent", "<=>")          \
+  _(TK_IDENT, "ident", "")                       \
+  _(TK_STRING, "string", "")                     \
+  _(TK_STRINGLITERAL, "string_literal", "")      \
+  _(TK_CONST, "const", "")                       \
+  _(TK_LIST, "list", "")                         \
+  _(TK_DICT, "dict", "")                         \
+  _(TK_OPTION, "option", "")                     \
+  _(TK_APPLY, "apply", "")                       \
+  _(TK_COMPREHENSION, "comprehension", "")       \
+  _(TK_RANGE_CONSTRAINT, "range_constraint", "") \
+  _(TK_PARAM, "param", "")                       \
+  _(TK_INFERRED, "inferred", "")                 \
+  _(TK_ACCESS, "access", "")                     \
+  _(TK_ASSIGN, "assign", "")                     \
+  _(TK_AUG_ASSIGN, "aug_assign", "")             \
+  _(TK_ATTRIBUTE, "attribute", "")               \
+  _(TK_IF, "if", "if")                           \
+  _(TK_ELSE, "else", "else")                     \
+  _(TK_ELIF, "elif", "elif")                     \
+  _(TK_WHILE, "while", "while")                  \
+  _(TK_EXPR_STMT, "expression statement", "")    \
+  _(TK_RETURN, "return", "return")               \
+  _(TK_IS, "is", "is")                           \
+  _(TK_ISNOT, "is not", "is not")                \
+  _(TK_NE, "ne", "!=")                           \
+  _(TK_EQ, "eq", "==")                           \
+  _(TK_LE, "le", "<=")                           \
+  _(TK_GE, "ge", ">=")                           \
+  _(TK_FLOOR_DIV, "floordiv", "//")              \
+  _(TK_IF_EXPR, "if", "")                        \
+  _(TK_TRUE, "True", "True")                     \
+  _(TK_FALSE, "False", "False")                  \
+  _(TK_NONE, "None", "None")                     \
+  _(TK_AND, "and", "and")                        \
+  _(TK_OR, "or", "or")                           \
+  _(TK_NOT, "not", "not")                        \
+  _(TK_LSHIFT, "<<", "<<")                       \
+  _(TK_RSHIFT, ">>", ">>")                       \
+  _(TK_CAST, "cast", "")                         \
+  _(TK_PLUS_EQ, "+=", "+=")                      \
+  _(TK_MINUS_EQ, "-=", "-=")                     \
+  _(TK_TIMES_EQ, "*=", "*=")                     \
+  _(TK_DIV_EQ, "/=", "/=")                       \
+  _(TK_MOD_EQ, "%=", "%=")                       \
+  _(TK_BIT_OR_EQ, "|=", "|=")                    \
+  _(TK_BIT_AND_EQ, "&=", "&=")                   \
+  _(TK_BIT_XOR_EQ, "^=", "^=")                   \
+  _(TK_LSHIFT_EQ, "<<=", "<<=")                  \
+  _(TK_RSHIFT_EQ, ">>=", ">>=")                  \
+  _(TK_POW_EQ, "**=", "**=")                     \
+  _(TK_GLOBAL, "global", "global")               \
+  _(TK_BUILT_IN, "built-in", "")                 \
+  _(TK_SUBSCRIPT, "subscript", "")               \
+  _(TK_VAR, "variable", "")                      \
+  _(TK_NOTHING, "nothing", "")                   \
+  _(TK_DICT_LITERAL, "dict-literal", "")         \
+  _(TK_LIST_LITERAL, "list-literal", "")         \
+  _(TK_TUPLE_LITERAL, "tuple-literal", "")       \
+  _(TK_FOR, "for", "for")                        \
+  _(TK_IN, "in", "in")                           \
+  _(TK_NOTIN, "not in", "not in")                \
+  _(TK_STARRED, "starred", "")                   \
+  _(TK_UNARY_MINUS, "unary minus", "")           \
+  _(TK_POW, "pow operator", "**")                \
+  _(TK_ARROW, "arrow", "->")                     \
+  _(TK_DECL, "decl", "")                         \
+  _(TK_SLICE_EXPR, "slice expr", "")             \
+  _(TK_TYPE_COMMENT, "type comment", "# type:")  \
+  _(TK_RAISE, "raise", "raise")                  \
+  _(TK_ASSERT, "assert", "assert")               \
+  _(TK_DOTS, "dots", "...")                      \
+  _(TK_LIST_COMP, "list comprehension", "")      \
+  _(TK_DICT_COMP, "dict comprehension", "")      \
+  _(TK_BREAK, "break", "break")                  \
+  _(TK_CONTINUE, "continue", "continue")         \
+  _(TK_DELETE, "del", "del")                     \
+  _(TK_PASS, "pass", "pass")                     \
+  _(TK_CLASS_DEF, "class", "class")              \
+  _(TK_IMPORT, "import", "import")               \
+  _(TK_WITH, "with", "with")                     \
+  _(TK_WITH_ITEM, "withitem", "")                \
+  _(TK_AS, "as", "as")                           \
+  _(TK_PROP, "property", "")                     \
+  _(TK_ELLIPSIS, "Ellipsis", "Ellipsis")         \
+  _(TK_NONE_TYPE, "NoneType", "NoneType")
+
+enum TokenKind {
+  // we use characters to represent themselves so skip all valid characters
+  // before
+  // assigning enum values to multi-char tokens.
+  TK_DUMMY_START = 256,
+#define DEFINE_TOKEN(tok, _, _2) tok,
+  TC_FORALL_TOKEN_KINDS(DEFINE_TOKEN)
+#undef DEFINE_TOKEN
+};
+
+TORCH_API std::string kindToString(int kind);
+TORCH_API int stringToKind(const std::string& str);
+
+// nested hash tables that indicate char-by-char what is a valid token.
+struct TokenTrie;
+using TokenTrieRef = std::unique_ptr<TokenTrie>;
+struct TokenTrie {
+  TokenTrie() = default;
+  void insert(const char* str, int tok) {
+    if (*str == '\0') {
+      AT_ASSERT(kind == 0);
+      kind = tok;
+      return;
+    }
+
+    for (size_t i = 0, e = child_chars.size(); i < e; ++i) {
+      if (child_chars[i] == *str) {
+        child_tries[i]->insert(str + 1, tok);
+        return;
+      }
+    }
+
+    child_chars.emplace_back(*str);
+    child_tries.emplace_back(std::make_unique<TokenTrie>());
+    child_tries.back()->insert(str + 1, tok);
+  }
+  int kind{0}; // 0 == invalid token
+
+  std::vector<char> child_chars;
+  std::vector<TokenTrieRef> child_tries;
+};
+
+// stuff that is shared against all TC lexers/parsers and is initialized only
+// once.
+struct TORCH_API SharedParserData {
+  SharedParserData() : head(new TokenTrie()) {
+    for (const char* c = valid_single_char_tokens; *c; c++) {
+      std::string str(1, *c);
+      head->insert(str.c_str(), *c);
+    }
+
+#define ADD_CASE(tok, _, tokstring)   \
+  if (*(tokstring) != '\0') {         \
+    head->insert((tokstring), (tok)); \
+  }
+    TC_FORALL_TOKEN_KINDS(ADD_CASE)
+#undef ADD_CASE
+  }
+
+  bool match(
+      StringCordView::Iterator pos,
+      bool continuation, // are we inside a scope where newlines don't count
+                         // (e.g. inside parens)
+      bool whitespace_token, // should we treat whitespace as a token
+      int* kind,
+      StringCordView::Iterator* start,
+      StringCordView::Iterator* end) {
+    *start = pos;
+    // skip whitespace
+    while (pos.has_next() && isblank(*pos)) {
+      ++pos;
+    }
+
+    // special handling
+    if (pos.has_next()) {
+      if (*pos == '#' && !isTypeComment(pos)) {
+        // skip comments
+        while (pos.has_next() && *pos != '\n')
+          ++pos;
+        // tail call, handle whitespace and more comments
+        return match(pos, continuation, whitespace_token, kind, start, end);
+      }
+      if (*pos == '\\') {
+        auto newiter = pos;
+        ++newiter;
+        if (newiter.has_next() && *newiter == '\n' && !whitespace_token) {
+          ++newiter;
+          return match(newiter, continuation, false, kind, start, end);
+        }
+      }
+      if (*pos == '\n') {
+        return match(++pos, continuation, !continuation, kind, start, end);
+      }
+    }
+    // we handle white space before EOF because in the case we have something
+    // like the following where we need to generate the dedent token if foo:
+    //   ...
+    // else:
+    //   pass
+    if (whitespace_token) {
+      *kind = !pos.has_next() ? TK_WHITESPACE_EOF : TK_WHITESPACE;
+      *end = pos;
+      return true;
+    }
+    if (!pos.has_next()) {
+      *kind = TK_EOF;
+      *start = pos;
+      *end = *start;
+      return true;
+    }
+    // invariant: the next token is not whitespace or newline
+    *start = pos;
+    // check for a valid number
+    size_t len = 0;
+    if (isNumber(pos.rest_line(), 0, &len)) {
+      *end = *start;
+      *end += len;
+      *kind = TK_NUMBER;
+      return true;
+    }
+    // check for string
+    if (isString(pos.rest_line(), 0, &len)) {
+      *kind = TK_STRINGLITERAL;
+      *end = *start;
+      *end += len;
+      return true;
+    }
+
+    // check for either an ident or a token
+    // ident tracks whether what we have scanned so far could be an identifier
+    // matched indicates if we have found any match.
+    bool matched = false;
+    bool ident = true;
+    TokenTrie* cur = head.get();
+    // for (size_t i = 0; pos + i < str.size() && (ident || cur != nullptr);
+    // i++)
+    for (size_t i = 0; pos.has_next() && (ident || cur != nullptr);
+         ++pos, ++i) {
+      ident = ident && validIdent(i, *pos);
+      if (ident) {
+        matched = true;
+        *end = pos.next_iter();
+        *kind = TK_IDENT;
+      }
+      // check for token second, so that e.g. 'max' matches the token TK_MAX
+      // rather the
+      // identifier 'max'
+      if (cur) {
+        const auto begin_it = cur->child_chars.begin();
+        const auto end_it = cur->child_chars.end();
+        const auto ch_it = std::find(begin_it, end_it, *pos);
+
+        cur = (ch_it == end_it) ? nullptr
+                                : cur->child_tries[ch_it - begin_it].get();
+
+        if (cur && cur->kind != 0) {
+          matched = true;
+          *end = pos.next_iter();
+          *kind = cur->kind;
+        }
+      }
+    }
+    return matched;
+  }
+
+  bool isUnary(int kind, int* prec);
+  bool isBinary(int kind, int* prec);
+  bool isRightAssociative(int kind) {
+    switch (kind) {
+      case '?':
+      case TK_POW:
+      case TK_IF:
+        return true;
+      default:
+        return false;
+    }
+  }
+
+ private:
+  bool validIdent(size_t i, char n) {
+    return isalpha(n) || n == '_' || (i > 0 && isdigit(n));
+  }
+
+  // 1. skip whitespace
+  // 2. handle comment or newline
+  //
+  bool isNumber(std::string_view str, size_t start, size_t* len) {
+    char first = str[start];
+    // strtod allows numbers to start with + or - or nan or inf
+    // http://en.cppreference.com/w/cpp/string/byte/strtof
+    // but we want only the number part, otherwise 1+3 will turn into two
+    // adjacent numbers in the lexer
+    if (first == '-' || first == '+' || isalpha(first))
+      return false;
+    const char* startptr = str.data() + start;
+    char* endptr = nullptr;
+    torch::jit::strtod_c(startptr, &endptr);
+    *len = endptr - startptr;
+    // check if the number is complex valued
+    // access is safe because string is assumed to be null terminated
+    if (endptr != nullptr && *endptr == 'j') {
+      *len += 1;
+    }
+    return *len > 0;
+  }
+
+  bool isCharCount(char c, std::string_view str, size_t start, int len) {
+    // count checks from [start, start + len)
+    return start + len <= str.size() &&
+        std::count(str.begin() + start, str.begin() + start + len, c) == len;
+  }
+
+  // python concatenates all adjacent strings "a" "b" == "ab"
+  // strings can be enclosed with 1 or 3 single or double quotes
+  // if enclosed with 3 quotes newlines are valid
+  // as elsewhere, backslash and new line should be ignored
+  bool isString(std::string_view str, size_t start, size_t* len) {
+    char quote = str[start];
+    if (quote != '\"' && quote != '\'')
+      return false;
+    int quote_len = isCharCount(quote, str, start, 3) ? 3 : 1;
+
+    // end is now set past the opening quotation marks
+    size_t end = start + quote_len;
+    while (end < str.size() && !isCharCount(quote, str, end, quote_len)) {
+      if (str[end] == '\n' && quote_len != 3) {
+        return false;
+      }
+      // handle escaped characters. advances past escaped quotation marks,
+      // escaped newlines and escaped backslashes
+      // multi-char escapes like \x1A are handled fine here because the
+      // remainder of the escape are valid string characters anyway
+      if (str[end] == '\\') {
+        end++;
+      }
+      end++;
+    }
+    // set length equal to the complete string including quotations
+    *len = end - start + quote_len;
+    // if end finished without going past the last character of the string than
+    // there is a match
+    return end < str.size();
+  }
+
+  bool isblank(int n) {
+    return isspace(n) && n != '\n';
+  }
+
+  bool isTypeComment(StringCordView::Iterator str_iter) {
+    std::string_view rest_line = str_iter.rest_line();
+    const std::string type_string = "# type:";
+    if (rest_line.size() < type_string.length()) {
+      return false;
+    }
+    auto match_string = rest_line.substr(0, type_string.size());
+    return match_string == type_string;
+  }
+
+  // Make an exception ignoring comments for type annotation comments
+  bool isTypeComment(const StringCordView& str, size_t pos) {
+    const std::string type_string = "# type:";
+    if (str.size() < pos + type_string.length()) {
+      return false;
+    }
+    auto match_string = str.substr(pos, type_string.size());
+    return match_string == type_string;
+  }
+
+  TokenTrieRef head;
+};
+
+TORCH_API SharedParserData& sharedParserData();
+
+struct Token {
+  int kind;
+  SourceRange range;
+  Token(int kind, SourceRange range) : kind(kind), range(std::move(range)) {}
+  std::string text() const {
+    return std::string(range.token_text());
+  }
+
+  std::string_view text_view() const {
+    return range.token_text();
+  }
+
+  std::string kindString() const {
+    return kindToString(kind);
+  }
+};
+
+struct Lexer {
+  explicit Lexer(std::shared_ptr<Source> source)
+      : source(std::move(source)), shared(sharedParserData()) {
+    auto first_indent = lexRaw(true);
+    indent_stack.push_back(first_indent.range.size());
+    lex();
+  }
+  // Return the current token, and then move to the next one
+  Token next() {
+    if (next_tokens.empty())
+      reportError("Lexer invariant violated: empty token queue");
+    Token r = std::move(next_tokens.front());
+    next_tokens.erase(next_tokens.begin());
+    if (next_tokens.empty()) {
+      lex();
+    }
+    return r;
+  }
+  // Skip the current token if it matches the given kind
+  bool nextIf(int kind) {
+    if (cur().kind != kind)
+      return false;
+    next();
+    return true;
+  }
+
+  [[noreturn]] void reportError(const std::string& what) {
+    reportError(what, cur());
+  }
+  [[noreturn]] void reportError(const std::string& what, const Token& t) {
+    std::stringstream ss;
+    ss << what << ":\n";
+    t.range.highlight(ss);
+    throw std::runtime_error(ss.str());
+  }
+  [[noreturn]] void expected(const std::string& what, const Token& t) {
+    std::stringstream ss;
+    ss << "expected " << what << " but found '" << t.kindString()
+       << "' here:\n";
+    t.range.highlight(ss);
+    throw std::runtime_error(ss.str());
+  }
+  [[noreturn]] void expected(const std::string& what) {
+    expected(what, cur());
+  }
+  // Check that the current token has a given kind, return the current token,
+  // and advance to the next one.
+  Token expect(int kind) {
+    if (cur().kind != kind) {
+      expected(kindToString(kind));
+    }
+    return next();
+  }
+  Token& lookahead() {
+    if (next_tokens.size() < 2) {
+      lex();
+    }
+    return next_tokens[1];
+  }
+  Token& cur() {
+    return next_tokens.front();
+  }
+
+ private:
+  void lex() {
+    auto r = lexRaw();
+    switch (r.kind) {
+      case '(':
+      case '[':
+      case '{':
+        nesting++;
+        break;
+      case ')':
+      case ']':
+      case '}':
+        nesting--;
+        break;
+      case TK_WHITESPACE:
+      case TK_WHITESPACE_EOF: {
+        const auto depth =
+            r.kind == TK_WHITESPACE_EOF ? indent_stack.front() : r.range.size();
+        // note: TK_WHITESPACE_EOF is whitespace right before the EOF token
+        // just like we allow the code to be indented to a particular initial
+        // indent level, we allow the final indent to be anything and set
+        // it back to the initial indent level. This allows the code to be
+        // put into string literals inside code without worrying about final
+        // whitespace
+        if (depth > indent_stack.back()) {
+          indent_stack.push_back(depth);
+          r.kind = TK_INDENT;
+        } else if (depth == indent_stack.back()) {
+          r.kind = TK_NEWLINE;
+        } else {
+          next_tokens.emplace_back(TK_NEWLINE, r.range);
+          while (indent_stack.back() != depth) {
+            indent_stack.pop_back();
+            next_tokens.emplace_back(TK_DEDENT, r.range);
+            if (indent_stack.empty()) {
+              reportError("invalid indent level " + std::to_string(depth), r);
+            }
+          }
+          return; // We've already queued the tokens
+        }
+      } break;
+      default:
+        break;
+    }
+    next_tokens.push_back(std::move(r));
+  }
+  Token lexRaw(bool whitespace_token = false) {
+    AT_ASSERT(source);
+    if (current == nullptr) {
+      AT_ASSERT(pos == 0);
+      current = std::make_unique<StringCordView::Iterator>(
+          source->text_str().begin());
+    }
+
+    StringCordView::Iterator start_iter = *current;
+    StringCordView::Iterator end_iter = *current;
+    int kind = 0;
+    if (!shared.match(
+            *current,
+            nesting > 0,
+            whitespace_token,
+            &kind,
+            &start_iter,
+            &end_iter)) {
+      expected(
+          "a valid token",
+          Token(
+              **current,
+              SourceRange(source, start_iter, start_iter.pos() + 1)));
+    }
+
+    auto t = Token(kind, SourceRange(source, start_iter, end_iter.pos()));
+    pos = end_iter.pos();
+    *current = end_iter;
+    return t;
+  }
+
+  std::shared_ptr<Source> source;
+  std::unique_ptr<StringCordView::Iterator> current;
+  size_t pos{0};
+  size_t nesting{0}; // depth of ( [ { nesting...
+  std::vector<size_t> indent_stack; // stack of indentation level of blocks
+  // Invariant: this should always contain at least a single element
+  std::vector<Token> next_tokens;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  SharedParserData& shared;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/mini_environment.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/mini_environment.h
new file mode 100644
index 0000000000000000000000000000000000000000..dfb7e43fe4b371b620fe1c1fd99ae675e71936da
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/mini_environment.h
@@ -0,0 +1,60 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Simple data structure for containing a type T in nested control blocks
+// Should only be used after initial compilation where type checking and
+// loads and stores are emitted
+
+template <typename T>
+struct MiniEnvironment {
+  MiniEnvironment(Block* b, std::shared_ptr<MiniEnvironment> next = nullptr)
+      : next(std::move(next)) {}
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::shared_ptr<MiniEnvironment<T>> next;
+
+  T findInThisFrame(const std::string& name) {
+    auto it = table.find(name);
+    if (it != table.end()) {
+      return it->second;
+    }
+    return nullptr;
+  }
+
+  T findInAnyFrame(const std::string& name) {
+    for (auto runner = this; runner; runner = runner->next.get()) {
+      if (auto r = runner->findInThisFrame(name)) {
+        return r;
+      }
+    }
+    return nullptr;
+  }
+
+  void setVar(const std::string& name, T value) {
+    table[name] = value;
+  }
+
+  std::vector<std::string> definedVariables() {
+    std::vector<std::string> result;
+    result.reserve(table.size());
+    for (auto& kv : table) {
+      result.push_back(kv.first);
+    }
+    std::sort(result.begin(), result.end());
+    return result;
+  }
+
+ private:
+  std::unordered_map<std::string, T> table;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/name_mangler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/name_mangler.h
new file mode 100644
index 0000000000000000000000000000000000000000..2cede9aaaffb8a0fe6c376845796614b8c7cf408
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/name_mangler.h
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/qualified_name.h>
+#include <torch/csrc/Export.h>
+
+namespace torch::jit {
+
+/**
+ * class NameMangler
+ *
+ * Utility to mangle qualified names in order to make them unique. We use this
+ * in various places where we to de-duplicate qualified names.
+ */
+class TORCH_API NameMangler {
+ public:
+  // Given a qualified name, return a mangled version that is guaranteed to be
+  // unique with respect to previous/future calls of `mangled()` on this name
+  // mangler instance.
+  c10::QualifiedName mangle(const c10::QualifiedName& name);
+
+ private:
+  size_t mangleIndex_ = 0;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parse_string_literal.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parse_string_literal.h
new file mode 100644
index 0000000000000000000000000000000000000000..16f32df8d9b12fb47d3f117019380aa3f4344b52
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parse_string_literal.h
@@ -0,0 +1,92 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/frontend/error_report.h>
+#include <torch/csrc/jit/frontend/lexer.h>
+#include <optional>
+
+namespace torch::jit {
+
+inline bool isCharCount(char c, const std::string& str, size_t start, int len) {
+  // count checks from [start, start + len)
+  return start + len <= str.size() &&
+      std::count(
+          str.begin() + static_cast<ptrdiff_t>(start),
+          str.begin() + static_cast<ptrdiff_t>(start + len),
+          c) == len;
+}
+
+inline std::optional<char> parseOctal(const std::string& str, size_t pos) {
+  //\xxx where x are 0-7
+  if (pos + 3 >= str.size())
+    return std::nullopt;
+  size_t c = 0;
+  for (size_t i = 1, b = 64; i < 4; ++i, b /= 8) {
+    auto d = str[pos + i];
+    if (d < '0' || d > '7')
+      return std::nullopt;
+    c += b * (d - '0');
+  }
+  if (c >= 256)
+    return std::nullopt;
+  return c;
+}
+
+inline std::string parseStringLiteral(
+    const SourceRange& range,
+    const std::string& str) {
+  size_t quote_len = isCharCount(str[0], str, 0, 3) ? 3 : 1;
+  auto ret_str = str.substr(quote_len, str.size() - quote_len * 2);
+  size_t pos = ret_str.find('\\');
+  while (pos != std::string::npos) {
+    // invariant: pos has to escape a character because it is a valid string
+    char c = ret_str[pos + 1];
+    size_t to_erase = 2;
+    switch (ret_str[pos + 1]) {
+      case '\\':
+      case '\'':
+      case '\"':
+      case '\n':
+        break;
+      case 'a':
+        c = '\a';
+        break;
+      case 'b':
+        c = '\b';
+        break;
+      case 'f':
+        c = '\f';
+        break;
+      case 'n':
+        c = '\n';
+        break;
+      case 'v':
+        c = '\v';
+        break;
+      case 't':
+        c = '\t';
+        break;
+      case 'x':
+        throw(ErrorReport(range) << "unsupported hex specifier");
+      case 'u':
+      case 'U':
+        throw(ErrorReport(range) << "unsupported unicode specifier");
+      default:
+        // octal value in format \nnn, n is [0-7]
+        if (auto v = parseOctal(ret_str, pos)) {
+          to_erase = 4;
+          c = *v;
+        } else {
+          throw(ErrorReport(range) << " ill formed octal specifier");
+        }
+    }
+    ret_str.replace(pos, to_erase, /* num copies */ 1, c);
+    pos = ret_str.find('\\', pos + 1);
+  }
+  return ret_str;
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parser.h
new file mode 100644
index 0000000000000000000000000000000000000000..77653ee4c4ab4dc00da00851d41d59344768295b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parser.h
@@ -0,0 +1,36 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/tree.h>
+#include <torch/csrc/jit/frontend/tree_views.h>
+#include <memory>
+
+namespace torch::jit {
+
+struct Decl;
+struct ParserImpl;
+struct Lexer;
+
+TORCH_API Decl mergeTypesFromTypeComment(
+    const Decl& decl,
+    const Decl& type_annotation_decl,
+    bool is_method);
+
+struct TORCH_API Parser {
+  explicit Parser(const std::shared_ptr<Source>& src);
+  TreeRef parseFunction(bool is_method);
+  TreeRef parseClass();
+  Decl parseTypeComment();
+  Expr parseExp();
+  Lexer& lexer();
+  ~Parser();
+
+ private:
+  std::unique_ptr<ParserImpl> pImpl;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parser_constants.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parser_constants.h
new file mode 100644
index 0000000000000000000000000000000000000000..b71bc27928ebf8667f1b6df5b9d4e4596aa70a23
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/parser_constants.h
@@ -0,0 +1,11 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::jit {
+static constexpr const char* valid_single_char_tokens =
+    "+-*/%@()[]:,={}><.?!&^|~";
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/resolver.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/resolver.h
new file mode 100644
index 0000000000000000000000000000000000000000..6be0d3b08423d6260e015246e0a36cd521a0eeab
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/resolver.h
@@ -0,0 +1,71 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/jit_type.h>
+#include <ATen/core/qualified_name.h>
+#include <torch/csrc/jit/frontend/sugared_value.h>
+
+namespace torch::jit {
+
+struct Resolver;
+using ResolverPtr = std::shared_ptr<Resolver>;
+
+/**
+ * class Resolver
+ *
+ * Represents an "outer environment" in which we an look up names and return
+ * a corresponding SugaredValue. This is used during compilation to resolve
+ * references to names which are not defined internal to the graph.
+ *
+ * Example: PythonResolver looks at the enclosing Python scope for `name`.
+ *
+ * NOTE: When adding methods, keep this an abstract class (i.e. all new methods
+ * should be purely virtual). Resist the urge to provide a default
+ * implementation; you should explicitly think about how each resolver would
+ * handle the method.
+ */
+struct Resolver {
+  virtual ~Resolver() = default;
+
+  // Resolve a given name to a SugaredValue. This takes the method `m` that the
+  // caller is currently constructing, since we may need to insert nodes into
+  // the graph to create a value.
+  virtual std::shared_ptr<SugaredValue> resolveValue(
+      const std::string& name,
+      GraphFunction& m,
+      const SourceRange& loc) {
+    return nullptr;
+  }
+
+  // Resolve `name` to a TypePtr.
+  virtual TypePtr resolveType(const std::string& name, const SourceRange& loc) {
+    return nullptr;
+  }
+};
+
+// A resolver that only understands "torch.foo()" lookups.
+struct NativeResolver : public Resolver {
+  std::shared_ptr<SugaredValue> resolveValue(
+      const std::string& name,
+      GraphFunction& m,
+      const SourceRange& loc) override {
+    if (name == "torch") {
+      return std::make_shared<BuiltinModule>("aten");
+    }
+    return nullptr;
+  }
+
+  TypePtr resolveType(const std::string& name, const SourceRange& loc)
+      override {
+    return nullptr;
+  }
+};
+
+inline std::shared_ptr<NativeResolver> nativeResolver() {
+  return std::make_shared<NativeResolver>();
+}
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/schema_matching.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/schema_matching.h
new file mode 100644
index 0000000000000000000000000000000000000000..c84708353f64269d594eb45c5819c2e293bdfe3f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/schema_matching.h
@@ -0,0 +1,73 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/named_value.h>
+
+#include <ATen/core/function_schema.h>
+
+namespace torch::jit {
+
+// Try to match a list of inputs and keyword 'attributes' to this
+// schema. Return the flat list of positional inputs to the call or
+// `std::nullopt` on failure (`failure_messages` contains a good error
+// report in this case)
+
+struct MatchedSchema {
+  std::vector<Value*> inputs;
+  std::vector<TypePtr> return_types;
+  c10::OptNameList return_field_names;
+  std::string schema_name;
+};
+
+TORCH_API bool isBlockListedSchema(const FunctionSchema& schema);
+
+TORCH_API MatchedSchema matchSchema(
+    const ::c10::FunctionSchema& schema,
+    const SourceRange& loc,
+    Graph& graph,
+    at::ArrayRef<NamedValue> args,
+    at::ArrayRef<NamedValue> kwargs,
+    const std::optional<NamedValue>& self = std::nullopt);
+
+TORCH_API std::pair<size_t, MatchedSchema> matchSchemas(
+    const std::vector<const ::c10::FunctionSchema*>& schemas,
+    const SourceRange& loc,
+    Graph& graph,
+    at::ArrayRef<NamedValue> args,
+    at::ArrayRef<NamedValue> kwargs,
+    const std::optional<NamedValue>& self = std::nullopt,
+    bool render_errors = false);
+
+TORCH_API bool convertibleToList(
+    const TypePtr& type,
+    const TypePtr& list_type_);
+
+TORCH_API std::string getFullSchemaName(const ::c10::FunctionSchema& schema);
+
+TORCH_API Value* emitBuiltinCall(
+    const SourceRange& loc,
+    Graph& graph,
+    Symbol name,
+    at::ArrayRef<NamedValue> args,
+    at::ArrayRef<NamedValue> kwargs,
+    const std::optional<NamedValue>& self = std::nullopt);
+
+TORCH_API std::optional<size_t> findInputWithName(
+    const std::string& name,
+    at::ArrayRef<NamedValue> kwargs,
+    bool is_aten = false);
+
+// applies implicit conversion from value trying to turn it into type
+// concrete_type it succeeds if the return_value->isSubtypeOf(concrete_type)
+TORCH_API Value* tryConvertToType(
+    const SourceRange& loc,
+    Graph& graph,
+    const TypePtr& concrete_type,
+    Value* value,
+    bool allow_conversions);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/schema_type_parser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/schema_type_parser.h
new file mode 100644
index 0000000000000000000000000000000000000000..3d148fe323c4bb54f582590376631af45c446984
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/schema_type_parser.h
@@ -0,0 +1,53 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/alias_info.h>
+#include <ATen/core/jit_type.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/FunctionRef.h>
+#include <torch/csrc/jit/frontend/lexer.h>
+
+namespace torch::jit {
+
+using TypePtr = c10::TypePtr;
+
+TORCH_API void registerOpaqueType(const std::string& type_name);
+TORCH_API void unregisterOpaqueType(const std::string& type_name);
+TORCH_API bool isRegisteredOpaqueType(const std::string& type_name);
+
+struct TORCH_API SchemaTypeParser {
+  TypePtr parseBaseType();
+  std::optional<c10::AliasInfo> parseAliasAnnotation();
+  std::pair<TypePtr, std::optional<c10::AliasInfo>> parseType();
+  std::tuple</*fake*/ TypePtr, /*real*/ TypePtr, std::optional<c10::AliasInfo>>
+  parseFakeAndRealType();
+  std::optional<at::ScalarType> parseTensorDType(const std::string& dtype);
+  TypePtr parseRefinedTensor();
+
+  SchemaTypeParser(
+      Lexer& L,
+      bool parse_complete_tensor_types,
+      bool allow_typevars)
+      : complete_tensor_types(parse_complete_tensor_types),
+        L(L),
+        allow_typevars_(allow_typevars) {}
+
+ private:
+  std::optional<bool> tryToParseRequiresGrad();
+  std::optional<c10::Device> tryToParseDeviceType();
+  void parseList(
+      int begin,
+      int sep,
+      int end,
+      c10::function_ref<void()> callback);
+
+  bool complete_tensor_types;
+  Lexer& L;
+  size_t next_id = 0;
+  bool allow_typevars_;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/script_type_parser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/script_type_parser.h
new file mode 100644
index 0000000000000000000000000000000000000000..aa4105637e51d13b451484ebac382c6f6e32037c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/script_type_parser.h
@@ -0,0 +1,58 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/resolver.h>
+#include <torch/csrc/jit/frontend/tree_views.h>
+
+namespace torch::jit {
+
+/**
+ * class ScriptTypeParser
+ *
+ * Parses expressions in our typed AST format (TreeView) into types and
+ * typenames.
+ */
+class TORCH_API ScriptTypeParser {
+ public:
+  explicit ScriptTypeParser() = default;
+  explicit ScriptTypeParser(ResolverPtr resolver)
+      : resolver_(std::move(resolver)) {}
+
+  c10::TypePtr parseTypeFromExpr(const Expr& expr) const;
+
+  std::optional<std::pair<c10::TypePtr, int32_t>> parseBroadcastList(
+      const Expr& expr) const;
+
+  c10::TypePtr parseType(const std::string& str);
+
+  FunctionSchema parseSchemaFromDef(const Def& def, bool skip_self);
+
+  c10::IValue parseClassConstant(const Assign& assign);
+
+ private:
+  c10::TypePtr parseTypeFromExprImpl(const Expr& expr) const;
+
+  std::optional<std::string> parseBaseTypeName(const Expr& expr) const;
+  at::TypePtr subscriptToType(
+      const std::string& typeName,
+      const Subscript& subscript) const;
+  std::vector<IValue> evaluateDefaults(
+      const SourceRange& r,
+      const std::vector<Expr>& default_types,
+      const std::vector<Expr>& default_exprs);
+  std::vector<Argument> parseArgsFromDecl(const Decl& decl, bool skip_self);
+
+  std::vector<Argument> parseReturnFromDecl(const Decl& decl);
+
+  ResolverPtr resolver_ = nullptr;
+
+  // Need to use `evaluateDefaults` in serialization
+  friend struct ConstantTableValue;
+  friend struct SourceImporterImpl;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/source_range.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/source_range.h
new file mode 100644
index 0000000000000000000000000000000000000000..f58ab00fac5e62175b160977891f89e3f1d17351
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/source_range.h
@@ -0,0 +1,608 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Exception.h>
+#include <optional>
+
+#include <algorithm>
+#include <iterator>
+#include <memory>
+#include <ostream>
+#include <sstream>
+#include <unordered_map>
+
+namespace torch::jit {
+
+class SourceRangeUnpickler;
+struct SourceRange;
+
+// A stringlike class backed by a vector of string_view
+// the string represented are logically the concatenation of  the string_views
+// This has advantage of not needing continues memory.
+struct TORCH_API StringCordView {
+  StringCordView();
+  StringCordView(const StringCordView&) = default;
+  StringCordView(StringCordView&&) noexcept = default;
+  StringCordView(
+      std::vector<std::string_view> inputs,
+      std::vector<std::shared_ptr<std::string>> ownerships);
+
+  StringCordView& operator=(const StringCordView&) = default;
+  StringCordView& operator=(StringCordView&&) noexcept = default;
+
+  size_t size() const {
+    return accumulated_sizes_.back();
+  }
+
+  size_t find(const std::string& tok, size_t start) const;
+  size_t find_regex(const std::string& tok, size_t start) const;
+  StringCordView substr(size_t start, size_t size) const;
+
+  char at(size_t index) const {
+    return *iter_for_pos(index);
+  }
+  char operator[](size_t index) const {
+    return at(index);
+  }
+
+  std::string str() const {
+    std::stringstream ss;
+    for (auto s : pieces_) {
+      ss << std::string(s);
+    }
+    return ss.str();
+  }
+
+  bool operator==(const std::string& rhs) const;
+
+  bool operator==(const StringCordView& rhs) const;
+
+  std::string_view piece(size_t index) const {
+    return pieces_[index];
+  }
+
+  // General-case iterator implementation.
+  struct IteratorImpl {
+    IteratorImpl(
+        const StringCordView* str,
+        size_t start_line,
+        size_t start_pos,
+        size_t size)
+        : line_(start_line), pos_(start_pos), str_(str), size_(size) {}
+    explicit IteratorImpl(const StringCordView* str)
+        : IteratorImpl(str, 0, 0, str->size()) {}
+
+    IteratorImpl() : IteratorImpl(nullptr, 0, 0, 0) {}
+
+    IteratorImpl(const IteratorImpl&) = default;
+    IteratorImpl(IteratorImpl&&) = default;
+    IteratorImpl& operator=(const IteratorImpl&) = default;
+    IteratorImpl& operator=(IteratorImpl&&) = default;
+
+    IteratorImpl& operator++() {
+      if (size_ == 0) {
+        return *this;
+      }
+      if ((pos_ + 1) < str_->pieces_[line_].size()) {
+        pos_++;
+      } else {
+        line_++;
+        pos_ = 0;
+      }
+      return *this;
+    }
+
+    IteratorImpl operator++(int) {
+      IteratorImpl prev(*this);
+      ++(*this);
+      return prev;
+    }
+
+    IteratorImpl next_iter() const {
+      IteratorImpl next(*this);
+      ++next;
+      return next;
+    }
+
+    IteratorImpl& operator+=(size_t num);
+
+    IteratorImpl operator+(size_t num) const {
+      IteratorImpl it(*this);
+      it += num;
+      return it;
+    }
+
+    bool operator==(const IteratorImpl& rhs) const {
+      if (!has_next() && !rhs.has_next()) {
+        return true;
+      }
+      return (str_ == rhs.str_) && (line_ == rhs.line_) && (pos_ == rhs.pos_);
+    }
+
+    bool operator!=(const IteratorImpl& rhs) const {
+      return !((*this) == rhs);
+    }
+    bool has_next() const {
+      return size_ > 0 && (line_ < str_->pieces_.size());
+    }
+
+    char operator*() const {
+      TORCH_INTERNAL_ASSERT(line_ < str_->pieces_.size());
+      TORCH_INTERNAL_ASSERT(pos_ < str_->pieces_[line_].size());
+      return str_->pieces_[line_].at(pos_);
+    }
+
+    // returns rest of the line of the current iterator
+    std::string_view rest_line() const {
+      if (line_ >= str_->pieces_.size()) {
+        return "";
+      }
+
+      std::string_view cur_line = str_->pieces_[line_];
+      return cur_line.substr(pos_, std::string::npos);
+    }
+
+    size_t pos() const {
+      if (size_ == 0) {
+        return 0;
+      }
+      return str_->accumulated_sizes_[line_] + pos_;
+    }
+
+   private:
+    size_t line_;
+    size_t pos_;
+    const StringCordView* str_;
+    size_t size_;
+    friend struct StringCordView;
+  };
+
+  // Either an IteratorImpl, or a simple std::string_view::iterator
+  // (which is faster) if possible.
+  struct Iterator {
+    Iterator() = default;
+
+    Iterator(
+        const StringCordView* str,
+        size_t start_line,
+        size_t start_pos,
+        size_t size)
+        : repr_(
+              str->pieces_.size() == 1
+                  ? repr_type(FastRepr(
+                        start_line ? str->pieces_[0].end()
+                                   : str->pieces_[0].begin() + start_pos,
+                        str))
+                  : repr_type(IteratorImpl(str, start_line, start_pos, size))) {
+    }
+
+    Iterator(const StringCordView* str) : Iterator(str, 0, 0, str->size()) {}
+
+    Iterator& operator++() {
+      if (auto* pit = std::get_if<IteratorImpl>(&repr_)) {
+        ++(*pit);
+      } else {
+        ++fast_repr().it;
+      }
+      return *this;
+    }
+
+    Iterator operator++(int) {
+      Iterator prev(*this);
+      ++(*this);
+      return prev;
+    }
+
+    Iterator next_iter() const {
+      Iterator next(*this);
+      ++next;
+      return next;
+    }
+
+    Iterator& operator+=(size_t num) {
+      if (auto* pit = std::get_if<IteratorImpl>(&repr_)) {
+        *pit += num;
+      } else {
+        fast_repr().it += num;
+      }
+      return *this;
+    }
+
+    Iterator operator+(size_t num) const {
+      Iterator it(*this);
+      it += num;
+      return it;
+    }
+
+    bool operator==(const Iterator& rhs) const {
+      return repr_ == rhs.repr_;
+    }
+
+    bool operator!=(const Iterator& rhs) const {
+      return repr_ != rhs.repr_;
+    }
+
+    bool has_next() const {
+      if (const auto* pit = std::get_if<IteratorImpl>(&repr_)) {
+        return pit->has_next();
+      } else {
+        return fast_repr().it != fast_repr().str->pieces_[0].end();
+      }
+    }
+
+    char operator*() const {
+      if (const auto* pit = std::get_if<IteratorImpl>(&repr_)) {
+        return **pit;
+      } else {
+        return *fast_repr().it;
+      }
+    }
+
+    std::string_view rest_line() const {
+      if (const auto* pit = std::get_if<IteratorImpl>(&repr_)) {
+        return pit->rest_line();
+      } else {
+        // NOTE: std::string_view(it, end) ctor wasn't added until C++20.
+        const auto fast_repr_end = fast_repr().str->pieces_[0].end();
+        if (fast_repr().it != fast_repr_end) {
+          return std::string_view(
+              &*fast_repr().it, fast_repr_end - fast_repr().it);
+        }
+        return std::string_view();
+      }
+    }
+
+    size_t pos() const {
+      if (const auto* pit = std::get_if<IteratorImpl>(&repr_)) {
+        return pit->pos();
+      } else {
+        return fast_repr().it - fast_repr().str->pieces_[0].begin();
+      }
+    }
+
+   private:
+    // When we have only one entry in pieces_ (importantly, such as
+    // when called from torch::Library::def during startup), we can
+    // skip extra complexity and just use string_view::iterator
+    // directly.
+    struct FastRepr {
+      std::string_view::iterator it;
+      const StringCordView* str;
+
+      FastRepr() : str(nullptr) {}
+
+      explicit FastRepr(
+          std::string_view::iterator it_,
+          const StringCordView* str_)
+          : it(it_), str(str_) {}
+
+      bool operator==(const FastRepr& rhs) const {
+        return it == rhs.it && str == rhs.str;
+      }
+
+      bool operator!=(const FastRepr& rhs) const {
+        return !operator==(rhs);
+      }
+    };
+    using repr_type = std::variant<FastRepr, IteratorImpl>;
+    repr_type repr_;
+
+    FastRepr& fast_repr() {
+      // -Oz refuses to inline std::get.
+      TORCH_INTERNAL_ASSERT_DEBUG_ONLY(std::holds_alternative<FastRepr>(repr_));
+      return *std::get_if<FastRepr>(&repr_);
+    }
+
+    const FastRepr& fast_repr() const {
+      // -Oz refuses to inline std::get.
+      TORCH_INTERNAL_ASSERT_DEBUG_ONLY(std::holds_alternative<FastRepr>(repr_));
+      return *std::get_if<FastRepr>(&repr_);
+    }
+  };
+
+  Iterator begin() const {
+    return Iterator(this, 0, 0, size());
+  }
+  Iterator end() const {
+    return Iterator(this, pieces_.size(), 0, 0);
+  }
+  Iterator iter_for_pos(size_t pos) const;
+
+ private:
+  IteratorImpl begin_impl() const {
+    return IteratorImpl(this, 0, 0, size());
+  }
+  IteratorImpl end_impl() const {
+    return IteratorImpl(this, pieces_.size(), 0, 0);
+  }
+  IteratorImpl iter_impl_for_pos(size_t pos) const;
+  std::vector<std::string_view> pieces_;
+  std::vector<size_t> accumulated_sizes_;
+  std::vector<std::shared_ptr<std::string>> owned_strings_;
+};
+
+// Source represents a code segment. It keeps track of:
+//  - text_view : the view into text of the code segment
+//  - filename (optional) : if present, represents the name of the file from
+//                          which the code segment originated.
+//  - starting_line_no : represents the line in the original file where the
+//                       code segment started.
+struct TORCH_API Source {
+  // Whether or not Source should copy the string passed in the constructor.
+  enum CopiesString { COPIES_STRING, DONT_COPY };
+
+  explicit Source(
+      std::string_view text_view,
+      std::optional<std::string> filename = std::nullopt,
+      size_t starting_line_no = 0,
+      std::shared_ptr<SourceRangeUnpickler> gen_ranges = nullptr,
+      CopiesString copies_str = COPIES_STRING)
+      : text_view_(create_text_view(copies_str, text_view)),
+        filename_(std::move(filename)),
+        starting_line_no_(starting_line_no),
+        gen_ranges_(std::move(gen_ranges)) {
+    calc_line_start_offsets();
+  }
+
+  explicit Source(
+      StringCordView str,
+      std::optional<std::string> filename = std::nullopt,
+      size_t starting_line_no = 0,
+      std::shared_ptr<SourceRangeUnpickler> gen_ranges = nullptr)
+      : text_view_(std::move(str)),
+        filename_(std::move(filename)),
+        starting_line_no_(starting_line_no),
+        gen_ranges_(std::move(gen_ranges)) {
+    calc_line_start_offsets();
+  }
+  // Given a line number (within source_), return the byte offset of the
+  // beginning of that line.
+  size_t offset_for_line(size_t line) const {
+    return line_starting_offsets_.at(line);
+  }
+
+  // Returns number of lines present.
+  size_t num_lines() const {
+    return line_starting_offsets_.size();
+  }
+
+  // Calculate the line (within the code segment) on which `offset` resides.
+  size_t lineno_for_offset(size_t offset) const {
+    auto iter = std::upper_bound(
+        line_starting_offsets_.begin(), line_starting_offsets_.end(), offset);
+    return iter - line_starting_offsets_.begin() - 1;
+  }
+
+  // Calculate the line (within the original source file, if present) on which
+  // `lineno` resides.
+  size_t lineno_to_source_lineno(size_t lineno) const {
+    if (filename_) {
+      return lineno + starting_line_no_;
+    } else {
+      return lineno;
+    }
+  }
+
+  StringCordView get_line(size_t lineno) const {
+    auto start = offset_for_line(lineno);
+    auto size = (lineno + 1) < num_lines() ? offset_for_line(lineno + 1) - start
+                                           : text_view_.size() - start;
+    return text_view_.substr(start, size);
+  }
+
+  const StringCordView& text_str() const {
+    return text_view_;
+  }
+
+  char char_at(size_t index) const {
+    return text_view_.at(index);
+  }
+
+  size_t size() const {
+    return text_view_.size();
+  }
+
+  std::optional<std::string>& filename() {
+    return filename_;
+  }
+
+  size_t starting_line_no() const {
+    return starting_line_no_;
+  }
+
+  std::optional<SourceRange> findSourceRangeThatGenerated(
+      const SourceRange& range);
+
+  ~Source() = default;
+
+ private:
+  void calc_line_start_offsets() {
+    line_starting_offsets_.clear();
+    line_starting_offsets_.push_back(0);
+    size_t pos = 0;
+    while ((pos = text_view_.find("\n", pos)) != std::string::npos) {
+      line_starting_offsets_.push_back(++pos);
+    }
+  }
+
+  static StringCordView create_text_view(
+      CopiesString copies_str,
+      std::string_view text_view) {
+    if (copies_str == COPIES_STRING) {
+      auto allocated_str =
+          std::make_shared<std::string>(text_view.data(), text_view.size());
+      return StringCordView({*allocated_str}, {allocated_str});
+    } else {
+      return StringCordView({text_view}, {});
+    }
+  }
+
+  StringCordView text_view_;
+
+  std::optional<std::string> filename_;
+  // If filename_ is not present, starting_line_no_ is don't care
+  size_t starting_line_no_;
+  // Starting offsets for lines into the source. e.g. line 0 starts at
+  // line_starting_offsets_[0], etc.
+  std::vector<size_t> line_starting_offsets_;
+
+  std::shared_ptr<SourceRangeUnpickler> gen_ranges_;
+};
+
+// A SourceRange is a reference to subset of a Source, specified by `start` and
+// `end` byte offsets into the source text.
+struct TORCH_API SourceRange {
+  SourceRange(std::shared_ptr<Source> source_view, size_t start_, size_t end_)
+      : source_view_(std::move(source_view)), start_(start_), end_(end_) {
+    if (source_view_) {
+      start_iter_ = source_view_->text_str().iter_for_pos(start_);
+    }
+  }
+
+  SourceRange() : source_view_(nullptr), start_(0), end_(0) {}
+
+  SourceRange(
+      std::shared_ptr<Source> source_view_,
+      StringCordView::Iterator start_iter,
+      size_t end_)
+      : source_view_(std::move(source_view_)),
+        start_(start_iter.pos()),
+        end_(end_),
+        start_iter_(start_iter) {}
+
+  const std::string_view token_text() const {
+    size_t size = end() - start();
+    return start_iter_.rest_line().substr(0, size);
+  }
+
+  const StringCordView text() const {
+    return source_view_->text_str().substr(start(), end() - start());
+  }
+  size_t size() const {
+    return end() - start();
+  }
+  static const size_t CONTEXT = 3;
+  void highlight(std::ostream& out) const;
+
+  // Customizable version of 'highlight' method.
+  void print_with_context(
+      std::ostream& out,
+      size_t context,
+      bool highlight,
+      const std::string& funcname) const;
+
+  const std::shared_ptr<Source>& source() const {
+    return source_view_;
+  }
+  size_t start() const {
+    return start_;
+  }
+  size_t end() const {
+    return end_;
+  }
+  std::string str() const {
+    std::stringstream ss;
+    highlight(ss);
+    return ss.str();
+  }
+
+  std::optional<std::tuple<std::string, size_t, size_t>> file_line_col() const {
+    if (!source_view_ || !source()->filename()) {
+      return std::nullopt;
+    }
+
+    auto lineno = source_view_->lineno_for_offset(start_);
+    auto col_offset = (int)start_ - (int)source_view_->offset_for_line(lineno);
+    // TODO: std::optional<>::value returns an rvalue ref so can't use it here??
+    return std::make_tuple<std::string, size_t, size_t>(
+        source_view_->filename().value_or(""),
+        source_view_->lineno_to_source_lineno(lineno),
+        (size_t)col_offset);
+  }
+
+  bool operator==(const SourceRange& rhs) const {
+    return start() == rhs.start() && end() == rhs.end() &&
+        source() == rhs.source();
+  }
+
+  bool operator!=(const SourceRange& rhs) const {
+    return !(*this == rhs);
+  }
+
+  std::optional<SourceRange> findSourceRangeThatGenerated() const {
+    if (!source_view_) {
+      return std::nullopt;
+    }
+    return source_view_->findSourceRangeThatGenerated(*this);
+  }
+
+ protected:
+  std::shared_ptr<Source> source_view_;
+
+ private:
+  size_t start_;
+  size_t end_;
+  StringCordView::Iterator start_iter_;
+};
+
+// OwnedSourceRange is just like a SourceRange except that it owns a `Source`
+// instead of `Source`. Thus OwnedSourceRange owns a copy of source text.
+struct OwnedSourceRange : public SourceRange {
+  explicit OwnedSourceRange(const SourceRange& source_range)
+      : SourceRange(source_range) {
+    const auto& source = source_range.source();
+    if (source) {
+      source_view_ = std::make_shared<Source>(
+          source->text_str().str(),
+          source->filename(),
+          source->starting_line_no());
+    }
+  }
+};
+
+struct TORCH_API SourceRangeHasher {
+ public:
+  size_t operator()(const torch::jit::SourceRange& key) const;
+};
+
+struct StackEntry {
+  std::string filename;
+  SourceRange range;
+};
+
+TORCH_API void format_stack_trace(
+    std::ostream& out,
+    const std::vector<StackEntry>& entries);
+
+inline std::ostream& operator<<(std::ostream& out, const SourceRange& range) {
+  range.highlight(out);
+  return out;
+}
+
+// A pair of (byte offset, SourceRange) describing a specific segment
+// of the output stream
+struct TaggedRange {
+  TaggedRange(size_t bytes, SourceRange range)
+      : bytes(bytes), range(std::move(range)) {}
+  size_t bytes;
+  SourceRange range;
+};
+using SourceRangeRecords = std::vector<TaggedRange>;
+using SourceRangeTagMap =
+    std::unordered_map<SourceRange, int64_t, SourceRangeHasher>;
+
+} // namespace torch::jit
+
+namespace std {
+template <>
+struct iterator_traits<torch::jit::StringCordView::Iterator> {
+  using value_type = char;
+  using difference_type = ptrdiff_t;
+  using pointer = char*;
+  using reference = char&;
+  using iterator_category = std::forward_iterator_tag;
+};
+} // namespace std
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/source_ref.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/source_ref.h
new file mode 100644
index 0000000000000000000000000000000000000000..c786cc523c9aa42e6f15afd81290e1e6619d0074
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/source_ref.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <functional>
+#include <memory>
+
+#include <ATen/core/ivalue.h>
+#include <c10/macros/Export.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+
+namespace torch::jit {
+
+/**
+ * SourceRef does two things:
+ *   1. Owns a Source object.
+ *   2. Serves as lookup key to the owned Source in associative containers, for
+ *      runtime data aggregation.
+ * We don't want to use std::shared_ptr<Source> directly because we want to
+ * support heteogeneous lookup, and also shared_ptr is an implementation detail
+ * which should be encapsulated.
+ */
+class TORCH_API SourceRef : public CustomClassHolder {
+ public:
+  explicit SourceRef(std::shared_ptr<Source> source_view)
+      : source_view_(std::move(source_view)) {}
+  bool operator==(const SourceRef& other) const {
+    return source_view_ == other.source_view_;
+  }
+  bool operator<(const Source& other) const {
+    return source_view_.get() < &other;
+  }
+  friend bool operator<(const Source& other, const SourceRef& self) {
+    return &other < self.source_view_.get();
+  }
+  bool operator<(const SourceRef& other) const {
+    return *this < *other.source_view_;
+  }
+  const Source* operator->() const {
+    return source_view_.get();
+  }
+
+ private:
+  std::shared_ptr<Source> source_view_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/strtod.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/strtod.h
new file mode 100644
index 0000000000000000000000000000000000000000..fefb4c36e6fd2682401ab299aa8a4fa9fcb6122b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/strtod.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Macros.h>
+
+namespace torch::jit {
+
+TORCH_API double strtod_c(const char* nptr, char** endptr);
+TORCH_API float strtof_c(const char* nptr, char** endptr);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/sugared_value.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/sugared_value.h
new file mode 100644
index 0000000000000000000000000000000000000000..24e3bcf8cb3130700b604ec8d4b7f31993215b74
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/sugared_value.h
@@ -0,0 +1,880 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include <ATen/core/symbol.h>
+#include <caffe2/serialize/versions.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/frontend/error_report.h>
+#include <torch/csrc/jit/frontend/schema_matching.h>
+#include <torch/csrc/jit/frontend/versioned_symbols.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+using SugaredValuePtr = std::shared_ptr<SugaredValue>;
+
+// The AST can contain nodes like `self`, `self.b` or `python_fn` that
+// are not first-class values in the graph representation, but instead
+// will be desugared based on how they are used in the AST.
+
+// SugaredValue is used to temporarily represent these values in a way
+// that separates their behavior from the AST -> IR converter itself.
+// This allows us to keep dependencies on python minimal.
+
+struct TORCH_API SugaredValue
+    : public std::enable_shared_from_this<SugaredValue> {
+  // what is this node? for error reporting (e.g. Module, python function)
+  virtual std::string kind() const = 0;
+
+  // what can we do with this thing?
+  // use it as a value e.g.  `this + 4`
+  virtual Value* asValue(const SourceRange& loc, GraphFunction& m) {
+    throw(ErrorReport(loc) << kind() << " cannot be used as a value");
+  }
+
+  // select an attribute on it, e.g. `this.field`
+  virtual std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) {
+    throw(ErrorReport(loc) << "attribute lookup is not defined on " << kind());
+  }
+
+  virtual bool hasAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) {
+    throw(ErrorReport(loc) << "attribute lookup is not defined on " << kind());
+  }
+
+  // assign an attribute on it, e.g. `this.field = newValue`
+  virtual void setAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field,
+      Value* newValue) {
+    throw(
+        ErrorReport(loc) << "attribute assignment is not defined on "
+                         << kind());
+  }
+
+  // use it as a vector of values, e.g. a tuple of values as return value from
+  // a method invocation
+  virtual std::vector<std::shared_ptr<SugaredValue>> asTuple(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::optional<size_t>& size_hint = {}) {
+    throw(ErrorReport(loc) << kind() << " cannot be used as a tuple");
+  }
+
+  // TODO @wconstab refactor to use ModuleValue::asTuple instead of new API
+  virtual SugaredValuePtr asTupleValue(
+      const SourceRange& loc,
+      GraphFunction& m) {
+    throw(ErrorReport(loc) << kind() << " cannot be used as a tuplevalue");
+  }
+
+  virtual std::vector<std::shared_ptr<SugaredValue>> asType(
+      const SourceRange& loc,
+      Method& m) {
+    throw(ErrorReport(loc) << kind() << " cannot be used as a type");
+  }
+
+  // call it like a function, e.g. `outputs = this(inputs)`
+  virtual std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      // note: names for args will be 'argument 0', 'argument 1', etc..
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) {
+    // n_binders is always set to the number of variables an expression is
+    // syntactically bound to:
+    //     a = foo() # 1 binder (note in this case the single binder might be a
+    //     tuple) a, * b = foo() # 1 binder a, b = foo() # 2 binders foo() # 0
+    //     binders
+    //
+    // In subexpressions, like bar() in foo(bar()), n_binders is always set to
+    // 1. n_binders is used as a hint to subexpressions to determine how many
+    // values they should return when that number is ambiguous statically. In
+    // particular it is currently used to decide how many tensors a call to a
+    // python function will return. It is only a hint, functions do not have to
+    // check that n_binders match the number of things they are returning, the
+    // assignment logic will do that anyway.
+
+    throw(ErrorReport(loc) << "cannot call a " << kind());
+  }
+
+  // This function is called when to convert a SugaredValue to its iterator.
+  // For example, when iterating through a Dict we iterate over its keys
+  virtual std::shared_ptr<SugaredValue> iter(
+      const SourceRange& loc,
+      GraphFunction& m) {
+    throw(ErrorReport(loc) << kind() << " cannot be used as an iterable");
+  }
+
+  // If we are iterating over a Sugared Value and it returns a value from this
+  // function, then we emit an unrolled loop over the variable. This allows us
+  // to support containers of Heterogeneous types, like Module Containers &
+  // Tuples
+  virtual std::optional<int64_t> staticLen() {
+    return std::nullopt;
+  }
+
+  // When iterating over this SugaredValue, should we emit the for loop as an
+  // unrolled loop.
+  bool shouldEmitUnrolled() {
+    return staticLen() != std::nullopt;
+  }
+
+  // return length of this thing, if not then it can't be iterated.
+  // If it does not have a statically-determinable length, then it cannot
+  // be iterated over with a modulelist. If it does it must return a constant
+  // Value *
+  virtual Value* len(const SourceRange& loc, GraphFunction& m) {
+    throw(
+        ErrorReport(loc) << "'" << kind() << "'" << " object is not iterable");
+  }
+
+  // expression for ith element for iterable value
+  virtual std::shared_ptr<SugaredValue> getitem(
+      const SourceRange& loc,
+      GraphFunction& m,
+      Value* idx,
+      TypePtr type_hint = nullptr) {
+    throw(
+        ErrorReport(loc) << "'" << kind() << "'"
+                         << " object is not subscriptable");
+  }
+
+  virtual ~SugaredValue() = default;
+};
+
+// most things in the environment are just simple value types
+// and not special python syntax sugar types
+struct TORCH_API SimpleValue : public SugaredValue {
+  SimpleValue(Value* value) : value_(value) {}
+  std::string kind() const override {
+    std::stringstream ss;
+    // NOLINTNEXTLINE(clang-analyzer-core.CallAndMessage)
+    ss << "value of type '" << value_->type()->annotation_str() << "'";
+    return ss.str();
+  }
+  Value* asValue(const SourceRange& range, GraphFunction& m) override {
+    return value_;
+  }
+  std::vector<std::shared_ptr<SugaredValue>> asTuple(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::optional<size_t>& size_hint = {}) override;
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  bool hasAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  void setAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field,
+      Value* newValue) override;
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      // note: names for args will be 'argument 0', 'argument 1', etc..
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+  std::shared_ptr<SugaredValue> iter(const SourceRange& loc, GraphFunction& m)
+      override;
+
+  Value* getValue() const {
+    return value_;
+  }
+
+  Value* len(const SourceRange& loc, GraphFunction& m) override;
+  SugaredValuePtr getitem(
+      const SourceRange& loc,
+      GraphFunction& m,
+      Value* idx,
+      TypePtr type_hint = nullptr) override;
+
+ private:
+  Value* value_;
+};
+
+struct TORCH_API BuiltinFunction : public SugaredValue {
+  BuiltinFunction(Symbol symbol, std::optional<NamedValue> self)
+      : symbol(symbol), self(std::move(self)) {}
+
+  // The symbol of the function (e.g. `aten::relu`).
+  Symbol symbol;
+
+  // if this is method, then this is the self argument.
+  std::optional<NamedValue> self;
+  std::string kind() const override {
+    return "builtin";
+  }
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+  // try to create this builtin but if it doesn't exist or the self argument
+  // cannot possibly match, then return nullptr. Use in situations where it is
+  // not clear if it is a valid builtin
+  static std::shared_ptr<BuiltinFunction> tryCreate(
+      Symbol symbol,
+      std::optional<NamedValue> self);
+};
+
+struct TORCH_API SugaredTupleValue : public SugaredValue {
+  explicit SugaredTupleValue(std::vector<std::shared_ptr<SugaredValue>> tup)
+      : tup_(std::move(tup)) {}
+
+  std::vector<std::shared_ptr<SugaredValue>> asTuple(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::optional<size_t>& size_hint = {}) override {
+    return tup_;
+  }
+
+  Value* asValue(const SourceRange& loc, GraphFunction& m) override {
+    std::vector<Value*> vec;
+    vec.reserve(tup_.size());
+    for (const auto& sv : tup_) {
+      vec.push_back(sv->asValue(loc, m));
+    }
+    Graph& g = *m.graph();
+    return g.insertNode(g.createTuple(vec))->output();
+  }
+
+  std::string kind() const override {
+    return "Tuple";
+  }
+
+  SugaredValuePtr getitem(
+      const SourceRange& loc,
+      GraphFunction& m,
+      Value* idx,
+      TypePtr type_hint = nullptr) override {
+    if (!(idx->type()->cast<IntType>() && toIValue(idx))) {
+      throw(
+          ErrorReport(loc)
+          << "Expected integer literal for index but got a variable or non-integer. "
+          << "ModuleList/Sequential indexing is only supported with integer literals. "
+          << "For example, 'i = 4; self.layers[i](x)' will fail because i is not a literal. "
+          << "Enumeration is supported, e.g. 'for index, v in enumerate(self): out = v(inp)'");
+    }
+    auto index = toIValue(idx)->toInt();
+    int64_t adj_index =
+        (index < 0) ? index + static_cast<int64_t>(tup_.size()) : index;
+    if (!(adj_index >= 0 && adj_index < static_cast<int64_t>(tup_.size()))) {
+      throw(
+          ErrorReport(loc) << "Index " << index << " out of range of length "
+                           << tup_.size());
+    }
+    return tup_.at(adj_index);
+  }
+
+  // This function is called when a SugaredValue is used to convert a
+  // SugaredValue to its iterator. For example, when iterating through a Dict we
+  // iterate over its keys
+  std::shared_ptr<SugaredValue> iter(const SourceRange& loc, GraphFunction& m)
+      override {
+    return shared_from_this();
+  }
+
+  // Because this is used to contain SugaredValues of Heterogeneous types,
+  // we define staticLen() so that when this is iterated over it is emitted
+  // as an unrolled loop.
+  std::optional<int64_t> staticLen() override {
+    return static_cast<int64_t>(tup_.size());
+  }
+
+  std::vector<std::shared_ptr<SugaredValue>> tup_;
+};
+
+struct TORCH_API BuiltinModule : public SugaredValue {
+  BuiltinModule(std::string name, std::optional<int64_t> version = std::nullopt)
+      : name(std::move(name)), version(version) {}
+
+  std::string kind() const override {
+    return "builtin module";
+  }
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override {
+    if (field == "autograd") {
+      // When referring torch.autograd, it is also considered to be a
+      // BuiltinModule and we will dispatch to the aten operators for the
+      // methods under its module.
+      return std::make_shared<BuiltinModule>("aten", version);
+    }
+
+    auto sym = Symbol::fromQualString(name + "::" + field);
+    return std::make_shared<BuiltinFunction>(sym, std::nullopt);
+  }
+
+ private:
+  std::string name;
+  // when we add operator versioning, emit this op as it existing at 'version'
+  // if not set, use the latest version
+  std::optional<int64_t> version;
+};
+
+// Represents a class, analogous to `int` or `dict`. Instances of classes,
+// like `1` or `{"foo": 5}`, are represented as SimpleValues
+struct TORCH_API ClassValue : public SugaredValue {
+  explicit ClassValue(ClassTypePtr type) : type_(std::move(type)) {}
+
+  // Call the type's constructor, as in:
+  //    n = Foo(constructor_arg)
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  std::string kind() const override {
+    return type_->str();
+  }
+
+  ClassTypePtr type_;
+};
+
+struct TORCH_API NamedTupleConstructor : public SugaredValue {
+  explicit NamedTupleConstructor(TupleTypePtr type) : type_(std::move(type)) {}
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+  std::string kind() const override {
+    return type_->str();
+  }
+
+  TupleTypePtr type_;
+};
+
+struct FunctionValue : public SugaredValue {
+  FunctionValue(Function* callee) : callees_({callee}) {}
+  FunctionValue(const StrongFunctionPtr& p)
+      : callees_({p.function_}), cu_(p.cu_) {}
+  FunctionValue(const std::vector<StrongFunctionPtr>& callees) {
+    for (const StrongFunctionPtr& callee : callees) {
+      cu_ = cu_ ? cu_ : callee.cu_;
+      TORCH_INTERNAL_ASSERT(callee.cu_ == cu_);
+      callees_.push_back(callee.function_);
+    }
+  }
+
+  std::string kind() const override {
+    return "function";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& f,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override {
+    std::vector<const FunctionSchema*> schemas;
+    for (Function* callee : callees_) {
+      try {
+        callee->ensure_defined();
+      } catch (const RecursiveMethodCallError&) {
+        throw(
+            ErrorReport(loc)
+            << " function '" << callee->name() << "' is called recursively. "
+            << "Recursive calls are not supported");
+      }
+      schemas.push_back(&callee->getSchema());
+    }
+    auto match = matchSchemas(schemas, loc, *f.graph(), args, kwargs);
+    Value* output =
+        f.graph()->insertFunctionCall(callees_[match.first], match.second);
+    output->node()->setSourceRange(loc);
+    return std::make_shared<SimpleValue>(output);
+  }
+
+  const std::vector<Function*>& callees() {
+    return callees_;
+  }
+
+ private:
+  std::vector<Function*> callees_;
+  // TODO holding this thing is creepy
+  std::shared_ptr<CompilationUnit> cu_;
+};
+
+struct TORCH_API ClosureValue : public SugaredValue {
+  ClosureValue(Value* value) : value_(value) {
+    TORCH_INTERNAL_ASSERT(value_->node()->kind() == prim::Closure);
+  }
+  std::string kind() const override {
+    return "closure";
+  }
+  Value* asValue(const SourceRange& range, GraphFunction& m) override {
+    return value_;
+  }
+  Value* value_;
+};
+
+// defines how a method obtained from a module/class/interface behaves in script
+struct MethodValue : public SugaredValue {
+  MethodValue(Value* self, std::vector<std::string> method_names)
+      : self_(self), method_names_(std::move(method_names)) {}
+  MethodValue(Value* self, std::string method_name)
+      : MethodValue(self, std::vector<std::string>({std::move(method_name)})) {}
+
+  std::string kind() const override {
+    return "method";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& f,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override {
+    std::vector<NamedValue> argsWithSelf = {self_};
+    argsWithSelf.insert(argsWithSelf.end(), args.begin(), args.end());
+    std::vector<const FunctionSchema*> schemas;
+    for (const std::string& method_name : method_names_) {
+      if (auto class_type = self_->type()->cast<ClassType>()) {
+        Function& method = class_type->getMethod(method_name);
+        try {
+          method.ensure_defined();
+        } catch (const RecursiveMethodCallError&) {
+          throw(
+              ErrorReport(loc)
+              << " method '" << method.name() << "' is called recursively. "
+              << "Recursive calls are not supported");
+        }
+        schemas.push_back(&method.getSchema());
+      } else if (auto interface_type = self_->type()->cast<InterfaceType>()) {
+        schemas.push_back(interface_type->getMethod(method_name));
+      } else {
+        TORCH_INTERNAL_ASSERT(
+            false, "method constructed that is not a class or interface");
+      }
+    }
+    auto match = matchSchemas(schemas, loc, *f.graph(), argsWithSelf, kwargs);
+    Value* output =
+        f.graph()->insertMethodCall(method_names_[match.first], match.second);
+    output->node()->setSourceRange(loc);
+    return std::make_shared<SimpleValue>(output);
+  }
+
+ private:
+  Value* self_;
+  std::vector<std::string> method_names_;
+};
+
+struct TORCH_API PrintValue : public SugaredValue {
+  std::string kind() const override {
+    return "print";
+  }
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+};
+
+// expressions like int(x)
+// these are the same as call prim::Int or equivalent except it
+// is a noop when the input is a subtype of 'type'
+struct TORCH_API CastValue : public BuiltinFunction {
+  CastValue(TypePtr type, c10::Symbol method)
+      : BuiltinFunction(method, std::nullopt), type_(std::move(type)) {}
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override {
+    if (args.size() == 1 && kwargs.empty()) {
+      auto len_op = std::make_shared<BuiltinFunction>(aten::len, std::nullopt);
+      auto gt_op = std::make_shared<BuiltinFunction>(aten::gt, std::nullopt);
+      auto zero = m.graph()->insertConstant(0);
+
+      auto v = args[0].value(*m.graph());
+      if (v->type()->isSubtypeOf(*type_)) {
+        return std::make_shared<SimpleValue>(v);
+      } else if (
+          *type_ == *BoolType::get() &&
+          (v->type()->isSubtypeOf(*AnyListType::get()) ||
+           v->type()->isSubtypeOf(*StringType::get()) ||
+           v->type()->cast<DictType>())) {
+        auto len = len_op->call(loc, m, {v}, {}, 1);
+        return gt_op->call(loc, m, {len->asValue(loc, m), zero}, {}, 1);
+      }
+    }
+    return BuiltinFunction::call(loc, m, args, kwargs, n_binders);
+  }
+
+ private:
+  TypePtr type_;
+};
+
+struct TORCH_API TensorCastValue : public SugaredValue {
+  TensorCastValue(at::ScalarType type, NamedValue self)
+      : dtype_(type), self_(std::move(self)) {}
+
+  std::string kind() const override {
+    return "Cast";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override {
+    TORCH_INTERNAL_ASSERT(args.empty() && kwargs.empty());
+    Value* dtype_const = m.graph()->insertConstant(dtype_, loc);
+    std::vector<NamedValue> kwargs_{
+        self_, NamedValue(loc, "dtype", dtype_const)};
+    Value* casted_val = m.graph()->insert(
+        /*opname=*/Symbol::fromQualString("aten::to"),
+        /*args=*/args,
+        /*kwargs=*/kwargs_,
+        /*range=*/loc);
+    return std::make_shared<SimpleValue>(casted_val);
+  }
+
+  at::ScalarType dtype_;
+  NamedValue self_;
+};
+
+// builtins operators and functions that call a method if it exists
+// on a class type, like 'len(x)' and 'x + y'
+struct TORCH_API MagicMethod : public SugaredValue {
+  MagicMethod(std::string desugared_name, SugaredValuePtr base)
+      : base_value_(std::move(base)),
+        desugared_name_(std::move(desugared_name)) {}
+
+  std::string kind() const override {
+    return desugared_name_;
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+ private:
+  SugaredValuePtr base_value_;
+  std::string desugared_name_;
+};
+
+// things that look like function applications, but
+// perform non-standard evaluation are represented
+// with SpecialFormValues, e.g.
+//   isinstance(x, int)
+//   fork(fn)
+//   annotate(int, 3)
+// The implementation of each value is handled by a case inside emitApplyExpr
+struct TORCH_API SpecialFormValue : public SugaredValue {
+  SpecialFormValue(Symbol form) : form_(form) {}
+  std::string kind() const override {
+    return form_.toUnqualString();
+  }
+  Symbol form() const {
+    return form_;
+  }
+  static std::shared_ptr<SpecialFormValue> create(Symbol form) {
+    return std::make_shared<SpecialFormValue>(form);
+  }
+
+ private:
+  Symbol form_;
+};
+
+struct TORCH_API LegacyTensorConstructor : public SpecialFormValue {
+  LegacyTensorConstructor(Symbol form, at::ScalarType dtype, at::Device device)
+      : SpecialFormValue(form), device_(device), dtype_(dtype) {}
+
+  static std::shared_ptr<LegacyTensorConstructor> create(
+      Symbol form,
+      at::ScalarType dtype,
+      at::Device device) {
+    return std::make_shared<LegacyTensorConstructor>(form, dtype, device);
+  }
+  at::ScalarType dtype() const {
+    return dtype_;
+  }
+
+ private:
+  at::Device device_;
+  at::ScalarType dtype_;
+};
+
+// matched against for special handling of range expressions
+struct TORCH_API RangeValue : SugaredValue {
+  RangeValue(
+      const SourceRange& loc,
+      GraphFunction& m,
+      std::vector<Value*> input,
+      std::optional<int64_t> static_len = std::nullopt);
+
+  std::string kind() const override {
+    return "range";
+  }
+  Value* len(const SourceRange& loc, GraphFunction& m) override;
+  SugaredValuePtr getitem(
+      const SourceRange& loc,
+      GraphFunction& m,
+      Value* idx,
+      TypePtr type_hint = nullptr) override;
+  std::shared_ptr<SugaredValue> iter(const SourceRange& loc, GraphFunction& m)
+      override;
+
+  // When Range is instantiated via enumerate(iterable_with_static_len),
+  // then it takes the static length of the iterable
+  std::optional<int64_t> staticLen() override {
+    return static_len_;
+  }
+
+ private:
+  Value* start_{};
+  Value* end_{};
+  Value* step_{};
+  // a flag to determine if it's a simple range() call with only end_ from
+  // arguments If true, we will not insert length calculation and index
+  // derivation nodes to simplify the graph and enable more possible
+  // optimizations
+  bool has_only_end_{};
+  std::optional<int64_t> static_len_;
+};
+
+// Specialized Tree structure to matched against for special handling
+// of builtin functions iterables expressions like zip(), enumerate(), etc.
+// zip and enumerate can be modeled as a tree of SimpleValue/RangeValue:
+//    zip(x, y) ->  (x, y) with tuple assignment to each loop target
+//    enumerate(x) -> (range(0, math.inf, 1), x)
+// So a complicated expression like zip(a, enumerate(b), range(0, 100)) will be:
+// (a, (range(0, math.inf, 1), b), range(0, 100))
+// We use those base iterables to fill in the loop information like
+// max_trip_count and set the value table for loop targets
+// Iterables can contain lists of SugaredValues like ModuleLists. If it
+// does, then we emit it unrolled and require that all values it contains
+// have a statically-determinable length.
+struct TORCH_API IterableTree : SugaredValue {
+  IterableTree() = default;
+  IterableTree(
+      const SourceRange& range,
+      GraphFunction& m,
+      at::ArrayRef<SugaredValuePtr> children) {
+    for (const auto& child : children) {
+      addChild(range, m, child);
+    }
+  }
+  std::string kind() const override {
+    return "iterabletree";
+  }
+
+  std::shared_ptr<SugaredValue> iter(const SourceRange& loc, GraphFunction& m)
+      override {
+    return shared_from_this();
+  }
+
+  void addChild(
+      const SourceRange& range,
+      GraphFunction& m,
+      const SugaredValuePtr& iter_value);
+
+  std::vector<SugaredValuePtr> get_children() {
+    return children_;
+  }
+
+  // If this iterable contains a ModuleList or Tuple, then it will have a
+  // static length, and we will emit it as an unrolled for loop.
+  std::optional<int64_t> staticLen() override {
+    return unroll_length_;
+  }
+
+  // given a IterableTree node, get all the base iterables/leaves under the
+  // IterableTree node. This enables
+  // us to get all the basic SugaredValues that contains valid loop information
+  // with len() and getitem()
+  std::vector<SugaredValuePtr> get_base_iterables();
+
+  Value* len(const SourceRange& loc, GraphFunction& m) override;
+  SugaredValuePtr getitem(
+      const SourceRange& loc,
+      GraphFunction& m,
+      Value* idx,
+      TypePtr type_hint = nullptr) override;
+
+ private:
+  std::optional<int64_t> unroll_length_ = std::nullopt;
+  std::vector<SugaredValuePtr> children_;
+};
+
+static inline std::vector<Value*> toValues(
+    Graph& g,
+    at::ArrayRef<NamedValue> nvs) {
+  return fmap(nvs, [&](const NamedValue& v) { return v.value(g); });
+}
+
+struct SimpleSelf : public Self {
+  explicit SimpleSelf(ClassTypePtr classType)
+      : Self(), classType_(std::move(classType)) {}
+  std::shared_ptr<SugaredValue> makeSugared(Value* v) const override {
+    v->setType(classType_);
+    return std::make_shared<SimpleValue>(v);
+  }
+  ClassTypePtr getClassType() const override {
+    return classType_;
+  }
+
+ private:
+  ClassTypePtr classType_;
+};
+
+// This is not a SimpleValue so it can not pass through the code paths that
+// expect a SimpleValue as a sugared value.
+struct TORCH_API ExceptionMessageValue : public SugaredValue {
+  explicit ExceptionMessageValue(
+      Value* value,
+      Value* qualified_class_name = nullptr)
+      : value_(value), qualified_class_name_(qualified_class_name) {}
+
+  std::string kind() const override {
+    return "exception message";
+  }
+
+  Value* getValue() {
+    return value_;
+  }
+
+  // qualified python class name
+  Value* getQualifiedClassName() {
+    return qualified_class_name_;
+  }
+
+ private:
+  Value* value_;
+  Value* qualified_class_name_;
+};
+
+struct TORCH_API ExceptionValue : public SugaredValue {
+  explicit ExceptionValue(std::string message) : message_(std::move(message)) {}
+
+  std::string kind() const override {
+    return "exception";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> /*attributes*/,
+      size_t /*n_binders*/) override {
+    auto exception_message = insertConstant(*m.graph(), message_ + ": ", loc);
+    for (auto& input : args) {
+      auto input_str = input.value(*m.graph());
+      if (!input_str->type()->isSubtypeOf(*StringType::get())) {
+        input_str =
+            emitBuiltinCall(loc, *m.graph(), aten::str, {input_str}, {});
+      }
+      exception_message = emitBuiltinCall(
+          loc, *m.graph(), aten::add, {exception_message, input_str}, {});
+    }
+    return std::make_shared<ExceptionMessageValue>(exception_message);
+  }
+
+  std::string message_;
+};
+
+struct TORCH_API SugaredEnumClass : public SugaredValue {
+  explicit SugaredEnumClass(EnumTypePtr enum_type)
+      : enum_type_(std::move(enum_type)) {}
+
+  std::string kind() const override {
+    return "EnumClass";
+  }
+
+  SugaredValuePtr attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  SugaredValuePtr iter(const SourceRange& loc, GraphFunction& m) override;
+
+ private:
+  EnumTypePtr enum_type_;
+};
+
+struct TORCH_API SliceValue : public SugaredValue {
+  explicit SliceValue(Value* start, Value* stop, Value* step)
+      : start_(start), stop_(stop), step_(step) {}
+
+  std::string kind() const override {
+    return "Python slice value";
+  }
+
+  Value* start() {
+    return start_;
+  }
+  Value* stop() {
+    return stop_;
+  }
+  Value* step() {
+    return step_;
+  }
+
+ private:
+  Value* start_;
+  Value* stop_;
+  Value* step_;
+};
+
+struct TORCH_API TorchCheckValue : public SugaredValue {
+  explicit TorchCheckValue() = default;
+
+  std::string kind() const override {
+    return "torch._check sugared value";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tracer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tracer.h
new file mode 100644
index 0000000000000000000000000000000000000000..a3383b06b939c78f776d408f991fa573548e699b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tracer.h
@@ -0,0 +1,418 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/Dimname.h>
+#include <ATen/core/class_type.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+#include <ATen/core/symbol.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/Export.h>
+
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/utils/variadic.h>
+
+#include <cstdint>
+#include <memory>
+#include <unordered_map>
+#include <vector>
+
+namespace torch::jit {
+struct Node;
+struct Value;
+struct Graph;
+struct Module;
+
+namespace tracer {
+
+using ::c10::ivalue::Shared;
+
+using ::c10::IValue;
+using ::c10::ivalue::Future;
+
+using ::c10::ArrayRef;
+using ::c10::TupleType;
+using ::c10::TupleTypePtr;
+using ::c10::ivalue::ConstantString;
+
+using torch::autograd::Variable;
+using variable_list = std::vector<Variable>;
+
+TORCH_API std::atomic<bool>& getTracerStateWarnMode();
+
+struct TORCH_API TracingState
+    : public std::enable_shared_from_this<TracingState> {
+  TracingState();
+  ~TracingState();
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::shared_ptr<Graph> graph;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  bool warn = getTracerStateWarnMode();
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  bool strict = true;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  bool force_outplace = false;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::function<std::string(const Variable& var)> lookup_var_name_fn =
+      [](const Variable& var) { return ""; };
+
+  void enterFrame() {
+    env_stack.emplace_back();
+  }
+
+  void leaveFrame() {
+    env_stack.pop_back();
+  }
+
+  void setValue(const IValue& v, Value* value);
+  void delValue(const IValue& var);
+  Value* getValue(const IValue& var);
+  Value* getOutput(const IValue& var, size_t i);
+  bool hasValue(const IValue& var) const;
+
+  Node* createNode(c10::Symbol op_name, size_t num_outputs);
+  void insertNode(Node* node);
+
+ private:
+  using WeakIValue = at::WeakIValue;
+
+  struct WeakIValueHasher {
+    size_t operator()(const WeakIValue& t) const {
+      return t.hash();
+    }
+  };
+
+  struct WeakIValueEq {
+    bool operator()(const WeakIValue& t1, const WeakIValue& t2) const {
+      return t1.isSameIdentity(t2);
+    }
+  };
+
+  using Frame =
+      std::unordered_map<WeakIValue, Value*, WeakIValueHasher, WeakIValueEq>;
+  std::vector<Frame> env_stack;
+};
+
+// This is meant to be used as a thread local place, where we can store extra
+// info that gets lost when we call into ATen from Python bindings. One example
+// for when this happens is when we get an IntArrayRef argument with e.g. sizes
+// for view. When tracing, those might be tensors, which let us encode extra
+// data dependencies, but once they get to the ATen call where we actually have
+// the tracing logic, they get converted into a raw IntArrayRef, and we loose
+// all information. To prevent this, we temporarily stash it in here.
+struct ArgumentStash {
+  struct IntArrayRefTrace : std::vector<Value*> {
+    IntArrayRefTrace(size_t size) : std::vector<Value*>(size, nullptr) {}
+  };
+
+  static bool empty() {
+    return stash.intlists.empty();
+  }
+
+  TORCH_API static void stashIntArrayRefElem(
+      const std::string& arg_name,
+      size_t size,
+      size_t idx,
+      const Variable& var);
+
+  static bool hasIntArrayRef(const std::string& arg_name) {
+    return stash.intlists.count(arg_name) > 0;
+  }
+
+  static IntArrayRefTrace popIntArrayRef(const std::string& arg_name) {
+    auto info = std::move(stash.intlists.at(arg_name));
+    stash.intlists.erase(arg_name);
+    return info;
+  }
+
+  // Value stashing: Use these methods to stash arguments which correspond
+  // to regular Value*'s in the graph. i.e. they don't require special
+  // handling like in the case of IntArrayRefs
+  TORCH_API static void stashValue(
+      const std::string& arg_name,
+      size_t idx,
+      const Variable& var,
+      const c10::TypePtr& type = nullptr);
+
+  static bool hasValue(const std::string& arg_name) {
+    return stash.values.count(arg_name) > 0;
+  }
+
+  static Value* popValue(const std::string& arg_name) {
+    auto info = stash.values.at(arg_name);
+    stash.values.erase(arg_name);
+    return info;
+  }
+
+ private:
+  static thread_local ArgumentStash stash;
+  std::unordered_map<std::string, IntArrayRefTrace> intlists;
+  std::unordered_map<std::string, Value*> values;
+};
+
+// Retrieve or set the current tracing state. Returns a nullptr if tracing is
+// disabled.
+TORCH_API const std::shared_ptr<TracingState>& getTracingState();
+TORCH_API void setTracingState(std::shared_ptr<TracingState> state);
+
+inline bool isTracing() {
+  return static_cast<bool>(getTracingState());
+}
+
+using warn_fn_type = void (*)(const std::string& msg);
+TORCH_API extern const char* WARN_PYTHON_DATAFLOW;
+TORCH_API extern const char* WARN_CONSTRUCTOR;
+TORCH_API extern const char* WARN_RESIZE;
+TORCH_API extern const char* STRICT_TRACER_MSG;
+TORCH_API void _do_warn(const char* _reason, const char* _kind);
+inline void warn(const char* _reason, const char* _kind = nullptr) {
+  if (const auto& state = getTracingState()) {
+    if (!state->warn)
+      return;
+    _do_warn(_reason, _kind);
+  }
+}
+TORCH_API void setWarn(warn_fn_type fn);
+
+struct TORCH_API NoWarn {
+  NoWarn() : state(getTracingState()) {
+    if (state) {
+      prev = state->warn;
+      state->warn = false;
+    }
+  }
+  ~NoWarn() {
+    if (state) {
+      state->warn = prev;
+    }
+  }
+  std::shared_ptr<TracingState> state;
+  bool prev{false};
+};
+
+struct WithNestedTracingFrame {
+  WithNestedTracingFrame() {
+    getTracingState()->enterFrame();
+  }
+
+  ~WithNestedTracingFrame() {
+    getTracingState()->leaveFrame();
+  }
+};
+TORCH_API void recordSourceLocation(Node* n);
+TORCH_API void setRecordSourceLocation(void (*v)(Node*));
+
+TORCH_API std::vector<StackEntry> pythonCallstack();
+TORCH_API void setPythonCallstack(std::vector<StackEntry> (*v)());
+
+// Having finished adding a new 'node' to the graph IR 'setValueTrace'
+// associates this node with an output variable, so that further operations
+// involving this variable know which node in the IR to reference.
+TORCH_API void setValueTrace(const IValue& v, Value* value);
+
+TORCH_API void delValueTrace(const IValue& var);
+
+TORCH_API std::function<void()> pauseTracing();
+
+TORCH_API Value* getValueTrace(const IValue& var);
+
+TORCH_API std::pair<std::shared_ptr<TracingState>, Stack> trace(
+    Stack inputs,
+    const std::function<Stack(Stack)>& traced_fn,
+    std::function<std::string(const Variable&)> var_name_lookup_fn,
+    bool strict = true,
+    bool force_outplace = false,
+    Module* self = nullptr,
+    const std::vector<std::string>& argument_names = {});
+
+TORCH_API void abandon();
+
+// NB: those serve both as an intermediate steps in addInputs below,
+// as well as the overloads that terminate template recursion
+TORCH_API void addInputs(Node* n, const char* name, int64_t value);
+TORCH_API void addInputs(Node* n, const char* name, const c10::SymInt& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    std::optional<int64_t> value);
+TORCH_API void addInputs(Node* n, const char* name, bool value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<bool>& value);
+TORCH_API void addInputs(Node* n, const char* name, double value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<double>& value);
+TORCH_API void addInputs(Node* n, const char* name, const at::Scalar& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<at::Scalar>& value);
+TORCH_API void addInputs(Node* n, const char* name, const at::Tensor& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<at::Tensor>& value);
+TORCH_API void addInputs(Node* n, const char* name, ArrayRef<int64_t> value);
+TORCH_API void addInputs(Node* n, const char* name, c10::SymIntArrayRef value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    std::optional<c10::SymInt> value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<ArrayRef<int64_t>>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const at::OptionalIntArrayRef& opt_value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const at::OptionalSymIntArrayRef& opt_value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    ArrayRef<at::Tensor> value,
+    bool allow_undefined = false);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::vector<at::Tensor>& value,
+    bool allow_undefined = false);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    at::ITensorListRef value,
+    bool allow_undefined = false);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const List<std::optional<at::Tensor>>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    ArrayRef<c10::intrusive_ptr<c10::ivalue::Object>> value,
+    const c10::ClassTypePtr& class_type);
+TORCH_API void addInputs(Node* n, const char* name, ArrayRef<double> value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<ArrayRef<double>>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::string_view value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<std::string_view>& value);
+TORCH_API void addInputs(Node* n, const char* name, at::Device value);
+TORCH_API void addInputs(Node* n, const char* name, c10::Stream stream);
+TORCH_API void addInputs(Node* n, const char* name, at::Layout value);
+TORCH_API void addInputs(Node* n, const char* name, at::ScalarType value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<at::ScalarType>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<at::Device>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<at::Layout>& value);
+TORCH_API void addInputs(Node* n, const char* name, at::MemoryFormat value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    std::optional<at::DimnameList> value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<at::MemoryFormat>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const std::optional<at::Generator>& value);
+
+inline void addInputs(
+    Node* n,
+    const char* name,
+    const std::vector<bool>& value) {
+  TORCH_CHECK(false, "Tracing a list of bool type is currently not supported!");
+}
+
+template <typename T>
+void addInputs(Node* n, const char* name, ArrayRef<T> value) {
+  TORCH_CHECK(
+      false, "Tracing a list of arbitrary type is currently not supported!");
+}
+template <typename K, typename V>
+void addInputs(
+    Node* n,
+    const char* name,
+    const std::unordered_map<K, V>& value) {
+  TORCH_CHECK(
+      false, "Tracing a dict of arbitrary types is currently not supported!");
+}
+
+template <size_t N>
+void addInputs(Node* n, const char* name, std::array<bool, N> value) {
+  throw std::runtime_error(
+      "Found an unsupported argument type in the JIT tracer. File a bug report.");
+}
+
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const c10::intrusive_ptr<c10::ivalue::Object>& obj);
+
+TORCH_API void ensureUniqueIfOutOfPlaced(
+    const char* name,
+    const at::Tensor& tensor);
+TORCH_API void ensureUniqueIfOutOfPlaced(
+    const char* name,
+    const std::optional<at::Tensor>& tensor);
+
+template <
+    typename T,
+    typename = std::enable_if_t<
+        (!std::is_convertible_v<std::decay_t<T>, at::TensorList> &&
+         !std::is_convertible_v<std::decay_t<T>, c10::List<at::Tensor>> &&
+         !std::is_convertible_v<std::decay_t<T>, at::Tensor> &&
+         !std::is_convertible_v<
+             std::decay_t<T>,
+             c10::intrusive_ptr<c10::ivalue::Object>>)>>
+void addOutput(Node* node, T&& /*unused*/) {
+  TORCH_CHECK(
+      false,
+      "Found an unsupported argument type ",
+      c10::demangle_type<T>(),
+      " in the JIT tracer. File a bug report.");
+}
+TORCH_API void addOutput(Node* node, const at::Tensor& tensor);
+TORCH_API void setOutput(Value* value, const at::Tensor& output);
+TORCH_API void addOutput(Node* node, const std::vector<at::Tensor>& list);
+TORCH_API void addOutput(Node* node, const c10::List<at::Tensor>& list);
+TORCH_API void addOutput(
+    Node* node,
+    const c10::intrusive_ptr<c10::ivalue::Object>& output);
+
+TORCH_API autograd::Variable getSizeOf(
+    const autograd::Variable& var,
+    int64_t dim);
+
+TORCH_API autograd::Variable getNumelOf(const autograd::Variable& var);
+
+} // namespace tracer
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tree.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tree.h
new file mode 100644
index 0000000000000000000000000000000000000000..99ae0ff06622e087a299ade7da2b39419aa840f6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tree.h
@@ -0,0 +1,227 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <unordered_map>
+#include <vector>
+
+#include <c10/util/Exception.h>
+#include <c10/util/SmallVector.h>
+#include <c10/util/intrusive_ptr.h>
+#include <torch/csrc/jit/frontend/lexer.h>
+
+namespace torch::jit {
+
+// Trees are used to represent all forms of TC IR, pre- and post-typechecking.
+// Rather than have a full class hierarchy for all TC statements, trees are a
+// slight variation of Lisp s-expressions. For instance, the expression a*b+1
+// is represented as:
+// (+ (* (ident a) (ident b)) (const 1))
+// Atoms like 'a', 'b', and '1' are represented by subclasses of Tree which
+// define stringValue(). Everything else is a Compound object, which has a
+// 'kind' that is a token from lexer.h's TokenKind enum. Single-character
+// operators like '+' are represented using the character itself (so, add.kind()
+// would be '+'). Each Compound object also contains a list of subtrees and is
+// associated with a SourceRange for error reporting.
+// Memory management of trees is done using intrusive_ptr.
+
+struct Tree;
+using TreeRef = c10::intrusive_ptr<Tree>;
+using TreeList = at::SmallVector<TreeRef, 4>;
+
+struct Tree : c10::intrusive_ptr_target {
+  Tree(int kind_) : kind_(kind_) {}
+  int kind() const {
+    return kind_;
+  }
+  virtual bool isAtom() const {
+    return true;
+  }
+  virtual const SourceRange& range() const {
+    TORCH_CHECK(false, "is an Atom");
+  }
+  virtual const std::string& stringValue() const {
+    TORCH_CHECK(false, "stringValue can only be called on TK_STRING");
+  }
+  virtual const TreeList& trees() const {
+    static const TreeList empty_trees = {};
+    return empty_trees;
+  }
+  const TreeRef& tree(size_t i) const {
+    return trees().at(i);
+  }
+  virtual TreeRef map(const std::function<TreeRef(TreeRef)>& fn) {
+    (void)fn;
+    c10::raw::intrusive_ptr::incref(this); // we are creating a new pointer
+                                           // from a raw `this` pointer
+                                           // so we need to bump the refcount
+                                           // to account for this ownership
+    return TreeRef::reclaim(this);
+  }
+  template <typename... Args>
+  void match(int k, Args&... args) const {
+    matchD(k, "unknown", 0, args...);
+  }
+  template <typename... Args>
+  void matchD(int k, const char* filename, int lineno, Args&... args) const {
+    std::initializer_list<TreeRef*> vars = {args...};
+    matchNumSubtreesD(k, filename, lineno, vars.size(), true);
+    size_t i = 0;
+    for (TreeRef* v : vars) {
+      *v = trees()[i++];
+    }
+  }
+  void matchNumSubtrees(int k, size_t expected_subtrees) {
+    return matchNumSubtreesD(k, "unknown", 0, expected_subtrees, false);
+  }
+  void matchNumSubtreesD(
+      int k,
+      const char* filename,
+      int lineno,
+      size_t expected_subtrees,
+      bool allow_more) const {
+    TORCH_CHECK(
+        kind() == k,
+        filename,
+        ":",
+        lineno,
+        ": expecting kind '",
+        kindToString(k),
+        "' but found '",
+        kindToString(kind()),
+        "'\n");
+    if (trees().size() < expected_subtrees ||
+        (!allow_more && trees().size() != expected_subtrees)) {
+      std::stringstream ss;
+      ss << filename << ':' << lineno << ": expected at least "
+         << expected_subtrees << " subtrees, but found only " << trees().size()
+         << '\n';
+      range().highlight(ss);
+      TORCH_CHECK(false, ss.str());
+    }
+  }
+  ~Tree() override = default;
+
+ private:
+  int kind_;
+};
+
+struct String : public Tree {
+  String(std::string value) : Tree(TK_STRING), value_(std::move(value)) {}
+  const std::string& stringValue() const override {
+    return value_;
+  }
+  template <typename... Args>
+  static TreeRef create(Args&&... args) {
+    return c10::make_intrusive<String>(std::forward<Args>(args)...);
+  }
+
+ private:
+  std::string value_;
+};
+
+static SourceRange mergeRanges(SourceRange c, const TreeList& others) {
+  for (const auto& t : others) {
+    if (t->isAtom())
+      continue;
+    size_t s = std::min(c.start(), t->range().start());
+    size_t e = std::max(c.end(), t->range().end());
+    c = SourceRange(c.source(), s, e);
+  }
+  return c;
+}
+
+struct Compound : public Tree {
+  Compound(int kind, SourceRange range)
+      : Tree(kind), range_(std::move(range)) {}
+  Compound(int kind, const SourceRange& range_, TreeList&& trees_)
+      : Tree(kind),
+        range_(mergeRanges(range_, trees_)),
+        trees_(std::move(trees_)) {}
+  const TreeList& trees() const override {
+    return trees_;
+  }
+  static TreeRef create(
+      int kind,
+      const SourceRange& range_,
+      TreeList&& trees_) {
+    return c10::make_intrusive<Compound>(kind, range_, std::move(trees_));
+  }
+  bool isAtom() const override {
+    return false;
+  }
+  TreeRef map(const std::function<TreeRef(TreeRef)>& fn) override {
+    TreeList ret;
+    for (auto& t : trees()) {
+      ret.push_back(fn(t));
+    }
+    return Compound::create(kind(), range(), std::move(ret));
+  }
+
+  const SourceRange& range() const override {
+    return range_;
+  }
+
+ private:
+  SourceRange range_;
+  TreeList trees_;
+};
+
+// tree pretty printer
+struct pretty_tree {
+  pretty_tree(const TreeRef& tree, size_t col = 40) : tree(tree), col(col) {}
+  const TreeRef& tree;
+  size_t col;
+  std::unordered_map<TreeRef, std::string> flat_strings;
+  const std::string& get_flat(const TreeRef& t) {
+    auto it = flat_strings.find(t);
+    if (it != flat_strings.end())
+      return it->second;
+
+    std::stringstream out;
+    switch (t->kind()) {
+      case TK_STRING:
+        out << t->stringValue();
+        break;
+      default:
+        out << '(' << kindToString(t->kind());
+        for (const auto& e : t->trees()) {
+          out << ' ' << get_flat(e);
+        }
+        out << ')';
+        break;
+    }
+    auto it_ = flat_strings.emplace(t, out.str());
+    return it_.first->second;
+  }
+  void print(std::ostream& out, const TreeRef& t, int indent) {
+    const std::string& s = get_flat(t);
+    if (indent + s.size() < col || t->isAtom()) {
+      out << s;
+      return;
+    }
+    std::string k = kindToString(t->kind());
+    out << '(' << k;
+    for (const auto& e : t->trees()) {
+      out << '\n' << std::string(indent + 2, ' ');
+      print(out, e, indent + 2);
+    }
+    out << ')';
+  }
+};
+
+static inline std::ostream& operator<<(std::ostream& out, pretty_tree t_) {
+  t_.print(out, t_.tree, 0);
+  return out << '\n';
+}
+
+static inline std::ostream& operator<<(std::ostream& out, const TreeRef& t) {
+  return out << pretty_tree(t);
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tree_views.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tree_views.h
new file mode 100644
index 0000000000000000000000000000000000000000..1dc386d938f69651b4f9dd3cf9d5a335531c63e1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/tree_views.h
@@ -0,0 +1,1285 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/frontend/error_report.h>
+#include <torch/csrc/jit/frontend/strtod.h>
+#include <torch/csrc/jit/frontend/tree.h>
+
+#include <c10/util/complex.h>
+#include <functional>
+#include <iostream>
+#include <string>
+#include <utility>
+
+namespace torch::jit {
+
+// clang-format off
+// TreeView provides a statically-typed way to traverse the tree, which should
+// be formed according to the grammar below.
+//
+// A few notes on types and their aliases:
+// - List<T> is really a Tree with kind TK_LIST and elements as subtrees
+// - Maybe<T> is really a Tree with kind TK_OPTION that has 0 or 1 subtree of type T
+// - Builtin types are: Ident (TK_IDENT), String (TK_STRING)
+//
+// Param = Param(Maybe<Expr> type, Ident name)                          TK_PARAM
+//
+// Decl  = Decl(List<Param> params, Maybe<Expr> return_type)            TK_DECL
+// Def   = Def(Ident name, Decl decl, List<Stmt> body)                  TK_DEF
+// ClassDef = ClassDef(Ident name,                                      TK_CLASS_DEF
+//                     Maybe<Expr> superclass,
+//                     List<Stmt> body)
+//
+// Stmt  = If(Expr cond, List<Stmt> true_body, List<Stmt> false_body)   TK_IF
+//       | For(List<Expr> targets, List<Expr> iters, List<Stmt> body)   TK_FOR
+//       | While(Expr cond, List<Stmt> body)                            TK_WHILE
+//       | Global(List<Ident> idents)                                   TK_GLOBAL
+//       -- NB: the only type of Expr's allowed on lhs are Var
+//          Or a tuple containing Var with an optional terminating Starred
+//       | Assign(Expr lhs, Maybe<Expr> rhs, Maybe<Expr> type)          TK_ASSIGN
+//       | AugAssign(Expr lhs, AugAssignKind aug_op, Expr rhs)          TK_AUG_ASSIGN
+//       | Return(List<Expr> values)                                    TK_RETURN
+//       | ExprStmt(List<Expr> expr)                                    TK_EXPR_STMT
+//       | Raise(Expr expr)                                             TK_RAISE
+//       | Def                                                          TK_DEF
+//       | With(List<WithItem> targets, List<Stmt> body)                TK_WITH
+//
+// Expr  = TernaryIf(Expr cond, Expr true_expr, Expr false_expr)        TK_IF_EXPR
+//       | BinOp(Expr lhs, Expr rhs)
+//       |     And                                                      TK_AND
+//       |     Or                                                       TK_OR
+//       |     Lt                                                       '<'
+//       |     Gt                                                       '>'
+//       |     Eq                                                       TK_EQ
+//       |     Le                                                       TK_LE
+//       |     Ge                                                       TK_GE
+//       |     Ne                                                       TK_NE
+//       |     Is                                                       TK_IS
+//       |     IsNot                                                    TK_ISNOT
+//       |     Add                                                      '+'
+//       |     Sub                                                      '-'
+//       |     Mul                                                      '*'
+//       |     Div                                                      '/'
+//       |     Mod                                                      '%'
+//       |     MatMult                                                  '@'
+//       |     Pow                                                      TK_POW
+//       | UnaryOp(Expr expr)
+//       |     Not                                                      TK_NOT
+//       |     USub                                                     '-'
+//       | Const(String value)                                          TK_CONST
+//       -- NB: x.name(y) is desugared into name(x, y)
+//       | Apply(Ident name, List<Expr> args, List<Attribute> kwargs)   TK_APPLY
+//       | Select(Expr value, Ident selector)                           '.'
+//       | Subscript(Expr value, List<Expr> subscript_exprs)            TK_SUBSCRIPT
+//       | SliceExpr(Maybe<Expr> start, Maybe<Expr> end)                TK_SLICE_EXPR
+//       | Var(Ident name)                                              TK_VAR
+//       | ListLiteral(List<Expr> inputs)                               TK_LIST_LITERAL
+//       | TupleLiteral(List<Expr> inputs)                              TK_TUPLE_LITERAL
+//       | Starred(Expr expr)                                           TK_STARRED
+//       | WithItem(Expr target, Maybe<Var> var)                        TK_WITH_ITEM
+// -- NB: only allowed expressions are Const or List(Const)
+//        (List as a value, not type constructor)
+// Attribute = Attribute(Ident name, Expr value)                        TK_ATTRIBUTE
+//
+// AugAssignKind =
+//            | Add()                                                   TK_PLUS_EQ
+//            | Sub()                                                   TK_MINUS_EQ
+//            | Mul()                                                   TK_TIMES_EQ
+//            | Div()                                                   TK_DIV_EQ
+//            | Mod()                                                   TK_MOD_EQ
+//
+
+// Each subclass of TreeView should provide:
+// 1. Constructor that takes a TreeRef, and checks that it's of the right type.
+// 2. Accessors that get underlying information out of the object. If they
+//    return subtrees, they should wrap them in appropriate views too.
+// 3. Static method 'create' that creates the underlying TreeRef object
+//    for every TreeRef kind that has a TreeView, the parser always uses
+//    (e.g.) Ident::create rather than Compound::Create, this means that
+//    changes to the structure of Ident are always made right here rather
+//    than both in the parser and in this code.
+// XXX: these structs should have no fields to prevent slicing when passing by value
+// clang-format on
+struct TreeView {
+  explicit TreeView(TreeRef tree) : tree_(std::move(tree)) {}
+  TreeRef tree() const {
+    return tree_;
+  }
+  const SourceRange& range() const {
+    return tree_->range();
+  }
+  operator TreeRef() const {
+    return tree_;
+  }
+  const TreeRef& get() const {
+    return tree_;
+  }
+  int kind() const {
+    return tree_->kind();
+  }
+  void dump() const {
+    std::cout << tree_;
+  }
+
+ protected:
+  const TreeRef& subtree(size_t i) const {
+    return tree_->trees().at(i);
+  }
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  TreeRef tree_;
+};
+
+template <typename T>
+struct ListIterator {
+  ListIterator(TreeList::const_iterator it) : it(it) {}
+  bool operator!=(const ListIterator& rhs) const {
+    return it != rhs.it;
+  }
+  bool operator==(const ListIterator& rhs) const {
+    return it == rhs.it;
+  }
+  T operator*() const {
+    return T(*it);
+  }
+  ListIterator& operator+=(std::ptrdiff_t n) {
+    it += n;
+    return *this;
+  }
+  ListIterator& operator++() {
+    ++it;
+    return *this;
+  }
+  ListIterator& operator--() {
+    --it;
+    return *this;
+  }
+
+ private:
+  TreeList::const_iterator it;
+};
+
+template <typename T>
+struct List : public TreeView {
+  using iterator = ListIterator<T>;
+  using const_iterator = ListIterator<T>;
+
+  List(const TreeRef& tree) : TreeView(tree) {
+    tree->match(TK_LIST);
+    // Iterate over list to temporarily instantiate Ts that will check the type
+    for (const T& elem : *this) {
+      (void)elem; // silence unused warning
+    }
+  }
+  iterator begin() const {
+    return iterator(tree_->trees().begin());
+  }
+  iterator end() const {
+    return iterator(tree_->trees().end());
+  }
+  bool empty() const {
+    return tree_->trees().begin() == tree_->trees().end();
+  }
+  T operator[](size_t i) const {
+    return T(subtree(i));
+  }
+  TreeRef map(const std::function<TreeRef(const T&)>& fn) {
+    return tree_->map([&](TreeRef v) { return fn(T(v)); });
+  }
+  static List create(const SourceRange& range, const std::vector<T>& subtrees) {
+    TreeList type_erased_sub{subtrees.begin(), subtrees.end()};
+    return List(Compound::create(TK_LIST, range, std::move(type_erased_sub)));
+  }
+  static List unsafeCreate(const SourceRange& range, TreeList&& subtrees) {
+    return List(Compound::create(TK_LIST, range, std::move(subtrees)));
+  }
+  size_t size() const {
+    return tree_->trees().size();
+  }
+};
+
+template <typename T>
+struct Maybe : public TreeView {
+  explicit Maybe(const TreeRef& tree) : TreeView(tree) {
+    tree_->match(TK_OPTION);
+    if (tree_->trees().size() > 1)
+      throw(ErrorReport(tree) << "Maybe trees can have at most one subtree");
+  }
+  /* implicit */ Maybe(const T& tree) : TreeView(tree) {}
+  bool present() const {
+    return tree_->trees().size() > 0;
+  }
+  T get() const {
+    return T(tree_->trees().at(0));
+  }
+  TreeRef map(const std::function<TreeRef(const T&)>& fn) {
+    return tree_->map([&](TreeRef v) { return fn(T(v)); });
+  }
+  static Maybe<T> create(const SourceRange& range) {
+    return Maybe<T>(Compound::create(TK_OPTION, range, {}));
+  }
+  static Maybe<T> create(const SourceRange& range, const T& value) {
+    return Maybe<T>(Compound::create(TK_OPTION, range, {value}));
+  }
+};
+
+struct Ident : public TreeView {
+  explicit Ident(const TreeRef& tree) : TreeView(tree) {
+    tree_->match(TK_IDENT);
+  }
+  const std::string& name() const {
+    return subtree(0)->stringValue();
+  }
+  static Ident create(const SourceRange& range, std::string name) {
+    return Ident(
+        Compound::create(TK_IDENT, range, {String::create(std::move(name))}));
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Base types (production LHS)
+////////////////////////////////////////////////////////////////////////////////
+
+struct Stmt : public TreeView {
+  explicit Stmt(const TreeRef& tree) : TreeView(tree) {
+    switch (tree->kind()) {
+      case TK_IF:
+      case TK_FOR:
+      case TK_WHILE:
+      case TK_GLOBAL:
+      case TK_ASSIGN:
+      case TK_AUG_ASSIGN:
+      case TK_RETURN:
+      case TK_EXPR_STMT:
+      case TK_RAISE:
+      case TK_ASSERT:
+      case TK_PASS:
+      case TK_BREAK:
+      case TK_DELETE:
+      case TK_CONTINUE:
+      case TK_DEF:
+      case TK_WITH:
+        return;
+      default:
+        throw(
+            ErrorReport(tree)
+            << kindToString(tree->kind()) << " is not a valid Stmt");
+    }
+  }
+};
+
+struct Expr : public TreeView {
+  explicit Expr(const TreeRef& tree) : TreeView(tree) {
+    switch (tree->kind()) {
+      case TK_IF_EXPR:
+      case TK_AND:
+      case TK_OR:
+      case '<':
+      case '>':
+      case TK_IS:
+      case TK_ISNOT:
+      case TK_EQ:
+      case TK_LE:
+      case TK_GE:
+      case TK_NE:
+      case '+':
+      case '-':
+      case TK_UNARY_MINUS:
+      case '~':
+      case '*':
+      case TK_STARRED:
+      case '/':
+      case '%':
+      case TK_NOT:
+      case TK_CONST:
+      case TK_STRINGLITERAL:
+      case TK_TRUE:
+      case TK_FALSE:
+      case TK_NONE:
+      case TK_NONE_TYPE:
+      case TK_CAST:
+      case TK_APPLY:
+      case '.':
+      case TK_SUBSCRIPT:
+      case TK_SLICE_EXPR:
+      case TK_VAR:
+      case TK_LIST_LITERAL:
+      case TK_TUPLE_LITERAL:
+      case TK_DICT_LITERAL:
+      case '@':
+      case TK_POW:
+      case TK_LSHIFT:
+      case TK_RSHIFT:
+      case TK_FLOOR_DIV:
+      case '&':
+      case '^':
+      case '|':
+      case TK_LIST_COMP:
+      case TK_DICT_COMP:
+      case TK_DOTS:
+      case TK_IN:
+      case TK_WITH_ITEM:
+        return;
+      default:
+        throw(
+            ErrorReport(tree)
+            << kindToString(tree->kind()) << " is not a valid Expr");
+    }
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Helper nodes (mostly for function arguments)
+////////////////////////////////////////////////////////////////////////////////
+
+struct Attribute : public TreeView {
+  explicit Attribute(const TreeRef& tree) : TreeView(tree) {
+    tree_->match(TK_ATTRIBUTE);
+  }
+  Ident name() const {
+    return Ident(subtree(0));
+  }
+  Expr value() const {
+    return Expr(subtree(1));
+  }
+  static Attribute create(
+      const SourceRange& range,
+      const Ident& name,
+      const TreeRef& value) {
+    return Attribute(Compound::create(TK_ATTRIBUTE, range, {name, value}));
+  }
+};
+
+struct Param : public TreeView {
+  explicit Param(const TreeRef& tree) : TreeView(tree) {
+    tree_->match(TK_PARAM);
+  }
+  static Param create(
+      const SourceRange& range,
+      const Ident& ident,
+      const Maybe<Expr>& type,
+      const Maybe<Expr>& def,
+      bool kwarg_only) {
+    TreeRef kwarg_only_tree =
+        Compound::create(kwarg_only ? TK_TRUE : TK_FALSE, range, {});
+    return Param(Compound::create(
+        TK_PARAM, range, {ident, type, def, std::move(kwarg_only_tree)}));
+  }
+  Ident ident() const {
+    return Ident(subtree(0));
+  }
+  Maybe<Expr> type() const {
+    return Maybe<Expr>(subtree(1));
+  }
+  Maybe<Expr> defaultValue() const {
+    return Maybe<Expr>(subtree(2));
+  }
+  bool kwarg_only() const {
+    return TK_TRUE == subtree(3)->kind();
+  }
+  Param withType(const Maybe<Expr>& typ) const {
+    return Param::create(range(), ident(), typ, defaultValue(), kwarg_only());
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Top level definitions
+////////////////////////////////////////////////////////////////////////////////
+
+struct Decl : public TreeView {
+  explicit Decl(const TreeRef& tree) : TreeView(tree) {
+    tree->match(TK_DECL);
+  }
+  List<Param> params() const {
+    return List<Param>(subtree(0));
+  }
+  Maybe<Expr> return_type() const {
+    return Maybe<Expr>(subtree(1));
+  }
+  static Decl create(
+      const SourceRange& range,
+      const List<Param>& params,
+      const Maybe<Expr>& return_type) {
+    return Decl(Compound::create(TK_DECL, range, {params, return_type}));
+  }
+};
+
+struct Def : public TreeView {
+  explicit Def(const TreeRef& tree) : TreeView(tree) {
+    tree->match(TK_DEF);
+  }
+  Def withName(std::string new_name) const {
+    auto new_ident = Ident::create(name().range(), std::move(new_name));
+    return create(range(), new_ident, decl(), statements());
+  }
+  Def withDecl(const Decl& decl) const {
+    return create(range(), name(), decl, statements());
+  }
+  Ident name() const {
+    return Ident(subtree(0));
+  }
+  Decl decl() const {
+    return Decl(subtree(1));
+  }
+  List<Stmt> statements() const {
+    return List<Stmt>(subtree(2));
+  }
+  static Def create(
+      const SourceRange& range,
+      const Ident& name,
+      const Decl& decl,
+      const List<Stmt>& stmts) {
+    return Def(Compound::create(TK_DEF, range, {name, decl, stmts}));
+  }
+};
+
+// Property represents a named attribute combined with a getter and setter
+// method to access and mutate that attribute.
+struct Property : public TreeView {
+  explicit Property(const TreeRef& tree) : TreeView(tree) {
+    tree->match(TK_PROP);
+  }
+  Ident name() const {
+    return Ident(subtree(0));
+  }
+  Def getter() const {
+    return Def(subtree(1));
+  }
+  Maybe<Def> setter() const {
+    return Maybe<Def>(subtree(2));
+  }
+  static Property create(
+      const SourceRange& range,
+      const Ident& name,
+      const Def& getter,
+      const Maybe<Def>& setter) {
+    return Property(Compound::create(TK_PROP, range, {name, getter, setter}));
+  }
+};
+
+struct Assign;
+
+struct ClassDef : public TreeView {
+  explicit ClassDef(const TreeRef& tree) : TreeView(tree) {
+    tree->match(TK_CLASS_DEF);
+  }
+  explicit ClassDef(TreeRef&& tree) : TreeView(std::move(tree)) {
+    tree_->match(TK_CLASS_DEF);
+  }
+  ClassDef withName(std::string new_name) const {
+    auto new_ident = Ident::create(name().range(), std::move(new_name));
+    return create(range(), new_ident, superclass(), body());
+  }
+  Ident name() const {
+    return Ident(subtree(0));
+  }
+  Maybe<Expr> superclass() const {
+    return Maybe<Expr>(subtree(1));
+  }
+  List<Stmt> body() const {
+    return List<Stmt>(subtree(2));
+  }
+  Maybe<List<Property>> properties() const {
+    return Maybe<List<Property>>(subtree(3));
+  }
+  Maybe<List<Assign>> assigns() const {
+    return Maybe<List<Assign>>(subtree(4));
+  }
+  static ClassDef create(
+      const SourceRange& range,
+      const Ident& name,
+      const Maybe<Expr>& superclass,
+      const List<Stmt>& body) {
+    return ClassDef(Compound::create(
+        TK_CLASS_DEF,
+        range,
+        {name,
+         superclass,
+         body,
+         Maybe<List<Property>>::create(range),
+         Maybe<List<Assign>>::create(range)}));
+  }
+  static ClassDef create(
+      const SourceRange& range,
+      const Ident& name,
+      const Maybe<Expr>& superclass,
+      const List<Stmt>& body,
+      const List<Property>& properties,
+      const List<Assign>& assigns);
+};
+
+TORCH_API std::vector<std::string> getUnresolvedClassAttributes(
+    const ClassDef& def);
+
+////////////////////////////////////////////////////////////////////////////////
+// Statements
+////////////////////////////////////////////////////////////////////////////////
+
+struct If : public Stmt {
+  explicit If(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_IF);
+  }
+  Expr cond() const {
+    return Expr(subtree(0));
+  }
+  List<Stmt> trueBranch() const {
+    return List<Stmt>(subtree(1));
+  }
+  List<Stmt> falseBranch() const {
+    return List<Stmt>(subtree(2));
+  }
+  If withNewBranches(
+      const List<Stmt>& true_branch,
+      const List<Stmt>& false_branch) const {
+    return create(range(), cond(), true_branch, false_branch);
+  }
+  static If create(
+      const SourceRange& range,
+      const Expr& cond,
+      const List<Stmt>& true_branch,
+      const List<Stmt>& false_branch) {
+    return If(
+        Compound::create(TK_IF, range, {cond, true_branch, false_branch}));
+  }
+};
+
+struct While : public Stmt {
+  explicit While(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_WHILE);
+  }
+  Expr cond() const {
+    return Expr(subtree(0));
+  }
+  List<Stmt> body() const {
+    return List<Stmt>(subtree(1));
+  }
+  static While create(
+      const SourceRange& range,
+      const Expr& cond,
+      const List<Stmt>& body) {
+    return While(Compound::create(TK_WHILE, range, {cond, body}));
+  }
+};
+
+struct For : public Stmt {
+  explicit For(const TreeRef& tree) : Stmt(tree) {
+    tree->match(TK_FOR);
+  }
+  List<Expr> targets() const {
+    return List<Expr>(subtree(0));
+  }
+  List<Expr> itrs() const {
+    return List<Expr>(subtree(1));
+  }
+  List<Stmt> body() const {
+    return List<Stmt>(subtree(2));
+  }
+  static For create(
+      const SourceRange& range,
+      const List<Expr>& targets,
+      const List<Expr>& itrs,
+      const List<Stmt>& body) {
+    return For(Compound::create(TK_FOR, range, {targets, itrs, body}));
+  }
+};
+
+// TODO: supports only single comprehension for now
+struct ListComp : public Expr {
+  explicit ListComp(const TreeRef& tree) : Expr(tree) {
+    tree->match(TK_LIST_COMP);
+  }
+  Expr elt() const {
+    return Expr(subtree(0));
+  }
+  Expr target() const {
+    return Expr(subtree(1));
+  }
+  Expr iter() const {
+    return Expr(subtree(2));
+  }
+  // TODO: no ifs for now
+  static ListComp create(
+      const SourceRange& range,
+      const Expr& elt,
+      const Expr& target,
+      const Expr& iter) {
+    return ListComp(Compound::create(TK_LIST_COMP, range, {elt, target, iter}));
+  }
+};
+
+// TODO: supports only single comprehension for now
+struct DictComp : public Expr {
+  explicit DictComp(const TreeRef& tree) : Expr(tree) {
+    tree->match(TK_DICT_COMP);
+  }
+  Expr key() const {
+    return Expr(subtree(0));
+  }
+  Expr value() const {
+    return Expr(subtree(1));
+  }
+  Expr target() const {
+    return Expr(subtree(2));
+  }
+  Expr iter() const {
+    return Expr(subtree(3));
+  }
+  // TODO: no ifs for now
+  static DictComp create(
+      const SourceRange& range,
+      const Expr& key,
+      const Expr& value,
+      const Expr& target,
+      const Expr& iter) {
+    return DictComp(
+        Compound::create(TK_DICT_COMP, range, {key, value, target, iter}));
+  }
+};
+
+struct Global : public Stmt {
+  explicit Global(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_GLOBAL);
+  }
+  List<Ident> names() {
+    return List<Ident>(subtree(0));
+  }
+  static Global create(const SourceRange& range, const List<Ident>& names) {
+    return Global(Compound::create(TK_GLOBAL, range, {names}));
+  }
+};
+
+struct AugAssignKind : public TreeView {
+  explicit AugAssignKind(const TreeRef& tree) : TreeView(tree) {
+    switch (tree->kind()) {
+      case '+':
+      case '-':
+      case '*':
+      case '/':
+      case '%':
+      case '|':
+      case '&':
+      case '^':
+      case TK_POW:
+      case TK_LSHIFT:
+      case TK_RSHIFT:
+        return;
+      default:
+        throw(ErrorReport(tree) << "is not a valid AugAssignKind");
+    }
+  }
+};
+
+// Augmented assignment, like "foo += bar"
+struct AugAssign : public Stmt {
+  explicit AugAssign(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_AUG_ASSIGN);
+  }
+  static AugAssign create(
+      const SourceRange& range,
+      const Expr& lhs,
+      const AugAssignKind& aug_op,
+      const Expr& rhs) {
+    return AugAssign(
+        Compound::create(TK_AUG_ASSIGN, range, {lhs, aug_op, rhs}));
+  }
+  Expr lhs() const {
+    return Expr(subtree(0));
+  }
+  int aug_op() const {
+    return subtree(1)->kind();
+  }
+  Expr rhs() const {
+    return Expr(subtree(2));
+  }
+};
+
+struct Assign : public Stmt {
+  explicit Assign(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_ASSIGN);
+  }
+  static Assign create(
+      const SourceRange& range,
+      const List<Expr>& lhs,
+      const Maybe<Expr>& rhs,
+      const Maybe<Expr>& type) {
+    return Assign(Compound::create(TK_ASSIGN, range, {lhs, rhs, type}));
+  }
+
+  List<Expr> lhs_list() const {
+    return List<Expr>(subtree(0));
+  }
+
+  Expr lhs() const {
+    const auto& li = lhs_list();
+    TORCH_INTERNAL_ASSERT(li.size() == 1);
+    return *li.begin();
+  }
+
+  Maybe<Expr> rhs() const {
+    return Maybe<Expr>(subtree(1));
+  }
+
+  Maybe<Expr> type() const {
+    return Maybe<Expr>(subtree(2));
+  }
+};
+
+struct Return : public Stmt {
+  explicit Return(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_RETURN);
+  }
+  Expr expr() const {
+    return Expr(subtree(0));
+  }
+  static Return create(const SourceRange& range, const Expr& value) {
+    return Return(Compound::create(TK_RETURN, range, {value}));
+  }
+};
+
+struct Raise : public Stmt {
+  explicit Raise(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_RAISE);
+  }
+  Expr expr() const {
+    return Expr(subtree(0));
+  }
+  static Raise create(const SourceRange& range, const Expr& expr) {
+    return Raise(Compound::create(TK_RAISE, range, {expr}));
+  }
+};
+
+struct Assert : public Stmt {
+  explicit Assert(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_ASSERT);
+  }
+  Expr test() const {
+    return Expr(subtree(0));
+  }
+  Maybe<Expr> msg() const {
+    return Maybe<Expr>(subtree(1));
+  }
+  static Assert create(
+      const SourceRange& range,
+      const Expr& test,
+      const Maybe<Expr>& msg) {
+    return Assert(Compound::create(TK_ASSERT, range, {test, msg}));
+  }
+};
+
+struct Pass : public Stmt {
+  explicit Pass(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_PASS);
+  }
+  static Pass create(const SourceRange& range) {
+    return Pass(Compound::create(TK_PASS, range, {}));
+  }
+};
+
+struct Dots : public Expr {
+  explicit Dots(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_DOTS);
+  }
+  static Dots create(const SourceRange& range) {
+    return Dots(Compound::create(TK_DOTS, range, {}));
+  }
+};
+
+struct Break : public Stmt {
+  explicit Break(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_BREAK);
+  }
+  static Break create(const SourceRange& range) {
+    return Break(Compound::create(TK_BREAK, range, {}));
+  }
+};
+
+struct Continue : public Stmt {
+  explicit Continue(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_CONTINUE);
+  }
+  static Continue create(const SourceRange& range) {
+    return Continue(Compound::create(TK_CONTINUE, range, {}));
+  }
+};
+
+struct ExprStmt : public Stmt {
+  explicit ExprStmt(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_EXPR_STMT);
+  }
+  Expr expr() {
+    return Expr(subtree(0));
+  }
+  static ExprStmt create(const SourceRange& range, const Expr& list) {
+    return ExprStmt(Compound::create(TK_EXPR_STMT, range, {list}));
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Expressions
+////////////////////////////////////////////////////////////////////////////////
+
+struct BinOp : public Expr {
+  explicit BinOp(const TreeRef& tree) : Expr(tree) {
+    switch (tree->kind()) {
+      case TK_AND:
+      case TK_OR:
+      case '<':
+      case '>':
+      case TK_IS:
+      case TK_ISNOT:
+      case TK_EQ:
+      case TK_LE:
+      case TK_GE:
+      case TK_NE:
+      case '+':
+      case '*':
+      case '/':
+      case '-':
+      case '@':
+      case TK_POW:
+      case TK_LSHIFT:
+      case TK_RSHIFT:
+      case '%':
+      case '&':
+      case '^':
+      case '|':
+      case TK_FLOOR_DIV:
+      case TK_IN:
+        if (tree->trees().size() != 2)
+          throw(
+              ErrorReport(tree)
+              << "BinOp expected 2 subtrees, found " << tree->trees().size());
+        return;
+      default:
+        throw(
+            ErrorReport(tree)
+            << kindToString(tree->kind()) << " is not a valid BinOp");
+    }
+  }
+  Expr lhs() const {
+    return Expr(subtree(0));
+  }
+  Expr rhs() const {
+    return Expr(subtree(1));
+  }
+  static BinOp create(
+      const SourceRange& range,
+      int kind,
+      const Expr& lhs,
+      const Expr& rhs) {
+    return BinOp(Compound::create(kind, range, {lhs, rhs}));
+  }
+};
+
+struct UnaryOp : public Expr {
+  explicit UnaryOp(const TreeRef& tree) : Expr(tree) {
+    switch (tree->kind()) {
+      case TK_UNARY_MINUS:
+      case '~':
+      case TK_NOT:
+        if (tree->trees().size() != 1)
+          throw(
+              ErrorReport(tree)
+              << "UnaryOp expected 1 subtree, found " << tree->trees().size());
+        return;
+      default:
+        throw(
+            ErrorReport(tree)
+            << kindToString(tree->kind()) << " is not a valid UnaryOp");
+    }
+  }
+  static UnaryOp create(const SourceRange& range, int kind, const Expr& expr) {
+    return UnaryOp(Compound::create(kind, range, {expr}));
+  }
+};
+
+struct Const : public Expr {
+  explicit Const(const TreeRef& tree) : Expr(tree) {
+    tree_->matchNumSubtrees(TK_CONST, 1);
+  }
+  bool isFloatingPoint() const {
+    if (isComplex())
+      return false;
+
+    bool is_inf = subtree(0)->stringValue() == "inf";
+    return is_inf ||
+        subtree(0)->stringValue().find_first_of(".eE") != std::string::npos;
+  }
+  bool isIntegral() const {
+    return !isFloatingPoint() && !isComplex();
+  }
+  bool isComplex() const {
+    return subtree(0)->stringValue().find_first_of('j') != std::string::npos;
+  }
+  int64_t asIntegral() const {
+    try {
+      return std::stoll(subtree(0)->stringValue(), nullptr, 0);
+    } catch (const std::out_of_range&) {
+      throw(
+          ErrorReport(range()) << "Integral constant out of range "
+                                  "(must fit in a signed 64 bit integer)");
+    }
+  }
+  double asFloatingPoint() const {
+    // We can't pass in nullptr as the dummy pointer gets dereferenced for
+    // Android version of strtod_c().
+    char* dummy = nullptr;
+    return torch::jit::strtod_c(subtree(0)->stringValue().c_str(), &dummy);
+  }
+  c10::complex<double> asComplex() const {
+    char* dummy = nullptr;
+    auto str = subtree(0)->stringValue();
+    // Complex numbers (a+bj, where a is non-zero) are parsed as an addition
+    // between float/int a and a complex number "bj". When a is 0, a complex
+    // number bj is created as above. So, while parsing the string, we don't
+    // have to worry about the real component of the complex number.
+    auto imag =
+        torch::jit::strtod_c(str.substr(0, str.size() - 1).c_str(), &dummy);
+    return c10::complex<double>(0, imag);
+  }
+  const std::string& text() const {
+    return subtree(0)->stringValue();
+  }
+  static Const create(const SourceRange& range, const std::string& value) {
+    return Const(Compound::create(TK_CONST, range, {String::create(value)}));
+  }
+};
+
+struct StringLiteral : public Expr {
+  explicit StringLiteral(const TreeRef& tree) : Expr(tree) {
+    tree_->matchNumSubtrees(TK_STRINGLITERAL, 1);
+  }
+  const std::string& text() const {
+    return subtree(0)->stringValue();
+  }
+  static StringLiteral create(
+      const SourceRange& range,
+      const std::string& value) {
+    return StringLiteral(
+        Compound::create(TK_STRINGLITERAL, range, {String::create(value)}));
+  }
+};
+
+struct Apply : public Expr {
+  explicit Apply(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_APPLY);
+  }
+  Expr callee() const {
+    return Expr(subtree(0));
+  }
+  List<Expr> inputs() const {
+    return List<Expr>(subtree(1));
+  }
+  List<Attribute> attributes() const {
+    return List<Attribute>(subtree(2));
+  }
+  static Apply create(
+      const SourceRange& range,
+      const Expr& callee,
+      const List<Expr>& inputs,
+      const List<Attribute>& attributes) {
+    return Apply(
+        Compound::create(TK_APPLY, range, {callee, inputs, attributes}));
+  }
+};
+
+struct Select : public Expr {
+  explicit Select(const TreeRef& tree) : Expr(tree) {
+    tree_->match('.');
+  }
+  Expr value() const {
+    return Expr(subtree(0));
+  }
+  Ident selector() const {
+    return Ident(subtree(1));
+  }
+  static Select create(
+      const SourceRange& range,
+      const Expr& value,
+      const Ident& selector) {
+    return Select(Compound::create('.', range, {value, selector}));
+  }
+};
+
+struct SliceExpr : public Expr {
+  explicit SliceExpr(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_SLICE_EXPR);
+  }
+  Maybe<Expr> start() const {
+    return Maybe<Expr>(subtree(0));
+  }
+  Maybe<Expr> end() const {
+    return Maybe<Expr>(subtree(1));
+  }
+  Maybe<Expr> step() const {
+    return Maybe<Expr>(subtree(2));
+  }
+  Expr startOr(int64_t alternative) const {
+    const auto startOption = start();
+    return startOption.present() ? startOption.get() : createInt(alternative);
+  }
+  Expr endOr(int64_t alternative) const {
+    const auto endOption = end();
+    return endOption.present() ? endOption.get() : createInt(alternative);
+  }
+  Expr stepOr(int64_t alternative) const {
+    const auto stepOption = step();
+    return stepOption.present() ? stepOption.get() : createInt(alternative);
+  }
+  static SliceExpr create(
+      const SourceRange& range,
+      const Maybe<Expr>& start,
+      const Maybe<Expr>& end,
+      const Maybe<Expr>& step) {
+    return SliceExpr(
+        Compound::create(TK_SLICE_EXPR, range, {start, end, step}));
+  }
+
+ private:
+  Expr createInt(int64_t value) const {
+    return Expr(Const::create(range(), std::to_string(value)));
+  }
+};
+
+struct Subscript : public Expr {
+  explicit Subscript(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_SUBSCRIPT);
+  }
+  Expr value() const {
+    return Expr(subtree(0));
+  }
+  List<Expr> subscript_exprs() const {
+    return List<Expr>(subtree(1));
+  }
+  static Subscript create(
+      const SourceRange& range,
+      const Expr& value,
+      const List<Expr>& subscript_exprs) {
+    auto whole_range = SourceRange(
+        range.source(), range.start(), subscript_exprs.range().end() + 1);
+    return Subscript(
+        Compound::create(TK_SUBSCRIPT, whole_range, {value, subscript_exprs}));
+  }
+};
+
+struct Var : public Expr {
+  explicit Var(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_VAR);
+  }
+  Ident name() const {
+    return Ident(subtree(0));
+  }
+  static Var create(const SourceRange& range, const Ident& name) {
+    return Var(Compound::create(TK_VAR, range, {name}));
+  }
+};
+
+// WithItem represents an item using with a WithStmt.
+struct WithItem : public Expr {
+  explicit WithItem(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_WITH_ITEM);
+  }
+
+  Expr target() const {
+    return Expr(subtree(0));
+  }
+
+  Maybe<Var> var() const {
+    return Maybe<Var>(subtree(1));
+  }
+
+  static WithItem create(
+      const SourceRange& range,
+      const Expr& target,
+      const Maybe<Var>& var) {
+    return WithItem(Compound::create(TK_WITH_ITEM, range, {target, var}));
+  }
+};
+
+// With represents a with statement consisting of a list of with items and a
+// body of statements.
+struct With : public Stmt {
+  explicit With(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_WITH);
+  }
+
+  List<WithItem> targets() const {
+    return List<WithItem>(subtree(0));
+  }
+
+  List<Stmt> body() const {
+    return List<Stmt>(subtree(1));
+  }
+
+  static With create(
+      const SourceRange& range,
+      const List<WithItem>& targets,
+      const List<Stmt>& body) {
+    return With(Compound::create(TK_WITH, range, {targets, body}));
+  }
+};
+
+struct TernaryIf : public Expr {
+  explicit TernaryIf(const TreeRef& tree) : Expr(tree) {
+    tree_->matchNumSubtrees(TK_IF_EXPR, 3);
+  }
+  Expr cond() const {
+    return Expr(subtree(0));
+  }
+  Expr true_expr() const {
+    return Expr(subtree(1));
+  }
+  Expr false_expr() const {
+    return Expr(subtree(2));
+  }
+  static TernaryIf create(
+      const SourceRange& range,
+      const Expr& cond,
+      const Expr& true_expr,
+      const Expr& false_expr) {
+    return TernaryIf(
+        Compound::create(TK_IF_EXPR, range, {cond, true_expr, false_expr}));
+  }
+};
+
+struct ListLiteral : public Expr {
+  explicit ListLiteral(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_LIST_LITERAL);
+  }
+  List<Expr> inputs() const {
+    return subtree(0);
+  }
+  static ListLiteral create(
+      const SourceRange& range,
+      const List<Expr>& inputs) {
+    return ListLiteral(Compound::create(TK_LIST_LITERAL, range, {inputs}));
+  }
+};
+
+struct TupleLiteral : public Expr {
+  explicit TupleLiteral(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_TUPLE_LITERAL);
+  }
+  List<Expr> inputs() const {
+    return subtree(0);
+  }
+  static TupleLiteral create(
+      const SourceRange& range,
+      const List<Expr>& inputs) {
+    return TupleLiteral(Compound::create(TK_TUPLE_LITERAL, range, {inputs}));
+  }
+};
+
+struct DictLiteral : public Expr {
+  explicit DictLiteral(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_DICT_LITERAL);
+  }
+  List<Expr> key_inputs() const {
+    return subtree(0);
+  }
+  List<Expr> value_inputs() const {
+    return subtree(1);
+  }
+  static DictLiteral create(
+      const SourceRange& range,
+      const List<Expr>& keys,
+      const List<Expr>& values) {
+    return DictLiteral(
+        Compound::create(TK_DICT_LITERAL, range, {keys, values}));
+  }
+};
+
+struct Starred : public Expr {
+  explicit Starred(const TreeRef& tree) : Expr(tree) {
+    tree_->match(TK_STARRED);
+  }
+  Expr expr() const {
+    return Expr(subtree(0));
+  }
+  static Starred create(const SourceRange& range, const Expr& expr) {
+    return Starred(Compound::create(TK_STARRED, range, {expr}));
+  }
+};
+
+struct Delete : public Stmt {
+  explicit Delete(const TreeRef& tree) : Stmt(tree) {
+    tree_->match(TK_DELETE);
+  }
+  List<Expr> targets() const {
+    return subtree(0);
+  }
+  static Delete create(const SourceRange& range, const List<Expr>& targets) {
+    return Delete(Compound::create(TK_DELETE, range, {targets}));
+  }
+};
+
+/*
+ * NOTE: transforming PEP 604 union into equivalent union type
+ *
+ * NOTE: Union[int, float] parses into:
+ * <EXPR> expr:(subscript
+ *  (variable (ident Union))
+ *  (list
+ *    (variable (ident int))
+ *    (variable (ident float))))
+ * <KIND> subscript
+ *
+ * NOTE: (int | float) parses into:
+ * <EXPR> expr:(|
+ *  (variable (ident int))
+ *  (variable (ident float)))
+ * <KIND> |
+ */
+
+inline void _flatten_pep604_union(
+    const torch::jit::Expr& node,
+    std::vector<torch::jit::Expr>* result) {
+  // flatten possibly nested union expressions like (int | (float | str))
+  // into a flat list of expressions like [int, float, str]
+  if (node.kind() == '|') {
+    auto as_binop = torch::jit::BinOp(node);
+    _flatten_pep604_union(as_binop.lhs(), result);
+    _flatten_pep604_union(as_binop.rhs(), result);
+  } else {
+    result->push_back(node);
+  }
+}
+
+inline std::vector<Expr> get_pep604_union_members(const Expr& node) {
+  std::vector<Expr> result;
+  _flatten_pep604_union(node, &result);
+  return result;
+}
+
+// Flattens a PEP 604 union into a classical union.
+// For example, ((x | y) | z) is transformed into Union[x, y, z].
+inline Expr pep604union_to_union(const Expr& expr) {
+  // noop if not a pep604 union
+  if (expr.kind() != '|')
+    return expr;
+
+  // In order to support unions with more than 2 operands ((x|y)|z), we need to
+  // recursively flatten the tree of | expressions.
+  auto members = get_pep604_union_members(expr);
+  auto synthesised_union = Subscript::create(
+      expr.range(),
+      Var::create(expr.range(), Ident::create(expr.range(), "Union")),
+      List<Expr>::create(expr.range(), members));
+#if defined(__clang__)
+  return std::move(synthesised_union);
+#else
+  return synthesised_union;
+#endif
+}
+
+} // namespace torch::jit
+
+namespace std {
+
+template <typename T>
+struct iterator_traits<torch::jit::ListIterator<T>>
+    : std::iterator_traits<torch::jit::TreeList::const_iterator> {};
+
+} // namespace std
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/versioned_symbols.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/versioned_symbols.h
new file mode 100644
index 0000000000000000000000000000000000000000..b8cb099cc27a3237256f2fa139e1973efc4f6ea7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/frontend/versioned_symbols.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <caffe2/serialize/versions.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/api/module.h>
+
+#include <cstdint>
+
+namespace torch::jit {
+// Maps the given symbol into an implementation of its behavior at the
+// given version.
+// See note [Versioned Symbols]
+TORCH_API Symbol
+get_symbol_for_version(const Symbol name, const uint64_t version);
+
+// Maps the given kind to the minimum version that supports it.
+// See note [Dynamic Versions and torch.jit.save vs. torch.save]
+TORCH_API uint64_t get_min_version_for_kind(const NodeKind& kind);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/alias_analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/alias_analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..598426bd6807d67c4e9c8e00441bb0680e0ca91a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/alias_analysis.h
@@ -0,0 +1,368 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/alias_info.h>
+#include <c10/util/flat_hash_map.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/type_hashing.h>
+#include <torch/csrc/jit/passes/create_functional_graphs.h>
+#include <torch/csrc/jit/passes/utils/memory_dag.h>
+
+namespace torch::jit {
+
+class ValueAndMemoryLocationSet;
+
+/**
+ * Alias analysis pass.
+ *
+ * This pass produces an AliasDb that contains aliasing and mutation
+ * information about the graph. Users can use this information to determine
+ * whether mutations to the graph are safe, i.e. they don't reorder/change
+ * nodes in a way that affects output.
+ *
+ * Every value with a mutable type (Tensors, Lists, Tuples, etc.) will be
+ * associated with one or more "alias sets". If two values share an alias set,
+ * that means they may alias, implying that a mutation to one value cannot be
+ * reordered past a use of the other. Only reordering two reads of an alias set
+ * is considered safe.
+ *
+ * There is a special alias set called the "wildcard set", which indicates that
+ * we're not sure what this value may alias. To be conservative, we consider the
+ * wildcard alias set as potentially aliasing any other wildcard value within
+ * the same type class. Whenever a value becomes contained by another value,
+ * such as when a Tensor is appended to a List[Tensor], the contained element
+ * becomes part of the wildcard set.
+ *
+ * Values that contain other mutable types, such as List[Tensor], are
+ * initialized as containing the Wildcard set for all contained mutable types.
+ *
+ * The AliasDb API references the idea of "mutable" vs "immutable"
+ * types. "Mutable" means that the object's value can change, while
+ * "immutable" means that the value is fixed. (For example, `List` is
+ * mutable, so you can add and delete elements from it. On the other
+ * hand, you can't modify a Tuple once you create it, making `Tuple` an
+ * immutable container.)
+ *
+ * `isFrozen` - if the Module is frozen then consider attributes as freshly
+ * created objects. Freezing API invokes alias analysis to check if they are
+ * mutated internally.
+ *
+ * `descendFunctionCalls` - recursively analyze function and method calls
+ * instead of conservative analysis. Generally analysis should be done after
+ * inlining so the implementation for recursive analysis is unoptimized.
+ */
+class AliasDb {
+ public:
+  TORCH_API explicit AliasDb(
+      std::shared_ptr<Graph> graphi,
+      bool isFrozen = false,
+      bool descendFunctionCalls = false);
+  TORCH_API ~AliasDb();
+
+  // There are limitations to what effects the alias analysis can track. Two
+  // kinds of nodes may have untracked effects:
+  // 1. Nodes that write to a value that may alias the graph inputs (since
+  //    the inputs can be used outside the graph).
+  // 2. Nodes that write to something in the wildcard set.
+  //
+  // These nodes are considered not safe to eliminate or mutate under any
+  // circumstances.
+  bool writesToWildcard(Node* n) const;
+
+  // Does `n` write to an alias of one of the values in `vs`?
+  // if `recurseBlocks` is true, consider writes on the nodes in `n`s sub-blocks
+  TORCH_API bool writesToAlias(Node* n, const ValueSet& vs) const;
+
+  // Does `n` write to any of the values in `vls`?
+  TORCH_API bool writesToAlias(Node* n, const ValueAndMemoryLocationSet& vls)
+      const;
+
+  TORCH_API ValueAndMemoryLocationSet getValueAndMemoryLocationSet() const;
+
+  // Does `a` and `b` potentially share a memory location or do either
+  // hold in memory any element that exists in the other
+  TORCH_API bool mayContainAlias(Value* a, Value* b) const;
+
+  TORCH_API bool mayContainAlias(Value* a, const at::ArrayRef<Value*> b) const;
+
+  // Do any values in group `a` share a memory location or hold in memory
+  // any element that exists in group `b`
+  TORCH_API bool mayContainAlias(
+      const at::ArrayRef<Value*> a,
+      const at::ArrayRef<Value*> b) const;
+
+  // Do `a` and `b` potentially share a memory location?
+  TORCH_API bool mayAlias(const Value* a, const Value* b) const;
+  // Do any values in group `a` potentially share a memory location with any
+  // value in group `b`? i.e. may they overlap?
+  TORCH_API bool mayAlias(const ValueSet& a, const ValueSet& b) const;
+
+  // Do any nodes write to an alias set input to `n`?
+  TORCH_API bool hasInputWriters(const Node* n) const;
+
+  // Do any nodes write to an alias set output by `n`?
+  TORCH_API bool hasOutputWriters(const Node* n) const;
+
+  // Do any nodes write to an alias set inputted/outputted by `n`?
+  TORCH_API bool hasWriters(const Node* n) const;
+
+  // Do any nodes write to `v`s memory location?
+  TORCH_API bool hasWriters(const Value* v) const;
+
+  // Is the operation in-place? i.e. doesn't write anywhere but locations it
+  // reads from.
+  TORCH_API bool isMutable(Node* n) const;
+
+  TORCH_API bool escapesScope(const at::ArrayRef<Value*>& vs) const;
+
+  // Is it safe to change whether `a` and `b` alias each other ?
+  TORCH_API bool safeToChangeAliasingRelationship(
+      const at::ArrayRef<Value*>& a,
+      const at::ArrayRef<Value*>& b) const;
+
+  // Move `n` (already in the graph) after `movePoint` in the topological order.
+  //
+  // Tries to preserve value dependencies, so other nodes might be moved. We
+  // make two guarantees about the postcondition of the node list:
+  //   - `n` is directly after `movePoint`.
+  //   - only nodes between `n` and `movePoint` have been moved.
+  //
+  // Returns `false` if it's impossible to move `n` after `MovePoint` without
+  // violating dependencies, otherwise executes the move and returns `true`
+  TORCH_API bool moveAfterTopologicallyValid(Node* n, Node* movePoint);
+  TORCH_API bool moveBeforeTopologicallyValid(Node* n, Node* movePoint);
+
+  bool couldMoveAfterTopologically(Node* n, Node* movePoint);
+  bool couldMoveBeforeTopologically(Node* n, Node* movePoint);
+
+  // For debugging: print alias db state to stdout
+  TORCH_API void dump() const;
+  TORCH_API std::string toString() const;
+
+  // Generates a DOT (www.graphviz.org) graph representation
+  //
+  // Returns `true` if the output file was successfully generated
+  //
+  // WARNING: The output dot file path can't include shell specific notations,
+  //  for example you can't use "~/temp/aliasdb.dot"
+  //  (instead, use "/home/user/temp/aliasdb.dot")
+  //
+  TORCH_API bool dumpToGraphvizFile(const char* filename) const;
+  TORCH_API std::string toGraphviz() const;
+
+  // Returns `true` if the given element is mutable or if it is a
+  // container type with an internal mutable element (e.g.
+  // `Tuple[int, Tensor]` has an internal mutable type `Tensor`, so
+  // it would be considered a "mutable type" in AliasDb)
+  static bool isMutableType(const Value* v);
+  static bool isMutableType(const TypePtr& type);
+
+  /**
+   * Mutation API
+   *
+   * These methods allow you to update AliasDb in-place if you are performing
+   * graph mutation.
+   *
+   * WARNING: These methods should be considered INTERNAL. They do not perform
+   * very many correctness checks, the user is responsible for making sure they
+   * are updating AliasDb correctly. `Lint()`ing the AliasDb can help with
+   * this.
+   */
+  // Copy `existing`s aliasing info to `new_value`, and remove `existing`.
+  TORCH_API void replaceWithNewValue(Value* existing, Value* new_value);
+  // Copy `from`s aliasing info to `to`.
+  TORCH_API void copyValue(Value* from, Value* to);
+  // Create a new `value` that does not alias anything else.
+  TORCH_API void createValue(const Value* value);
+
+  // Enable more precise treatment of prim::TupleConstruct.
+  void enablePreciseTupleContainerAnalysis();
+
+  friend struct MutationRemover;
+  friend class ValueAndMemoryLocationSet;
+
+ private:
+  // Helper for topologically-safe node moves.
+  class WorkingSet;
+  enum class MoveSide { BEFORE, AFTER };
+  bool tryMove(Node* toMove, Node* movePoint, MoveSide moveSide, bool dryRun);
+  void move(Node* toMove, Node* movePoint, MoveSide moveSide);
+  bool isBeforeOrAfter(const Node* n, MoveSide moveSide) const;
+
+  bool isMutableTypeInternal(const Value* v) const;
+  bool isMutableTypeInternal(const TypePtr& type) const;
+
+  /**
+   * Write and read internal API
+   */
+  // Get all the values that `n` writes to.
+  // NOTE: this only returns values directly written to, not aliases thereof
+  //
+  // if `recurseBlocks` is true, gather writes on the nodes in `n`s sub-blocks
+  MemoryLocations getWrites(Node* n) const;
+  void getWritesImpl(Node* n, MemoryLocations& ret) const;
+  // Register the fact that `n` writes to `v`.
+  void registerWrite(const Value* v, Node* n, bool writeToContained = false);
+  // Get all the values that `n` reads from.
+  // if `recurseBlocks` is true, gather reads on the nodes in `n`s sub-blocks
+  MemoryLocations getReads(Node* n) const;
+  void getReadsImpl(Node* n, MemoryLocations& ret) const;
+  MemoryLocations getMemoryLocations(Value* v) const;
+
+  /**
+   * Wildcard methods
+   */
+  // Register `v` as a wildcard value.
+  std::optional<Element*> setWildcard(const Value* v);
+
+  // Is this a value which will not alias?
+  bool nonAliasingValue(const Value* elem) const;
+
+  /**
+   * Special analysis methods
+   */
+  void analyze(const std::shared_ptr<Graph>& graph);
+  void analyze(Block* block);
+  void analyze(Node* node);
+  void analyzeImpl(Node* node);
+  void analyzeIf(Node* node);
+  void analyzeLoop(Node* node);
+  void analyzeSubgraph(Node* node, const std::shared_ptr<Graph>& subgraph);
+  void analyzeSubgraph(Node* node);
+  void analyzeCreator(Node* node);
+  void analyzeExtractor(Node* node);
+  void analyzeChunk(Node* node);
+  void analyzeBroadcastingChunk(Node* node);
+  void analyzeFork(Node* node);
+  void analyzeWait(Node* node);
+  void analyzeAwaitable(Node* node);
+  void analyzeAwaitableWait(Node* node);
+  void analyzeRpcAsync(Node* node);
+  void analyzeBatchNorm(Node* node);
+  void analyzeInstanceNorm(Node* node);
+  void analyzeGradOf(Node* node);
+  void analyzeSetAttr(Node* node);
+  void analyzeConservative(Node* node);
+  void analyzeContainerConstruct(Node* node);
+  bool tryRegisteredAnalysis(Node* node);
+
+  /**
+   * Alias manipulation methods
+   */
+  void makeAllAlias(const std::vector<Value*>& values);
+  void makePointerTo(const Value* value, const Value* to);
+  TORCH_API void addToContainedElements(
+      const Value* element,
+      const Value* container);
+  void mapAliases(at::ArrayRef<Value*> to, at::ArrayRef<Value*> from);
+  void giveFreshAlias(
+      const Value* value,
+      bool add_wildcard_to_contained_elems = true);
+  Element* getOrCreateElement(const Value* value);
+
+  const AliasTypeSet* mapTypeToAliasTypeSetPtr(const TypePtr& type) const;
+  bool functionalNonEscapingListUse(const Use& use) const;
+  bool functionalNonEscapingTupleUse(const Use& use) const;
+
+  std::shared_ptr<Graph> graph_;
+
+  // If the Module is frozen then consider attributes as freshly created
+  // objects. Freezing API invokes alias analysis to check if they are mutated
+  // internally.
+  bool isFrozen_;
+
+  bool descend_function_calls_;
+  std::unordered_map<Graph*, std::vector<std::shared_ptr<Graph>>>
+      function_call_copies_;
+
+  // The points-to graph that stores aliasing relationships
+  std::unique_ptr<MemoryDAGBuilder> memoryDAGBuilder_;
+  std::unique_ptr<MemoryDAG> memoryDAG_;
+
+  // Mapping of values to MemoryDAG elements
+  ska::flat_hash_map<const Value*, Element*> elementMap_;
+  // All wildcard Elements (one for each unique mutable type)
+  ska::flat_hash_map<TypePtr, Element*, HashType, EqualType> wildcardIndex_;
+  Element* getWildcard(const TypePtr& type) const;
+  std::optional<Element*> tryGetOrCreateWildcard(const TypePtr& type);
+  void addContainedTypesToFreshElement(
+      Element* container_elem,
+      const AliasTypeSet& mut_types);
+  void pointUnionTypeElementToAllContainedTypes(
+      Element* container_elem,
+      const AliasTypeSet& mut_types);
+
+  std::vector<Element*> getElements(at::ArrayRef<Value*> vs) const;
+  bool mayAliasWildcard(const Value* v) const;
+  bool mayAliasWildcard(const at::ArrayRef<Value*> vs) const;
+  bool hasWriters(const at::ArrayRef<Value*>& values) const;
+
+  // Cached mapping of type ptrs to their mutable types
+  mutable ska::flat_hash_map<TypePtr, AliasTypeSet> mapped_mutable_types_;
+
+  /**
+   * State for tracking write info.
+   */
+  // Write registry where the analysis can record the writes as it sees them.
+  // This information is later denormalized into various caches to improve query
+  // efficiency.
+  struct WriteRegistry;
+  std::unique_ptr<WriteRegistry> writeRegistry_;
+
+  // Map of nodes to the memory locations that they write to
+  using TWriteIndex = ska::flat_hash_map<Node*, MemoryLocations>;
+  std::optional<TWriteIndex> writeIndex_;
+  // Collection of all memory locations that are written to.
+  std::optional<MemoryLocations> writtenToLocationsIndex_;
+  void buildWrittenToLocationsIndex();
+
+  std::unordered_set<const Value*> wildcards_;
+
+  std::string getElementName(const Element* e) const;
+
+  friend void Lint(const AliasDb* db);
+};
+
+// Helper check that invariants over AliasDb are maintained.
+// Useful if you are using the AliasDb mutation API and want to check you did
+// the right thing.
+TORCH_API void Lint(const AliasDb* db);
+
+/**
+ * ValueAndMemoryLocationSet
+ *
+ * A insert-only set of values which also maintains a MemoryLocations bitset
+ * of the memory locations that the values alias. It is insert-only. It
+ * should be constructed by calling aliasDb.getValueAndMemoryLocationSet().
+ *
+ * WARNING:
+ *  * The AliasDb must not be mutated after construction of a
+ *    ValueAndMemoryLocationsSet, or else the MemoryLocations stored in the
+ *    ValueAndMemoryLocationSet will no longer be accurate.
+ *  * A ValueAndMemoryLocationsSet is tied to an instance of AliasDb but
+ *    does not own the AliasDb. It is the user's responsibility to ensure
+ *    that the AliasDb outlives the ValuesAndMemoryLocationsSet.
+ *
+ * The use case for this is to be able to implement writesToAlias
+ * more efficiently for a set of values.
+ */
+class ValueAndMemoryLocationSet {
+ public:
+  TORCH_API void insert(Value* v);
+  TORCH_API ValueSet& getValueSet();
+
+  friend class AliasDb;
+
+ private:
+  ValueAndMemoryLocationSet(const AliasDb* db) : aliasDb_(db) {}
+
+  const AliasDb* aliasDb_;
+  ValueSet valueSet_;
+  MemoryLocations memoryLocations_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/attributes.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/attributes.h
new file mode 100644
index 0000000000000000000000000000000000000000..35dfe7a87642f60500ce791979e3f9d610e84994
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/attributes.h
@@ -0,0 +1,185 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/Tensor.h>
+#include <string>
+#include <vector>
+
+#include <ATen/core/jit_type_base.h>
+#include <ATen/core/symbol.h>
+
+#include <torch/csrc/Export.h>
+
+namespace torch::jit {
+
+using ::c10::Symbol;
+
+constexpr int max_tensor_display_size = 10;
+
+enum class AttributeKind {
+  f,
+  fs,
+  c,
+  cs,
+  i,
+  is,
+  s,
+  ss,
+  t,
+  ts,
+  g,
+  gs,
+  ty,
+  tys,
+  ival
+};
+static inline const char* toString(AttributeKind kind) {
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+  static constexpr const char* names[] = {
+      "f",
+      "c",
+      "cs",
+      "fs",
+      "i",
+      "is",
+      "s",
+      "ss",
+      "t",
+      "ts",
+      "g",
+      "gs",
+      "ty",
+      "tys",
+      "ival"};
+  AT_ASSERT(size_t(kind) < sizeof(names) / sizeof(*names));
+  return names[int(kind)];
+}
+
+struct AttributeValue {
+  AttributeValue(Symbol name) : name(name) {}
+  using Ptr = std::unique_ptr<AttributeValue>;
+  Symbol name;
+  virtual AttributeKind kind() const = 0;
+  virtual Ptr clone() const = 0;
+  virtual ~AttributeValue() = default;
+};
+
+template <typename T, AttributeKind Kind>
+struct ScalarAttributeValue : public AttributeValue {
+  using ConstructorType = T;
+  using ValueType = T;
+  ScalarAttributeValue(Symbol name, ConstructorType value_)
+      : AttributeValue(name), value_(std::move(value_)) {}
+  ValueType& value() {
+    return value_;
+  }
+  Ptr clone() const override {
+    return Ptr(new ScalarAttributeValue(name, value_));
+  }
+  AttributeKind kind() const override {
+    return Kind;
+  }
+
+ private:
+  ValueType value_;
+};
+
+template <typename T, AttributeKind Kind>
+struct VectorAttributeValue : public AttributeValue {
+  using ConstructorType = std::vector<T>;
+  using ValueType = std::vector<T>;
+  VectorAttributeValue(Symbol name, ConstructorType value_)
+      : AttributeValue(name), value_(std::move(value_)) {}
+  ValueType& value() {
+    return value_;
+  }
+  AttributeKind kind() const override {
+    return Kind;
+  }
+  std::unique_ptr<AttributeValue> clone() const override {
+    auto copy = value_;
+    return Ptr(new VectorAttributeValue(name, std::move(copy)));
+  }
+
+ private:
+  ValueType value_;
+};
+
+using ComplexAttr =
+    ScalarAttributeValue<c10::complex<double>, AttributeKind::c>;
+using ComplexValsAttr =
+    VectorAttributeValue<c10::complex<double>, AttributeKind::cs>;
+using FloatAttr = ScalarAttributeValue<double, AttributeKind::f>;
+using FloatsAttr = VectorAttributeValue<double, AttributeKind::fs>;
+using IntAttr = ScalarAttributeValue<int64_t, AttributeKind::i>;
+using IntsAttr = VectorAttributeValue<int64_t, AttributeKind::is>;
+using StringAttr = ScalarAttributeValue<std::string, AttributeKind::s>;
+using StringsAttr = VectorAttributeValue<std::string, AttributeKind::ss>;
+using TensorAttr = ScalarAttributeValue<at::Tensor, AttributeKind::t>;
+using TensorsAttr = VectorAttributeValue<at::Tensor, AttributeKind::ts>;
+using TypeAttr = ScalarAttributeValue<c10::TypePtr, AttributeKind::ty>;
+using TypesAttr = VectorAttributeValue<c10::TypePtr, AttributeKind::tys>;
+using IValueAttr = ScalarAttributeValue<at::IValue, AttributeKind::ival>;
+
+struct Graph;
+
+// We special case Graph attributes like this because we want to ensure that
+// Graph::copy() is called when we clone() these attributes.
+struct TORCH_API GraphAttr : public AttributeValue {
+  using ConstructorType = std::shared_ptr<Graph>;
+  using ValueType = std::shared_ptr<Graph>;
+  GraphAttr(Symbol name, ConstructorType value_)
+      : AttributeValue(name), value_(std::move(value_)) {}
+  ValueType& value() {
+    return value_;
+  }
+  Ptr clone() const override;
+  AttributeKind kind() const override {
+    return AttributeKind::g;
+  }
+
+ private:
+  std::shared_ptr<Graph> value_;
+};
+
+struct TORCH_API GraphsAttr : public AttributeValue {
+  using ConstructorType = std::vector<std::shared_ptr<Graph>>;
+  using ValueType = std::vector<std::shared_ptr<Graph>>;
+  GraphsAttr(Symbol name, ConstructorType value_)
+      : AttributeValue(name), value_(std::move(value_)) {}
+  ValueType& value() {
+    return value_;
+  }
+  AttributeKind kind() const override {
+    return AttributeKind::gs;
+  }
+  std::unique_ptr<AttributeValue> clone() const override;
+
+ private:
+  ValueType value_;
+};
+
+struct IRAttributeError : public std::exception {
+  IRAttributeError(Symbol name, bool defined) {
+    std::stringstream ss;
+    // NOLINTNEXTLINE(bugprone-branch-clone)
+    if (!defined) {
+      ss << "required keyword attribute '" << name.toUnqualString()
+         << "' is undefined";
+    } else {
+      ss << "required keyword attribute '" << name.toUnqualString()
+         << "' has the wrong type";
+    }
+    msg = ss.str();
+  }
+  const char* what() const noexcept override {
+    return msg.c_str();
+  }
+
+ private:
+  std::string msg;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/constants.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/constants.h
new file mode 100644
index 0000000000000000000000000000000000000000..f7fd7f01304ca0fff41a52a4f1de7fd7b1e9ca6d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/constants.h
@@ -0,0 +1,65 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/ir/scope.h>
+
+// helpers for handling constants in the IR
+// - create constant nodes from ints, floats, complex, intlist, Tensors, and
+// other types
+// - implement primitive constant ops.
+
+namespace torch::jit {
+
+using ::c10::IValue;
+
+struct Graph;
+struct Value;
+
+// thrown when insertConstant cannot encode the IValue into a graph
+struct TORCH_API constant_not_supported_error : public std::runtime_error {
+  using runtime_error::runtime_error;
+};
+
+TORCH_API Value* insertConstant(
+    Graph& g,
+    const IValue& val,
+    std::optional<SourceRange> loc = std::nullopt,
+    std::optional<ScopePtr> scope = std::nullopt);
+
+// note: prefer g.insertConsant(val, loc) which does exactly the same thing
+// this function is only declared/defined here because its implementation is
+// closely related to the implementation of prim::Constant that is also in
+// constants.cpp.
+//
+// returns a std::nullopt if the IValue kind cannot be inserted as a constant
+TORCH_API std::optional<Value*> tryInsertConstant(
+    Graph& g,
+    const IValue& val,
+    std::optional<SourceRange> loc = std::nullopt,
+    std::optional<ScopePtr> scope = std::nullopt);
+
+////////////////////////////////////////////////////////////////////////////////
+// Helper for retrieving constants
+////////////////////////////////////////////////////////////////////////////////
+
+// attempt to convert a (possibly constant) Value* into an interpreter value
+// (IValue). returns std::nullopt if the Value* was not constant
+TORCH_API std::optional<IValue> toIValue(const Value* v);
+
+// if a value is a constant then try to turn into type T using the
+// same rules as the interpreter
+template <typename T>
+std::optional<T> constant_as(const Value* v) {
+  if (auto ivalue = toIValue(v)) {
+    return ivalue->to<T>();
+  }
+  return std::nullopt;
+}
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/graph_node_list.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/graph_node_list.h
new file mode 100644
index 0000000000000000000000000000000000000000..e1ffa81efea513b2046ea9f86acc495b1f48e45c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/graph_node_list.h
@@ -0,0 +1,204 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/Exception.h>
+
+namespace torch::jit {
+
+// Intrusive doubly linked lists with sane reverse iterators.
+// The header file is named generic_graph_node_list.h because it is ONLY
+// used for Graph's Node lists, and if you want to use it for other
+// things, you will have to do some refactoring.
+//
+// At the moment, the templated type T must support a few operations:
+//
+//  - It must have a field: T* next_in_graph[2] = { nullptr, nullptr };
+//    which are used for the intrusive linked list pointers.
+//
+//  - It must have a method 'destroy()', which removes T from the
+//    list and frees a T.
+//
+// In practice, we are only using it with Node and const Node.  'destroy()'
+// needs to be renegotiated if you want to use this somewhere else.
+//
+// Regardless of the iteration direction, iterators always physically point
+// to the element they logically point to, rather than
+// the off-by-one behavior for all standard library reverse iterators like
+// std::list.
+
+// The list is includes two sentinel nodes, one at the beginning and one at the
+// end with a circular link between them. It is an error to insert nodes after
+// the end sentinel node but before the beginning node:
+
+// Visualization showing only the next() links:
+//  HEAD -> first -> second  -> ... -> last -> TAIL
+//   ^------------------------------------------
+
+// Visualization showing only the prev() links:
+//  HEAD <- first <- second  <- ... <- last <- TAIL
+//   ------------------------------------------^
+
+static constexpr int kNextDirection = 0;
+static constexpr int kPrevDirection = 1;
+
+template <typename T>
+struct generic_graph_node_list;
+
+template <typename T>
+struct generic_graph_node_list_iterator;
+
+struct Node;
+using graph_node_list = generic_graph_node_list<Node>;
+using const_graph_node_list = generic_graph_node_list<const Node>;
+using graph_node_list_iterator = generic_graph_node_list_iterator<Node>;
+using const_graph_node_list_iterator =
+    generic_graph_node_list_iterator<const Node>;
+
+template <typename T>
+struct generic_graph_node_list_iterator {
+  generic_graph_node_list_iterator() : cur(nullptr), d(kNextDirection) {}
+  generic_graph_node_list_iterator(T* cur, int d) : cur(cur), d(d) {}
+  generic_graph_node_list_iterator(
+      const generic_graph_node_list_iterator& rhs) = default;
+  generic_graph_node_list_iterator(
+      generic_graph_node_list_iterator&& rhs) noexcept = default;
+  generic_graph_node_list_iterator& operator=(
+      const generic_graph_node_list_iterator& rhs) = default;
+  generic_graph_node_list_iterator& operator=(
+      generic_graph_node_list_iterator&& rhs) noexcept = default;
+  T* operator*() const {
+    return cur;
+  }
+  T* operator->() const {
+    return cur;
+  }
+  generic_graph_node_list_iterator& operator++() {
+    AT_ASSERT(cur);
+    cur = cur->next_in_graph[d];
+    return *this;
+  }
+  generic_graph_node_list_iterator operator++(int) {
+    generic_graph_node_list_iterator old = *this;
+    ++(*this);
+    return old;
+  }
+  generic_graph_node_list_iterator& operator--() {
+    AT_ASSERT(cur);
+    cur = cur->next_in_graph[reverseDir()];
+    return *this;
+  }
+  generic_graph_node_list_iterator operator--(int) {
+    generic_graph_node_list_iterator old = *this;
+    --(*this);
+    return old;
+  }
+
+  // erase cur without invalidating this iterator
+  // named differently from destroy so that ->/. bugs do not
+  // silently cause the wrong one to be called.
+  // iterator will point to the previous entry after call
+  void destroyCurrent() {
+    T* n = cur;
+    cur = cur->next_in_graph[reverseDir()];
+    n->destroy();
+  }
+  generic_graph_node_list_iterator reverse() {
+    return generic_graph_node_list_iterator(cur, reverseDir());
+  }
+
+ private:
+  int reverseDir() {
+    return d == kNextDirection ? kPrevDirection : kNextDirection;
+  }
+  T* cur;
+  int d; // direction 0 is forward 1 is reverse, see next_in_graph
+};
+
+template <typename T>
+struct generic_graph_node_list {
+  using iterator = generic_graph_node_list_iterator<T>;
+  using const_iterator = generic_graph_node_list_iterator<const T>;
+  generic_graph_node_list_iterator<T> begin() {
+    return generic_graph_node_list_iterator<T>(head->next_in_graph[d], d);
+  }
+  generic_graph_node_list_iterator<const T> begin() const {
+    return generic_graph_node_list_iterator<const T>(head->next_in_graph[d], d);
+  }
+  generic_graph_node_list_iterator<T> end() {
+    return generic_graph_node_list_iterator<T>(head->next_in_graph[!d], d);
+  }
+  generic_graph_node_list_iterator<const T> end() const {
+    return generic_graph_node_list_iterator<const T>(
+        head->next_in_graph[!d], d);
+  }
+  generic_graph_node_list_iterator<T> rbegin() {
+    return reverse().begin();
+  }
+  generic_graph_node_list_iterator<const T> rbegin() const {
+    return reverse().begin();
+  }
+  generic_graph_node_list_iterator<T> rend() {
+    return reverse().end();
+  }
+  generic_graph_node_list_iterator<const T> rend() const {
+    return reverse().end();
+  }
+  generic_graph_node_list reverse() {
+    return generic_graph_node_list(head->next_in_graph[!d], !d);
+  }
+  const generic_graph_node_list reverse() const {
+    return generic_graph_node_list(head->next_in_graph[!d], !d);
+  }
+  T* front() {
+    return head->next_in_graph[d];
+  }
+  const T* front() const {
+    return head->next_in_graph[d];
+  }
+  T* back() {
+    return head->next_in_graph[!d];
+  }
+  const T* back() const {
+    return head->next_in_graph[!d];
+  }
+  generic_graph_node_list(T* head, int d) : head(head), d(d) {}
+
+ private:
+  T* head; // both head and tail are sentinel nodes
+           // the first real node is head->next_in_graph[d]
+           // the tail sentinel is head->next_in_graph[!d]
+  int d;
+};
+
+template <typename T>
+static inline bool operator==(
+    generic_graph_node_list_iterator<T> a,
+    generic_graph_node_list_iterator<T> b) {
+  return *a == *b;
+}
+
+template <typename T>
+static inline bool operator!=(
+    generic_graph_node_list_iterator<T> a,
+    generic_graph_node_list_iterator<T> b) {
+  return *a != *b;
+}
+
+} // namespace torch::jit
+
+namespace std {
+
+template <typename T>
+struct iterator_traits<torch::jit::generic_graph_node_list_iterator<T>> {
+  using difference_type = int64_t;
+  using value_type = T*;
+  using pointer = T**;
+  using reference = T*&;
+  using iterator_category = bidirectional_iterator_tag;
+};
+
+} // namespace std
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/graph_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/graph_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..fe3c83f45e44792baf30e582f6cdbada415afc23
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/graph_utils.h
@@ -0,0 +1,28 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <vector>
+
+namespace torch::jit {
+
+TORCH_API TypePtr getTensorType(const at::Tensor& t, bool complete);
+
+TORCH_API TypePtr inferShapeAndTypeForInput(
+    TypePtr input_type,
+    Stack::const_iterator& s_iter,
+    const Stack::const_iterator& s_iter_end,
+    bool complete);
+
+TORCH_API void setInputTensorTypes(
+    Graph& g,
+    const Stack& stack,
+    bool complete,
+    const std::vector<int>& param_count_list = {});
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/ir.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/ir.h
new file mode 100644
index 0000000000000000000000000000000000000000..6fe375e6ad1b0d05f0e38cc1b697f1a3eb285301
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/ir.h
@@ -0,0 +1,1828 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/attributes.h>
+#include <torch/csrc/jit/ir/graph_node_list.h>
+#include <torch/csrc/jit/ir/named_value.h>
+#include <torch/csrc/jit/ir/scope.h>
+#include <torch/csrc/jit/runtime/operator.h>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/utils/python_stub.h>
+#include <torch/csrc/utils/schema_info.h>
+
+#include <ATen/Utils.h>
+#include <ATen/core/Tensor.h>
+#include <ATen/core/dynamic_type.h>
+#include <ATen/core/enum_type.h>
+#include <ATen/core/functional.h>
+#include <ATen/core/interned_strings.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/Exception.h>
+#include <optional>
+
+#include <functional>
+#include <iosfwd>
+#include <unordered_set>
+#include <vector>
+
+// Forward declare, the real meat is in python_ir.cpp
+template <class T>
+class THPPointer;
+using THPObjectPtr = THPPointer<PyObject>;
+using pyobj_list = std::vector<THPObjectPtr>;
+
+namespace torch::jit {
+namespace utils {
+TORCH_API std::string getNodesModuleHierarchy(const Node& n);
+} // namespace utils
+class AliasDb;
+
+using ::c10::Argument;
+using ::c10::FunctionSchema;
+using ::c10::Symbol;
+
+using ::c10::ivalue::Shared;
+
+using ::c10::IValue;
+using ::c10::ivalue::Future;
+
+using ::c10::ivalue::ConstantString;
+
+#define C10_USING(T) using ::c10::T;
+C10_FORALL_TYPES(C10_USING)
+#undef C10_USING
+
+#define C10_USING(T) using ::c10::T##Ptr;
+C10_FORALL_TYPES(C10_USING)
+#undef C10_USING
+
+using ::c10::Type;
+using ::c10::TypeEnv;
+using ::c10::TypePtr;
+
+using ::c10::getTypePtr;
+using ::c10::MatchTypeReturn;
+using ::c10::TypeKind;
+
+using ::c10::fmap;
+
+namespace prim = ::c10::prim;
+namespace attr = ::c10::attr;
+namespace aten = ::c10::aten;
+namespace cuda = ::c10::cuda;
+
+struct Function;
+struct GraphFunction;
+struct MatchedSchema;
+
+// A Graph represents one "function" of computation.
+// It uses a simple ownership model where the graph owns all the nodes inside
+// it. All references inside the graph are raw pointers. Destroying the Graph
+// will invalidate any pointers to nodes in the graph.
+struct Graph;
+
+// Node is the base class of the IR graph. It represents one computation
+// and dependencies on a list of Values. The "prim-ops", so to speak.
+struct Node;
+
+// A Value represents an input or output to node that is either a
+// Tensor or an opaque Handle object, as determined by type().
+struct Value;
+
+TORCH_API std::ostream& operator<<(std::ostream& out, const Graph& g);
+TORCH_API std::ostream& operator<<(std::ostream& out, const Node& n);
+
+// A list of nodes, with inputs and outputs
+struct Block;
+
+// Each use is represented by this type, see 'Node::uses()'
+// 'user' is the consumer of the value, 'offset' is the index into
+// 'user's input this where the producers will be found.
+struct Use {
+  Use(Node* user, size_t offset) : user(user), offset(offset) {}
+  Node* user;
+  size_t offset;
+
+  bool operator==(const Use& b) {
+    return user == b.user && offset == b.offset;
+  }
+};
+
+// Note [User node does not uniquely identify use]
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// A while back, we wrote some code manipulating uses that looked like this:
+//
+//    for (auto& use : used_val->uses_) {
+//      if (use.user == this_node) {
+//        use.offset += 1;
+//        break;
+//      }
+//    }
+//
+// This code is trying to find a particular use (our node's use) to update it.
+// However, it's wrong: there may be *multiple* uses of a value %x in a node,
+// as might be the case in this IR:
+//
+//    %y = Add %x %x
+//
+// In this case, there are two uses of %x whose user is the node 'Add %x %x'.
+// So, "use induced by this node" is not a well-formed concept.
+//
+// If you are looking for "use induced by an input", it's best to use
+// findUseForInput() to get it.
+
+// the list types are intentionally simple, but we type-def
+// them here so if we need to change them, refactoring will be easier
+using node_list = std::vector<Node*>;
+using value_list = std::vector<Value*>;
+using use_list = std::vector<Use>;
+template <typename T>
+using ArrayRef = at::ArrayRef<T>;
+using NodeKind = Symbol;
+using topo_position_t = int64_t;
+using ValueSet = std::unordered_set<const Value*>;
+
+struct OperatorSet;
+template <typename T>
+struct OperatorMap;
+
+// This is a wrapper to allow invalidating the Python object
+// safely when the C++ object for a Node/Value/Block is deleted
+// like much of graph, it isn't safe for different threads to
+// access the same graph
+template <typename T>
+struct Wrap {
+  explicit Wrap(T* p) : elem(p) {}
+  void clear() {
+    if (clear_cb) {
+      clear_cb(elem);
+    }
+    elem = nullptr;
+  }
+  T* elem;
+  void (*clear_cb)(void*){nullptr};
+};
+
+struct Value {
+  AT_DISALLOW_COPY_AND_ASSIGN(Value);
+  Value(Node* node_, size_t offset_);
+
+ private:
+  friend struct Node;
+  friend struct Graph;
+  Node* node_;
+  size_t offset_;
+  size_t unique_ = 0; // unique id
+  use_list uses_;
+  std::string unique_name_;
+  TypePtr type_;
+  // a managing wrapper for Python to allow invalidation
+  std::shared_ptr<Wrap<Value>> wrap_;
+
+ public:
+  Value* setType(TypePtr type);
+  TORCH_API void inferTypeFrom(const at::Tensor& output);
+  TORCH_API void inferTypeFrom(
+      const c10::intrusive_ptr<c10::ivalue::Object>& output);
+  const TypePtr& type() const {
+    AT_ASSERT(type_ != nullptr);
+    return type_;
+  }
+  bool requires_grad() const {
+    return type()->requires_grad();
+  }
+  bool isCompleteTensor() const {
+    if (auto pt = type()->cast<TensorType>()) {
+      return pt->isComplete();
+    }
+    return false;
+  }
+  TORCH_API bool mustBeNone() const;
+  TORCH_API bool mustNotBeNone() const;
+  size_t unique() const {
+    return unique_;
+  }
+  bool hasDebugName() const {
+    return !unique_name_.empty();
+  }
+  static bool isValidName(const std::string& name);
+  TORCH_API Value* setDebugName(const std::string& name);
+  std::string debugName() const {
+    if (hasDebugName()) {
+      return unique_name_;
+    }
+    return std::to_string(unique());
+  }
+  TORCH_API std::string debugNameBase() const;
+  Node* node() {
+    return node_;
+  }
+  size_t offset() const {
+    return offset_;
+  }
+  void setOffset(size_t offset) {
+    offset_ = offset;
+  }
+  const Node* node() const {
+    return node_;
+  }
+
+  /**
+   * @warning NEVER pass raw pointer of smart pointer managed Graph to Python.
+   * Check #87343 for details.
+   */
+  Graph* owningGraph();
+  const Graph* owningGraph() const;
+  // TODO: make this more const correct
+  const use_list& uses() const {
+    return uses_;
+  }
+
+  bool hasUses() const {
+    return !uses().empty();
+  }
+
+  TORCH_API void replaceFirstUseWith(Value* newValue);
+
+  // Replaces all uses of this value with 'newValue'.
+  //
+  // Given:   %3 = f(%1, %2)
+  //          %4 = g(%3)
+  //          %5 = h(%3, %3)
+  // Execute: %3.replaceAllUsesWith(%6)
+  // Result:  %3 = f(%1, %2)
+  //          %4 = g(%6)
+  //          %5 = h(%6, %6)
+  TORCH_API void replaceAllUsesWith(Value* newValue);
+
+  // Replaces all uses of this value with 'newValue' after 'node'.
+  // Given:   %3 = f(%1, %2)
+  //          %4 = g(%3)
+  //          %5 = inplace_(%3)
+  //          %6 = h(%3, %3)
+  // Execute: %3.replaceAllUsesAfterNodeWith(%5.node(), %5)
+  // Result:  %3 = f(%1, %2)
+  //          %4 = g(%3)
+  //          %5 = inplace_(%3)
+  //          %6 = h(%5, %5)
+  // XXX: does not check scoping legality, consider using
+  // replaceAllUsesDominatedByNodeWith
+  TORCH_API void replaceAllUsesAfterNodeWith(const Node* node, Value* newValue);
+
+  // Replaces all uses of this value with 'newValue' that are dominated by
+  // 'node'. Given:
+  // x = op(...).
+  // if cond:
+  //    z = foo(..)
+  //    bar(x)
+  // else:
+  //    print(x)
+  // x.replaceAllUsesDominatedByNodeWith(foo, z) would replace bar(x)
+  // but not print(x) because print is not dominated by foo.
+  // replaceAllUsesAfterNode does not check domination, so in this example
+  // it would produce invalid IR.
+  TORCH_API void replaceAllUsesDominatedByNodeWith(
+      const Node* node,
+      Value* newValue);
+
+  TORCH_API Value* copyMetadata(Value* from);
+
+  TORCH_API std::shared_ptr<Wrap<Value>> wrap() {
+    if (!wrap_) {
+      wrap_ = std::make_shared<Wrap<Value>>(this);
+    }
+    return wrap_;
+  }
+
+  virtual ~Value() {
+    if (wrap_) {
+      wrap_->clear();
+    }
+  }
+};
+
+struct TORCH_API Node {
+  AT_DISALLOW_COPY_AND_ASSIGN(Node);
+  friend struct Graph;
+  friend struct Block;
+  friend struct Value;
+  friend graph_node_list;
+  friend const_graph_node_list;
+  friend graph_node_list_iterator;
+  friend const_graph_node_list_iterator;
+
+ private:
+  const NodeKind kind_;
+  std::vector<Value*> inputs_;
+  std::vector<Value*> outputs_;
+  // subblocks
+  std::vector<Block*> blocks_;
+  Graph* graph_;
+  Block* owning_block_{nullptr};
+  std::optional<SourceRange> source_range_;
+  ScopePtr scope_;
+  std::optional<InlinedCallStackPtr> callstack_;
+  // Assumes FunctionSchemas are persistent, so we don't manage their lifetime.
+  // This field is effective a cache that's populated on attribute lookups and
+  // invalidated every time we perform an operation that could potentially
+  // change the schema. note: mutable because schema_ is effectively a cache
+  mutable const Operator* op_{nullptr};
+  topo_position_t topo_position_ = 0;
+  // a managing wrapper for Python to allow invalidation
+  std::shared_ptr<Wrap<Node>> wrap_;
+  // Stores the full schema name, if the operator is historic
+  // When the operator is deprecated or the name of the operator
+  // is changed, we need to rely on this name
+  // to retrieve old schemas to successfully apply upgraders
+  // for this operator.
+  std::optional<std::string> historic_schema_name_ = std::nullopt;
+
+ protected:
+  Node(Graph* graph_, NodeKind kind_); // defined after graph
+ public:
+  // Each Node but Return/Param Nodes are associated with exactly one
+  // place in the Node list of the Graph. The Graph itself is a circular
+  // doubly-linked list. The Return Node is used as the sentinel for the
+  // "beginning"/"end" of the list. This means that you can tell when
+  // you've traversed the entire list without means worrying about null
+  // pointers. `next_in_graph[0]` is the pointer to the next Node, while
+  // `next_in_graph[1]` is the pointer to the previous Node. The
+  // linked list is implemented as an array to allow the same iterator
+  // class for forward and reversed Node lists. Taken together, this
+  // list also represents a topological sort of the Nodes in the Graph.
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-non-private-member-variables-in-classes,modernize-avoid-c-arrays)
+  Node* next_in_graph[2] = {nullptr, nullptr};
+
+  std::shared_ptr<Wrap<Node>> wrap() {
+    if (!wrap_) {
+      wrap_ = std::make_shared<Wrap<Node>>(this);
+    }
+    return wrap_;
+  }
+
+  const std::optional<std::string> getHistoricSchemaName() {
+    return historic_schema_name_;
+  }
+
+  void setHistoricSchemaName(const std::string& name) {
+    historic_schema_name_ = name;
+  }
+
+  Node*& next() {
+    return next_in_graph[kNextDirection];
+  }
+  Node*& prev() {
+    return next_in_graph[kPrevDirection];
+  }
+  Node* const& next() const {
+    return next_in_graph[kNextDirection];
+  }
+  Node* const& prev() const {
+    return next_in_graph[kPrevDirection];
+  }
+
+  NodeKind kind() const {
+    return kind_;
+  }
+  Node* setSourceRange(SourceRange r) {
+    source_range_ = std::move(r);
+    return this;
+  }
+  SourceRange sourceRange() const;
+
+  /**
+   * @warning NEVER pass raw pointer of smart pointer managed Graph to Python.
+   * Check #87343 for details.
+   */
+  Graph* owningGraph() {
+    return graph_;
+  }
+  const Graph* owningGraph() const {
+    return graph_;
+  }
+  Block* owningBlock() {
+    return owning_block_;
+  }
+  const Block* owningBlock() const {
+    return owning_block_;
+  }
+  ScopePtr scope() {
+    return scope_;
+  }
+  void setScope(ScopePtr scope) {
+    scope_ = std::move(scope);
+  }
+  std::string scopeName() const {
+    if (!scope_) {
+      return "";
+    }
+    return scope_->namesFromRoot();
+  }
+
+  // Copies the source range, scope and callstack from another node.
+  Node* copyMetadata(Node* from) {
+    this->setSourceRange(from->sourceRange());
+    this->setScope(from->scope());
+    if (auto cs = from->callstack()) {
+      this->setCallStack(*cs);
+    }
+    return this;
+  }
+
+  std::optional<InlinedCallStackPtr> callstack() const {
+    return callstack_;
+  }
+  void setCallStack(InlinedCallStackPtr cs) {
+    callstack_ = std::move(cs);
+  }
+
+  // NB: This returns an ArrayRef; that means that it will
+  // get invalidated if you resize inputs (e.g., using addInput)
+  // We can't return a std::vector<Node*>& because there's no
+  // way to soundly cast to std::vector<const Node*> (an insane
+  // implementation of std::vector could make this representationally
+  // different.)
+  at::ArrayRef<Value*> inputs() {
+    return inputs_;
+  }
+  at::ArrayRef<const Value*> inputs() const {
+    // Vectors are not convertible in const-ness of elements, but
+    // raw pointers are.
+    return {inputs_.data(), inputs_.size()};
+  }
+  // NB: This returns an ArrayRef; that means that it will
+  // get invalidated if you resize inputs (e.g., using addInput)
+  // We can't return a std::vector<Node*>& because there's no
+  // way to soundly cast to std::vector<const Node*> (an insane
+  // implementation of std::vector could make this representationally
+  // different.)
+  at::ArrayRef<Value*> outputs() {
+    return outputs_;
+  }
+  at::ArrayRef<const Value*> outputs() const {
+    // Vectors are not convertible in const-ness of elements, but
+    // raw pointers are.
+    return {outputs_.data(), outputs_.size()};
+  }
+  Value* output(size_t i) const {
+    return outputs_.at(i);
+  }
+  bool hasUses() const {
+    for (auto o : outputs()) {
+      if (!o->uses().empty()) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  void replaceAllUsesWith(Node* n);
+
+  // replaces `this` with a new node with the same inputs and outputs
+  // but a new node symbol. does not destroy `this`
+  Node* replaceWithNewSymbol(Symbol new_symbol);
+
+  // Checks if this node is dominated by `dominator` which means that
+  // `dominator` will always be executed before `this` and `dominator`
+  // is in scope of `this.
+  bool isDominatedBy(const Node* dominator) const;
+
+  // lots of things like chunk have a single input or single output, so we have
+  // a helper to make accessing it easier
+  Value* input() {
+    AT_ASSERT(inputs_.size() == 1);
+    return inputs_.at(0);
+  }
+  Value* output() {
+    AT_ASSERT(outputs_.size() == 1);
+    return outputs_.at(0);
+  }
+  const Value* output() const {
+    AT_ASSERT(outputs_.size() == 1);
+    return outputs_.at(0);
+  }
+  const Value* input() const {
+    AT_ASSERT(inputs_.size() == 1);
+    return inputs_.at(0);
+  }
+  // Access a particular input.  This is a checked index.
+  Value* input(size_t i) const {
+    return inputs_.at(i);
+  }
+
+  bool hasNamedInput(const std::string& unqualName) const;
+  Value* namedInput(const std::string& unqualName) const;
+  Value* namedInput(Symbol name) const;
+
+  std::optional<IValue> get(Symbol name) const;
+
+  template <typename T>
+  std::optional<T> get(Symbol name) const {
+    if (auto v = get(name)) {
+      return v->template to<T>();
+    }
+    return std::nullopt;
+  }
+
+  // Returns true if the value of input name is statically known
+  bool is_constant(Symbol name) const {
+    return static_cast<bool>(get(name));
+  }
+  bool mustBeNone() const;
+
+  bool isNondeterministic() const;
+  bool hasSideEffects() const;
+
+  // instructions lowered by the interpreter and not run in the optimized graph
+  bool notExecutedOp() const {
+    return kind_ == prim::Constant || kind_ == prim::profile ||
+        kind_ == prim::profile_ivalue;
+  }
+
+  // Graphs
+
+  // Note [Topological invariant]
+  // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+  // We always maintain an up-to-date topological ordering of all nodes via
+  // the next()/prev() links.  All transformations to graphs must preserve
+  // this topological ordering: for example, it is only valid to 'addInput'
+  // with an input which is topologically before the current node.
+  //
+  // Usually, it is obvious whether or not topological order is maintained;
+  // for example, if you are adding nodes to the end of the topsort, it's
+  // impossible for them to refer to inputs that are not in the topsort.
+  // If it is not obvious, please comment accordingly.
+
+  // Add 'node' as an input to 'this' at the end of existing
+  // arguments.  Returns the added node for ease of chaining.
+  //
+  // Given:   %3 = f(%1, %2)
+  // Execute: %3.addInput(%4)
+  // Result:  %3 = f(%1, %2, %4)
+  Value* addInput(Value* value);
+
+  // Add 'value' as an input to 'this' at the specified position in the
+  // arguments. Returns the added value for ease of chaining.
+  Value* insertInput(size_t i, Value* value);
+
+  // Replace the input of 'this' at position 'i' with
+  // 'newValue', returning the old node.
+  //
+  // Given:   %3 = f(%1, %2)
+  // Execute: %3.replaceInput(1, %4)
+  // Result:  %3 = f(%1, %4)
+  Value* replaceInput(size_t i, Value* newValue);
+
+  // Replace all occurrences of 'from' in the inputs of this
+  // node with 'to'. Corresponds to llvm's replaceUsesOfWith.
+  //
+  // Given:   %3 = f(%1, %2, %1)
+  // Execute: %3.replaceInputWith(%1, %4)
+  // Result:  %3 = f(%4, %2, %4)
+  void replaceInputWith(Value* from, Value* to);
+
+  Value* addOutput();
+
+  Value* insertOutput(size_t i);
+
+  void eraseOutput(size_t i);
+
+  Block* addBlock();
+  void eraseBlock(size_t i);
+
+  // Each Node can have a list of subblocks. These are used to define structured
+  // nested control flow operators such as If and Loop.
+  // The meaning of a block is specific to the kind of node it is in, but
+  // all blocks share these semantics:
+  // * Nested lexical scoping: If a node 'Parent' has a subblock which contains
+  //   a node 'Child', Child can use any value that was in scope for the Parent
+  //   node in addition to any values defined before 'Child' in the subblock.
+  // * The list of inputs to the block are in scope for the duration of the
+  //   block
+  // * the outputs of the Parent node are not in scope for the subblocks
+  // Typically the inputs to a block that represents control flow act as
+  // as the equivalents phi-nodes in standard SSA form,
+  // defining a new Value to represent any term that has multiple
+  // definitions depending on how control flowed. Outputs of the node containing
+  // control flow serve a similar purpose defining new values for variables
+  // that would have different definitions depending on which way control
+  // flowed.
+
+  at::ArrayRef<Block*> blocks() {
+    return blocks_;
+  }
+  at::ArrayRef<const Block*> blocks() const {
+    // Vectors are not convertible in const-ness of elements, but
+    // raw pointers are.
+    return {blocks_.data(), blocks_.size()};
+  }
+
+  // Is 'this' before 'n' in the topological order?
+  bool isBefore(const Node* n) const;
+
+  // Is 'this' after 'n' in the topological order?
+  bool isAfter(const Node* n) const;
+
+  // Insert unattached 'this' node before 'n' in the topological order.
+  // Returns this (for chaining).
+  //
+  // Given:   %3 = f(%1, %2)
+  //          %4 = g(%3)
+  // and unattached: %5 = h(%1)
+  // Execute: %5.insertBefore(%4)
+  // Result:  %3 = f(%1, %2)
+  //          %5 = h(%1)
+  //          %4 = g(%3)
+  Node* insertBefore(Node* n);
+
+  // Insert unattached 'this' node after 'n' in the topological order.
+  // Returns this (for chaining).
+  //
+  // Given: %3 = f(%1, %2)
+  //        %4 = g(%3)
+  // and unattached: %5 = h(%1)
+  // Execute: %5.insertAfter(%4)
+  // Result:  %3 = f(%1, %2)
+  //          %4 = g(%3)
+  //          %5 = h(%1)
+  Node* insertAfter(Node* n);
+
+  // Move 'this' (already in the graph) after 'n' in the topological order.
+  //
+  // NOTE: Does not check that value dependencies are preserved, see
+  //   AliasDb::moveAfterTopologicallyValid
+  //
+  // Given: %2 = f(%1)
+  //        %3 = g(%1)
+  // Execute: %2.moveAfter(%3)
+  // Result: %3 = g(%1)
+  //         %2 = f(%1)
+  //
+  void moveAfter(Node* n);
+
+  // Move a node 'n' (already in the graph) before 'this' in the topological
+  // order.
+  //
+  // NOTE: Does not check that value dependencies are preserved, see
+  //   AliasDb::moveBeforeTopologicallyValid
+  //
+  // Given: %2 = f(%1)
+  //        %3 = g(%1)
+  // Execute: %3.moveBefore(%2)
+  // Result: %3 = g(%1)
+  //         %2 = f(%1)
+  void moveBefore(Node* n);
+
+  // Remove the input at 'i' from this node.
+  //
+  // WARNING: This is O(n) in the number of inputs, so avoid repeatedly calling
+  // removeInput.
+  //
+  // Given: %3 = f(%1, %2)
+  // Execute: %3.removeInput(1)
+  // Result: %3 = f(%1)
+  void removeInput(size_t i);
+
+  // Remove all inputs from a node.
+  //
+  // Given: %3 = f(%1, %2)
+  // Execute: %3.removeAllInputs()
+  // Result: %3 = f()
+  void removeAllInputs();
+
+  // Remove all outputs from a node.
+  //
+  // Given: %1, %2 = f()
+  // Execute:removeAllInputs()
+  // Result: = f()
+  void removeAllOutputs();
+
+  // Rearrange the ordering of inputs or outputs of a node
+  // Given: %3 = f(%1, %2)
+  // Execute: %3.permuteInputs({1, 0})
+  // Result: %3 = f(%2, %1)
+  // Each index must appear exactly once
+  void permuteInputs(const std::vector<size_t>& new_inputs);
+  void permuteOutputs(const std::vector<size_t>& new_inputs);
+
+  // iterators of the node list starting at this node
+  // useful for resuming a search starting at this node
+  inline graph_node_list_iterator iterator() {
+    return {this, 0};
+  }
+  inline graph_node_list_iterator reverseIterator() {
+    return iterator().reverse();
+  }
+  inline const_graph_node_list_iterator iterator() const {
+    return {this, 0};
+  }
+  inline const_graph_node_list_iterator reverseIterator() const {
+    return iterator().reverse();
+  }
+
+  // Remove 'this' from the instruction list and deallocate it.
+  //
+  // Invariant: no outputs of 'this' may have any uses.
+  //
+  // Given: %2 = f(%1)
+  //        %3 = g(%1)
+  // Execute: %2.destroy()
+  // Result: %3 = g(%1)
+  void destroy();
+
+  // Dynamically cast this node to the subclass indicated by the
+  // template variable, returning nullptr if the cast is invalid..
+  //
+  // Example usage: if(auto s = n.cast<Select>()) { ... }
+  template <typename T>
+  T* cast() {
+    if (T::Kind == kind()) {
+      return static_cast<T*>(this);
+    }
+    return nullptr;
+  }
+  template <typename T>
+  const T* cast() const {
+    if (T::Kind == kind()) {
+      return static_cast<const T*>(this);
+    }
+    return nullptr;
+  }
+
+  template <typename T>
+  T* expect() {
+    TORCH_CHECK(
+        T::Kind == kind(),
+        "expected a ",
+        T::Kind.toDisplayString(),
+        " but found a ",
+        kind().toDisplayString());
+    return static_cast<T*>(this);
+  }
+
+  bool matches(const FunctionSchema& schema) const;
+
+  // XXX: this function is meant to be used with string literals only!
+  bool matches(
+      const char* signature_literal,
+      at::ArrayRef<Symbol> const_inputs = {}) const;
+
+  bool isMemberOf(const OperatorSet& os) const;
+  template <typename T>
+  bool isMemberOf(const OperatorMap<T>& om) const {
+    auto it = om.map.find(kind());
+    if (it == om.map.end()) {
+      return false;
+    }
+    for (auto& op : it->second) {
+      if (matches(op.first->schema())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  const FunctionSchema& schema() const;
+  const FunctionSchema* maybeSchema() const;
+  const Operator& getOperator() const;
+  Operation getOperation() const;
+
+  const Operator* maybeOperator() const;
+
+  void dump() const;
+
+  std::ostream& print(
+      std::ostream& out,
+      size_t level,
+      std::vector<const Node*>* groups,
+      bool print_source_locations = true,
+      bool print_attributes = true,
+      bool print_scopes = true,
+      bool print_body = true) const;
+
+  virtual ~Node() {
+    if (wrap_) {
+      wrap_->clear();
+    }
+  }
+
+  // Methods for accessing attributes
+  Node* copyAttributes(const Node& rhs) {
+    values_.clear();
+    for (const AVPtr& i : rhs.values_) {
+      values_.push_back(i->clone());
+    }
+    return this;
+  }
+  bool hasAttribute(Symbol name) const {
+    AT_ASSERT(name.is_attr());
+    return findAttr(name, false) != values_.end();
+  }
+  bool hasAttributeS(const std::string& name) const {
+    return hasAttribute(Symbol::attr(name));
+  }
+  AttributeKind kindOf(Symbol name) const {
+    AT_ASSERT(name.is_attr());
+    return (*findAttr(name, true))->kind();
+  }
+  AttributeKind kindOfS(const std::string& name) const {
+    return kindOf(Symbol::attr(name));
+  }
+  Node* removeAttribute(Symbol name) {
+    AT_ASSERT(name.is_attr());
+    values_.erase(findAttr(name, true));
+    return this;
+  }
+  Node* removeAttributeS(const std::string& name) {
+    return removeAttribute(Symbol::attr(name));
+  }
+  bool hasAttributes() const {
+    return !values_.empty();
+  }
+  size_t numAttributes() const {
+    return values_.size();
+  }
+  // The names are returned in order, since name actually is the index.
+  std::vector<Symbol> attributeNames() const {
+    std::vector<Symbol> names;
+    names.reserve(values_.size());
+    for (const AVPtr& a : values_) {
+      names.push_back(a->name);
+    }
+    return names;
+  }
+  std::vector<const char*> attributeNamesS() const {
+    std::vector<const char*> names;
+    names.reserve(values_.size());
+    for (const AVPtr& a : values_) {
+      names.push_back(a->name.toUnqualString());
+    }
+    return names;
+  }
+
+#define CREATE_ACCESSOR(Kind, method)                           \
+  Node* method##_(Symbol name, Kind##Attr::ConstructorType v) { \
+    return setAttr<Kind##Attr>(                                 \
+        name, std::forward<Kind##Attr::ConstructorType>(v));    \
+  }                                                             \
+  const Kind##Attr::ValueType& method(Symbol name) const {      \
+    return getAttr<Kind##Attr>(name);                           \
+  }
+
+  CREATE_ACCESSOR(Float, f)
+  CREATE_ACCESSOR(Complex, c)
+  CREATE_ACCESSOR(Floats, fs)
+  CREATE_ACCESSOR(ComplexVals, cs)
+  CREATE_ACCESSOR(String, s)
+  CREATE_ACCESSOR(Strings, ss)
+  CREATE_ACCESSOR(Int, i)
+  CREATE_ACCESSOR(Ints, is)
+  CREATE_ACCESSOR(Graph, g)
+  CREATE_ACCESSOR(Graphs, gs)
+  CREATE_ACCESSOR(Type, ty)
+  CREATE_ACCESSOR(Types, tys)
+  CREATE_ACCESSOR(IValue, ival)
+
+#undef CREATE_ACCESSOR
+
+  // Our Graphs are not very const-correct, so we need to allow returning
+  // non-const references too
+  GraphAttr::ValueType& g(Symbol name) {
+    return getAttr<GraphAttr>(name);
+  }
+
+  // does not use CREATE_ACCESSOR because we need additional asserts
+  Node* t_(Symbol name, TensorAttr::ConstructorType v) {
+    return setAttr<TensorAttr>(
+        name, std::forward<TensorAttr::ConstructorType>(v));
+  }
+  const TensorAttr::ValueType& t(Symbol name) const {
+    return getAttr<TensorAttr>(name);
+  }
+
+  Node* ts_(Symbol name, TensorsAttr::ConstructorType v) {
+    return setAttr<TensorsAttr>(
+        name, std::forward<TensorsAttr::ConstructorType>(v));
+  }
+  const TensorsAttr::ValueType& ts(Symbol name) const {
+    return getAttr<TensorsAttr>(name);
+  }
+
+  Block* findCommonAncestorBlockWith(Node* n);
+
+  size_t blocksFromGraphBlock();
+
+ private:
+  void printAttrValue(std::ostream& out, const Symbol& name) const;
+  void printAttributes(std::ostream& out, bool ignore_subgraph) const;
+
+  template <typename T>
+  Node* setAttr(Symbol name, typename T::ConstructorType v) {
+    AT_ASSERT(name.is_attr());
+    auto it = findAttr(name, false);
+    auto nv = AVPtr(new T(name, std::forward<typename T::ConstructorType>(v)));
+    // NOLINTNEXTLINE(bugprone-branch-clone)
+    if (it == values_.end()) {
+      values_.push_back(std::move(nv));
+    } else {
+      *it = std::move(nv);
+    }
+    return this;
+  }
+  template <typename T>
+  typename T::ValueType& getAttr(Symbol name) const {
+    AT_ASSERT(name.is_attr());
+    auto it = findAttr(name, true);
+    auto* child = dynamic_cast<T*>(it->get());
+    if (child == nullptr) {
+      throw IRAttributeError(name, true);
+    }
+    return child->value();
+  }
+  using AVPtr = AttributeValue::Ptr;
+  // NB: For determinism, we use a vector rather than a hash map.  This does
+  // mean that lookups are O(n), so you shouldn't use Attributes to store
+  // a big pile of messages.
+  std::vector<AVPtr> values_;
+  std::vector<AVPtr>::iterator findAttr(Symbol name, bool required) {
+    AT_ASSERT(name.is_attr());
+    auto it = std::find_if(values_.begin(), values_.end(), [&](const AVPtr& v) {
+      return v->name == name;
+    });
+    if (required && it == values_.end()) {
+      throw IRAttributeError(name, false);
+    }
+    AT_ASSERT(!required || it != values_.end());
+    return it;
+  }
+  std::vector<AVPtr>::const_iterator findAttr(Symbol name, bool required)
+      const {
+    AT_ASSERT(name.is_attr());
+    auto it = std::find_if(values_.begin(), values_.end(), [&](const AVPtr& v) {
+      return v->name == name;
+    });
+    if (required && it == values_.end()) {
+      throw IRAttributeError(name, false);
+    }
+    AT_ASSERT(!required || it != values_.end());
+    return it;
+  }
+
+  enum class MoveSide { BEFORE, AFTER };
+  bool isBeforeOrAfter(const Node* n, MoveSide moveSide) const;
+
+  std::pair<Value*, const Argument&> findInput(Symbol name);
+  // Lookup iterator in use list of _input i_ that corresponds to its use of
+  // _this_
+  use_list::iterator findUseForInput(size_t i);
+
+  // remove the use of input i, this sets input i to nullptr, but
+  // is only used internally to Node before setting it to a new value
+  // or erasing the entry from the list.
+  Value* dropInput(size_t i);
+
+  bool inBlockList() const {
+    if (next() == nullptr) {
+      AT_ASSERT(prev() == nullptr);
+    }
+    return next() != nullptr;
+  }
+
+  void removeFromList();
+  void lint() const;
+
+  void assignTopoPosition();
+
+ protected:
+  // subclasses must override
+  // this function is used by createClone to initialize a new version
+  // of a node in another graph. It should allocate a new instance of the same
+  // concrete type as 'this', but in graph 'g' which might be different
+  // than graph_
+  virtual Node* allocNewInstance(Graph* g) {
+    return new Node(g, kind());
+  }
+  // create a copy of all properties of Node s into this.
+  // subclasses should extend if they have additional information to copy.
+  // 'this' will be allocated with s->allocNewInstance(g) so it should have
+  // the same concrete type as 's'
+  virtual void cloneFrom(Node* s);
+};
+
+struct Block {
+  friend struct Node;
+  friend struct Graph;
+
+  AT_DISALLOW_COPY_AND_ASSIGN(Block);
+  TORCH_API Block(Graph* graph_, Node* node_);
+
+  at::ArrayRef<Value*> inputs() {
+    return input_->outputs();
+  }
+  at::ArrayRef<const Value*> inputs() const {
+    const auto& inputs = input_->outputs();
+    return {inputs.data(), inputs.size()};
+  }
+  at::ArrayRef<Value*> outputs() {
+    return output_->inputs();
+  }
+  at::ArrayRef<const Value*> outputs() const {
+    return static_cast<const Node*>(output_)->inputs();
+  }
+  graph_node_list nodes() {
+    return {input_, kNextDirection};
+  }
+  const_graph_node_list nodes() const {
+    return {input_, kNextDirection};
+  }
+  Node* return_node() {
+    return output_;
+  }
+  const Node* return_node() const {
+    return output_;
+  }
+  Node* param_node() {
+    return input_;
+  }
+  const Node* param_node() const {
+    return input_;
+  }
+  /**
+   * @warning NEVER pass raw pointer of smart pointer managed Graph to Python.
+   * Check #87343 for details.
+   */
+  Graph* owningGraph() {
+    return graph_;
+  }
+  const Graph* owningGraph() const {
+    return graph_;
+  }
+  Node* owningNode() {
+    return owning_node_;
+  }
+  const Node* owningNode() const {
+    return owning_node_;
+  }
+
+  Value* addInput(const std::string& name = "") {
+    Value* v = input_->addOutput();
+    v->setDebugName(name);
+    return v;
+  }
+  Value* insertInput(size_t i, const std::string& name = "") {
+    Value* v = input_->insertOutput(i);
+    v->setDebugName(name);
+    return v;
+  }
+  void eraseInput(size_t i) {
+    input_->eraseOutput(i);
+  }
+  void removeAllInputs() {
+    input_->removeAllOutputs();
+  }
+  size_t registerOutput(Value* v) {
+    output_->addInput(v);
+    return outputs().size() - 1;
+  }
+  size_t insertOutput(size_t i, Value* n) {
+    output_->insertInput(i, n);
+    return i;
+  }
+  void eraseOutput(size_t i) {
+    output_->removeInput(i);
+  }
+  void removeAllOutputs() {
+    output_->removeAllInputs();
+  }
+
+  void replaceOutput(size_t i, Value* n) {
+    output_->replaceInput(i, n);
+  }
+  void permuteOutputs(const std::vector<size_t>& new_inputs) {
+    output_->permuteInputs(new_inputs);
+  }
+  void permuteInputs(const std::vector<size_t>& new_inputs) {
+    input_->permuteOutputs(new_inputs);
+  }
+
+  Node* appendNode(Node* n) {
+    AT_ASSERT(n->graph_ == graph_ && !n->inBlockList());
+    n->insertBefore(output_);
+    return n;
+  }
+  Node* prependNode(Node* n) {
+    AT_ASSERT(n->graph_ == graph_ && !n->inBlockList());
+    n->insertAfter(input_);
+    return n;
+  }
+
+  // clone all inputs, nodes, and outputs from src and append them
+  // to the inputs, nodes, and outputs of this block
+  // value_map is used whenever a node in src references a free variable
+  // in src to look up its corresponding value
+  TORCH_API void cloneFrom(Block* src, std::function<Value*(Value*)> value_map);
+  TORCH_API void remapTypes(const std::function<TypePtr(TypePtr)>& type_map);
+
+  TORCH_API std::shared_ptr<Wrap<Block>> wrap() {
+    if (!wrap_) {
+      wrap_ = std::make_shared<Wrap<Block>>(this);
+    }
+    return wrap_;
+  }
+
+  virtual ~Block() {
+    if (wrap_) {
+      wrap_->clear();
+    }
+  }
+
+  void clear() {
+    removeAllOutputs();
+    for (auto it = nodes().rbegin(); it != nodes().rend(); it++) {
+      it.destroyCurrent();
+    }
+    removeAllInputs();
+  }
+
+ private:
+  void reIndexTopology();
+
+  // get rid of all nodes
+  // destroys in reverse order so that uses internal to this block
+  // do not have to be removed before you can destroy the block
+  void destroy();
+
+  Graph* const graph_;
+  // holds outputs in a way that can be reflected
+  // as a Use object
+  // also used as the beginning/end of the circular node list to avoid
+  // having corner cases where the list is empty.
+  Node* const output_;
+  Node* const input_;
+  Node* const
+      owning_node_; // either the node that has this block or nullptr for root
+  // a managing wrapper for Python to allow invalidation
+  std::shared_ptr<Wrap<Block>> wrap_;
+};
+
+struct Graph : std::enable_shared_from_this<Graph> {
+  AT_DISALLOW_COPY_AND_ASSIGN(Graph);
+  friend struct Node;
+  friend struct Value;
+  friend struct Block;
+
+ private:
+  // only used to keep track of allocated nodes
+  // actual representation of Graph is done with
+  // inputs, outputs, nodes
+
+  std::unordered_set<const Node*> all_nodes;
+  std::unordered_set<const Value*> all_values;
+  std::unordered_set<const Block*> all_blocks;
+  size_t next_unique_{0};
+
+  std::unordered_map<std::string, Value*> unique_names_;
+  // name_base_suffix tracks largest suffix currently used by all names sharing
+  // same name_base. Key of this map is name_base, value is largest suffix
+  // numeric value.
+  std::unordered_map<std::string, size_t> name_base_suffix_;
+
+  ScopePtr current_scope_;
+
+  Block* const block_;
+  // when insertNode() is called, the node is inserted before this node
+  // by default this is set to append to the top level block
+  Node* insert_before_;
+  int64_t predicted_insert_count_ = 0;
+
+  std::optional<size_t> op_version_;
+
+ public:
+  Graph(ScopePtr scope_root = c10::make_intrusive<Scope>())
+      : current_scope_(std::move(scope_root)),
+        block_(new Block(this, nullptr)),
+        insert_before_(return_node()) {}
+
+  at::ArrayRef<Value*> inputs() {
+    return block_->inputs();
+  }
+  at::ArrayRef<const Value*> inputs() const {
+    const Block& block = *block_;
+    return block.inputs();
+  }
+  at::ArrayRef<Value*> outputs() {
+    return block_->outputs();
+  }
+  at::ArrayRef<const Value*> outputs() const {
+    const Block& block = *block_;
+    return block.outputs();
+  }
+  graph_node_list nodes() {
+    return block_->nodes();
+  }
+  const_graph_node_list nodes() const {
+    const Block& block = *block_;
+    return block.nodes();
+  }
+  Node* param_node() {
+    return block_->param_node();
+  }
+  const Node* param_node() const {
+    return block_->param_node();
+  }
+  Node* return_node() {
+    return block_->return_node();
+  }
+  const Node* return_node() const {
+    return block_->return_node();
+  }
+  const std::unordered_map<std::string, Value*>& debugNames() const {
+    return unique_names_;
+  }
+
+  TORCH_API void push_scope(const std::string& scope_name);
+  TORCH_API void pop_scope();
+
+  ScopePtr current_scope() {
+    return current_scope_;
+  }
+
+  void set_op_version(std::optional<size_t> version) {
+    op_version_ = version;
+  }
+
+  std::optional<size_t> get_op_version() {
+    return op_version_;
+  }
+
+  void set_current_scope(ScopePtr scope) {
+    current_scope_ = std::move(scope);
+  }
+
+  Value* addInput(const std::string& name = "") {
+    return block_->addInput(name);
+  }
+  Value* insertInput(size_t i, const std::string& name = "") {
+    return block_->insertInput(i, name);
+  }
+  void eraseInput(size_t i) {
+    block_->eraseInput(i);
+  }
+  size_t registerOutput(Value* n) {
+    return block_->registerOutput(n);
+  }
+  void eraseOutput(size_t i) {
+    block_->eraseOutput(i);
+  }
+
+  TORCH_API Node* create(NodeKind kind, size_t num_outputs = 1);
+  TORCH_API Node* create(
+      NodeKind kind,
+      ArrayRef<Value*> inputs,
+      size_t num_outputs = 1);
+
+  TORCH_API Node* createNone();
+  TORCH_API Node* createAutogradZero();
+  TORCH_API Node* createUninitialized(TypePtr typ);
+  TORCH_API Node* createWithSubgraph(Symbol kind);
+  TORCH_API Node* createDifferentiableSubgraph();
+  TORCH_API Node* createTuple(
+      at::ArrayRef<Value*> values,
+      TupleTypePtr optional_named_tuple = nullptr);
+  TORCH_API Node* createTupleUnpack(Value* v);
+  TORCH_API Node* createTupleIndex(
+      Value* tup,
+      Value* idx,
+      const TypePtr& output_type);
+  TORCH_API Node* createTupleSlice(
+      Value* tup,
+      int64_t beg,
+      int64_t step_size,
+      int64_t num_values);
+  TORCH_API Node* createEnumName(Value* e);
+  TORCH_API Node* createEnumValue(Value* e);
+  TORCH_API Node* createList(
+      const TypePtr& contained_type,
+      at::ArrayRef<Value*> values);
+  TORCH_API Node* createListUnpack(Value* v, size_t size);
+  TORCH_API Node* createDict(
+      const TypePtr& key_type,
+      const TypePtr& value_type,
+      at::ArrayRef<Value*> keys,
+      at::ArrayRef<Value*> values);
+  TORCH_API Node* createNumToTensor(Value* value);
+  TORCH_API Node* createObject(const ClassTypePtr& type);
+  TORCH_API Node* createSetAttr(
+      Value* obj,
+      const std::string& field,
+      Value* newValue);
+  TORCH_API Node* createGetAttr(Value* obj, const std::string& field);
+  Value* insertGetAttr(Value* obj, const std::string& field) {
+    return insertNode(createGetAttr(obj, field))->output();
+  }
+  TORCH_API Node* createStore(const std::string& name, Value* v);
+  TORCH_API Node* createLoad(const std::string& name, const TypePtr& type);
+  TORCH_API Node* createIsInstance(Value* v, at::ArrayRef<TypePtr> types);
+
+  TORCH_API Value* insertUncheckedCast(Value* v, TypePtr type);
+
+  // Insert a ToList operator with argument \p v and output type \p type.
+  // \returns the output of the operation.
+  TORCH_API Value* insertToList(Value* v, TypePtr type);
+
+  TORCH_API Value* insertFunctionCall(
+      Function* callee,
+      const MatchedSchema& matched);
+  TORCH_API Value* insertMethodCall(
+      std::string method_name,
+      const MatchedSchema& matched);
+
+  // Note: defined in python_ir.cpp and can be used only in python extension
+  Node* createPythonOp(
+      THPObjectPtr&& pyobj,
+      const std::string& cconv,
+      pyobj_list&& scalar_args);
+  // clone n, making a new node in _this_ graph.
+  // use value_map to translate inputs of n to inputs of the cloned node
+  // if copy_blocks is false, it will not recursively clone the nested blocks
+  // this node contains.
+  TORCH_API Node* createClone(
+      Node* n,
+      const std::function<Value*(Value*)>& value_map,
+      bool copy_blocks = true);
+
+  // Insert constant IValue into the graph.
+  TORCH_API Value* insertConstant(
+      const IValue& val,
+      std::optional<SourceRange> loc = std::nullopt,
+      std::optional<ScopePtr> scope = std::nullopt);
+
+  // Schema-driven insert:
+  // This inserts a node into the graph with inputs determined from args and
+  // kwargs using Python argument matching rules, and checks that the op matches
+  // a known schema.
+  //
+  // If this node successfully completes, it guarantees the node
+  // is a correctly-formed invocation of opname
+  TORCH_API Value* insert(
+      Symbol opname,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs = {},
+      const std::optional<SourceRange>& range = {});
+
+  Node* appendNode(Node* n) {
+    return block_->appendNode(n);
+  }
+
+  Node* prependNode(Node* n) {
+    return block_->prependNode(n);
+  }
+
+  // insert before insert_before_ node
+  // initialized to insert at the end of the top level block
+  // can be changed with setInsertPoint()
+  Node* insertNode(Node* n) {
+    AT_ASSERT(
+        insert_before_->inBlockList() &&
+        "insert point node is no longer in a block list");
+    return n->insertBefore(insert_before_);
+  }
+  // set where nodes are inserted to append to the end of this block
+  void setInsertPoint(Block* b) {
+    AT_ASSERT(b->owningGraph() == this);
+    setInsertPoint(b->return_node());
+  }
+  // set where nodes are inserted to insert _before_ this node
+  // for implementation simplicity we only support inserting before a node for
+  // now
+  void setInsertPoint(Node* n) {
+    AT_ASSERT(n->owningGraph() == this && n->inBlockList());
+    insert_before_ = n;
+    predicted_insert_count_ = 0;
+  }
+  Node* insertPoint() {
+    return insert_before_;
+  }
+
+  // the top level block
+  Block* block() {
+    return block_;
+  }
+  const Block* block() const {
+    return block_;
+  }
+
+  // Checks well-formedness and invariants of graph
+  TORCH_API void lint() const;
+  // for use in debugger
+  TORCH_API void dump() const;
+
+  TORCH_API ~Graph();
+
+  TORCH_API std::string toString(bool print_source_locations = true) const;
+
+  TORCH_API std::ostream& print(
+      std::ostream& out,
+      bool print_source_locations = true) const;
+
+  friend TORCH_API std::ostream& operator<<(std::ostream& out, const Graph& g);
+
+  TORCH_API std::shared_ptr<Graph> copy();
+  TORCH_API std::unique_ptr<Graph> copyUnique();
+  TORCH_API void remapTypes(const std::function<TypePtr(TypePtr)>& type_map);
+
+ private:
+  friend TORCH_API void Lint(const AliasDb* db);
+  TORCH_API void freeNode(Node* n);
+  TORCH_API void freeValue(Value* v);
+  TORCH_API void freeBlock(Block* b);
+  void cloneFrom(Graph& src);
+};
+
+/** \brief An utility class for setting temporary insertion points.
+ *
+ * When an object of this class is created, it stores the current insertion
+ * point, sets the new one, and restores the original insertion point when the
+ * object is destroyed.
+ */
+struct WithInsertPoint {
+  WithInsertPoint(Node* n) : prev_(n->owningGraph()->insertPoint()) {
+    n->owningGraph()->setInsertPoint(n);
+  }
+  WithInsertPoint(Block* b) : WithInsertPoint(b->return_node()) {}
+
+  ~WithInsertPoint() {
+    prev_->owningGraph()->setInsertPoint(prev_);
+  }
+
+ private:
+  Node* prev_;
+};
+
+/** \brief An utility class for setting temporary scopes.
+ *
+ * When an object of this class is created, it stores the current scope, sets
+ * the new one, and restores the original scope when the object is destroyed.
+ */
+struct WithCurrentScope {
+  WithCurrentScope(Graph& g, ScopePtr scope)
+      : graph_(&g), prev_scope_(g.current_scope()) {
+    g.set_current_scope(std::move(scope));
+  }
+  ~WithCurrentScope() {
+    graph_->set_current_scope(prev_scope_);
+  }
+
+ private:
+  Graph* graph_;
+  ScopePtr prev_scope_;
+};
+
+inline Value::Value(Node* node_, size_t offset_)
+    : node_(node_),
+      offset_(offset_),
+      unique_(node_->graph_->next_unique_++),
+      type_(TensorType::get()) {
+  node_->graph_->all_values.emplace(this);
+}
+
+inline Value* Value::setType(TypePtr type) {
+  AT_ASSERT(type);
+  if (auto dyn = type->castRaw<c10::DynamicType>()) {
+    type = dyn->fallback();
+  }
+  type_ = std::move(type);
+  for (Use& use : uses_) {
+    use.user->op_ = nullptr;
+  }
+  return this;
+}
+
+inline Graph* Value::owningGraph() {
+  return node()->owningGraph();
+}
+
+inline const Graph* Value::owningGraph() const {
+  return node()->owningGraph();
+}
+
+/************* All nodes not required to be defined before Graph **************/
+struct ProfileOp : public Node {
+  static const Symbol Kind;
+  ProfileOp(Graph* graph, std::function<void(std::vector<IValue>&)> callback)
+      : Node(graph, ::c10::prim::profile), callback_(std::move(callback)) {}
+
+  void cloneFrom(Node* other_) override;
+  Node* allocNewInstance(Graph* g) override;
+
+  const std::function<void(std::vector<IValue>&)>& getCallback() const {
+    return callback_;
+  }
+
+  void setCallback(std::function<void(std::vector<IValue>&)> callback) {
+    callback_ = std::move(callback);
+  }
+
+  bool hasSeenTensor() const {
+    return has_seen_tensor_;
+  }
+
+  void setHasSeenTensor(bool has_seen_tensor) {
+    has_seen_tensor_ = has_seen_tensor;
+  }
+
+ private:
+  std::function<void(std::vector<IValue>&)> callback_;
+  bool has_seen_tensor_ = false;
+};
+
+struct TORCH_API ProfileIValueOp : public Node {
+  static const Symbol Kind;
+  ProfileIValueOp(
+      Graph* graph,
+      std::function<void(std::vector<IValue>&)> callback)
+      : Node(graph, ::c10::prim::profile_ivalue),
+        callback_(std::move(callback)) {}
+
+  void cloneFrom(Node* other_) override;
+  Node* allocNewInstance(Graph* g) override;
+
+  const std::function<void(std::vector<IValue>&)>& getCallback() const {
+    return callback_;
+  }
+
+  void setCallback(std::function<void(std::vector<IValue>&)> callback) {
+    callback_ = std::move(callback);
+  }
+
+ private:
+  std::function<void(std::vector<IValue>&)> callback_;
+};
+
+// execute a Python function, used for Ops we can't optimize but that we want to
+// optimize around
+//
+// Note: actual implementation (ConcretePythonOp) is defined in python_ir.cpp
+// which is not included in libtorch.so. We still include some bits and pieces
+// of PythonOp here to enable writing simple passes generically. In general,
+// python-aware bits need to be moved to the descendant classes.
+struct TORCH_API PythonOp : public Node {
+  using Node::Node;
+
+  virtual std::string name() const = 0;
+  virtual void writeScalars(std::ostream& out) const = 0;
+  void cloneFrom(Node* other_) override = 0;
+  Node* allocNewInstance(Graph* g) override = 0;
+  // recover the autograd.Function instance, if this PythonOp's function
+  // was originally SomeFunction.apply
+  // used in ONNX for discovering symbolics
+  virtual std::optional<THPObjectPtr> autogradFunction() const = 0;
+
+  virtual void lint_python() const = 0;
+};
+
+TORCH_API void LintGraph(const std::shared_ptr<Graph>& graph);
+
+TORCH_API at::ArrayRef<Value*> createTupleUnpack(Value* v);
+
+/** Insert graph \p CALLEE into graph \p G using \p INPUTS as input values.
+ * The insertion happens at the current insertion point.
+ * Optionally, one can also pass \p VALUE_MAP to get a map between \p CALLEE
+ * values and their cloned copies in \p G.
+ */
+TORCH_API std::vector<Value*> insertGraph(
+    Graph& g,
+    Graph& callee,
+    ArrayRef<Value*> inputs);
+TORCH_API std::vector<Value*> insertGraph(
+    Graph& g,
+    Graph& callee,
+    ArrayRef<Value*> inputs,
+    std::unordered_map<Value*, Value*>& value_map);
+
+/** Insert function \p CALLEE after node \p TO_REPLACE, remove the node and
+ * replace all its uses with corresponding outputs of the inserted function.
+ * This asserts that the number of outputs of the original node and the
+ * graph are the same.
+ */
+TORCH_API std::vector<Value*> inlineCallTo(
+    Node* to_replace,
+    GraphFunction* callee,
+    bool use_graph = true);
+
+TORCH_API std::vector<Value*> inlineCallTo(
+    Node* to_replace,
+    GraphFunction* callee,
+    Graph* callee_graph);
+
+/** If there is only one value in \p OUTPUTS and its kind is Tuple, insert a
+ * tuple unpack node and return the resulting values.
+ */
+TORCH_API std::vector<Value*> unpackOutputs(const std::vector<Value*>& outputs);
+
+TORCH_API std::vector<Node*> findAllNodes(Graph& g, Symbol kind, bool recurse);
+TORCH_API std::vector<Node*> findAllNodes(Block& b, Symbol kind, bool recurse);
+TORCH_API std::vector<Node*> findAllNodes(
+    at::ArrayRef<Block*> a,
+    Symbol kind,
+    bool recurse);
+
+struct TORCH_API OperatorSet {
+  OperatorSet(std::initializer_list<const char*> sig_literals);
+  std::vector<std::shared_ptr<Operator>> getOps() const;
+  void insert(std::initializer_list<const char*> sig_literals);
+
+ private:
+  friend struct Node;
+  std::unordered_map<Symbol, std::vector<std::shared_ptr<Operator>>> ops;
+};
+
+template <typename T>
+struct OperatorMap {
+  // Type aliasing
+  using OpMapType = typename std::pair<std::shared_ptr<Operator>, T>;
+  using ValueType = std::vector<OpMapType>;
+  using MapType = std::unordered_map<Symbol, ValueType>;
+
+  OperatorMap() = default;
+  explicit OperatorMap(
+      std::initializer_list<std::pair<std::shared_ptr<Operator>, T>> init) {
+    insert(init);
+  }
+  explicit OperatorMap(std::initializer_list<std::pair<const char*, T>> init) {
+    insert(init);
+  }
+
+  void insert(const std::shared_ptr<Operator>& op, T val) {
+    // Remove if exists before insert
+    erase(op);
+    map[Symbol::fromQualString(op->schema().name())].emplace_back(
+        std::make_pair(op, val));
+  }
+
+  void insert(const OperatorSet& op_set, T val) {
+    for (auto& op : op_set.getOps()) {
+      insert(op, val);
+    }
+  }
+
+  void insert(
+      std::initializer_list<std::pair<std::shared_ptr<Operator>, T>> v) {
+    for (auto& el : v) {
+      insert(el.first, el.second);
+    }
+  }
+
+  void insert(std::initializer_list<std::pair<const char*, T>> v) {
+    for (auto& el : v) {
+      insert(getOperatorForLiteral(el.first), el.second);
+    }
+  }
+
+  void erase(const std::shared_ptr<Operator>& op) {
+    auto it = map.find(Symbol::fromQualString(op->schema().name()));
+    if (it == map.end()) {
+      return;
+    }
+    for (auto vit = it->second.begin(); vit != it->second.end(); ++vit) {
+      if (vit->first->schema() == op->schema()) {
+        it->second.erase(vit);
+        break;
+      }
+    }
+    if (it->second.size() == 0) {
+      map.erase(Symbol::fromQualString(op->schema().name()));
+    }
+  }
+
+  bool contains(const Operator& op) const {
+    const auto it = map.find(Symbol::fromQualString(op.schema().name()));
+    if (it == map.end()) {
+      return false;
+    }
+    for (auto vit = it->second.begin(); vit != it->second.end(); ++vit) {
+      if (vit->first->schema() == op.schema()) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  bool contains(const Node* n) const {
+    return n->maybeOperator() && contains(n->getOperator());
+  }
+
+  std::optional<T> find(const Operator& op) {
+    const auto it = map.find(Symbol::fromQualString(op.schema().name()));
+    if (it == map.end()) {
+      return std::nullopt;
+    }
+    for (auto vit = it->second.begin(); vit != it->second.end(); ++vit) {
+      if (vit->first->schema() == op.schema()) {
+        return vit->second;
+      }
+    }
+    return std::nullopt;
+  }
+
+  // TODO: return iterator
+  std::vector<OpMapType> getAllKeysAndValues() const {
+    std::vector<OpMapType> keys_values;
+    keys_values.reserve(map.size());
+    for (auto& symbol_mapping : map) {
+      auto& vec = symbol_mapping.second;
+      for (auto& pair : vec) {
+        keys_values.push_back(pair);
+      }
+    }
+    return keys_values;
+  }
+
+ private:
+  friend struct Node;
+  MapType map;
+};
+
+template <typename T>
+struct FunctionSchemaMap {
+  // Type aliasing
+  using FuncSchemaMapType = typename std::pair<FunctionSchema, T>;
+  using ValueType = std::vector<FuncSchemaMapType>;
+  using MapType = std::unordered_map<Symbol, ValueType>;
+
+  FunctionSchemaMap() = default;
+  void insert(const FunctionSchema& schema, T val) {
+    // Remove if exists before insert
+    erase(schema);
+    map[Symbol::fromQualString(schema.name())].emplace_back(
+        std::make_pair(schema, val));
+  }
+
+  void erase(const FunctionSchema& schema) {
+    auto it = map.find(Symbol::fromQualString(schema.name()));
+    if (it == map.end()) {
+      return;
+    }
+    for (auto vit = it->second.begin(); vit != it->second.end(); ++vit) {
+      if (vit->first == schema) {
+        it->second.erase(vit);
+        break;
+      }
+    }
+    if (it->second.size() == 0) {
+      map.erase(Symbol::fromQualString(schema.name()));
+    }
+  }
+
+  bool contains(const FunctionSchema& schema) const {
+    const auto it = map.find(Symbol::fromQualString(schema.name()));
+    if (it == map.end()) {
+      return false;
+    }
+    for (auto vit = it->second.begin(); vit != it->second.end(); ++vit) {
+      if (vit->first->schema() == schema) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  std::optional<T> find(const FunctionSchema& schema) const {
+    const auto it = map.find(Symbol::fromQualString(schema.name()));
+    if (it == map.end()) {
+      return std::nullopt;
+    }
+    for (auto vit = it->second.begin(); vit != it->second.end(); ++vit) {
+      if (vit->first == schema) {
+        return vit->second;
+      }
+    }
+    return std::nullopt;
+  }
+
+  // TODO: return iterator
+  std::vector<FuncSchemaMapType> getAllKeysAndValues() const {
+    std::vector<FuncSchemaMapType> keys_values;
+    keys_values.reserve(map.size());
+    for (auto& symbol_mapping : map) {
+      auto& vec = symbol_mapping.second;
+      for (auto& pair : vec) {
+        keys_values.push_back(pair);
+      }
+    }
+    return keys_values;
+  }
+
+ private:
+  friend struct Node;
+  MapType map;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/ir_views.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/ir_views.h
new file mode 100644
index 0000000000000000000000000000000000000000..371d7c9def5748af1c17c5e48c7ad4f4a8cf8dee
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/ir_views.h
@@ -0,0 +1,167 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/irange.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+struct IfView {
+  explicit IfView(Node* node) : node_(node) {
+    AT_ASSERT(node->kind() == ::c10::prim::If);
+  }
+  Value* cond() const {
+    return node_->input(0);
+  }
+  Block* thenBlock() const {
+    return node_->blocks().at(0);
+  }
+  Block* elseBlock() const {
+    return node_->blocks().at(1);
+  }
+  ArrayRef<Value*> thenOutputs() const {
+    return thenBlock()->outputs();
+  }
+  ArrayRef<Value*> elseOutputs() const {
+    return elseBlock()->outputs();
+  }
+  ArrayRef<Value*> outputs() const {
+    return node_->outputs();
+  }
+  Node* node() const {
+    return node_;
+  }
+  operator Node*() const {
+    return node_;
+  }
+
+  void permuteOutputs(const std::vector<size_t>& new_output_order) {
+    node_->permuteOutputs(new_output_order);
+    thenBlock()->permuteOutputs(new_output_order);
+    elseBlock()->permuteOutputs(new_output_order);
+  }
+
+ private:
+  Node* node_;
+};
+
+struct LoopView {
+  explicit LoopView(Node* node) : node_(node) {
+    AT_ASSERT(
+        node->kind() == ::c10::prim::Loop || node->kind() == ::c10::onnx::Loop);
+  }
+  Block* bodyBlock() const {
+    return node_->blocks().at(0);
+  }
+  Value* cond() const {
+    return node_->input(0);
+  }
+  Value* maxTripCount() const {
+    return node_->input(0);
+  }
+  Value* inputCond() const {
+    return node_->input(1);
+  }
+  Value* nextCond() const {
+    return bodyBlock()->outputs().at(0);
+  }
+  Value* currentTripCount() const {
+    return bodyBlock()->inputs().at(0);
+  }
+  ArrayRef<Value*> carriedInputs() const {
+    // skip trip count and cond
+    return node_->inputs().slice(2);
+  }
+  ArrayRef<Value*> carriedInputsWithCond() const {
+    // skip trip count and cond
+    return node_->inputs().slice(1);
+  }
+  ArrayRef<Value*> carriedOutputs() const {
+    return node_->outputs();
+  }
+  ArrayRef<Value*> bodyCarriedInputs() const {
+    // skip trip count and cond
+    return bodyBlock()->inputs().slice(1);
+  }
+  ArrayRef<Value*> bodyCarriedOutputs() const {
+    return bodyBlock()->outputs().slice(1);
+  }
+  Node* node() const {
+    return node_;
+  }
+  operator Node*() const {
+    return node_;
+  }
+
+  void permuteLoopCarried(const std::vector<size_t>& new_output_order) {
+    node_->permuteOutputs(new_output_order);
+    // skip trip count and cond
+    node_->permuteInputs(adjustIndices(2, new_output_order));
+    auto adjusted_block_order = adjustIndices(1, new_output_order);
+    bodyBlock()->permuteOutputs(adjusted_block_order);
+    bodyBlock()->permuteInputs(adjusted_block_order);
+  }
+
+  void replaceMaxTripCount(Value* new_max_trip_count) {
+    node_->replaceInput(0, new_max_trip_count);
+  }
+  void replaceInputCondition(Value* new_input_condition) {
+    node_->replaceInput(1, new_input_condition);
+  }
+
+  // our way of encoding loops makes them difficult to turn back into python
+  // syntax. we have to check properties of the condition and trip count inputs
+  // to figure out which one it initially was. ModifiedLoops are not directly
+  // mappable to either For or While
+  enum LoopType { While, For, ModifiedLoop };
+
+  LoopType loopType() {
+    auto trip_count = toIValue(maxTripCount());
+    auto cond_input = toIValue(inputCond());
+    auto cond_next = toIValue(nextCond());
+
+    bool condition_is_always_true =
+        cond_input && cond_input->toBool() && cond_next && cond_next->toBool();
+    bool trip_count_is_specified = !trip_count || // trip is not a constant
+        trip_count->toInt() !=
+            std::numeric_limits<int64_t>::max() || // it is a constant but not
+                                                   // the default one
+        !currentTripCount()
+             ->uses()
+             .empty(); // it is actually being used in the body.
+
+    if (condition_is_always_true) {
+      // if the trip count was not specified this was a user-written while True:
+      return trip_count_is_specified ? For : While;
+    } else {
+      if (trip_count_is_specified) {
+        return ModifiedLoop;
+      }
+      return While;
+    }
+  }
+
+ private:
+  Node* node_;
+
+  // adjust index_ordering by adding indices 0 - thorough adjust, and
+  // incrementing all existing inputs by adjust
+  static std::vector<size_t> adjustIndices(
+      size_t adjust,
+      const std::vector<size_t>& index_ordering) {
+    std::vector<size_t> adjusted;
+    adjusted.reserve(adjust + index_ordering.size());
+    for (const auto i : c10::irange(adjust)) {
+      adjusted.push_back(i);
+    }
+    for (auto index : index_ordering) {
+      adjusted.push_back(index + adjust);
+    }
+    return adjusted;
+  }
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/irparser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/irparser.h
new file mode 100644
index 0000000000000000000000000000000000000000..b6455d7585e8626c36d00fd6026a905ace5ffbfd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/irparser.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <optional>
+#include <string>
+#include <unordered_map>
+
+#include <torch/csrc/Export.h>
+
+namespace torch::jit {
+
+struct Graph;
+struct Value;
+
+// \brief Parse IR from \p STR constructing the corresponding IR in\ GRAPH.
+// if parse_tensor_constants is true will construct empty tensors
+// for Tensor constants with random or uninitialized contents, otherwise will
+// throw
+TORCH_API void parseIR(
+    const std::string& str,
+    torch::jit::Graph* graph,
+    bool parse_tensor_constants = false);
+
+/** \brief Parse IR from \p STR constructing the corresponding IR in\ GRAPH.
+ *
+ * \p VMAP is filled with String to Value pairs allowing to index Values in the
+ * newly created graph by their name in the original IR string.
+ * if parse_tensor_constants is true will construct empty tensors
+ * for Tensor constants with random or uninitialized contents, otherwise will
+ * throw
+ */
+TORCH_API void parseIR(
+    const std::string& str,
+    torch::jit::Graph* graph,
+    std::unordered_map<std::string, Value*>& vmap,
+    bool parse_tensor_constants = false);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/named_value.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/named_value.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ce6ea7c0200fef7461ca8ad69ecb255d2ae75b7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/named_value.h
@@ -0,0 +1,86 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/ir/constants.h>
+#include <torch/csrc/utils/variadic.h>
+
+namespace torch::jit {
+
+struct Value;
+
+/**
+ * A value with optional extra name and location information. Used during
+ * schema matching to provide extra error information and resolve kwargs.
+ */
+struct NamedValue {
+  NamedValue(const SourceRange& loc, const std::string& name, Value* value)
+      : loc_(loc), name_(name), value_(value) {}
+  NamedValue(const SourceRange& loc, Value* value) : loc_(loc), value_(value) {}
+
+  /* implicit */ NamedValue(Value* value) : value_(value) {}
+  NamedValue(const std::string& name, Value* value)
+      : name_(name), value_(value) {}
+
+  /* implicit */ NamedValue(IValue value) : ivalue_(std::move(value)) {}
+
+  NamedValue(const std::string& name, IValue value)
+      : name_(name), ivalue_(std::move(value)) {}
+
+  template <
+      typename T,
+      typename = std::enable_if_t<
+          (!std::is_same_v<std::decay_t<T>, NamedValue> &&
+           !std::is_same_v<std::decay_t<T>, Value*> &&
+           !std::is_same_v<std::decay_t<T>, IValue>)>>
+  // NOLINTNEXTLINE(bugprone-forwarding-reference-overload)
+  NamedValue(T&& t) : NamedValue(IValue(std::forward<T>(t))) {}
+
+  template <
+      typename T,
+      typename = std::enable_if_t<
+          (!std::is_same_v<std::decay_t<T>, Value*> &&
+           !std::is_same_v<std::decay_t<T>, IValue>)>>
+  NamedValue(const std::string& name, T&& t)
+      : NamedValue(name, IValue(std::forward<T>(t))) {}
+
+  SourceRange locOr(const SourceRange& backup_location) const {
+    if (!loc_)
+      return backup_location;
+    return loc();
+  }
+
+  // note: this will insert a constant node into the graph at the current
+  // insert point if this NamedValue is actually a constant
+  Value* value(Graph& g) const {
+    if (!value_)
+      return insertConstant(
+          g, ivalue_); // use insertConstant to remove need to include ir.h here
+    return value_;
+  }
+
+  const std::string& name() const {
+    AT_ASSERT(name_);
+    return *name_;
+  }
+
+  const SourceRange& loc() const {
+    AT_ASSERT(loc_);
+    return *loc_;
+  }
+
+  at::TypePtr type() const;
+
+ private:
+  std::optional<SourceRange> loc_;
+  std::optional<std::string> name_;
+  Value* value_{nullptr};
+  // only valid if value_ == nullptr;
+  IValue ivalue_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/node_hashing.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/node_hashing.h
new file mode 100644
index 0000000000000000000000000000000000000000..dd49b94db895020a8ae806ec5cd8cd537b6a7784
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/node_hashing.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+struct TORCH_API HashNode {
+  size_t operator()(const Node* k) const;
+};
+
+struct TORCH_API EqualNode {
+  bool operator()(const Node* lhs, const Node* rhs) const;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/scope.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/scope.h
new file mode 100644
index 0000000000000000000000000000000000000000..aa087d0169e4d4571119acef7d7e45b1ca3dc1ff
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/scope.h
@@ -0,0 +1,223 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/jit_type.h>
+#include <ATen/core/symbol.h>
+#include <c10/util/intrusive_ptr.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <optional>
+#include <unordered_map>
+
+namespace torch::jit {
+struct ModuleInstanceInfo;
+constexpr size_t kModuleInstanceInfo = 2;
+
+namespace utils {
+std::string get_module_info(const ModuleInstanceInfo& module_instance_info);
+} // namespace utils
+
+// Scope is a node of a trie that represents the tree of nested scopes.
+// Individual scopes are pushed and popped from Graph, which holds a
+// pointer to the current scope. Each Node in Graph holds a pointer
+// to the scope that was current when the node was created.
+// The trie never needs to shrink, it only grows until it is disposed
+// of when Graph is deallocated. Hence, pointers to scopes held by nodes
+// will always be valid as long as Graph is alive.
+struct Scope;
+using ScopePtr = c10::intrusive_ptr<Scope>;
+using c10::Symbol;
+
+struct TORCH_API Scope : public c10::intrusive_ptr_target {
+ private:
+  ScopePtr parent_;
+  Symbol name_;
+  ScopePtr intrusive_from_this();
+
+ public:
+  Scope();
+
+  Scope(ScopePtr parent, Symbol name);
+
+  ScopePtr push(Symbol name);
+
+  ScopePtr parent();
+
+  bool isRoot() const;
+
+  bool isBlank() const;
+
+  ScopePtr getRoot();
+
+  size_t getDepth();
+
+  Symbol name() const;
+
+  std::string namesFromRoot(const std::string& separator = "/") const;
+};
+
+struct Function;
+struct InlinedCallStack;
+
+/**
+ * ModuleInstanceInfo is a structure to include the module type and instance
+ * name. It also provide public methods to get the pointer to module type and
+ * instance name.
+ *
+ * This structure is mainly used as a private member in InlinedCallStack, such
+ * that one can follow the callstack to find the relevant module hierarchy.
+ */
+struct ModuleInstanceInfo {
+ private:
+  c10::ClassTypePtr module_type_{nullptr};
+  std::string instance_name_;
+
+ public:
+  ModuleInstanceInfo() = default;
+  ModuleInstanceInfo(c10::ClassTypePtr module_type, std::string instance_name);
+  c10::ClassTypePtr class_type() {
+    return module_type_;
+  }
+  c10::ClassTypePtr class_type() const {
+    return module_type_;
+  }
+  std::string instance_name() const {
+    return instance_name_;
+  }
+
+  bool operator==(const ModuleInstanceInfo& rhs) const {
+    return (class_type() == rhs.class_type()) &&
+        (instance_name() == rhs.instance_name());
+  }
+};
+
+/**
+ * InlinedCallStack is an element in a list representing callstack of functions
+ * that have been inlined.
+ *
+ * Each such element holds info about the current callsite (Function and
+ * SourceRange) and a pointer to the next element in the list. The last element
+ * in the list represents the innermost function that was inlined.
+ *
+ * For instance, if a node has a callstack
+ *    [foo, source_range1] -> [bar, source_range2]
+ * it means that this node was originally from function 'bar' that was called
+ * at 'source_range2' in function 'foo' that was called in the current function
+ * at 'source_range1'.
+ *
+ * If a node did not come from any inlined function, its callstack will be
+ * empty.
+ *
+ * The callstack lists only grow, we never remove elements from them, which
+ * allows us to reuse same elements in different lists. For instance, if we
+ * inline function 'bar' to 'foo' and then inline 'foo' to two functions 'ham'
+ * and 'baz', the callstacks would look like:
+ *
+ *  [baz, source_range3]  --
+ *                           \
+ *                             --> [foo, source_range1] -> [bar, source_range2]
+ *                           /
+ *  [ham, source_range4]  --
+ */
+using InlinedCallStackPtr = c10::intrusive_ptr<InlinedCallStack>;
+using InlinedCallStackEntry =
+    std::tuple<Function*, SourceRange, std::optional<ModuleInstanceInfo>>;
+
+struct TORCH_API InlinedCallStack : public c10::intrusive_ptr_target {
+ private:
+  std::optional<InlinedCallStackPtr> callee_;
+  Function* fn_;
+  // Reason for fn_name_ even though we have fn_
+  // Serialized callstack is used in circustmances where InlinedCallstack
+  // cannot be constructed during runtime, e.g. mobile runtime or
+  // delegated backends.
+  // Since in those cases we do not have Function* we store function name
+  // fn_name does not give you access to the same information that Function*
+  // does, however in mobile/delegated backend runtime we use InlindedCallStack
+  // for exception stack and for that purpose fn_name_ suffices.
+  const std::string fn_name_;
+  SourceRange source_range_;
+  InlinedCallStackPtr intrusive_from_this();
+  std::optional<ModuleInstanceInfo> module_instance_info_;
+
+ public:
+  // Constructor for a leaf callstack node.
+  InlinedCallStack(Function* fn, SourceRange source_range);
+
+  // Constructor for a leaf callstack node.
+  InlinedCallStack(
+      Function* fn,
+      SourceRange source_range,
+      std::optional<ModuleInstanceInfo> module_instance_info);
+
+  // Constructor for a leaf callstack node.
+  InlinedCallStack(
+      Function* fn,
+      SourceRange source_range,
+      std::optional<ModuleInstanceInfo> module_instance_info,
+      std::string& function_name);
+
+  // Constructor for an inner callstack node.
+  InlinedCallStack(
+      InlinedCallStackPtr callee,
+      Function* fn,
+      SourceRange source_range);
+
+  InlinedCallStack(
+      InlinedCallStackPtr callee,
+      Function* fn,
+      SourceRange source_range,
+      std::optional<ModuleInstanceInfo> module_instance_info);
+
+  InlinedCallStack(
+      InlinedCallStackPtr callee,
+      Function* fn,
+      SourceRange source_range,
+      std::optional<ModuleInstanceInfo> module_instance_info,
+      std::string& function_name);
+
+  // Return next element in the callstack list.
+  std::optional<InlinedCallStackPtr> callee() const;
+
+  // Return module instance associated with the current element.
+  std::optional<ModuleInstanceInfo> module_instance() const;
+
+  // Returns the source range of the node
+  SourceRange source_range() const;
+
+  Function* function() const;
+
+  const std::string& function_name() const;
+
+  // Return callstack as a vector of [Function, SourceRange] pairs.
+  std::vector<InlinedCallStackEntry> vec();
+
+  void setCallee(std::optional<InlinedCallStackPtr> /*callee*/);
+
+  bool operator==(const InlinedCallStack& rhs) const {
+    // No need to compare fn_, since source_range equivalence check
+    // should suffice.
+    return (module_instance().has_value() ==
+            rhs.module_instance().has_value()) &&
+        (module_instance().has_value() &&
+         module_instance().value() == rhs.module_instance().value()) &&
+        callee() == rhs.callee() && source_range() == rhs.source_range();
+  }
+
+  bool operator!=(const InlinedCallStack& rhs) const {
+    return !(*this == rhs);
+  }
+};
+
+// {source range, node name, InlinedCallStack}
+// We store node name because same debug info will be used for
+// profiling as well, so we need to know op names as well.
+using DebugInfoTuple =
+    std::tuple<SourceRange, std::string, InlinedCallStackPtr>;
+constexpr size_t kDebugInfoTupleSourceRangeIndex{0};
+constexpr size_t kDebugInfoTupleNodeNameIndex{1};
+constexpr size_t kDebugInfoTupleInlinedCSIndex{2};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/subgraph_matcher.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/subgraph_matcher.h
new file mode 100644
index 0000000000000000000000000000000000000000..968db93a3545b124169de30c379234997c85aaad
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/subgraph_matcher.h
@@ -0,0 +1,77 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace torch::jit {
+
+/**
+ * \brief A structure describing a match of a pattern in a graph.
+ *
+ * The structure contains an anchor node, from which the match was found, and
+ * match-maps for nodes and values. A match-map specifies the correspondence
+ * between nodes in the pattern graph (match-map keys) with nodes in the actual
+ * graph (match-map values). We keep such maps for both nodes and values.
+ */
+struct Match {
+  Node* anchor;
+  std::unordered_map<const Node*, Node*> nodes_map;
+  std::unordered_map<const Value*, Value*> values_map;
+};
+
+/**
+ * \brief Find all matches of a \p PATTERN in a \p GRAPH.
+ *
+ * The function returns a vector of match-descriptors (see description of
+ * `struct Match`).
+ *
+ * Matching rules:
+ *  - Pattern graph must contain a single block.
+ *  - Matched subgraphs do not span across different blocks.
+ *  - No uses outside the match are allowed, except for Param and Return nodes.
+ *  Basically, we're matching hammocks, not arbitrary subgraphs.
+ *  - The pattern graph must return only one value (i.e. it must have a single
+ *  node leading to return).
+ *  - Nodes that are not used in computation of the return value in the pattern
+ * graph are ignored during matching (IOW, we're essentially performing DCE on
+ * the pattern).
+ *  - Pattern graph nodes cannot alias. TODO: the check not implemented yet.
+ *  - Aliasing nodes in the graph cannot constitute a match (i.e. through all
+ * found matches, no nodes in the subgraph alias with each other). TODO: check
+ * not implemented yet.
+ *  - The matcher will not mutate either the pattern graph or the matched graph.
+ * The matched graph is taken as non-const so that Match may contain non-const
+ * pointers.  This enables clients of this API to use Match to drive mutations.
+ *
+ * Note [Multi-output Patterns]
+ * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+ * Subgraph matcher provides limited support for multi-output patterns. With a
+ * single output pattern, a single scan through the graph is sufficient to
+ * find all the matches: given a starting node (an "anchor"), we can
+ * deterministically check whether a pattern matches a subgraph corresponding to
+ * this anchor node. For a general case of multi-output patterns, we would have
+ * N anchors, which would result in M^N comparisons (M is the size of the
+ * graph). Clearly this is computationally prohibitive.
+ *
+ * To overcome this, we impose some constraints on the multi-output patterns
+ * that we accept. We require that checking whether the pattern matches a
+ * subgraph would still be fully determined by a single node in the graph. To
+ * achieve this, we designate the first output in the pattern as the "main"
+ * output and assume that we can traverse up from this node to match the
+ * entire pattern.
+ *
+ * Corrolary 1: the order of outputs in the pattern matters!
+ * Corollary 2: patterns cannot contain any nodes not participating in the main
+ * output computation.
+ */
+std::vector<Match> TORCH_API
+findPatternMatches(const Graph& pattern, Graph& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/type_hashing.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/type_hashing.h
new file mode 100644
index 0000000000000000000000000000000000000000..74734be6bea23bce7d2f7de595422a0132f7ad51
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/ir/type_hashing.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+struct TORCH_API HashType {
+  size_t operator()(const TypePtr& type) const;
+  size_t operator()(const c10::ConstTypePtr& type) const;
+};
+
+struct EqualType {
+  bool operator()(const TypePtr& a, const TypePtr& b) const;
+  bool operator()(const c10::ConstTypePtr& a, const c10::ConstTypePtr& b) const;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/jit_log.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/jit_log.h
new file mode 100644
index 0000000000000000000000000000000000000000..593de26f9216a32d09f0c159f2d125d71363aaff
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/jit_log.h
@@ -0,0 +1,132 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/StringUtil.h>
+#include <torch/csrc/Export.h>
+#include <memory>
+#include <ostream>
+#include <string>
+#include <unordered_map>
+
+// `TorchScript` offers a simple logging facility that can enabled by setting an
+// environment variable `PYTORCH_JIT_LOG_LEVEL`.
+
+// Logging is enabled on a per file basis. To enable logging in
+// `dead_code_elimination.cpp`, `PYTORCH_JIT_LOG_LEVEL` should be
+// set to `dead_code_elimination.cpp` or, simply, to `dead_code_elimination`
+// (i.e. `PYTORCH_JIT_LOG_LEVEL=dead_code_elimination`).
+
+// Multiple files can be logged by separating each file name with a colon `:` as
+// in the following example,
+// `PYTORCH_JIT_LOG_LEVEL=dead_code_elimination:guard_elimination`
+
+// There are 3 logging levels available for your use ordered by the detail level
+// from lowest to highest.
+
+// * `GRAPH_DUMP` should be used for printing entire graphs after optimization
+// passes
+// * `GRAPH_UPDATE` should be used for reporting graph transformations (i.e.
+// node deletion, constant folding, etc)
+// * `GRAPH_DEBUG` should be used for providing information useful for debugging
+//   the internals of a particular optimization pass or analysis
+
+// The default logging level is `GRAPH_DUMP` meaning that only `GRAPH_DUMP`
+// statements will be enabled when one specifies a file(s) in
+// `PYTORCH_JIT_LOG_LEVEL`.
+
+// `GRAPH_UPDATE` can be enabled by prefixing a file name with an `>` as in
+// `>alias_analysis`.
+// `GRAPH_DEBUG` can be enabled by prefixing a file name with an `>>` as in
+// `>>alias_analysis`.
+// `>>>` is also valid and **currently** is equivalent to `GRAPH_DEBUG` as there
+// is no logging level that is higher than `GRAPH_DEBUG`.
+
+namespace torch::jit {
+
+struct Node;
+struct Graph;
+
+enum class JitLoggingLevels {
+  GRAPH_DUMP = 0,
+  GRAPH_UPDATE,
+  GRAPH_DEBUG,
+};
+
+TORCH_API std::string get_jit_logging_levels();
+
+TORCH_API void set_jit_logging_levels(std::string level);
+
+TORCH_API void set_jit_logging_output_stream(std::ostream& out_stream);
+
+TORCH_API std::ostream& get_jit_logging_output_stream();
+
+TORCH_API std::string getHeader(const Node* node);
+
+TORCH_API std::string log_function(const std::shared_ptr<Graph>& graph);
+
+TORCH_API ::torch::jit::JitLoggingLevels jit_log_level();
+
+// Prefix every line in a multiline string \p IN_STR with \p PREFIX.
+TORCH_API std::string jit_log_prefix(
+    const std::string& prefix,
+    const std::string& in_str);
+
+TORCH_API std::string jit_log_prefix(
+    ::torch::jit::JitLoggingLevels level,
+    const char* fn,
+    int l,
+    const std::string& in_str);
+
+TORCH_API bool is_enabled(
+    const char* cfname,
+    ::torch::jit::JitLoggingLevels level);
+
+TORCH_API std::ostream& operator<<(
+    std::ostream& out,
+    ::torch::jit::JitLoggingLevels level);
+
+#define JIT_LOG(level, ...)                                         \
+  if (is_enabled(__FILE__, level)) {                                \
+    ::torch::jit::get_jit_logging_output_stream()                   \
+        << ::torch::jit::jit_log_prefix(                            \
+               level, __FILE__, __LINE__, ::c10::str(__VA_ARGS__)); \
+  }
+
+// tries to reconstruct original python source
+#define SOURCE_DUMP(MSG, G)                       \
+  JIT_LOG(                                        \
+      ::torch::jit::JitLoggingLevels::GRAPH_DUMP, \
+      MSG,                                        \
+      '\n',                                       \
+      ::torch::jit::log_function(G));
+// use GRAPH_DUMP for dumping graphs after optimization passes
+#define GRAPH_DUMP(MSG, G) \
+  JIT_LOG(                 \
+      ::torch::jit::JitLoggingLevels::GRAPH_DUMP, MSG, '\n', (G)->toString());
+// use GRAPH_UPDATE for reporting graph transformations (i.e. node deletion,
+// constant folding, CSE)
+#define GRAPH_UPDATE(...) \
+  JIT_LOG(::torch::jit::JitLoggingLevels::GRAPH_UPDATE, __VA_ARGS__);
+// use GRAPH_DEBUG to provide information useful for debugging a particular opt
+// pass
+#define GRAPH_DEBUG(...) \
+  JIT_LOG(::torch::jit::JitLoggingLevels::GRAPH_DEBUG, __VA_ARGS__);
+// use GRAPH_EXPORT to export a graph so that the IR can be loaded by a script
+#define GRAPH_EXPORT(MSG, G)                       \
+  JIT_LOG(                                         \
+      ::torch::jit::JitLoggingLevels::GRAPH_DEBUG, \
+      MSG,                                         \
+      "\n<GRAPH_EXPORT>\n",                        \
+      (G)->toString(),                             \
+      "</GRAPH_EXPORT>");
+
+#define GRAPH_DUMP_ENABLED \
+  (is_enabled(__FILE__, ::torch::jit::JitLoggingLevels::GRAPH_DUMP))
+#define GRAPH_UPDATE_ENABLED \
+  (is_enabled(__FILE__, ::torch::jit::JitLoggingLevels::GRAPH_UPDATE))
+#define GRAPH_DEBUG_ENABLED \
+  (is_enabled(__FILE__, ::torch::jit::JitLoggingLevels::GRAPH_DEBUG))
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/jit_opt_limit.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/jit_opt_limit.h
new file mode 100644
index 0000000000000000000000000000000000000000..59a2d591837d8867032fedb709ee7090b850bc29
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/jit_opt_limit.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <string>
+#include <unordered_map>
+
+// `TorchScript` offers a simple optimization limit checker
+// that can be configured through environment variable `PYTORCH_JIT_OPT_LIMIT`.
+// The purpose is to limit how many optimization you can make per pass.
+// This is useful for debugging any passes.
+
+// Opt limit checker is enabled on a per file basis (hence per pass). For
+// example, in `constant_propagation.cpp`, `PYTORCH_JIT_OPT_LIMIT` should be set
+// to `constant_propagation=<opt_limit>` or, simply, to
+// `constant_propagation=<opt_limit>` where <opt_limit> is the number of
+// optimizations you want to make for the pass. (i.e.
+// `PYTORCH_JIT_OPT_LIMIT="constant_propagation=<opt_limit>"`).
+
+// Multiple files can be configured by separating each file name with a colon
+// `:` as in the following example,
+// `PYTORCH_JIT_OPT_LIMIT="constant_propagation=<opt_limit>:dead_code_elimination=<opt_limit>"`
+
+// You can call opt limiter by calling JIT_OPT_ALLOWED. It will return true if
+// we haven't reached the optimization limit yet. Otherwise, it will return
+// false. Typical usage:
+
+// if (!JIT_OPT_ALLOWED) {
+//     GRAPH_DUMP(...); //supplied from jit_log
+//     return;
+// }
+
+namespace torch::jit {
+
+TORCH_API bool opt_limit(const char* pass_name);
+
+#define JIT_OPT_ALLOWED opt_limit(__FILE__)
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/code.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/code.h
new file mode 100644
index 0000000000000000000000000000000000000000..62d7b160ebd14f5f2ded616347f012db37b44701
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/code.h
@@ -0,0 +1,39 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <vector>
+
+#include <ATen/core/ivalue.h>
+#include <ATen/core/operator_name.h>
+#include <torch/csrc/jit/runtime/instruction.h>
+
+namespace torch::jit::mobile {
+
+using Stack = std::vector<c10::IValue>;
+using DebugHandle = int64_t;
+
+class Function;
+
+struct Code {
+  std::vector<Instruction> instructions_;
+  std::vector<DebugHandle> debug_handles_;
+  std::vector<c10::OperatorName> op_names_;
+  std::vector<int> operator_input_sizes_;
+  std::vector<std::function<void(Stack&)>> operators_;
+  std::vector<c10::IValue> constants_;
+  std::vector<c10::TypePtr> types_;
+  // TODO After we actually export CALL instructions we can remove this.
+  // We may need a two-stage importing scheme, where we firstly construct all
+  // function objects, and then append referenced function pointers. This could
+  // be done in parseMethods().
+  std::vector<mobile::Function*> functions_;
+  size_t register_size_ = 0; // Aggregated output size.
+  // initialized means operators_ array is filled with operators
+  bool initialized = false;
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/backport.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/backport.h
new file mode 100644
index 0000000000000000000000000000000000000000..c43fd3542c3b9430702f5dbe4067462f36900ffa
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/backport.h
@@ -0,0 +1,33 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Export.h>
+#include <istream>
+
+namespace torch::jit {
+
+TORCH_API bool _backport_for_mobile(
+    std::istream& in,
+    std::ostream& out,
+    const int64_t to_version);
+
+TORCH_API bool _backport_for_mobile(
+    std::istream& in,
+    const std::string& output_filename,
+    const int64_t to_version);
+
+TORCH_API bool _backport_for_mobile(
+    const std::string& input_filename,
+    std::ostream& out,
+    const int64_t to_version);
+
+TORCH_API bool _backport_for_mobile(
+    const std::string& input_filename,
+    const std::string& output_filename,
+    const int64_t to_version);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/backport_manager.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/backport_manager.h
new file mode 100644
index 0000000000000000000000000000000000000000..1a06d441ac992af2c5c1236d53d8b3398d19849e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/backport_manager.h
@@ -0,0 +1,53 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <unordered_map>
+
+namespace c10 {
+struct IValue;
+}
+
+namespace caffe2::serialize {
+class PyTorchStreamWriter;
+} // namespace caffe2::serialize
+
+namespace torch::jit {
+
+/*
+BackportManager manages a list of backport from n to n-1 function, and provides
+function to check if a specific function exists.
+*/
+class BackportManager final {
+ public:
+  bool hasBytecodeBackportFunction(const int64_t from_version) const;
+
+  std::unordered_map<
+      int64_t,
+      std::function<std::stringstream(std::stringstream&)>>&
+  bytecodeBackportFunctions() const;
+
+  bool backport(
+      std::istream& oss,
+      caffe2::serialize::PyTorchStreamWriter& final_writer,
+      int64_t from_version,
+      int64_t to_version) const;
+
+  BackportManager(BackportManager const&) = delete;
+  BackportManager& operator=(BackportManager const&) = delete;
+  BackportManager();
+
+ private:
+  // Registry of backport functions.
+  void registerBytecodeBackportFunction(
+      const int64_t from_version,
+      const std::function<std::stringstream(std::stringstream&)>&
+          backport_function);
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/model_compatibility.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/model_compatibility.h
new file mode 100644
index 0000000000000000000000000000000000000000..1469432f1e1a9e99e3b770e06afb0b281a63d301
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/model_compatibility.h
@@ -0,0 +1,109 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <c10/macros/Export.h>
+#include <torch/csrc/jit/mobile/compatibility/runtime_compatibility.h>
+
+#include <istream>
+#include <memory>
+#include <unordered_map>
+#include <vector>
+
+namespace caffe2::serialize {
+class PyTorchStreamReader;
+class ReadAdapterInterface;
+} // namespace caffe2::serialize
+
+namespace torch::jit {
+
+// The family of methods below to get bytecode version from a model
+// Throws if not passed in a well formed model
+TORCH_API uint64_t _get_model_bytecode_version(std::istream& in);
+
+TORCH_API uint64_t _get_model_bytecode_version(const std::string& filename);
+
+TORCH_API uint64_t _get_model_bytecode_version(
+    const std::shared_ptr<caffe2::serialize::ReadAdapterInterface>& rai);
+
+uint64_t _get_model_bytecode_version(
+    const std::vector<c10::IValue>& bytecode_ivalues);
+
+// The family of methods below to get the operator version from a model
+// Throws if not passed in a well formed model
+TORCH_API uint64_t _get_model_operator_version(std::istream& in);
+
+TORCH_API uint64_t _get_model_operator_version(const std::string& filename);
+
+TORCH_API uint64_t _get_model_operator_version(
+    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);
+
+// Utility Functions
+std::vector<c10::IValue> get_bytecode_ivalues(
+    caffe2::serialize::PyTorchStreamReader& reader);
+
+c10::IValue readArchive(
+    const std::string& archive_name,
+    caffe2::serialize::PyTorchStreamReader& stream_reader);
+
+bool check_zip_file(
+    const std::shared_ptr<caffe2::serialize::ReadAdapterInterface>& rai);
+
+// The family of methods below to get the root ops and information from a model
+TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::istream& in);
+
+TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    const std::string& filename);
+
+TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
+    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);
+
+// The family of methods below to get contained types from a model
+// Throws if not passed in a well formed model
+TORCH_API std::unordered_set<std::string> _get_mobile_model_contained_types(
+    std::istream& in);
+
+TORCH_API std::unordered_set<std::string> _get_mobile_model_contained_types(
+    const std::string& filename);
+
+TORCH_API std::unordered_set<std::string> _get_mobile_model_contained_types(
+    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);
+
+std::unordered_set<std::string> _get_mobile_model_contained_types(
+    const std::vector<c10::IValue>& bytecode_ivalues);
+
+// The family of methods below return the compatibility information of a model
+struct ModelCompatibilityInfo {
+  uint64_t bytecode_version;
+  std::unordered_map<std::string, OperatorInfo> operator_info;
+  std::unordered_set<std::string> type_table;
+  uint64_t operator_version;
+
+  // Factory Methods
+  static TORCH_API ModelCompatibilityInfo get(std::istream& in);
+  static TORCH_API ModelCompatibilityInfo get(const std::string& filename);
+  static TORCH_API ModelCompatibilityInfo
+  get(std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);
+};
+
+enum ModelCompatibilityStatus {
+  OK = 1,
+  ERROR = 2,
+};
+
+struct ModelCompatCheckResult {
+  ModelCompatibilityStatus status;
+  std::vector<std::string> errors;
+};
+// Takes in information about a runtime and a model and returns if the two are
+// compatible with one another.
+TORCH_API ModelCompatCheckResult is_compatible(
+    RuntimeCompatibilityInfo runtime_info,
+    const ModelCompatibilityInfo& model_info);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/runtime_compatibility.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/runtime_compatibility.h
new file mode 100644
index 0000000000000000000000000000000000000000..dbe8f38ef1b12e115e666a5ab111801d6ba080cc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/compatibility/runtime_compatibility.h
@@ -0,0 +1,47 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Export.h>
+#include <optional>
+
+#include <memory>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace torch::jit {
+
+// Struct storing metadata of an operator that can be useful for versioning
+struct OperatorInfo {
+  // The number of arguments within the schema of the op
+  std::optional<int> num_schema_args;
+};
+
+struct RuntimeCompatibilityInfo {
+  std::pair<uint64_t, uint64_t> min_max_supported_bytecode_version;
+  std::unordered_map<std::string, OperatorInfo> operator_info;
+  std::unordered_set<std::string> supported_types;
+  std::pair<uint64_t, uint64_t> min_max_supported_operator_versions;
+
+  // Factory Method
+  static TORCH_API RuntimeCompatibilityInfo get();
+};
+
+TORCH_API uint64_t _get_runtime_bytecode_version();
+
+TORCH_API std::pair<uint64_t, uint64_t> _get_runtime_bytecode_min_max_versions();
+
+TORCH_API std::pair<uint64_t, uint64_t>
+_get_runtime_operators_min_max_versions();
+
+TORCH_API std::unordered_map<std::string, OperatorInfo>
+_get_runtime_ops_and_info();
+
+TORCH_API std::unordered_set<std::string> _get_mobile_supported_types();
+
+TORCH_API std::unordered_set<std::string> _get_loaded_custom_classes();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/debug_info.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/debug_info.h
new file mode 100644
index 0000000000000000000000000000000000000000..4c601d6c123343d9114d9c1be35e5779e692ac39
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/debug_info.h
@@ -0,0 +1,60 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/flat_hash_map.h>
+#include <caffe2/serialize/inline_container.h>
+#include <torch/csrc/jit/api/compilation_unit.h>
+#include <torch/csrc/jit/ir/scope.h>
+#include <torch/csrc/jit/serialization/source_range_serialization.h>
+
+namespace torch::jit {
+/*
+ * MobileDebugTable:
+ * Deserializes debug_pkl and callstack_map records from PT model's zip archive
+ * and stores them in a map of debug handles to DebugInfoPair. Debug handles are
+ * unique per model and runtime, be in lite interpreter or delegate, an
+ * exception of BackendRuntimeException should raised using debug handles.
+ * getSourceDebugString method is responsible for translating debug
+ * handles to correspond debug information.
+ * This debug information includes stack trace of model level source code and
+ * module hierarchy where the exception occurred.
+ */
+class MobileDebugTable {
+ public:
+  MobileDebugTable() = default;
+  MobileDebugTable(
+      std::unique_ptr<caffe2::serialize::PyTorchStreamReader>& reader,
+      const std::shared_ptr<CompilationUnit>& cu);
+
+  template <typename It>
+  MobileDebugTable(It begin, It end) : callstack_ptr_map_(begin, end) {}
+
+  std::string getSourceDebugString(
+      const int64_t debug_handle,
+      const std::string& top_module_type_name = "ModuleTypeUnknown") const;
+  std::string getSourceDebugString(
+      const std::vector<int64_t>& debug_handles,
+      const std::string& top_module_type_name = "ModuleTypeUnknown") const;
+  std::string getModuleHierarchyInfo(
+      const int64_t debug_handle,
+      const std::string& top_module_type_name = "ModuleTypeUnknown") const;
+  std::string getModuleHierarchyInfo(
+      const std::vector<int64_t>& debug_handles,
+      const std::string& top_module_type_name = "ModuleTypeUnknown") const;
+
+  const ska::flat_hash_map<int64_t, DebugInfoTuple>& getCallStackPtrMap()
+      const {
+    return callstack_ptr_map_;
+  }
+
+ private:
+  std::pair<std::string, std::string> getSourceDebugModuleHierarchyInfo(
+      const std::vector<int64_t>& debug_handles,
+      const std::string& top_module_type_name = "ModuleTypeUnknown") const;
+  ska::flat_hash_map<int64_t, DebugInfoTuple> callstack_ptr_map_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/file_format.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/file_format.h
new file mode 100644
index 0000000000000000000000000000000000000000..69ce02b0f4e16eef79f60162358de2b2540ef973
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/file_format.h
@@ -0,0 +1,200 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <cerrno>
+#include <cstddef>
+#include <cstring>
+#include <fstream>
+#include <istream>
+#include <memory>
+
+#include <c10/core/CPUAllocator.h>
+#include <c10/core/impl/alloc_cpu.h>
+#include <caffe2/serialize/read_adapter_interface.h>
+
+#if defined(HAVE_MMAP)
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+#endif
+
+/**
+ * @file
+ *
+ * Helpers for identifying file formats when reading serialized data.
+ *
+ * Note that these functions are declared inline because they will typically
+ * only be called from one or two locations per binary.
+ */
+
+namespace torch::jit {
+
+/**
+ * The format of a file or data stream.
+ */
+enum class FileFormat {
+  UnknownFileFormat = 0,
+  FlatbufferFileFormat,
+  ZipFileFormat,
+};
+
+/// The size of the buffer to pass to #getFileFormat(), in bytes.
+constexpr size_t kFileFormatHeaderSize = 8;
+constexpr size_t kMaxAlignment = 16;
+
+/**
+ * Returns the likely file format based on the magic header bytes in @p header,
+ * which should contain the first bytes of a file or data stream.
+ */
+// NOLINTNEXTLINE(facebook-hte-NamespaceScopedStaticDeclaration)
+static inline FileFormat getFileFormat(const char* data) {
+  // The size of magic strings to look for in the buffer.
+  static constexpr size_t kMagicSize = 4;
+
+  // Bytes 4..7 of a Flatbuffer-encoded file produced by
+  // `flatbuffer_serializer.h`. (The first four bytes contain an offset to the
+  // actual Flatbuffer data.)
+  static constexpr std::array<char, kMagicSize> kFlatbufferMagicString = {
+      'P', 'T', 'M', 'F'};
+  static constexpr size_t kFlatbufferMagicOffset = 4;
+
+  // The first four bytes of a ZIP file.
+  static constexpr std::array<char, kMagicSize> kZipMagicString = {
+      'P', 'K', '\x03', '\x04'};
+
+  // Note that we check for Flatbuffer magic first. Since the first four bytes
+  // of flatbuffer data contain an offset to the root struct, it's theoretically
+  // possible to construct a file whose offset looks like the ZIP magic. On the
+  // other hand, bytes 4-7 of ZIP files are constrained to a small set of values
+  // that do not typically cross into the printable ASCII range, so a ZIP file
+  // should never have a header that looks like a Flatbuffer file.
+  if (std::memcmp(
+          data + kFlatbufferMagicOffset,
+          kFlatbufferMagicString.data(),
+          kMagicSize) == 0) {
+    // Magic header for a binary file containing a Flatbuffer-serialized mobile
+    // Module.
+    return FileFormat::FlatbufferFileFormat;
+  } else if (std::memcmp(data, kZipMagicString.data(), kMagicSize) == 0) {
+    // Magic header for a zip file, which we use to store pickled sub-files.
+    return FileFormat::ZipFileFormat;
+  }
+  return FileFormat::UnknownFileFormat;
+}
+
+/**
+ * Returns the likely file format based on the magic header bytes of @p data.
+ * If the stream position changes while inspecting the data, this function will
+ * restore the stream position to its original offset before returning.
+ */
+// NOLINTNEXTLINE(facebook-hte-NamespaceScopedStaticDeclaration)
+static inline FileFormat getFileFormat(std::istream& data) {
+  FileFormat format = FileFormat::UnknownFileFormat;
+  std::streampos orig_pos = data.tellg();
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+  std::array<char, kFileFormatHeaderSize> header;
+  data.read(header.data(), header.size());
+  if (data.good()) {
+    format = getFileFormat(header.data());
+  }
+  data.seekg(orig_pos, data.beg);
+  return format;
+}
+
+/**
+ * Returns the likely file format based on the magic header bytes of the file
+ * named @p filename.
+ */
+// NOLINTNEXTLINE(facebook-hte-NamespaceScopedStaticDeclaration)
+static inline FileFormat getFileFormat(const std::string& filename) {
+  std::ifstream data(filename, std::ifstream::binary);
+  return getFileFormat(data);
+}
+
+// NOLINTNEXTLINE(facebook-hte-NamespaceScopedStaticDeclaration)
+static void file_not_found_error() {
+  std::stringstream message;
+  message << "Error while opening file: ";
+  if (errno == ENOENT) {
+    message << "no such file or directory" << '\n';
+  } else {
+    message << "error no is: " << errno << '\n';
+  }
+  TORCH_CHECK(false, message.str());
+}
+
+// NOLINTNEXTLINE(facebook-hte-NamespaceScopedStaticDeclaration)
+static inline std::tuple<std::shared_ptr<char>, size_t> get_file_content(
+    const char* filename) {
+#if defined(HAVE_MMAP)
+  int fd = open(filename, O_RDONLY);
+  if (fd < 0) {
+    // failed to open file, chances are it's no such file or directory.
+    file_not_found_error();
+  }
+  struct stat statbuf{};
+  fstat(fd, &statbuf);
+  size_t size = statbuf.st_size;
+  void* ptr = mmap(nullptr, statbuf.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
+  close(fd);
+  auto deleter = [statbuf](char* ptr) { munmap(ptr, statbuf.st_size); };
+  std::shared_ptr<char> data(reinterpret_cast<char*>(ptr), deleter);
+#else
+  FILE* f = fopen(filename, "rb");
+  if (f == nullptr) {
+    file_not_found_error();
+  }
+  fseek(f, 0, SEEK_END);
+  size_t size = ftell(f);
+  fseek(f, 0, SEEK_SET);
+  // make sure buffer size is multiple of alignment
+  size_t buffer_size = (size / kMaxAlignment + 1) * kMaxAlignment;
+  std::shared_ptr<char> data(
+      static_cast<char*>(c10::alloc_cpu(buffer_size)), c10::free_cpu);
+  auto nread = fread(data.get(), size, 1, f);
+  TORCH_CHECK(nread == 1, "Failed to read file: ", filename);
+  fclose(f);
+#endif
+  return std::make_tuple(data, size);
+}
+
+// NOLINTNEXTLINE(facebook-hte-NamespaceScopedStaticDeclaration)
+static inline std::tuple<std::shared_ptr<char>, size_t> get_stream_content(
+    std::istream& in) {
+  // get size of the stream and reset to orig
+  std::streampos orig_pos = in.tellg();
+  in.seekg(orig_pos, std::ios::end);
+  const long size = in.tellg();
+  in.seekg(orig_pos, in.beg);
+
+  // read stream
+  // NOLINT make sure buffer size is multiple of alignment
+  size_t buffer_size = (size / kMaxAlignment + 1) * kMaxAlignment;
+  std::shared_ptr<char> data(
+      static_cast<char*>(c10::alloc_cpu(buffer_size)), c10::free_cpu);
+  in.read(data.get(), size);
+
+  // reset stream to original position
+  in.seekg(orig_pos, in.beg);
+  return std::make_tuple(data, size);
+}
+
+// NOLINTNEXTLINE(facebook-hte-NamespaceScopedStaticDeclaration)
+static inline std::tuple<std::shared_ptr<char>, size_t> get_rai_content(
+    caffe2::serialize::ReadAdapterInterface* rai) {
+  size_t buffer_size = (rai->size() / kMaxAlignment + 1) * kMaxAlignment;
+  std::shared_ptr<char> data(
+      static_cast<char*>(c10::alloc_cpu(buffer_size)), c10::free_cpu);
+  rai->read(
+      0, data.get(), rai->size(), "Loading ReadAdapterInterface to bytes");
+  return std::make_tuple(data, buffer_size);
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/flatbuffer_loader.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/flatbuffer_loader.h
new file mode 100644
index 0000000000000000000000000000000000000000..fc11a2f54808cb206d0c531ba8b23ee93f755990
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/flatbuffer_loader.h
@@ -0,0 +1,132 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <istream>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <ATen/core/ivalue.h>
+#include <c10/core/Device.h>
+#include <c10/macros/Macros.h>
+#include <torch/csrc/jit/mobile/module.h>
+#include <optional>
+
+/**
+ * Defines the public API for loading flatbuffer-serialized mobile modules.
+ * Note that this header must not include or depend on flatbuffer-defined
+ * types, to avoid leaking those details to PyTorch clients.
+ */
+
+namespace torch::jit {
+
+/// All non-copied data pointers provided to `parse_and_initialize_*` functions
+/// must be aligned to this boundary. Since the Module will point directly into
+/// the data, this alignment is necessary to ensure that certain types/structs
+/// are properly aligned.
+constexpr size_t kFlatbufferDataAlignmentBytes = 16;
+
+/// Maps file names to file contents.
+using ExtraFilesMap = std::unordered_map<std::string, std::string>;
+
+// On high level, to produce a Module from a file on disk, we need to go
+// through the follow steps:
+// 1. Read: Read the file from disk -> memory
+// 2. Deserialize: Parse the bytes to produce some in memory manipulable
+//    structure
+// 3. Module initialization: Produce mobile::Module out of the structure
+//    produced in 2.
+// Under this context, the structure described in 2. is the flatbuffer-defined
+// type mobile::serialization::Module. However, this step/type is not visible in
+// the public API.
+
+// Parse a mobile::Module from raw bytes.
+//
+// This function does steps 2+3 described above.
+//
+// Does not take ownership of `data`; if you want it to take ownership, see the
+// shared_ptr overload of this function.
+//
+// If should_copy_tensor_memory is true, then the returned module will NOT have
+// references to `data`, so `data` can be freed immediately.
+//
+// If should_copy_tensor_memory is false, then returned module will have tensors
+// that points inside of `data`; the caller will need to make sure that `data`
+// outlives the returned Module. Also, `data` must be aligned to
+// kFlatbufferDataAlignmentBytes.
+TORCH_API mobile::Module parse_and_initialize_mobile_module(
+    void* data,
+    size_t size, // of `data`, in bytes.
+    std::optional<at::Device> device = std::nullopt,
+    ExtraFilesMap* extra_files = nullptr,
+    bool should_copy_tensor_memory = false);
+
+// Parse a mobile::Module from raw bytes.
+//
+// This function does steps 2+3 described above.
+//
+// The returned Module holds a reference to `data`, which must be aligned to
+// kFlatbufferDataAlignmentBytes.
+//
+// If you do not want the Module to hold a reference to `data`, see the raw
+// pointer overload of this function.
+TORCH_API mobile::Module parse_and_initialize_mobile_module(
+    std::shared_ptr<char> data,
+    size_t size, // of `data`, in bytes.
+    std::optional<at::Device> device = std::nullopt,
+    ExtraFilesMap* extra_files = nullptr);
+
+// Parse a mobile::Module from raw bytes, also returning JIT-related metadata.
+//
+// This is the same as parse_and_initialize_mobile_module() except that it also
+// extracts JIT source files and constants. Can be used to construct a
+// jit::Module.
+TORCH_API mobile::Module parse_and_initialize_mobile_module_for_jit(
+    void* data,
+    size_t size, // of `data`, in bytes.
+    ExtraFilesMap& jit_sources,
+    std::vector<IValue>& jit_constants,
+    std::optional<at::Device> device = std::nullopt,
+    ExtraFilesMap* extra_files = nullptr);
+
+// Load a mobile::Module from a filepath.
+//
+// This function does steps 1+2+3 described above.
+//
+// We need to have this as a convenience because Python API will need to wrap
+// this. C++ clients should use one of the versions of
+// parse_and_initialize_mobile_module() so they can manage the raw data more
+// directly.
+TORCH_API mobile::Module load_mobile_module_from_file(
+    const std::string& filename,
+    std::optional<at::Device> device = std::nullopt,
+    ExtraFilesMap* extra_files = nullptr);
+
+TORCH_API uint64_t get_bytecode_version(std::istream& in);
+TORCH_API uint64_t get_bytecode_version(const std::string& filename);
+TORCH_API uint64_t get_bytecode_version_from_bytes(char* flatbuffer_content);
+
+TORCH_API mobile::ModuleInfo get_module_info_from_flatbuffer(
+    char* flatbuffer_content);
+
+// The methods below are less efficient because it need to read the stream in
+// its entirety to a buffer
+TORCH_API mobile::Module load_mobile_module_from_stream_with_copy(
+    std::istream& in,
+    std::optional<at::Device> device = std::nullopt,
+    ExtraFilesMap* extra_files = nullptr);
+
+TORCH_API mobile::Module parse_flatbuffer_no_object(
+    std::shared_ptr<char> data,
+    size_t size,
+    std::optional<at::Device> device);
+
+// no op, TODO(qihan) delete
+TORCH_API bool register_flatbuffer_loader();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/frame.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/frame.h
new file mode 100644
index 0000000000000000000000000000000000000000..72238214ca28691f490ab7f805ea1160ba0be072
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/frame.h
@@ -0,0 +1,54 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstddef>
+
+#include <torch/csrc/jit/mobile/code.h>
+#include <optional>
+
+namespace torch::jit::mobile {
+
+class Frame {
+ public:
+  explicit Frame(const Code& code) : code_(code) {}
+  const Code& getCode() const {
+    return code_;
+  }
+
+  void step() {
+    pc_++;
+  }
+
+  void jump(size_t n) {
+    pc_ += n;
+  }
+
+  size_t getPC() const {
+    return pc_;
+  }
+
+  const Instruction& getInstruction() const {
+    return code_.instructions_.at(pc_);
+  }
+
+  std::optional<int64_t> getDebugHandle() const {
+    return getDebugHandle(pc_);
+  }
+
+  std::optional<int64_t> getDebugHandle(size_t pc) const {
+    if (pc >= code_.debug_handles_.size()) {
+      return {};
+    }
+    return code_.debug_handles_[pc];
+  }
+
+ private:
+  const Code& code_;
+  size_t pc_{0};
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/function.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/function.h
new file mode 100644
index 0000000000000000000000000000000000000000..82543310a83fd4827b24329d8e4b2af1646c1fbb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/function.h
@@ -0,0 +1,91 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <vector>
+
+#include <ATen/core/function.h>
+#include <ATen/core/function_schema.h>
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/mobile/code.h>
+
+namespace torch::jit {
+enum OpCode : uint8_t;
+struct Instruction;
+struct OperatorString;
+
+namespace mobile {
+
+class TORCH_API Function : public torch::jit::Function {
+ public:
+  explicit Function(c10::QualifiedName name);
+  Function(
+      c10::QualifiedName name,
+      Code code,
+      std::optional<c10::FunctionSchema> schema);
+  void run(Stack& stack) override;
+  at::IValue operator()(Stack& stack);
+  void ensure_defined() override {}
+  size_t num_inputs() const override;
+  const c10::QualifiedName& qualname() const override;
+  bool call(
+      Stack& /*unused*/,
+      c10::function_ref<void(const mobile::Code&)> /*f*/ /*unused*/) override;
+
+  // NOTE: the APIs below is dangerous: if you call append_instruction with
+  // dbg_handle and then call it without; then the dbg_handle will become
+  // misaligned. Therefore only use ONE variant at time.
+  void append_instruction(OpCode op, int64_t X, int64_t N, int64_t dbg_handle);
+  void append_instruction(OpCode op, int64_t X, int64_t N);
+  void append_operator(
+      const std::string& name,
+      const std::string& overload_name,
+      const std::optional<int>& num_specified_args);
+  void append_constant(const c10::IValue& constant);
+  void append_type(const c10::TypePtr& type);
+  void append_function(mobile::Function& func);
+
+  void set_register_size(size_t size);
+
+  int64_t get_debug_handle(size_t pc) const;
+  const Code& get_code() const;
+  Code& get_code();
+
+  torch::jit::Function& setSchema(c10::FunctionSchema schema) override;
+  bool hasSchema() const;
+  const c10::FunctionSchema& getSchema() const override;
+
+  // Returns the debug handle corresponding to where the execution
+  // is halted due to exception.
+  // If no corresponding debug handle is found then -1 is returned.
+  const std::vector<int64_t>& getExceptionDebugHandles() const;
+  static Function& registerFunc(
+      const std::string& qualified_name,
+      const std::vector<Instruction>& instructions,
+      const std::vector<c10::IValue>& constants,
+      const std::vector<c10::TypePtr>& types,
+      const size_t register_size);
+
+  // if not initialize, initialize by loading operators.
+  // return true of all op loaded, return false if some op is not found
+  // in the current runtime. Then, the ops that did not found will be filled
+  // in unsupported_op_names
+  bool initialize_operators(bool should_check_operators);
+
+ private:
+  c10::QualifiedName name_;
+  Code code_;
+  std::optional<c10::FunctionSchema> schema_; // (byte-code version 4+)
+};
+
+std::optional<std::function<void(Stack&)>> makeOperatorFunction(
+    const c10::OperatorName& opname,
+    std::optional<int> num_specified_args);
+
+TORCH_API std::string operator_str(const c10::OperatorName& opname);
+
+} // namespace mobile
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import.h
new file mode 100644
index 0000000000000000000000000000000000000000..fc3abe4f36c2f03d9dd85cc5a5318d2280973539
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import.h
@@ -0,0 +1,113 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/mobile/module.h>
+#include <torch/csrc/jit/mobile/parse_operators.h>
+
+#include <istream>
+#include <memory>
+
+#include <caffe2/serialize/file_adapter.h>
+
+namespace torch::jit {
+using caffe2::serialize::ReadAdapterInterface;
+using ExtraFilesMap = std::unordered_map<std::string, std::string>;
+
+constexpr const char* kArchiveNameBytecode = "bytecode";
+constexpr const char* kArchiveNameConstants = "constants";
+constexpr const char* kArchiveNameVersion = "version";
+
+// The family of methods below load a serialized Mobile Module
+// into a mobile::Module object.
+TORCH_API mobile::Module _load_for_mobile(
+    std::istream& in,
+    std::optional<at::Device> device,
+    ExtraFilesMap& extra_file,
+    uint64_t module_load_options = kDefaultMobileLoadOptions);
+
+TORCH_API mobile::Module _load_for_mobile(
+    const std::string& filename,
+    std::optional<at::Device> device,
+    ExtraFilesMap& extra_files);
+
+TORCH_API mobile::Module _load_for_mobile(
+    std::unique_ptr<ReadAdapterInterface> rai,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    uint64_t module_load_options = kDefaultMobileLoadOptions);
+
+TORCH_API mobile::Module _load_for_mobile(
+    const std::string& filename,
+    std::optional<at::Device> device,
+    ExtraFilesMap& extra_files,
+    uint64_t module_load_options);
+
+TORCH_API mobile::Module _load_for_mobile(
+    std::istream& in,
+    std::optional<at::Device> device = std::nullopt);
+
+TORCH_API mobile::Module _load_for_mobile(
+    const std::string& filename,
+    std::optional<at::Device> device = std::nullopt);
+
+TORCH_API mobile::Module _load_for_mobile(
+    std::unique_ptr<ReadAdapterInterface> rai,
+    std::optional<c10::Device> device = std::nullopt);
+
+/**
+ * Load only the contents of the "extra/" files whose names are
+ * passed in the map (extra_files). Populate the corresponding values
+ * with the contents of those files. Do not attempt to load the entire
+ * model, and stop once the extra files have been extracted.
+ *
+ * This API is needed to be able to load GPU models on linux CPU
+ * machines and extract only the extra files so that we can inspect
+ * the metadata that was added to the .ptl archive when it was
+ * generated.
+ *
+ */
+void _load_extra_only_for_mobile(
+    const std::string& filename,
+    std::optional<at::Device> device,
+    ExtraFilesMap& extra_files);
+
+// Currently used by both mobile/import.cpp and model_compatibility.cpp.
+// Should be removed after model_compatibility.cpp start using simplified
+// version type_resolver and obj_loader.
+at::TypePtr resolveTypeNameMobile(
+    const c10::QualifiedName& qn,
+    const std::shared_ptr<CompilationUnit>& compilation_unit);
+c10::StrongTypePtr typeResolverMobile(
+    const c10::QualifiedName& qn,
+    const std::shared_ptr<CompilationUnit>& compilation_unit);
+c10::intrusive_ptr<c10::ivalue::Object> objLoaderMobile(
+    const at::StrongTypePtr& type,
+    const at::IValue& input,
+    mobile::CompilationUnit& mobile_compilation_unit);
+
+// Given a reader, which has access to a model file,
+// return true if there exists tensors in `bytecode` archive
+bool isTensorInBytecodeArchive(
+    caffe2::serialize::PyTorchStreamReader& stream_reader);
+
+namespace mobile {
+
+/**
+ * Given a torch::jit::mobile::Module, return a set of operator names
+ * (with overload name) that are used by any method in this mobile
+ * Mobile. This method runs through the bytecode for all methods
+ * in the specified model (module), and extracts all the root
+ * operator names. Root operators are operators that are called
+ * directly by the model (as opposed to non-root operators, which
+ * may be called transitively by the root operators).
+ *
+ */
+TORCH_API std::set<std::string> _export_operator_list(
+    torch::jit::mobile::Module& module);
+
+} // namespace mobile
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import_data.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import_data.h
new file mode 100644
index 0000000000000000000000000000000000000000..4d1f3e5f6a81b9c0429e342967e8b7af409fda67
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import_data.h
@@ -0,0 +1,41 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/TensorBase.h>
+#include <c10/core/Device.h>
+#include <torch/csrc/jit/mobile/module.h>
+#include <optional>
+
+#include <istream>
+#include <map>
+#include <string>
+
+namespace torch::jit {
+
+/**
+ * Loads named parameters from the serialized data in @p in.
+ *
+ * Calls #TORCH_CHECK() if the data format is not recognized.
+ */
+TORCH_API std::map<std::string, at::Tensor> _load_parameters(
+    std::istream& in,
+    std::optional<at::Device> device = std::nullopt);
+
+/**
+ * Loads named parameters from the serialized data in @p filename.
+ *
+ * Calls #TORCH_CHECK() if the data format is not recognized.
+ */
+TORCH_API std::map<std::string, at::Tensor> _load_parameters(
+    const std::string& filename,
+    std::optional<at::Device> device = std::nullopt);
+
+// NOTE: Please prefer using _load_parameters over using the function below.
+TORCH_API std::map<std::string, at::Tensor> mobile_module_to_parameter_map(
+    const mobile::Module& module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import_export_common.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import_export_common.h
new file mode 100644
index 0000000000000000000000000000000000000000..4463396fa43d270f25ad747d3403bcbcdfec5ef7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/import_export_common.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+/**
+ * @file
+ * Declarations shared between import_data.cpp and export_data.cpp
+ */
+
+namespace torch::jit::mobile::internal {
+/**
+ * The name of the mobile::Module attribute which contains saved parameters, as
+ * a Dict of names to Tensors. Only used for Flatbuffer serialization.
+ */
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+constexpr char kSavedParametersAttributeName[] = "data";
+} // namespace torch::jit::mobile::internal
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/interpreter.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/interpreter.h
new file mode 100644
index 0000000000000000000000000000000000000000..62442d6530b50a871fc2097e0c0076aba849b0c2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/interpreter.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <vector>
+
+#include <torch/csrc/jit/mobile/code.h>
+#include <torch/csrc/jit/mobile/frame.h>
+
+namespace torch::jit::mobile {
+
+struct InterpreterState {
+  TORCH_API explicit InterpreterState(const Code& code);
+  TORCH_API bool run(Stack& stack);
+
+ private:
+  void enterFrame(const Code& /*code*/);
+  void leaveFrame();
+  void saveExceptionDebugHandles();
+  void callFunction(torch::jit::Function& f, Stack& stack);
+
+  c10::IValue& reg(size_t reg);
+  std::vector<c10::IValue> registers_;
+  std::vector<Frame> frames_;
+};
+
+const std::vector<DebugHandle>& getInterpretersExceptionDebugHandles();
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/method.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/method.h
new file mode 100644
index 0000000000000000000000000000000000000000..86d2dcd7ced9de9965943cc417d7dc8dafde09a6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/method.h
@@ -0,0 +1,46 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/mobile/function.h>
+
+namespace torch::jit::mobile {
+
+class Module;
+
+struct TORCH_API Method {
+  Method(const Module* owner, Function* function);
+
+  void run(Stack& stack) const;
+  void run(Stack&& stack) const {
+    run(stack);
+  }
+
+  c10::IValue operator()(std::vector<c10::IValue> stack) const;
+
+  const std::string& name() const {
+    return function_->name();
+  }
+
+  int64_t get_debug_handle(size_t pc) const {
+    return function_->get_debug_handle(pc);
+  }
+
+  Function& function() const {
+    return *function_;
+  }
+
+ private:
+  // Methods are uniquely owned by a single module.
+  // This raw pointer allows referencing the module
+  const Module* owner_;
+
+  // Underlying unbound function
+  Function* function_;
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.h
new file mode 100644
index 0000000000000000000000000000000000000000..76ecd8a5ebea3a3121e156bbca3f29e7edfe8e12
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/record_function.h>
+#include <c10/util/Synchronized.h>
+#include <map>
+#include <set>
+#include <string>
+
+namespace torch::jit::mobile {
+
+/* The BuildFeatureTracer class handles the attachment and removal of a
+ * recording callback that traces the invocation of code that handles executing
+ * generic build features.
+ *
+ * You can get the set of used build features using
+ * getBuildFeatures().
+ *
+ * Note: This class is not thread safe or re-entrant, and should not be used
+ * across multiple threads of execution.
+ *
+ */
+struct BuildFeatureTracer final {
+  at::CallbackHandle handle_;
+  /* These are the custom class names (constant
+   * character string) which shows up in code.
+   */
+  typedef std::set<std::string> build_feature_type;
+
+  BuildFeatureTracer();
+  static c10::Synchronized<build_feature_type>& getBuildFeatures();
+
+  ~BuildFeatureTracer() {
+    at::removeCallback(handle_);
+  }
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/CustomClassTracer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/CustomClassTracer.h
new file mode 100644
index 0000000000000000000000000000000000000000..4a88312b2f5872e5edf81dfa8a3006b2a1b8b21f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/CustomClassTracer.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/record_function.h>
+#include <c10/util/Synchronized.h>
+#include <map>
+#include <set>
+#include <string>
+
+namespace torch::jit::mobile {
+
+/* The CustomClassTracer class handles the attachment and removal of a recording
+ * callback that traces the invocation of code that handles loading custom
+ * classes on mobile.
+ *
+ * You can get the set of used custom classes using
+ * getLoadedClasses().
+ *
+ * Note: This class is not thread safe or re-entrant, and should not be used
+ * across multiple threads of execution.
+ *
+ */
+struct CustomClassTracer final {
+  at::CallbackHandle handle_;
+  /* These are the custom class names (constant
+   * character string) which shows up in code.
+   */
+  typedef std::set<std::string> custom_classes_type;
+
+  CustomClassTracer();
+  static c10::Synchronized<custom_classes_type>& getLoadedClasses();
+
+  ~CustomClassTracer() {
+    at::removeCallback(handle_);
+  }
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/KernelDTypeTracer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/KernelDTypeTracer.h
new file mode 100644
index 0000000000000000000000000000000000000000..e0ddb8c3aaea93ec3242aa3d1b86ec72890aec2e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/KernelDTypeTracer.h
@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/record_function.h>
+#include <c10/util/Synchronized.h>
+#include <map>
+#include <set>
+#include <string>
+
+namespace torch::jit::mobile {
+/* The KernelDTypeTracer class handles the attachment and removal of a recording
+ * callback that traces the invocation of code that handles specific dtypes in
+ * kernel function implementations that are tagged with specific tags.
+ *
+ * You can get the set of kernel tags and the dtypes using
+ * getCalledKernelTags().
+ *
+ * Note: This class is not thread safe or re-entrant, and should not be used
+ * across multiple threads of execution.
+ *
+ */
+struct KernelDTypeTracer final {
+  at::CallbackHandle handle_;
+  /* The key of the map below (std::string) is the kernel tag name (constant
+   * character string) which shows up in code. The value part of type
+   * std::set<std::string> is the collection of dtypes for which we need to
+   * generate code for the said kernel tag.
+   */
+  typedef std::map<std::string, std::set<std::string>> kernel_tags_type;
+
+  KernelDTypeTracer();
+  static c10::Synchronized<kernel_tags_type>& getCalledKernelTags();
+
+  ~KernelDTypeTracer() {
+    at::removeCallback(handle_);
+  }
+};
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/MobileModelRunner.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/MobileModelRunner.h
new file mode 100644
index 0000000000000000000000000000000000000000..2bb7593ca65735a8b8c21acba3d61bfdaa17760e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/MobileModelRunner.h
@@ -0,0 +1,151 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <mutex>
+#include <sstream>
+
+#include <torch/csrc/autograd/grad_mode.h>
+#include <torch/csrc/jit/mobile/import.h>
+#include <torch/csrc/jit/mobile/module.h>
+#include <torch/csrc/jit/serialization/export.h>
+#include <torch/script.h>
+
+namespace torch::jit::mobile {
+
+class MobileModelRunner {
+  std::shared_ptr<torch::jit::mobile::Module> module_;
+
+ public:
+  explicit MobileModelRunner(std::string const& file_path) {
+    module_ = std::make_shared<torch::jit::mobile::Module>(
+        torch::jit::_load_for_mobile(file_path));
+  }
+
+  MobileModelRunner(
+      std::string const& file_path,
+      uint64_t module_load_options) {
+    std::unordered_map<std::string, std::string> extra_files;
+    module_ = std::make_shared<torch::jit::mobile::Module>(
+        torch::jit::_load_for_mobile(
+            file_path,
+            at::Device(at::DeviceType::CPU, 0),
+            extra_files,
+            module_load_options));
+  }
+
+  MobileModelRunner(std::stringstream oss) {
+    module_ = std::make_shared<torch::jit::mobile::Module>(
+        torch::jit::_load_for_mobile(oss, at::Device(at::DeviceType::CPU, 0)));
+  }
+
+  /**
+   * Returns true if the list of operators passed in has a Metal GPU operator,
+   * and false otherwise.
+   *
+   */
+  static bool set_has_metal_gpu_operators(std::set<std::string> const& op_list);
+
+  /**
+   * Fetches the set of root operators in the file "extra/mobile_info.json"
+   * within the .ptl archive at location file_path.
+   *
+   * An exception is thrown if:
+   *
+   * 1. The file at file_path does not exist, or
+   * 2. The contents of extra/mobile_info.json is not a JSON, or
+   * 3. The file extra/mobile_info.json does not exist, or
+   * 4. The JSON is malformed in some way and the operator list can not be
+   * extracted correctly.
+   *
+   */
+  static std::set<std::string> get_operators_from_mobile_info_json(
+      std::string const& file_path);
+
+  static std::vector<std::vector<at::IValue>> ivalue_to_bundled_inputs(
+      const c10::IValue& bundled_inputs);
+
+  static std::unordered_map<std::string, std::string>
+  ivalue_to_bundled_inputs_map(const c10::IValue& bundled_inputs);
+
+  /**
+   * Fetches all the bundled inputs of the loaded mobile model.
+   *
+   * A bundled input itself is of type std::vector<at::IValue> and the
+   * elements of this vector<> are the arguments that the "forward"
+   * method of the model accepts. i.e. each of the at::IValue is a
+   * single argument to the model's "forward" method.
+   *
+   * The outer vector holds a bundled input. For models with bundled
+   * inputs, the outer most vector will have size > 0.
+   */
+  std::vector<std::vector<at::IValue>> get_all_bundled_inputs();
+
+  /**
+   * Fetches all the bundled inputs for all functions of the loaded mobile
+   * model.
+   *
+   * The mapping is from 'function_names' eg 'forward' to bundled inputs for
+   * that function
+   *
+   * A bundled input itself is of type std::vector<at::IValue> and the
+   * elements of this vector<> are the arguments that the corresponding
+   * method of the model accepts. i.e. each of the at::IValue in the entry
+   * for forward is a single argument to the model's "forward" method.
+   *
+   * The outer vector of each value holds a bundled input. For models with
+   * bundled inputs, the outer most vector will have size > 0.
+   */
+  std::unordered_map<std::string, std::vector<std::vector<at::IValue>>>
+  get_many_functions_bundled_inputs();
+
+  /**
+   * Returns true if a model possesses get_bundled_inputs_functions_and_info()
+   */
+  bool has_new_style_bundled_inputs() const {
+    return module_->find_method("get_bundled_inputs_functions_and_info") !=
+        std::nullopt;
+  }
+
+  /**
+   * For each tensor in bundled inputs, call the user-provided function 'func'.
+   */
+  void for_each_tensor_in_bundled_inputs(
+      std::function<void(const ::at::Tensor&)> const& func);
+
+  /**
+   * Get the root operators directly called by this model's Bytecode.
+   */
+  std::set<std::string> get_root_operators() {
+    return torch::jit::mobile::_export_operator_list(*module_);
+  }
+
+  /**
+   * Runs the model against all of the provided inputs using the model's
+   * "forward" method. Returns an std::vector<at::IValue>, where each element
+   * of the returned vector is one of the return values from calling forward().
+   */
+  std::vector<at::IValue> run_with_inputs(
+      std::vector<std::vector<at::IValue>> const& bundled_inputs);
+
+  /**
+   * Runs the model against all of the provided inputs for all the specified
+   * function. Returns an std::vector<at::IValue>, where each element
+   * of the returned vector is one of the return values from calling the
+   * method named "function_name" on this model.
+   */
+  std::vector<at::IValue> run_with_inputs(
+      const std::string& function_name,
+      std::vector<std::vector<at::IValue>> const& bundled_inputs) const;
+
+  /**
+   * Attempts to run all functions in the passed in list if they exist. All
+   * funcs should require no args
+   */
+  void run_argless_functions(const std::vector<std::string>& functions);
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/OperatorCallTracer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/OperatorCallTracer.h
new file mode 100644
index 0000000000000000000000000000000000000000..2182156d87b9d2bcf7c0b40d7b9c5d49296d7e69
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/OperatorCallTracer.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/record_function.h>
+#include <c10/util/Synchronized.h>
+
+namespace torch::jit::mobile {
+/* The OperatorCallTracer class handles the attachment and removal of a
+ * recording callback that traces invocation of ATen (and other) PyTorch
+ * operators that get called via the Dispatcher.
+ *
+ * You can get the set of operators that were called (op_name.overload_name)
+ * using getCalledOperators().
+ *
+ * Note: This class is not thread safe or re-entrant, and should not be used
+ * across multiple threads of execution.
+ *
+ */
+struct OperatorCallTracer final {
+  at::CallbackHandle handle_;
+
+  OperatorCallTracer();
+
+  static c10::Synchronized<std::set<std::string>>& getCalledOperators() {
+    static c10::Synchronized<std::set<std::string>> called_operators_;
+    return called_operators_;
+  }
+
+  ~OperatorCallTracer() {
+    at::removeCallback(handle_);
+  }
+};
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/TensorUtils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/TensorUtils.h
new file mode 100644
index 0000000000000000000000000000000000000000..f437d4b8ee6ea875a75c5856f6be898d3792b7fa
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/TensorUtils.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+
+namespace torch::jit::mobile {
+/**
+ * Recursively scan the IValue object, traversing lists, tuples, dicts, and stop
+ * and call the user provided callback function 'func' when a Tensor is found.
+ */
+void for_each_tensor_in_ivalue(
+    const ::c10::IValue& iv,
+    std::function<void(const ::at::Tensor&)> const& func);
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/TracerRunner.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/TracerRunner.h
new file mode 100644
index 0000000000000000000000000000000000000000..ddef55b7bc52711df4dce97cfa82e28338d2acdb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/model_tracer/TracerRunner.h
@@ -0,0 +1,44 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <set>
+#include <string>
+#include <vector>
+
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.h>
+#include <torch/csrc/jit/mobile/model_tracer/CustomClassTracer.h>
+#include <torch/csrc/jit/mobile/model_tracer/KernelDTypeTracer.h>
+
+namespace torch::jit::mobile {
+
+const std::vector<std::string> always_included_traced_ops = {
+    // The following are called from setup sections.
+    "aten::resize_",
+    "aten::slice.Tensor",
+};
+
+struct TracerResult {
+  std::set<std::string> root_ops;
+  std::set<std::string> traced_operators;
+  KernelDTypeTracer::kernel_tags_type called_kernel_tags;
+  CustomClassTracer::custom_classes_type loaded_classes;
+  BuildFeatureTracer::build_feature_type build_features;
+  std::set<std::string> enabled_backends;
+};
+
+/**
+ * Trace a single model and return the TracerResult.
+ */
+TracerResult trace_run(const std::string& input_module_path);
+
+/**
+ * Trace multiple models and return the TracerResult.
+ */
+TracerResult trace_run(const std::vector<std::string>& input_module_paths);
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/module.h
new file mode 100644
index 0000000000000000000000000000000000000000..e7b69b620dc180efe9303ee22baefedc0e22de20
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/module.h
@@ -0,0 +1,198 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/jit/mobile/debug_info.h>
+#include <torch/csrc/jit/mobile/function.h>
+#include <torch/csrc/jit/mobile/method.h>
+#include <torch/csrc/jit/mobile/quantization.h>
+
+#include <utility>
+
+namespace torch::jit::mobile {
+using Stack = std::vector<c10::IValue>;
+
+// A CompilationUnit object is the one that gets executed by the lite
+// interpreter.
+//
+// A CompilationUnit object contains a list of Method Objects. These are methods
+// that appear in the original PyTorch Model. These method correspond to Python
+// member functions of the Model class.
+//
+// Methods in turn contain a Function, and a back-pointer to the Module that
+// owns this Method instance.
+//
+// A Function contains a Code Object (code_) which is defined in interpreter.h
+//
+// A Code object contains the following:
+//
+// std::vector<Instruction> instructions_;
+// std::vector<c10::OperatorName> op_names_;
+// std::vector<std::function<void(Stack&)>> operators_;
+// std::vector<c10::IValue> constants_;
+// std::vector<c10::TypePtr> types_;
+// size_t register_size_; // Aggregated output size.
+//
+class CompilationUnit {
+ public:
+  void register_function(std::unique_ptr<Function> fn);
+  std::vector<std::unique_ptr<Function>>& methods() {
+    return methods_;
+  }
+  const std::vector<std::unique_ptr<Function>>& methods() const {
+    return methods_;
+  }
+  Function* find_function(const c10::QualifiedName& qn);
+  const Function* find_function(const c10::QualifiedName& qn) const;
+
+  void unsafeRemoveFunction(const int64_t index) {
+    methods_.erase(methods_.begin() + index);
+  }
+
+ private:
+  std::vector<std::unique_ptr<Function>> methods_;
+};
+
+// A Torch Mobile Module is a representation of the model (trained in case
+// of inference). A Mobile Module contains
+//
+// 1. data (object_)
+// 2. metadata (optional) about the model (metadata_ from the metadata.pkl
+//    file added after training)
+// 3. Compilation Unit (cu_)
+//
+class TORCH_API Module {
+ public:
+  Module(
+      c10::intrusive_ptr<c10::ivalue::Object> object,
+      std::shared_ptr<CompilationUnit> cu)
+      : object_(std::move(object)), cu_(std::move(cu)) {}
+  Module() = default;
+  Method get_method(const std::string& method_name) const;
+  template <typename... Types>
+  c10::IValue run_method(const std::string& method_name, Types&&... args) {
+    return get_method(method_name)({IValue(std::forward<Types>(args))...});
+  }
+  c10::IValue forward(std::vector<c10::IValue> inputs) {
+    return get_method("forward")(std::move(inputs));
+  }
+  std::optional<Method> find_method(const std::string& basename) const;
+
+  const std::string name() const {
+    return object_->name();
+  }
+  const std::vector<at::IValue>& slots() const {
+    return object_->slots();
+  }
+  const c10::intrusive_ptr<c10::ivalue::Object> _ivalue() const {
+    return object_;
+  }
+  const std::vector<at::Tensor> parameters() const;
+  const std::map<std::string, at::Tensor> named_parameters() const;
+  std::string get_forward_method_debug_info(int64_t debug_handle) const;
+  std::string getModuleHierarchy(const int64_t debug_handle) const;
+  std::string getCallStack(const int64_t debug_handle) const;
+  /// Enables "training" mode.
+  void train(bool on = true);
+  /// Calls train(false) to enable "eval" mode.
+  void eval() {
+    train(/*on=*/false);
+  }
+  /// True if the module is in training mode.
+  bool is_training() const;
+  const std::unordered_map<std::string, std::string> getMetadata() const {
+    return metadata_;
+  }
+  void setMetadata(
+      const std::unordered_map<std::string, std::string>& metadata) {
+    metadata_ = metadata;
+  }
+  const std::vector<Method> get_methods() const;
+
+  c10::IValue attr(const std::string& name, c10::IValue or_else) const {
+    if (auto r = object_->type()->findAttributeSlot(name)) {
+      return object_->getSlot(*r);
+    }
+    if (auto r = object_->type()->findConstantSlot(name)) {
+      return object_->type()->getConstant(*r);
+    }
+    return or_else;
+  }
+
+  void setDebugTable(MobileDebugTable&& debug_table) {
+    debug_table_ = std::move(debug_table);
+  }
+  const MobileDebugTable& getDebugTable() const {
+    return debug_table_;
+  }
+
+  void setHasDebugHandles(bool has_debug_handles) {
+    has_debug_handles_ = has_debug_handles;
+  }
+
+  bool hasDebugHandles() const {
+    return has_debug_handles_;
+  }
+
+  const CompilationUnit& compilation_unit() const {
+    return *cu_;
+  }
+
+  void set_delete_memory(std::shared_ptr<char> delete_mem) {
+    mem_to_delete_ = std::move(delete_mem);
+  }
+
+  void set_min_operator_version(int64_t version) {
+    min_operator_version_ = version;
+  }
+
+  int64_t min_operator_version() const {
+    return min_operator_version_;
+  }
+
+  void set_bytecode_version(int64_t version) {
+    bytecode_version_ = version;
+  }
+
+  int64_t bytecode_version() const {
+    return bytecode_version_;
+  }
+
+ private:
+  friend class quantization::PTQQuanizationHelper;
+
+  bool compareMethodSchemas(
+      const std::string& name_1,
+      const std::string& name_2);
+
+  void unsafeRemoveMethod(const std::string& basename);
+
+  void unsafeCopyMethod(
+      const std::string& new_method_name,
+      const Function& to_be_copied);
+
+  c10::intrusive_ptr<c10::ivalue::Object> object_;
+  std::unordered_map<std::string, std::string> metadata_;
+  std::shared_ptr<CompilationUnit> cu_;
+  MobileDebugTable debug_table_;
+  bool has_debug_handles_ = false;
+  int64_t min_operator_version_ = 4;
+  int64_t bytecode_version_ = 4;
+
+  // Extra handle for the module to delete when itself is deleted
+  std::shared_ptr<char> mem_to_delete_;
+};
+
+struct TORCH_API ModuleInfo {
+  uint64_t bytecode_version;
+  uint64_t operator_version;
+  std::unordered_map<std::string, int> opname_to_num_args;
+  std::unordered_set<std::string> function_names;
+  std::unordered_set<std::string> type_names;
+};
+TORCH_API ModuleInfo get_module_info(const mobile::Module& module);
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/aot_compiler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/aot_compiler.h
new file mode 100644
index 0000000000000000000000000000000000000000..f82f4a1e31fba6f037daf313a81d73cbba78d6e6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/aot_compiler.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/mobile/nnc/context.h>
+
+namespace torch::jit::mobile::nnc {
+
+// Performs Ahead Of Time compilation of a given method in a model
+// returning the compiled function and LLVM assembly code
+TORCH_API std::pair<std::unique_ptr<Function>, const std::string> aotCompile(
+    const std::string& method_name,
+    std::shared_ptr<Graph>& subgraph,
+    const std::vector<std::vector<int64_t>>& sizes,
+    const std::vector<at::ScalarType>& types,
+    const std::string& kernel_func_name = "func");
+
+} // namespace torch::jit::mobile::nnc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/context.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/context.h
new file mode 100644
index 0000000000000000000000000000000000000000..5ac9e2dad7a137faa0ad3b6840af03d90e533477
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/context.h
@@ -0,0 +1,228 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <ATen/core/ivalue.h>
+#include <c10/core/ScalarType.h>
+
+namespace torch::jit::mobile::nnc {
+
+// Specify the requirements on an input tensor.
+// TODO: support input tensor with dynamic shape (PR #54982)
+struct TORCH_API InputSpec {
+  InputSpec() = default;
+
+  // Deserialize the spec from an IValue.
+  explicit InputSpec(const c10::IValue& value);
+
+  // Serialize the spec into an IValue.
+  [[nodiscard]] c10::IValue serialize() const;
+
+  // Check whether the input tensor adheres to the spec.
+  [[nodiscard]] bool validate(const at::Tensor& input) const;
+
+  std::vector<int64_t> sizes_;
+  c10::ScalarType dtype_{c10::ScalarType::Undefined};
+};
+
+// Specify the sizes/dtype/... of output tensor to preallocate the output.
+// TODO: support the case where kernel allocates output tensors dynamically.
+struct TORCH_API OutputSpec {
+  OutputSpec() = default;
+
+  // Deserialize the spec from an IValue.
+  explicit OutputSpec(const c10::IValue& value);
+
+  // Serialize the spec into an IValue.
+  [[nodiscard]] c10::IValue serialize() const;
+
+  // Allocate an output tensor in accordance with the spec.
+  [[nodiscard]] at::Tensor allocate() const;
+
+  std::vector<int64_t> sizes_;
+  c10::ScalarType dtype_{c10::ScalarType::Undefined};
+  std::optional<double> qscale_;
+  std::optional<int64_t> qzero_;
+};
+
+// Hold the temporary buffers / states needed during the execution.
+struct TORCH_API ExecutionState {
+  ExecutionState() = default;
+  ExecutionState(const ExecutionState&) = delete;
+  ExecutionState(ExecutionState&&) = default;
+  ExecutionState& operator=(const ExecutionState&) = delete;
+  ExecutionState& operator=(ExecutionState&&) = default;
+
+  // Preallocated buffers needed by the NNC kernel.
+  std::vector<c10::DataPtr> preallocations_;
+
+  // The NNC kernel expects the following arguments layout:
+  //   input tensor 1
+  //   ...
+  //   input tensor INPUT_NUM
+  //   output tensor 1
+  //   ...
+  //   output tensor OUTPUT_NUM
+  //   parameter tensor 1
+  //   ...
+  //   parameter tensor PARAM_NUM
+  //   temporary buffer 1
+  //   ...
+  //   temporary buffer BUFFER_NUM
+  std::vector<void*> arguments_;
+};
+
+// Specify how to allocate temporary buffers at initialization.
+struct TORCH_API MemoryPlan {
+  MemoryPlan() = default;
+
+  explicit MemoryPlan(const c10::IValue& value);
+
+  [[nodiscard]] c10::IValue serialize() const;
+
+  void allocate(ExecutionState* state) const;
+
+  std::vector<int64_t> buffer_sizes_;
+};
+
+// Location of a symbolic shape among dimensions of the inputs
+struct TORCH_API SymbolicShapePosition {
+  SymbolicShapePosition() = default;
+  SymbolicShapePosition(int64_t input_idx, int64_t dim_idx)
+      : input_idx_(input_idx), dim_idx_(dim_idx) {}
+
+  int64_t input_idx_;
+  int64_t dim_idx_;
+};
+
+// Represents a compiled NNC function which has a 1-1 correspondence with a
+// `Method` (e.g. `forward`). It's similar as torch::jit::mobile::Function.
+class TORCH_API Function {
+ public:
+  explicit Function() = default;
+
+  // Deserialize from an IValue that is generated by the 'serialize()' method.
+  explicit Function(const c10::IValue& value);
+
+  // Serialize into an IValue.
+  c10::IValue serialize() const;
+
+  // Execute the compiled NNC function.
+  c10::impl::GenericList run(const c10::impl::GenericList& inputs) const;
+
+  // The name of the function as specified in the model code.
+  c10::QualifiedName name() const {
+    return name_;
+  }
+
+  void set_name(const c10::QualifiedName& name) {
+    name_ = name;
+  }
+
+  // The unique id of the generated NNC kernel corresponding to the function.
+  const std::string& nnc_kernel_id() const {
+    return nnc_kernel_id_;
+  }
+
+  void set_nnc_kernel_id(const std::string& name) {
+    nnc_kernel_id_ = name;
+  }
+
+  // The parameters (e.g. weights / bias tensors) to be passed to the generated
+  // NNC kernel.
+  const c10::impl::GenericList& parameters() const {
+    return parameters_;
+  }
+
+  void set_parameters(const c10::impl::GenericList& parameters) {
+    parameters_ = parameters;
+  }
+
+  const std::vector<InputSpec>& input_specs() const {
+    return input_specs_;
+  }
+
+  void set_input_specs(const std::vector<InputSpec>& input_specs) {
+    input_specs_ = input_specs;
+  }
+
+  const std::vector<OutputSpec>& output_specs() const {
+    return output_specs_;
+  }
+
+  void set_output_specs(const std::vector<OutputSpec>& output_specs) {
+    output_specs_ = output_specs;
+  }
+
+  const MemoryPlan& memory_plan() const {
+    return memory_plan_;
+  }
+
+  void set_memory_plan(const MemoryPlan& memory_plan) {
+    memory_plan_ = memory_plan;
+  }
+
+  const std::vector<SymbolicShapePosition>& sym_shape_positions() const {
+    return sym_shape_positions_;
+  }
+
+  void set_sym_shape_positions(
+      const std::vector<SymbolicShapePosition>& sym_shape_pos) {
+    sym_shape_positions_ = sym_shape_pos;
+  }
+
+ private:
+  void init_execution_state() const;
+
+  c10::QualifiedName name_;
+  std::string nnc_kernel_id_;
+  c10::impl::GenericList parameters_{at::AnyType::get()};
+  std::vector<InputSpec> input_specs_;
+  std::vector<OutputSpec> output_specs_;
+  std::vector<SymbolicShapePosition> sym_shape_positions_;
+  MemoryPlan memory_plan_;
+  mutable std::unique_ptr<ExecutionState> execution_state_;
+};
+
+// CompilationUnit consists of a set of compiled NNC functions. It has a 1-1
+// correspondence with a `Module`.
+// It's similar as torch::jit::mobile::CompilationUnit.
+class TORCH_API CompilationUnit {
+ public:
+  CompilationUnit() = default;
+  CompilationUnit(const CompilationUnit&) = delete;
+  CompilationUnit(CompilationUnit&&) = default;
+  CompilationUnit& operator=(const CompilationUnit&) = delete;
+  CompilationUnit& operator=(CompilationUnit&&) = default;
+
+  // Deserialize from an IValue that is generated by the 'serialize()' method.
+  explicit CompilationUnit(const c10::IValue& value);
+
+  // Serialize all registered functions into an IValue. The IValue will be save
+  // into the compiled TorchScript model file ahead-of-time on the host, and
+  // will be deserialized at runtime on the target device.
+  [[nodiscard]] c10::IValue serialize() const;
+
+  // Execute a registered function.
+  [[nodiscard]] c10::impl::GenericList run(
+      const c10::QualifiedName& function_name,
+      const c10::impl::GenericList& inputs) const;
+
+  // Register a function to the compilation unit.
+  void register_function(std::unique_ptr<Function> fn);
+
+ private:
+  [[nodiscard]] Function* find_function(const c10::QualifiedName& qn) const;
+
+  std::unordered_map<c10::QualifiedName, std::unique_ptr<Function>> functions_;
+};
+
+} // namespace torch::jit::mobile::nnc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/registry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/registry.h
new file mode 100644
index 0000000000000000000000000000000000000000..d6579ce1450ccae6fb8e58a273f944fff53844ee
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/nnc/registry.h
@@ -0,0 +1,45 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/Exception.h>
+#include <c10/util/Registry.h>
+
+namespace torch::jit::mobile::nnc {
+
+using nnc_kernel_function_type = int(void**);
+
+struct TORCH_API NNCKernel {
+  virtual ~NNCKernel() = default;
+  virtual int execute(void** /* args */) = 0;
+};
+
+TORCH_DECLARE_REGISTRY(NNCKernelRegistry, NNCKernel);
+
+#define REGISTER_NNC_KERNEL(id, kernel, ...)     \
+  extern "C" {                                   \
+  nnc_kernel_function_type kernel;               \
+  }                                              \
+  struct NNCKernel_##kernel : public NNCKernel { \
+    int execute(void** args) override {          \
+      return kernel(args);                       \
+    }                                            \
+  };                                             \
+  C10_REGISTER_TYPED_CLASS(NNCKernelRegistry, id, NNCKernel_##kernel);
+
+namespace registry {
+
+inline bool has_nnc_kernel(const std::string& id) {
+  return NNCKernelRegistry()->Has(id);
+}
+
+inline std::unique_ptr<NNCKernel> get_nnc_kernel(const std::string& id) {
+  return NNCKernelRegistry()->Create(id);
+}
+
+} // namespace registry
+
+} // namespace torch::jit::mobile::nnc
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/observer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/observer.h
new file mode 100644
index 0000000000000000000000000000000000000000..205d8095dbc9e3e078d8de4d48c4dd8b45df8b73
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/observer.h
@@ -0,0 +1,117 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/ThreadLocalDebugInfo.h>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace torch {
+
+class MobileDebugInfo : public c10::DebugInfoBase {
+ public:
+  const std::string& getModelName() {
+    return model_name_;
+  }
+
+  void setModelName(const std::string& model_name) {
+    model_name_ = model_name;
+  }
+
+  const std::string& getMethodName() {
+    return method_name_;
+  }
+
+  void setMethodName(const std::string& method_name) {
+    method_name_ = method_name;
+  }
+
+  size_t getOpIdx() {
+    return op_idx_;
+  }
+
+  void setOpIdx(size_t op_idx) {
+    op_idx_ = op_idx;
+  }
+
+ private:
+  std::string model_name_;
+  std::string method_name_;
+  // TODO: Kimish
+  // If we launch a thread such as for at::launch, interepter continuation
+  // and if the caching allocator is enabled in the base thread
+  // then, in order to propagate this information, that is caching allocator
+  // is enabled, across thread boundaries we can use the mechanism provided
+  // by ThreadLocalDebugInfo
+  // Once the thread local MobileDebugInfo is accessible in the launched
+  // thread, it can be accessed in that thread and that thread can set
+  // its own thread local CachingAllocatorInfo.
+  // However, we cannot expect every launched thread to extract and set
+  // its own thread local copy of CachingAllocatorInfo.
+  // But this can be done in lite interpreter, where in the run method
+  // it can do info =
+  // c10::ThreadLocalDebugInfo::get(c10::DebugInfoKind::MOBILE_RUNTIME_INFO))
+  // .get_caching_allocator_info();
+  // GetThreadLocalCachingAllocatorInfo() = info;
+  // Other option is to have MobileDebugInfo itself be the place where thread
+  // local copy of CachingAllocatorInfo is stored. Then
+  // DefaultMobileCPUAllocator inspects this to decide if to use
+  // CachingAllocator. However, current lite interpreter does not support FORK,
+  // thus from the run method of lite interpreter we are not really gonna launch
+  // another instance of lite interpreter in a different thread. So for now not
+  // getting bothered about passing CachingAllocatorInfo across thread
+  // boundaries. c10::CachingAllocatorInfo caching_allocator_info;
+  size_t op_idx_ = 0;
+};
+
+class MobileModuleObserver {
+ public:
+  virtual ~MobileModuleObserver() = default;
+
+  virtual void onEnterRunMethod(const int32_t /*unused*/) {}
+  virtual void onExitRunMethod(
+      const std::unordered_map<std::string, std::string>& /*unused*/,
+      const std::string& /*unused*/,
+      const int32_t /*unused*/) {}
+  virtual void onFailRunMethod(
+      const std::unordered_map<std::string, std::string>& /*unused*/,
+      const std::string& /*unused*/,
+      const int32_t /*unused*/,
+      const char* /*unused*/) {}
+  virtual void onEnterLoadModel(const int32_t /*unused*/) {}
+  virtual void onExitLoadModel(
+      const int32_t /*unused*/,
+      const std::unordered_map<std::string, std::string>& /*unused*/) {
+  } // key: filename, value: file content
+  virtual void onFailLoadModel(
+      const int32_t /*unused*/,
+      const char* /*unused*/) {}
+  virtual void onFailLoadModel(
+      const int32_t /*unused*/,
+      const char* /*unused*/,
+      const std::unordered_map<std::string, std::string>& /*unused*/) {}
+  virtual std::vector<std::string> getDefaultExtraFiles() = 0;
+  virtual std::unordered_map<std::string, std::string> processMetadataFromExtra(
+      const std::unordered_map<std::string, std::string>&) = 0;
+};
+
+class MobileObserverConfig {
+ public:
+  void setModuleObserver(std::unique_ptr<MobileModuleObserver> reporter) {
+    module_observer_ = std::move(reporter);
+  }
+  MobileModuleObserver* getModuleObserver() {
+    return module_observer_.get();
+  }
+
+ private:
+  std::unique_ptr<MobileModuleObserver> module_observer_;
+};
+
+MobileObserverConfig& observerConfig();
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/parse_bytecode.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/parse_bytecode.h
new file mode 100644
index 0000000000000000000000000000000000000000..00ada45c7a8011bec75f73a57a888e72365b6736
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/parse_bytecode.h
@@ -0,0 +1,26 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/mobile/function.h>
+
+namespace torch::jit::mobile {
+using c10::IValue;
+TORCH_API void parseInstructions(
+    const std::string& function_name,
+    c10::ivalue::TupleElements&& ins_list,
+    c10::ivalue::TupleElements& debug_handles_m_tuple,
+    mobile::Function* function);
+TORCH_API void parseConstants(
+    const c10::ivalue::TupleElements& consts_list,
+    mobile::Function* function);
+TORCH_API void parseTypes(
+    const c10::ivalue::TupleElements& types_list,
+    mobile::Function* function);
+TORCH_API void parseRegisterSize(size_t rsize, mobile::Function* function);
+TORCH_API void applyUpgrader(
+    mobile::Function* function,
+    uint64_t operator_version);
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/parse_operators.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/parse_operators.h
new file mode 100644
index 0000000000000000000000000000000000000000..92c24de81af3da6b3df59a43b581fbbf3176efe2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/parse_operators.h
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/mobile/function.h>
+
+namespace torch::jit {
+using c10::IValue;
+
+enum MobileModuleLoadOptions {
+  OPERATOR_CHECK = 1,
+  // PARSE_ALL_EXTRA_FILE_MAPS is used to gate for ExtraFileMaps to pull all
+  // files automatically without explicit entries mapping. Refer to PR for a
+  // detail: https://github.com/pytorch/pytorch/pull/99747
+  PARSE_ALL_EXTRA_FILE_MAPS = 2,
+};
+
+const uint64_t kDefaultMobileLoadOptions =
+    MobileModuleLoadOptions::OPERATOR_CHECK;
+
+namespace mobile {
+
+TORCH_API void parseOperators(
+    c10::ivalue::TupleElements&& ops_list,
+    const uint64_t& module_load_options,
+    mobile::Function* function);
+} // namespace mobile
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/prim_ops_registery.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/prim_ops_registery.h
new file mode 100644
index 0000000000000000000000000000000000000000..f6f297f966cad2201f0e85cc15a869e75eb2f618
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/prim_ops_registery.h
@@ -0,0 +1,33 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <functional>
+#include <vector>
+
+namespace torch::jit::mobile {
+
+using Stack = std::vector<c10::IValue>;
+
+void registerPrimOpsFunction(
+    const std::string& name,
+    const std::function<void(Stack&)>& fn);
+
+bool hasPrimOpsFn(const std::string& name);
+
+std::function<void(Stack&)>& getPrimOpsFn(const std::string& name);
+
+class prim_op_fn_register {
+ public:
+  prim_op_fn_register(
+      const std::string& name,
+      const std::function<void(Stack&)>& fn) {
+    registerPrimOpsFunction(name, fn);
+  }
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/profiler_edge.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/profiler_edge.h
new file mode 100644
index 0000000000000000000000000000000000000000..f1aaa2288fc6bbddca0ff7869c2d228322d11208
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/profiler_edge.h
@@ -0,0 +1,120 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/autograd/profiler_kineto.h>
+#include <torch/csrc/jit/mobile/module.h>
+
+namespace torch::jit::mobile {
+
+// If we dont have kineto available then edge profiler does not
+// work since it relies on Kineto
+#ifdef USE_KINETO
+class TORCH_API KinetoEdgeCPUProfiler {
+ public:
+  // This profiler only profiles KINETO events
+  // No GPU_FALLBACK or NVTX
+  /*
+   * @param m is the instance of mobile Module which is being profiled.
+   *        Note that this implies that KinetoEdgeCPUProfiler can be used
+   *        to profile specific Module (see usage below), unliked ProfilerKineto
+   *        which can profile pytorch runtime in arbitrary scope.
+   * @param fname is the name of the file to which chrome trace is written.
+   * @param report_input_shapes: whether to record shapes of op's inputs.
+   * @param with_stack: whether to record model's python stacktrace for the op.
+   * @param with_flops: whether to report flops corresponding to the op.
+   * @param with_modules: whether to report original python module
+   *        hierarchy to which the op belongs.
+   * @param events
+   * @param adjust_vulkan_timestamps: whether to adjust vulkan timestamps from
+   *        query pool to align with cpu event times
+   *
+   * Usage pattern for this profiler must be as follows:
+   *
+   * {
+   *   KinetoEdgeCPUProfiler(m, filename, args);
+   *   m.forward(...);
+   * }
+   *
+   * The reason being that KinetoEdgeCPUProfiler has a dependency on Module
+   * and thus it must not outlive it.
+   *
+   * Thus, when KinetoEdgeCPUProfiler is used as RAII to do profiling
+   * within certain scope. In that scope, the captured reference to
+   * Module will outlive KinetoEdgeCPUProfiler. This is guaranteed because
+   * KinetoEdgeCPUProfiler must be constructed later than Module, on stack.
+   *
+   * An example of the anti-pattern and wrong usage is:
+   *
+   * std::shared_ptr<KinetoMobileCPUProfiler> profiler(m, filename, args);
+   * m.forward(...);
+   *
+   * Since KinetoEdgeCPUProfiler object would then be constructed on heap
+   * with its lifetime managed manually or via smart pointers.
+   */
+  KinetoEdgeCPUProfiler(
+      const torch::jit::mobile::Module& m,
+      const std::string& fname,
+      const bool report_input_shapes = false,
+      const bool profile_memory = false,
+      const bool with_stack = false,
+      const bool with_flops = false,
+      const bool with_modules = false,
+      std::vector<std::string> events = {},
+      const bool adjust_vulkan_timestamps = false);
+
+  const std::unique_ptr<torch::autograd::profiler::ProfilerResult>&
+  disableProfiler();
+  const std::unique_ptr<torch::autograd::profiler::ProfilerResult>&
+  getProfilerResult();
+  void recordBackendEvent(
+      const int64_t start_time_us,
+      const int64_t end_time_us,
+      const int64_t debug_handle,
+      const std::string& event_name,
+      const std::string& backend_name);
+  void recordBackendMemoryEvent(
+      void* ptr,
+      int64_t alloc_size,
+      size_t total_allocated,
+      size_t total_reserved,
+      c10::Device device);
+
+  ~KinetoEdgeCPUProfiler();
+
+ private:
+  /*
+   * We store a reference to Module to make such dependency explicit, since
+   * a Module reference is already stored in a functor.
+   */
+  const mobile::Module& m_;
+  std::string trace_file_name_;
+  std::unique_ptr<torch::autograd::profiler::ProfilerResult> profiler_result_;
+};
+
+TORCH_API KinetoEdgeCPUProfiler* getCurrentEdgeProfiler();
+
+#define RECORD_BACKEND_EVENT_TO_EDGE_PROFILER(                               \
+    start_time_us, end_time_us, debug_handle, event_name, backend_name)      \
+  if (mobile::getCurrentEdgeProfiler()) {                                    \
+    mobile::getCurrentEdgeProfiler()->recordBackendEvent(                    \
+        start_time_us, end_time_us, debug_handle, event_name, backend_name); \
+  }
+
+#define RECORD_BACKEND_MEMORY_EVENT_TO_EDGE_PROFILER(              \
+    ptr, alloc_size, total_allocated, total_reserved, device)      \
+  if (mobile::getCurrentEdgeProfiler()) {                          \
+    mobile::getCurrentEdgeProfiler()->recordBackendMemoryEvent(    \
+        ptr, alloc_size, total_allocated, total_reserved, device); \
+  }
+#else
+
+#define RECORD_BACKEND_EVENT_TO_EDGE_PROFILER( \
+    start_time_us, end_time_us, debug_handle, event_name, backend_name)
+
+#define RECORD_BACKEND_MEMORY_EVENT_TO_EDGE_PROFILER( \
+    ptr, alloc_size, total_allocated, total_reserved, device)
+#endif
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/promoted_prim_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/promoted_prim_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..bf507cc088421b2317f2a9381cdbb54b275bb65d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/promoted_prim_ops.h
@@ -0,0 +1,66 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/mobile/prim_ops_registery.h>
+#include <torch/csrc/jit/mobile/register_ops_common_utils.h>
+
+namespace torch::jit {
+
+void tupleIndex(Stack& stack);
+
+void raiseException(Stack& stack);
+
+void is(Stack& stack);
+
+void unInitialized(Stack& stack);
+
+void isNot(Stack& stack);
+
+void aten_format(Stack& stack);
+
+void size(Stack& stack);
+
+void sym_size(Stack& stack);
+
+void sym_size_int(Stack& stack);
+
+void sym_stride_int(Stack& stack);
+
+void sym_numel(Stack& stack);
+
+void sym_storage_offset(Stack& stack);
+
+void sym_stride(Stack& stack);
+
+void device(Stack& stack);
+
+void device_with_index(Stack& stack);
+
+void dtype(Stack& stack);
+
+void layout(Stack& stack);
+
+void toPrimDType(Stack& stack);
+
+void dim(Stack& stack);
+
+void _not(Stack& stack);
+
+void boolTensor(Stack& stack);
+
+void toList(Stack& stack);
+
+void numToTensorScalar(Stack& stack);
+
+void isCuda(Stack& stack);
+
+void numToTensorBool(Stack& stack);
+
+void dictIndex(Stack& stack);
+
+void raiseExceptionWithMessage(Stack& stack);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/quantization.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/quantization.h
new file mode 100644
index 0000000000000000000000000000000000000000..fb0157c680ab760cdba5bfcacb825233721d5b63
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/quantization.h
@@ -0,0 +1,39 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Export.h>
+#include <string>
+
+namespace torch::jit::mobile {
+class Module;
+namespace quantization {
+/*
+ * Device side PTQ API.
+ * Once the model has been prepared for quantization on server side, such model
+ * is sent to device. On device side the model is further trained. At the end of
+ * the training, before the model is readied for inference, we need to quantize
+ * the model.
+ * Usage of this API is as follows.
+ * PTQQuanizationHelper ptq_helper;
+ * ptq_helper.quantize_dynamic(m, "forward");
+ * Args:
+ * m: Captured by reference, an instance of mobile::Module. This module will be
+ * mutated in place to replace its <method_name> method with quantized
+ * equivalent. method:name: Name of the method to be quantized. AOT preparation
+ * for quantization must also have been done for this method. Returns: In place
+ * mutated `m` whose size should be smaller due to weight quantization and whose
+ * <method_name> method should use quantized ops
+ */
+class TORCH_API PTQQuanizationHelper {
+ public:
+  PTQQuanizationHelper() = default;
+  void quantize_dynamic(
+      torch::jit::mobile::Module& m,
+      const std::string& method_name);
+};
+} // namespace quantization
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/register_ops_common_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/register_ops_common_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..d6cce1434bb52f9656373fd6f31702928cd06a48
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/register_ops_common_utils.h
@@ -0,0 +1,58 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Context.h>
+#include <ATen/NativeFunctions.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/jit/runtime/jit_exception.h>
+#include <torch/csrc/jit/runtime/vararg_functions.h>
+
+namespace torch::jit {
+
+inline void noop(Stack& n) {}
+
+int64_t normalizeIndex(int64_t idx, int64_t list_size);
+
+// reference function THPVariable_to in python_variable_methods.cpp
+[[maybe_unused]] static at::Tensor to_dispatch(
+    at::Tensor self,
+    std::optional<at::Device> device,
+    std::optional<at::ScalarType> scalarType,
+    bool non_blocking,
+    bool copy) {
+  if (device && device->is_cuda()) {
+    at::globalContext().lazyInitDevice(c10::DeviceType::CUDA);
+  }
+  if (!device && !scalarType && !copy) {
+    return self;
+  } else if (!device) {
+    return self.to(*scalarType, non_blocking, copy);
+  } else if (!scalarType) {
+    return self.to(*device, non_blocking, copy);
+  } else {
+    return self.to(*device, *scalarType, non_blocking, copy);
+  }
+}
+
+// Convert the tensor pointed to by \p data to a nested list. \p dim is the
+// number of dimensions in the tensor and \p cur_dim is the dimension being
+// processed by the current invocation. \p ty is the expected output IR type of
+// the operation. \p is the scalar type of \p data. \p sizes and \p strides are
+// the sizes and strides of the tensor operand and \p element_size is the size
+// in bytes of one tensor element.
+IValue tensorToListRecursive(
+    char* data,
+    int64_t cur_dim,
+    int64_t num_tensor_dims,
+    at::TypePtr ty,
+    at::ScalarType scalar_ty,
+    at::IntArrayRef sizes,
+    at::IntArrayRef strides,
+    size_t element_size);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/export_data.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/export_data.h
new file mode 100644
index 0000000000000000000000000000000000000000..bf8613ed2be04eed0ce6f5c45306a9c5f30667f4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/export_data.h
@@ -0,0 +1,56 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/mobile/module.h>
+
+namespace torch::jit {
+
+/**
+ * Serializes the provided tensor map to the provided stream.
+ *
+ * @param[in] map The tensors to serialize.
+ * @param[in] out The stream to write the serialized data to.
+ * @param[in] use_flatbuffer If true, use Flatbuffers to serialize the data.
+ *     If false, use Pickle.
+ */
+TORCH_API void _save_parameters(
+    const std::map<std::string, at::Tensor>& map,
+    std::ostream& out,
+    bool use_flatbuffer = false);
+
+/**
+ * Serializes the provided tensor map to a file.
+ *
+ * @param[in] map The tensors to serialize.
+ * @param[in] filename The stem of the file name to write to. If
+ *     @p use_flatbuffer is false, the extension ".pkl" will be appended. If
+ *     @p use_flatbuffer is true, the extension ".ff" will be appended.
+ * @param[in] use_flatbuffer If true, use Flatbuffers to serialize the data.
+ *     If false, use Pickle.
+ */
+TORCH_API void _save_parameters(
+    const std::map<std::string, at::Tensor>& map,
+    const std::string& filename,
+    bool use_flatbuffer = false);
+
+namespace mobile {
+
+// NOTE: Please prefer using _save_parameters directly over using the 2
+// functions below.
+TORCH_API mobile::Module tensor_dict_to_mobile(
+    const c10::Dict<std::string, at::Tensor>& dict);
+
+c10::Dict<std::string, at::Tensor> tensor_map_to_dict(
+    const std::map<std::string, at::Tensor>& map);
+
+} // namespace mobile
+
+extern void (*_save_mobile_module_to)(
+    const mobile::Module& module,
+    const std::function<size_t(const void*, size_t)>& writer_func);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/optim/sgd.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/optim/sgd.h
new file mode 100644
index 0000000000000000000000000000000000000000..563cafb5e7fc25f911d6b66b065db6d66cdd7a50
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/optim/sgd.h
@@ -0,0 +1,130 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/arg.h>
+#include <torch/types.h>
+
+#include <utility>
+#include <vector>
+
+namespace torch::jit::mobile {
+
+class SGDParamState {
+  TORCH_ARG(torch::Tensor, momentum_buffer);
+
+ public:
+  std::unique_ptr<SGDParamState> clone() const {
+    return std::make_unique<SGDParamState>(
+        static_cast<const SGDParamState&>(*this));
+  }
+  friend bool operator==(const SGDParamState& lhs, const SGDParamState& rhs);
+};
+
+struct TORCH_API SGDOptions {
+  /* implicit */ SGDOptions(double lr);
+  TORCH_ARG(double, lr);
+  TORCH_ARG(double, momentum) = 0;
+  TORCH_ARG(double, dampening) = 0;
+  TORCH_ARG(double, weight_decay) = 0;
+  TORCH_ARG(bool, nesterov) = false;
+
+ public:
+  std::unique_ptr<SGDOptions> clone() const {
+    return std::make_unique<SGDOptions>(static_cast<const SGDOptions&>(*this));
+  }
+  TORCH_API friend bool operator==(
+      const SGDOptions& lhs,
+      const SGDOptions& rhs);
+};
+
+/// Stores parameters in the param_group and stores a pointer to the SGDOptions
+class TORCH_API SGDParamGroup {
+ public:
+  // NOTE: In order to store `SGDParamGroup` in a `std::vector`, it has to be
+  // copy-constructible.
+  SGDParamGroup(const SGDParamGroup& param_group)
+      : params_(param_group.params()),
+        options_(
+            param_group.has_options() ? param_group.options().clone()
+                                      : nullptr) {}
+  SGDParamGroup& operator=(const SGDParamGroup& param_group) {
+    this->params_ = param_group.params();
+    this->options_ =
+        param_group.has_options() ? param_group.options().clone() : nullptr;
+    return *this;
+  }
+  /* implicit */ SGDParamGroup(std::vector<Tensor> params)
+      : params_(std::move(params)) {}
+  SGDParamGroup(std::vector<Tensor> params, std::unique_ptr<SGDOptions> options)
+      : params_(std::move(params)), options_(std::move(options)) {}
+
+  bool has_options() const;
+  SGDOptions& options();
+  const SGDOptions& options() const;
+  void set_options(std::unique_ptr<SGDOptions> options);
+  std::vector<Tensor>& params();
+  const std::vector<Tensor>& params() const;
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::vector<Tensor> params_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::unique_ptr<SGDOptions> options_;
+};
+
+class TORCH_API SGD {
+ public:
+  explicit SGD(
+      const std::vector<torch::jit::mobile::SGDParamGroup>& param_groups,
+      SGDOptions defaults)
+      : defaults_(std::make_unique<SGDOptions>(defaults)) {
+    for (const auto& param_group : param_groups) {
+      add_param_group(param_group);
+    }
+    TORCH_CHECK(defaults.lr() >= 0, "Invalid learning rate: ", defaults.lr());
+    TORCH_CHECK(
+        defaults.momentum() >= 0,
+        "Invalid momentum value: ",
+        defaults.momentum());
+    TORCH_CHECK(
+        defaults.weight_decay() >= 0,
+        "Invalid weight_decay value: ",
+        defaults.weight_decay());
+    TORCH_CHECK(
+        !defaults.nesterov() ||
+            (defaults.momentum() > 0 && defaults.dampening() == 0),
+        "Nesterov momentum requires a momentum and zero dampening");
+  }
+
+  explicit SGD(std::vector<Tensor> params, SGDOptions defaults)
+      : SGD({SGDParamGroup(std::move(params))}, defaults) {}
+
+  /// Adds the given param_group to the optimizer's param_group list.
+  void add_param_group(const SGDParamGroup& param_group);
+
+  ~SGD() = default;
+
+  using LossClosure = std::function<Tensor()>;
+  /// A loss function closure, which is expected to return the loss value.
+  torch::Tensor step(const LossClosure& closure = nullptr);
+
+  /// Zeros out the gradients of all parameters.
+  void zero_grad();
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::vector<SGDParamGroup> param_groups_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  ska::flat_hash_map<void*, std::unique_ptr<SGDParamState>> state_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::unique_ptr<SGDOptions> defaults_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::vector<Tensor> params_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::unique_ptr<SGDOptions> options_;
+};
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/random.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/random.h
new file mode 100644
index 0000000000000000000000000000000000000000..d6f2389163c4ddc89fd7be4e60fffe84a0936176
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/random.h
@@ -0,0 +1,55 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/data/samplers/base.h>
+#include <torch/types.h>
+
+#include <cstddef>
+#include <vector>
+
+namespace torch::serialize {
+class OutputArchive;
+class InputArchive;
+} // namespace torch::serialize
+
+namespace torch::jit::mobile {
+
+/// A lighter `Sampler` that returns indices randomly and cannot be
+/// serialized.
+class TORCH_API RandomSampler : public torch::data::samplers::Sampler<> {
+ public:
+  /// Constructs a `RandomSampler` with a size and dtype for the stored indices.
+  ///
+  /// The constructor will eagerly allocate all required indices, which is the
+  /// sequence `0 ... size - 1`. `index_dtype` is the data type of the stored
+  /// indices. You can change it to influence memory usage.
+  explicit RandomSampler(int64_t size, Dtype index_dtype = torch::kInt64);
+
+  ~RandomSampler() override;
+
+  /// Resets the `RandomSampler` to a new set of indices.
+  void reset(std::optional<size_t> new_size = std::nullopt) override;
+
+  /// Returns the next batch of indices.
+  std::optional<std::vector<size_t>> next(size_t batch_size) override;
+
+  /// Serializes the `RandomSampler` to the `archive`.
+  void save(serialize::OutputArchive& archive) const override;
+
+  /// Deserializes the `RandomSampler` from the `archive`.
+  void load(serialize::InputArchive& archive) override;
+
+  /// Returns the current index of the `RandomSampler`.
+  size_t index() const noexcept;
+
+ private:
+  at::Tensor indices_;
+  int64_t index_ = 0;
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/sequential.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/sequential.h
new file mode 100644
index 0000000000000000000000000000000000000000..88d9b6e13b16da6610218b8c66e193f51b14291f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/train/sequential.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/data/samplers/base.h>
+#include <torch/types.h>
+
+#include <cstddef>
+#include <vector>
+
+namespace torch::serialize {
+class OutputArchive;
+class InputArchive;
+} // namespace torch::serialize
+
+namespace torch::jit::mobile {
+
+/// A lighter `Sampler` that returns indices sequentially and cannot be
+/// serialized.
+class TORCH_API SequentialSampler : public torch::data::samplers::Sampler<> {
+ public:
+  /// Creates a `SequentialSampler` that will return indices in the range
+  /// `0...size - 1`.
+  explicit SequentialSampler(size_t size);
+
+  /// Resets the `SequentialSampler` to zero.
+  void reset(std::optional<size_t> new_size = std::nullopt) override;
+
+  /// Returns the next batch of indices.
+  std::optional<std::vector<size_t>> next(size_t batch_size) override;
+
+  /// Not supported for mobile SequentialSampler
+  void save(serialize::OutputArchive& archive) const override;
+
+  /// Not supported for mobile SequentialSampler
+  void load(serialize::InputArchive& archive) override;
+
+  /// Returns the current index of the `SequentialSampler`.
+  size_t index() const noexcept;
+
+ private:
+  size_t size_;
+  size_t index_{0};
+};
+
+} // namespace torch::jit::mobile
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/type_parser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/type_parser.h
new file mode 100644
index 0000000000000000000000000000000000000000..42f3ce8374e378d77092f2f3f8102b73ede5ec78
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/type_parser.h
@@ -0,0 +1,59 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/dynamic_type.h>
+#include <ATen/core/jit_type.h>
+#include <unordered_set>
+
+namespace c10 {
+
+class TORCH_API TypeParser {
+ public:
+  explicit TypeParser(std::string pythonStr);
+  explicit TypeParser(std::vector<std::string>& pythonStrs);
+
+  TypePtr parse();
+  std::vector<TypePtr> parseList();
+  static const std::unordered_set<std::string>& getNonSimpleType();
+  static const std::unordered_set<std::string>& getCustomType();
+  std::unordered_set<std::string> getContainedTypes();
+
+ private:
+  TypePtr parseNamedTuple(const std::string& qualified_name);
+  TypePtr parseCustomType();
+  TypePtr parseTorchbindClassType();
+  TypePtr parseNonSimple(const std::string& token);
+
+  void expect(const char* s);
+  void expectChar(char c);
+  template <typename T>
+  TypePtr parseSingleElementType();
+
+  void lex();
+
+  std::string next();
+  std::string_view nextView();
+  void advance();
+  [[nodiscard]] std::string_view cur() const;
+
+  std::string pythonStr_;
+  size_t start_;
+  std::string_view next_token_;
+
+  // Used for parsing string list
+  std::vector<std::string> pythonStrs_;
+  std::unordered_map<std::string, c10::TypePtr> str_type_ptr_map_;
+
+  // Store all contained types when parsing a string
+  std::unordered_set<std::string> contained_types_;
+};
+
+TORCH_API TypePtr parseType(const std::string& pythonStr);
+
+TORCH_API std::vector<TypePtr> parseType(std::vector<std::string>& pythonStr);
+
+} // namespace c10
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/upgrader_mobile.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/upgrader_mobile.h
new file mode 100644
index 0000000000000000000000000000000000000000..d3ba135e4634af9c4ff33e7bfe72002e7c12166b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/mobile/upgrader_mobile.h
@@ -0,0 +1,44 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue_inl.h>
+
+#include <torch/csrc/jit/mobile/code.h>
+#include <torch/csrc/jit/mobile/function.h>
+#include <torch/csrc/jit/serialization/import_export_functions.h>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace torch::jit {
+struct Instruction;
+struct Upgrader {
+  int min_version;
+  int max_version;
+  std::string upgrader_name;
+  int index;
+};
+
+// From operator_versions.yaml
+TORCH_API const std::unordered_map<std::string, std::vector<Upgrader>>
+getOperatorVersionMapForMobile();
+
+struct OperatorString {
+  const std::string name;
+  const std::string overload_name;
+  const std::optional<int> num_specified_args;
+};
+
+struct ByteCodeFunctionWithOperator {
+  mobile::Function& function;
+  std::vector<OperatorString> operators;
+};
+
+TORCH_API const std::vector<ByteCodeFunctionWithOperator>&
+getUpgraderBytecodeList();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/upgraders.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/upgraders.h
new file mode 100644
index 0000000000000000000000000000000000000000..91d14a9bc272762a69b8011a2d052f30186d3e2e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/upgraders.h
@@ -0,0 +1,51 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/macros/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+
+namespace torch::jit {
+
+class UpgradersMap {
+ public:
+  void set_content(
+      std::unordered_map<std::string, std::shared_ptr<Graph>>&& content);
+  int count();
+  const std::unordered_map<std::string, std::shared_ptr<Graph>>& get_content();
+  bool is_populated();
+  // THESE METHODS ARE ONLY USED FOR TESTING PURPOSES
+  void test_only_set_content(
+      const std::unordered_map<std::string, std::string>& content);
+  void test_only_remove_content(
+      const std::unordered_map<std::string, std::string>& content);
+
+ private:
+  std::unordered_map<std::string, std::shared_ptr<Graph>> content_;
+  std::mutex lock;
+  bool isPopulated = false;
+};
+
+TORCH_API void populate_upgraders_map(
+    std::unordered_map<std::string, std::shared_ptr<Graph>>&& content);
+
+TORCH_API int get_upgraders_map_size();
+
+TORCH_API bool is_upgraders_map_populated();
+
+TORCH_API const std::unordered_map<std::string, std::shared_ptr<Graph>>&
+dump_upgraders_map();
+
+// THESE TWO METHODS BELOW ARE ONLY USED FOR TESTING
+TORCH_API void test_only_populate_upgraders(
+    const std::unordered_map<std::string, std::string>& content);
+
+TORCH_API void test_only_remove_upgraders(
+    const std::unordered_map<std::string, std::string>& content);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/upgraders_entry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/upgraders_entry.h
new file mode 100644
index 0000000000000000000000000000000000000000..840b357f2ba31ad0dff6d0bee1a6e89429f06786
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/upgraders_entry.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/macros/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <string>
+#include <unordered_map>
+
+namespace torch::jit {
+
+TORCH_API void populate_upgraders_graph_map();
+
+TORCH_API std::unordered_map<std::string, std::shared_ptr<Graph>>
+generate_upgraders_graph();
+
+TORCH_API std::unordered_map<std::string, std::string> get_upgraders_entry_map();
+
+std::shared_ptr<Graph> create_upgrader_graph(
+    const std::string& upgrader_name,
+    const std::string& upgrader_body);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..53a588552db965a1d3e9e660ed42131088a9e349
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/utils.h
@@ -0,0 +1,56 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/macros/Export.h>
+#include <torch/csrc/jit/operator_upgraders/version_map.h>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace torch::jit {
+
+struct UpgraderRange {
+  int min_version;
+  int max_version;
+};
+
+// Given a list of upgrader entries for a single operator
+// and the model version for that operator, find a valid
+// upgrader.
+TORCH_API std::optional<UpgraderEntry> findUpgrader(
+    const std::vector<UpgraderEntry>& upgraders_for_schema,
+    size_t current_version);
+
+// Utility methods to find if the operator is up-to-date
+// based on all registered upgraders for this operator.
+// This can be different from the current server version
+// because the implementation of this operator could have
+// been consistent for many later version bumps.
+TORCH_API bool isOpCurrentBasedOnUpgraderEntries(
+    const std::vector<UpgraderEntry>& upgraders_for_schema,
+    size_t current_version);
+
+TORCH_API bool isOpSymbolCurrent(
+    const std::string& name,
+    size_t current_version);
+
+// Returns the possible old schemas for the operator that
+// doesn't exist anymore. This can be true for deprecated
+// operators. Since name is always a symbol name, there
+// can be multiple schemas for different overloads.
+TORCH_API std::vector<std::string> loadPossibleHistoricOps(
+    const std::string& name,
+    std::optional<size_t> version);
+
+TORCH_API uint64_t getMaxOperatorVersion();
+
+// Returns the list of min and max version numbers of the operators
+// that an upgrader `x` support for all upgraders for op `foo`
+TORCH_API std::vector<UpgraderRange> getUpgradersRangeForOp(
+    const std::string& name);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/version_map.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/version_map.h
new file mode 100644
index 0000000000000000000000000000000000000000..f3db22b1003ec561d576c34996655d0b775f2e1c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/operator_upgraders/version_map.h
@@ -0,0 +1,38 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/macros/Export.h>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace torch::jit {
+
+struct UpgraderEntry {
+  int bumped_at_version;
+  std::string upgrader_name;
+  std::string old_schema;
+};
+
+// Toggle the behaviour of calculating version for the module.
+// If this is true, we calculate solely based on upgraders
+// If this is false, we calculate it based on historic per op version map
+TORCH_API void calculate_package_version_based_on_upgraders(bool val);
+
+TORCH_API bool get_version_calculator_flag();
+
+TORCH_API const std::unordered_map<std::string, std::vector<UpgraderEntry>>&
+get_operator_version_map();
+
+TORCH_API void test_only_add_entry(
+    const std::string& op_name,
+    UpgraderEntry entry);
+
+TORCH_API void test_only_remove_entry(const std::string& op_name);
+
+TORCH_API void test_only_reset_flag();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/add_if_then_else.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/add_if_then_else.h
new file mode 100644
index 0000000000000000000000000000000000000000..db13ccdcac467d59bf24e34166099a4dca804a7c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/add_if_then_else.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API bool AddIfThenElseOp(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/annotate_warns.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/annotate_warns.h
new file mode 100644
index 0000000000000000000000000000000000000000..b535fbda72ca74110ef2c39d7557f2a51d96a7a9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/annotate_warns.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void AnnotateWarns(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/autocast.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/autocast.h
new file mode 100644
index 0000000000000000000000000000000000000000..26076b23ead34ab9e6ca1715e06f674d129dda0f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/autocast.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void Autocast(const std::shared_ptr<Graph>& graph);
+
+TORCH_API bool setAutocastMode(bool value);
+TORCH_API bool autocastEnabled();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/bailout_graph.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/bailout_graph.h
new file mode 100644
index 0000000000000000000000000000000000000000..f93bf5f0265a27a1e1f8c6b3cb0a14fc7db06594
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/bailout_graph.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <list>
+#include <vector>
+
+namespace torch::jit {
+
+// Replaces prim::Guard nodes with prim::BailOut nodes and
+// computes sets of inputs needed to resume execution at
+// bailout points
+TORCH_API void InsertBailOuts(std::shared_ptr<Graph> graph);
+
+// Builds a bailout graph into `target` (which is an empty graph)
+// for a given bailout point `bailout_index`
+// from the original graph `orig` (the original unoptimized graph)
+// BailOut graphs allow Interpreter to resume
+// execution of the (un/de)optimized graph (i.e.
+// a graph that doesn't rely on any assumptions derived from
+// on profiling information) from a given BailOut point
+// should any of the assumptions fail for an actual input.
+TORCH_API std::shared_ptr<Graph> BuildBailOutGraphFrom(
+    int64_t bailout_index,
+    const std::shared_ptr<Graph>& orig,
+    const std::shared_ptr<Graph>& target);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/batch_mm.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/batch_mm.h
new file mode 100644
index 0000000000000000000000000000000000000000..d69d2d0735619683efdfe554841f24ea9607f27d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/batch_mm.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void BatchMM(std::shared_ptr<Graph>& graph);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/canonicalize.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/canonicalize.h
new file mode 100644
index 0000000000000000000000000000000000000000..3dca00139a05a7f1f5bebba6d8097ac10206bed7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/canonicalize.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API std::shared_ptr<Graph> Canonicalize(
+    const std::shared_ptr<Graph>& graph,
+    bool keep_unique_names = true);
+
+TORCH_API void CanonicalizeOutputs(std::shared_ptr<Graph>& graph);
+
+TORCH_API std::optional<const Use> firstOrLastUse(Value* v, bool find_first);
+
+TORCH_API bool isBeforeOrAfter(
+    const Use& a,
+    const Use& b,
+    bool checking_before);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/canonicalize_graph_fuser_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/canonicalize_graph_fuser_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..c9bfa164240df394dec717fa0a8d74c59ef551c8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/canonicalize_graph_fuser_ops.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void CanonicalizeOps(const std::shared_ptr<Graph>& graph);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/check_strict_fusion.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/check_strict_fusion.h
new file mode 100644
index 0000000000000000000000000000000000000000..e815d2db7b4d00f99f1a5c02e0f7377988f2df61
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/check_strict_fusion.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void CheckStrictFusion(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/clear_profiling.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/clear_profiling.h
new file mode 100644
index 0000000000000000000000000000000000000000..7df9c419fbcbf31340d5b3d4ad137f89b0ba0216
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/clear_profiling.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void unprofileGraphInputs(const std::shared_ptr<Graph>& graph);
+TORCH_API void unprofileBlock(Block* start_block);
+// Unprofiles all the node outputs in a block.
+
+TORCH_API void ClearProfilingInformation(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/clear_undefinedness.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/clear_undefinedness.h
new file mode 100644
index 0000000000000000000000000000000000000000..134a50a2c6b521735b7dab5e1225394dda1dfee0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/clear_undefinedness.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Undefinedness makes argument matching fail for regular tensor operations
+// if 1+ arguments are undefined or possibly undefined tensors.
+// Technically, undefined tensors are **not** tensors as the regular tensor
+// operations do not know how to handle them.
+// However, in practice, there are guards and conversion operators that
+// **always** gate regular operations if undefined tensors may be present
+// Eventually, we would love to move to the world where we use optionals
+// in lieu of undefined tensors.
+// When this happens, this pass will be removed
+TORCH_API void ClearUndefinedness(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/common_subexpression_elimination.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/common_subexpression_elimination.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ba9168870d9bb1db49aaee14ba5b94e90f5e7a8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/common_subexpression_elimination.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API bool EliminateCommonSubexpression(
+    const std::shared_ptr<Graph>& graph);
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/concat_opt.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/concat_opt.h
new file mode 100644
index 0000000000000000000000000000000000000000..2bb17068b66a9ea42362a7779808341319eca31e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/concat_opt.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Eliminates common inputs among `aten::cat` ops.
+TORCH_API bool EliminateConcatCommonInputs(const std::shared_ptr<Graph>& graph);
+
+// Expands `aten::cat` ops into `aten::copy` ops and eliminates redudancies
+// in the buffers used for concatenation if possible.
+TORCH_API void ExpandConcatAndEliminateRedundancy(
+    const std::shared_ptr<Graph>& graph);
+
+TORCH_API bool CombineConcats(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/constant_pooling.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/constant_pooling.h
new file mode 100644
index 0000000000000000000000000000000000000000..7fcaf321f353f1a89d8e2ec7170c55463a58871c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/constant_pooling.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void ConstantPooling(const std::shared_ptr<Graph>& graph);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/constant_propagation.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/constant_propagation.h
new file mode 100644
index 0000000000000000000000000000000000000000..bd25c7c3851924b85195d5c4fca6555acd38ad2f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/constant_propagation.h
@@ -0,0 +1,35 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Runs constant propagation on all objects unless ignore_custom_classes is
+// specified as true, in which case user defined classes are skipped.  This is
+// useful to prevent early fusion of packing operations, which end up lowering
+// away information about their constructors (e.g. packed::linear_clamp_prepack
+// and prepacked::conv2d_clamp_prepack)
+// Returns True if the pass made a change to the graph
+TORCH_API bool ConstantPropagation(
+    std::shared_ptr<Graph>& graph,
+    bool ignore_custom_classes = false);
+
+// runs constant propagation only on ops that have non-aliasing inputs & outputs
+// Returns True if the pass made a change to the graph
+TORCH_API bool ConstantPropagationImmutableTypes(std::shared_ptr<Graph>& graph);
+
+// Runs the node if its inputs are constants. Callers of this function must
+// make their own determination if constant prop is appropriate - for example
+// non-deterministic ops or ops with side effects.  If ignore_custom_classes is
+// specified, nodes that output user defined classes are not run.
+TORCH_API std::optional<Stack> runNodeIfInputsAreConstant(
+    const Node* node,
+    bool ignore_custom_classes = false,
+    AliasDb* db = nullptr);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/create_autodiff_subgraphs.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/create_autodiff_subgraphs.h
new file mode 100644
index 0000000000000000000000000000000000000000..fb4824c8fb8c190e781fb035ef32c023926aefc4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/create_autodiff_subgraphs.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <cstddef>
+
+namespace torch::jit {
+
+// insert GraphExecutor nodes that group together
+// subgraphs that are differentiable by the jit's autodiff passes
+// threshold - minimum number of nodes that will appear in a block
+// returns all differentiable blocks that have been found
+TORCH_API std::vector<Node*> CreateAutodiffSubgraphs(
+    const std::shared_ptr<Graph>& graph,
+    size_t threshold = 2);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/create_functional_graphs.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/create_functional_graphs.h
new file mode 100644
index 0000000000000000000000000000000000000000..ad3c2e3c9778ea4930c131a0739532a6720bea78
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/create_functional_graphs.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void CreateFunctionalGraphs(const std::shared_ptr<Graph>& graph);
+
+TORCH_API void InlineFunctionalGraphs(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dbr_quantization/remove_redundant_aliases.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dbr_quantization/remove_redundant_aliases.h
new file mode 100644
index 0000000000000000000000000000000000000000..5d5ec2c402a79d8b201241becdb491efe5b61596
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dbr_quantization/remove_redundant_aliases.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+
+namespace torch::jit {
+
+// This function replaces instances of
+//
+//   %b = aten::alias(%a)
+//   %c = foo(%b)
+//
+// with
+//
+//   %c = foo(%a)
+//
+// on the module forward, if it's safe to do so.
+TORCH_API Module DBRQuantRemoveRedundantAliases(Module& module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dead_code_elimination.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dead_code_elimination.h
new file mode 100644
index 0000000000000000000000000000000000000000..0768c102d94460bd159d4ff0304c06671a4068f7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dead_code_elimination.h
@@ -0,0 +1,45 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// If given a top-level graph, DCE will construct do alias analysis that allows
+// for "smarter" dead code elimination (we will eliminate mutable ops if we can
+// prove the mutated values are not used). Otherwise, we will not allow DCE to
+// eliminate mutable ops.
+//
+// So, prefer to use the graph version if you can.
+enum class DCESideEffectPolicy : uint8_t {
+  // default behavior: dead code elimination will check if a node has side
+  // effects
+  // and not delete it if it does.
+  DONT_DELETE_NODES_WITH_SIDE_EFFECTS,
+  // with this flag, dead code elimination will not check if a node has side
+  // effects and treat nodes with side effects like any other node,
+  // i.e. delete them if their outputs aren't used anywhere.
+  ALLOW_DELETING_NODES_WITH_SIDE_EFFECTS
+};
+
+TORCH_API void EliminateDeadCode(
+    const std::shared_ptr<Graph>& graph,
+    DCESideEffectPolicy sideEffectPolicy =
+        DCESideEffectPolicy::DONT_DELETE_NODES_WITH_SIDE_EFFECTS);
+TORCH_API void EliminateDeadCode(
+    Block* block,
+    bool recurse = true,
+    DCESideEffectPolicy sideEffectPolicy =
+        DCESideEffectPolicy::DONT_DELETE_NODES_WITH_SIDE_EFFECTS);
+
+// Invoke the user-provided callback on all live values before deleting anything
+TORCH_API void EliminateDeadCode(
+    Block* block,
+    std::function<void(const std::unordered_set<const Value*>&)> cb,
+    DCESideEffectPolicy sideEffectPolicy =
+        DCESideEffectPolicy::DONT_DELETE_NODES_WITH_SIDE_EFFECTS);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/decompose_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/decompose_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..eca49af30188b48f5a3f5ca3d3b9a787c112b5a6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/decompose_ops.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void DecomposeOps(std::shared_ptr<Graph>& graph);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/device_type_analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/device_type_analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..da2918b45f0c8ae32bab087241ef188af1615043
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/device_type_analysis.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+struct Graph;
+
+// Propagates Device type info throughout the given graph.
+TORCH_API bool DeviceTypePropagation(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dtype_analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dtype_analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..2a35b85ac435cfaaed3a45b4603f73949532ba32
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/dtype_analysis.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <memory>
+
+namespace torch::jit {
+struct Graph;
+
+// Propagate tensor properties (e.g., dtype, device, is_contiguous, layout)
+// propagation on all tensor objects. Currently, we only support dtype
+// propagation
+TORCH_API bool DtypePropagation(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/eliminate_no_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/eliminate_no_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..b54a0f597ae04bd74b089d4d63e06d1820480a9a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/eliminate_no_ops.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Remove ops that do nothing on the forward pass (like aten::detach).
+// This pass is invoked as a part of freeze_module.
+// This function also takes a set of custom ops to eliminate. All ops in this
+// set must take their output as their first input, i.e. x = f(x, ...)
+TORCH_API bool EliminateNoOps(
+    std::shared_ptr<Graph>& graph,
+    std::unordered_set<c10::Symbol> custom_ops = {});
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/erase_number_types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/erase_number_types.h
new file mode 100644
index 0000000000000000000000000000000000000000..a07403371551c7fc51de7e6158a7e4f52406fd98
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/erase_number_types.h
@@ -0,0 +1,26 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Erase NumberType information. This is necessary for and only used in
+// exporting to ONNX. This pass ensures that no remaining Values have
+// NumberType types, replacing them with tensors.
+// The following things are done to erase NumberType info:
+// - NumberType outputs are changed to DynamicType.
+// - prim::Constant nodes which are numbers get changed into 0-dim tensors of
+//   the corresponding type
+// - prim::TensorToNum, aten::Float, aten::Int and prim::NumToTensor nodes
+//   are erased.
+//
+// The pass assumes that DCE will be called sometime after.
+TORCH_API void EraseNumberTypes(const std::shared_ptr<Graph>& graph);
+TORCH_API void EraseNumberTypesOnBlock(Block* block);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fixup_trace_scope_blocks.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fixup_trace_scope_blocks.h
new file mode 100644
index 0000000000000000000000000000000000000000..e6a5a15dc1d6428ff0776412c2eb860ad924c586
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fixup_trace_scope_blocks.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Directly after tracing, we have an ill-formed graph with blocks inserted.
+// Example:
+//
+// graph(%self : ClassType<Module>,
+//       %input.1 : Float(3, 4)):
+//   %1 : ClassType<Module> = prim::GetAttr[name="relu1"](%self)
+//   %2 : ClassType<Module> = prim::GetAttr[name="relu2"](%self)
+//   %3 : ClassType<Module> = prim::GetAttr[name="rrr"](%2)
+//    = prim::TracedModuleForward[scope="__module.relu1"]()
+//     block0():
+//       %input : Float(3, 4) = aten::relu(%input.1),
+//       -> ()
+//    = prim::TracedModuleForward[scope="__module.relu2"](),
+//     block0():
+//        = prim::TracedModuleForward[scope="__module.relu2.rrr"](),
+//         block0():
+//           %6 : Float(3, 4) = aten::relu(%input),
+//           -> ()
+//       -> ()
+//   return (%6)
+//
+// In this pass, we:
+//   1) Lift Value defs to as high of a scope as needed to ensure that
+//      they dominate all their uses. For example, `input` in the above
+//      graph needs to be lifted to the top-level block so that its use
+//      in the second `relu` operator is dominated.
+//   2) Lambda lift the blocks. This ensures that all values used within
+//      each scope have their defs captured.
+//   3) Convert the scope blocks into methods on their respective Modules,
+//      and convert TracedModuleForward nodes to CallMethod nodes into those
+//      methods.
+//
+//  Then, we'll have a well-formed graph with proper method calls.
+TORCH_API void FixupTraceScopeBlocks(
+    std::shared_ptr<Graph>& graph,
+    Module* self);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fold_conv_bn.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fold_conv_bn.h
new file mode 100644
index 0000000000000000000000000000000000000000..ecdd7d39540262c49af085a7d4f4a0b5a4ccb125
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fold_conv_bn.h
@@ -0,0 +1,40 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+
+namespace torch::jit {
+
+/** \brief Fold Conv2d-BatchNorm2d into Conv2d in all methods of this
+ * module and all its submodules, forward is included by default.
+ *
+ * The weight and bias of the Conv2d are correspondingly updated. Should only be
+ * used on modules in eval mode.
+ */
+TORCH_API Module FoldConvBatchNorm(const Module& module);
+
+struct TORCH_API ConvBNParameters {
+  at::Tensor conv_w;
+  at::Tensor conv_b;
+  at::Tensor bn_rm;
+  at::Tensor bn_rv;
+  double bn_eps = 0.0;
+  at::Tensor bn_w;
+  at::Tensor bn_b;
+};
+
+/**
+ * Given the current weight and bias tensors of a Conv module and parameters
+ * of the BatchNorm module we're folding with, compute the updated values
+ * for the weight and bias.
+ *
+ * The function is basically copied from torch/nn/utils/fusion.py
+ */
+TORCH_API std::tuple<at::Tensor, at::Tensor> computeUpdatedConvWeightAndBias(
+    const ConvBNParameters& p);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fold_linear_bn.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fold_linear_bn.h
new file mode 100644
index 0000000000000000000000000000000000000000..9725255f4efa62ef1a2a66b80b5c88da307f57b7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fold_linear_bn.h
@@ -0,0 +1,32 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+
+namespace torch::jit {
+
+struct TORCH_API LinearBNParameters {
+  at::Tensor linear_w;
+  at::Tensor linear_b;
+  at::Tensor bn_rm;
+  at::Tensor bn_rv;
+  double bn_eps = 0.0;
+  at::Tensor bn_w;
+  at::Tensor bn_b;
+};
+
+/**
+ * Given the current weight and bias tensors of a Linear module and parameters
+ * of the BatchNorm module we're folding with, compute the updated values
+ * for the weight and bias.
+ *
+ * The function is basically copied from torch/nn/utils/fusion.py
+ */
+TORCH_API std::tuple<at::Tensor, at::Tensor> computeUpdatedLinearWeightAndBias(
+    const LinearBNParameters& p);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/freeze_module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/freeze_module.h
new file mode 100644
index 0000000000000000000000000000000000000000..9281743a4801dd11ee383ac4e893fe561c190a1b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/freeze_module.h
@@ -0,0 +1,39 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/** \brief This file defines freezing Torchscript module API.
+ *
+ * This API has python-binding and can be invoked directly or as a part of
+ * general optimization pipeline.
+ */
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+/** \brief Freeze Module, i.e., Assume all attributes are constants.
+ *
+ * Freezing module is a functionality that allows the JIT to internalize
+ * immutable attributes. Combined with inlining, the module is aggressively
+ * optimized and significant overhead is optimized away. The freezeModule API
+ * produces a cloned frozen module.
+ */
+
+namespace torch::jit {
+
+TORCH_API Module freeze_module(
+    const Module& module,
+    std::vector<std::string> preservedAttrs = std::vector<std::string>(),
+    bool freezeInterfaces = true,
+    bool preserveParameters = false);
+
+// Clone-free version of freeze_module. This modifies the module inplace.
+// Use this version to avoid extra memory usage incurred by cloning the module.
+TORCH_API void freeze_module_inplace(
+    Module* module,
+    std::vector<std::string> preservedAttrs = std::vector<std::string>(),
+    bool freezeInterfaces = true,
+    bool preserveParameters = false);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_concat_linear.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_concat_linear.h
new file mode 100644
index 0000000000000000000000000000000000000000..c8dc43e1f264dd6379de56370041caa2f122adf8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_concat_linear.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Concats multiple linear ops with the same Tensor input
+// into a single linear op.
+TORCH_API bool FrozenConcatLinear(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_conv_add_relu_fusion.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_conv_add_relu_fusion.h
new file mode 100644
index 0000000000000000000000000000000000000000..c0ec86b5e84b5346e2ba002234b745d123d15edc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_conv_add_relu_fusion.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API extern std::function<void(std::shared_ptr<Graph>&)>&
+getFuseFrozenConvAddReluImpl();
+
+TORCH_API void FuseFrozenConvAddRelu(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_conv_folding.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_conv_folding.h
new file mode 100644
index 0000000000000000000000000000000000000000..7be1df9bb211fe4d2d895c81bf26000e12bee60f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_conv_folding.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Fuses Convolution -> Batchnorm into a single Convolution by
+// folding batchnorm weights into conv weights.
+// This pass only works on Frozen Graphs; otherwise it is a No-Op.
+TORCH_API bool FoldFrozenConvBatchnorm(std::shared_ptr<Graph>& graph);
+
+// Fuses Convolution -> Add/Sub into a single Convolution by
+// folding add constant tensor into conv weights.
+// This pass only works on Frozen Graphs; otherwise it is a No-Op.
+TORCH_API bool FoldFrozenConvAddOrSub(std::shared_ptr<Graph>& graph);
+
+// Fuses Convolution -> Mul/Div into a single Convolution by
+// folding add constant tensor into conv weights.
+// This pass only works on Frozen Graphs; otherwise it is a No-Op.
+TORCH_API bool FoldFrozenConvMulOrDiv(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_graph_optimizations.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_graph_optimizations.h
new file mode 100644
index 0000000000000000000000000000000000000000..8c435cbebf000189967373afe992f1c633e7af4d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_graph_optimizations.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+/** \brief Runs a set of Optimizations that Optimize Frozen Graphs
+ *
+ * Currently this set of optimizations is:
+ * - FoldFrozenConvBatchnorm
+ * - FoldFrozenConvAddOrSub
+ * - FoldFrozenConvMulOrDiv
+ * - FoldFrozenLinearBatchnorm
+ */
+
+namespace torch::jit {
+
+TORCH_API void OptimizeFrozenGraph(
+    std::shared_ptr<Graph>& graph,
+    bool optimize_numerics = true);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_linear_folding.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_linear_folding.h
new file mode 100644
index 0000000000000000000000000000000000000000..9615bff1c4f477064c89e9991d4e14cc68f19f63
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_linear_folding.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Fuses Linear -> BatchNormNd into a single Linear by
+// folding batchnorm weights into linear weights.
+// This pass only works on Frozen Graphs; otherwise it is a No-Op.
+TORCH_API bool FoldFrozenLinearBatchnorm(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_linear_transpose.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_linear_transpose.h
new file mode 100644
index 0000000000000000000000000000000000000000..c1e7d52f78130c983f1ba63b8f46911a29c29cdf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_linear_transpose.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Transposes the weight matrix for frozen linear modules.
+// and converts it into a matmul
+TORCH_API bool FrozenLinearTranspose(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_ops_to_mkldnn.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_ops_to_mkldnn.h
new file mode 100644
index 0000000000000000000000000000000000000000..3a6dc2b906a9bcd1236c6a516a2b3c4bba439a90
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/frozen_ops_to_mkldnn.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Converts operators & their parameters to mkldnn if it is profitable
+// Currently encompassing Conv2d and Conv3d, and Linear
+// Op must be in float32 and mkldnn must be built
+// This pass only works on frozen graph
+TORCH_API void ConvertFrozenOpsToMKLDNN(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fuse_linear.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fuse_linear.h
new file mode 100644
index 0000000000000000000000000000000000000000..ad13f50ea137822d1e38617e7ba17ddea80f76c5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fuse_linear.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/** \brief Fusing linear patterns as single at::linear for easier pattern
+ * matching in later passes
+ */
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+/** \brief Match the at::linear pattern and fuse it into a single at::linear
+ * This pass fuse the addmm or matmul + add generated by JIT back to linear
+ * This pass can be deleted once the JIT can emit the aten::linear in the future
+ */
+TORCH_API void FuseLinear(std::shared_ptr<Graph>& graph);
+
+/** Swap functional linear CallFunctions to aten::linear
+ */
+TORCH_API void SwapFunctionalLinear(std::shared_ptr<Graph>& graph);
+/** Swap all functional linear CallFunctions in module
+ */
+TORCH_API void SwapFunctionalLinear(Module& module);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fuse_relu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fuse_relu.h
new file mode 100644
index 0000000000000000000000000000000000000000..749a6487b8fa93f818437d3089ce95aa05d8f4c2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/fuse_relu.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+TORCH_API void FuseAddRelu(script::Module& module);
+TORCH_API void FuseAddRelu(std::shared_ptr<Graph>& graph);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/graph_fuser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/graph_fuser.h
new file mode 100644
index 0000000000000000000000000000000000000000..eb1c5de6f88efff396eb2e756febb0cda5291212
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/graph_fuser.h
@@ -0,0 +1,40 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API bool canFuseOnCPULegacy();
+TORCH_API void overrideCanFuseOnCPULegacy(bool value);
+
+// NB: Be sure to run DCE before fusion, because dead instructions
+// can prevent fusion opportunities from being exploited.
+// On Windows will noop, NYI
+TORCH_API void FuseGraph(
+    std::shared_ptr<Graph>& graph,
+    bool strict_fuser_check = false);
+
+// \brief Custom fusion pass using a node-level callback to
+// determine the inclusion of nodes in a subgraph.
+//
+// This helper omits aliased inputs and fusion across control flow
+// boundaries.
+//
+// \arg graph The graph to be modified in-place
+// \arg is_fusable A callback run on each fusable node in the graph.
+// \arg kind The label given to the resultant fused subgraph
+// \arg arg_limit The maximum number of args the resultant fused subgraph
+//                should have.  Note: This will likely develop into a general
+//                post condition on the fused subgraph.
+TORCH_API void CustomFuseGraph(
+    std::shared_ptr<Graph>& graph,
+    const std::function<bool(Node*)>& is_fusable,
+    Symbol kind,
+    size_t arg_limit = std::numeric_limits<size_t>::max());
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/graph_rewrite_helper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/graph_rewrite_helper.h
new file mode 100644
index 0000000000000000000000000000000000000000..011d55177772914c920e3070dce107ede8d50e47
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/graph_rewrite_helper.h
@@ -0,0 +1,55 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/irparser.h>
+#include <torch/csrc/jit/ir/subgraph_matcher.h>
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+
+namespace torch::jit::graph_rewrite_helper {
+
+std::string getFuncName(Value* func_value);
+Value* getValue(
+    const std::string& name,
+    const std::unordered_map<const Value*, Value*>& match_vmap,
+    const std::unordered_map<std::string, Value*>& vmap);
+std::optional<IValue> getIValue(
+    const std::string& name,
+    const std::unordered_map<const Value*, Value*>& match_vmap,
+    const std::unordered_map<std::string, Value*>& vmap);
+TORCH_API void replaceConvolutionWithAtenConv(std::shared_ptr<Graph>& graph);
+
+bool isClampFusable(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+// This struct contains a compiled IR patterns slated for use in the
+// findPatternMatches function. The struct encapsulates the common
+// information from parseIR that is used in conjunction with the
+// pattern matching facility. A const instance of this struct can
+// also be stored away to cache the compiled IR pattern and reduce
+// runtime cost
+struct PatternInfo {
+  std::string pattern_string;
+  std::unique_ptr<Graph> pattern_graph;
+  std::unordered_map<std::string, Value*> vmap;
+  std::vector<MatchFilter> filters;
+
+  static PatternInfo parse_from_str(
+      std::string pattern_string,
+      const std::vector<MatchFilter>& filters = {}) {
+    PatternInfo rv{
+        std::move(pattern_string),
+        std::make_unique<Graph>(),
+        decltype(vmap){},
+        filters};
+    parseIR(rv.pattern_string, rv.pattern_graph.get(), rv.vmap);
+    return rv;
+  }
+};
+
+} // namespace torch::jit::graph_rewrite_helper
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/guard_elimination.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/guard_elimination.h
new file mode 100644
index 0000000000000000000000000000000000000000..79661a9d5d5333e2c2af2669745b8ef6976d53b6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/guard_elimination.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <list>
+#include <vector>
+
+namespace torch::jit {
+
+TORCH_API void EliminateRedundantGuards(std::shared_ptr<Graph> graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/hoist_conv_packed_params.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/hoist_conv_packed_params.h
new file mode 100644
index 0000000000000000000000000000000000000000..bab51295e51dfe2a069fa18c8c351e7ccf82e0a2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/hoist_conv_packed_params.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+void HoistConvPackedParams(script::Module& m);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_autodiff_subgraphs.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_autodiff_subgraphs.h
new file mode 100644
index 0000000000000000000000000000000000000000..38c7253d947a70304412967356fc137a9dea2f1c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_autodiff_subgraphs.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API bool canRunWithAutograd(Node* node);
+
+TORCH_API void InlineAutodiffSubgraphs(
+    std::shared_ptr<Graph>& graph,
+    size_t threshold = 5);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_fork_wait.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_fork_wait.h
new file mode 100644
index 0000000000000000000000000000000000000000..a4fbf517439474178cc726152b08de73052eabcf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_fork_wait.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Inline Fork and Wait calls. This is used, for example, in ONNX export, where
+// we do not support the explicit parallelism structures and would rather
+// just have a flat graph. This inlines the forked section in the fork()
+// callsite and replaces uses of the result of wait() calls with the values
+// produced from the (now-inlined) forked section.
+TORCH_API void InlineForkWait(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_forked_closures.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_forked_closures.h
new file mode 100644
index 0000000000000000000000000000000000000000..125a5681c5cd6f4d9181e6e49982cce61970134d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inline_forked_closures.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void inlineForkedClosures(std::shared_ptr<Graph>& to_clean);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inliner.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inliner.h
new file mode 100644
index 0000000000000000000000000000000000000000..9de6df90451d490af3d52244d0e782bf1a655f03
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inliner.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Inline function and method calls.
+TORCH_API void Inline(Graph& graph);
+
+TORCH_API GraphFunction* tryToGraphFunction(Node* n);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inplace_check.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inplace_check.h
new file mode 100644
index 0000000000000000000000000000000000000000..c45c0a27ed634e4540164cfaf0aeac99a35a0e08
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/inplace_check.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void CheckInplace(std::shared_ptr<Graph>& graph);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/insert_guards.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/insert_guards.h
new file mode 100644
index 0000000000000000000000000000000000000000..fe477baa97fb17cb074ea1a35f53bf7671ab4b10
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/insert_guards.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <list>
+#include <vector>
+
+namespace torch::jit {
+
+TORCH_API void InsertGuards(std::shared_ptr<Graph> graph);
+
+TORCH_API void RemoveProfilingNodes(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/integer_value_refinement.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/integer_value_refinement.h
new file mode 100644
index 0000000000000000000000000000000000000000..8a692daee3b8f2074855ab6b25fceb73ba270c38
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/integer_value_refinement.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// return true if graph is modified
+TORCH_API bool RefineIntegerValues(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lift_closures.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lift_closures.h
new file mode 100644
index 0000000000000000000000000000000000000000..978c14cde21f4995073a6c795ada4c040e6e6661
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lift_closures.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void liftClosures(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/liveness.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/liveness.h
new file mode 100644
index 0000000000000000000000000000000000000000..574a0be83f81d2bbd32c5fcbde3d1e5972e54536
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/liveness.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+#include <c10/util/sparse_bitset.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <list>
+#include <unordered_map>
+#include <vector>
+
+namespace torch::jit {
+
+using SparseBitVector = ::c10::SparseBitVector<256>;
+
+// BuildLivenessSets computes "bailout" liveness which is equivalent to
+// "{LIVE_IN} or {GEN}" or "{LIVE_OUT} - {KILL}"
+TORCH_API std::unordered_map<Node*, std::vector<Value*>> BuildLivenessSets(
+    std::shared_ptr<Graph> graph);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/loop_unrolling.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/loop_unrolling.h
new file mode 100644
index 0000000000000000000000000000000000000000..e6c064f7b1870062022e497223535d9853fd2698
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/loop_unrolling.h
@@ -0,0 +1,39 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// return true if graph is modified
+TORCH_API bool UnrollLoops(std::shared_ptr<Graph>& graph);
+
+// Only unrolls constant loops. Will unroll them regardless of loop block size
+TORCH_API bool UnrollConstantLoops(std::shared_ptr<Graph>& graph);
+
+TORCH_API Node* PeelLoop(Node* n, size_t times);
+
+// return true if graph is modified
+TORCH_API bool PeelProfilingLoops(const std::shared_ptr<Graph>& graph);
+
+struct TORCH_API LoopsPeeler {
+  LoopsPeeler(std::function<bool(Node* n)> callback, size_t num_iterations = 1)
+      : callback_(std::move(callback)), num_iterations_(num_iterations) {}
+
+  bool run(const std::shared_ptr<Graph>& graph);
+
+ private:
+  void collectLoop(Node* n);
+  void collectLoops(Block* block);
+  void peelLoops();
+
+  std::function<bool(Node* n)> callback_ = nullptr;
+  Node* in_loop_ = nullptr;
+  std::list<Node*> loops_to_peel_;
+  size_t num_iterations_ = 1;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_grad_of.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_grad_of.h
new file mode 100644
index 0000000000000000000000000000000000000000..cf4d238e1355513649139a43deba9173723726cc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_grad_of.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// This pass removes 'grad_of' nodes, replacing them with conditionals of
+// the form:
+// if any_defined(inputs):
+//  outputs = <original_computation>
+// else:
+//  outputs = undefineds
+TORCH_API void LowerGradOf(Graph& g);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_graph.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_graph.h
new file mode 100644
index 0000000000000000000000000000000000000000..d1e6338cd1bcc70c475a825e58f56b8bd1a12367
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_graph.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+using ModulePtr = c10::intrusive_ptr<c10::ivalue::Object>;
+
+// Given a graph with of a method which first argument is %self, lower it to a
+// graph where all attributes accesses are replaced with explicit inputs of the
+// graph (rather than results of prim::GetAttr executed on %self).
+//
+// Returns a tuple (graph, parameters) where the last module.parameters.size()
+// inputs to the graph are the trainable parameters used in this method. The
+// remaining inputs are the true inputs to the function.
+TORCH_API std::pair<std::shared_ptr<Graph>, std::vector<IValue>> LowerGraph(
+    Graph& graph,
+    const ModulePtr& self);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_tuples.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_tuples.h
new file mode 100644
index 0000000000000000000000000000000000000000..5358b4f68e1a774da4b05370656e367fa7ea130b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/lower_tuples.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// removes tuples where TupleConstruct and TupleUnpack are matched
+// but leaves tuples in place across if statements, loops, and as inputs/outputs
+TORCH_API void LowerSimpleTuples(const std::shared_ptr<Graph>& graph);
+
+// removes _all_ tuples and raises an error if some cannot be removed
+// this is used by ONNX to ensure there are not tuples before conversion,
+// but will not work on graphs whose inputs contain tuples.
+TORCH_API void LowerAllTuples(const std::shared_ptr<Graph>& graph);
+
+TORCH_API void LowerSimpleTuples(Block* block);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/metal_rewrite.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/metal_rewrite.h
new file mode 100644
index 0000000000000000000000000000000000000000..e26b4ff635d69d87ae8694b09acc61b3abc6f4f8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/metal_rewrite.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <string>
+#include <vector>
+
+namespace torch::jit {
+TORCH_API void metalInsertPrePackedOps(std::shared_ptr<Graph>& graph);
+TORCH_API void metalInsertPrePackedOps(script::Module& module);
+TORCH_API void metalFusePrePackedConvWithClamp(script::Module& module);
+TORCH_API void metalFoldPrePackingOps(script::Module& module);
+TORCH_API script::Module metalOptimizeForMobile(
+    const script::Module& module,
+    const std::vector<std::string>& preserved_methods);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/mkldnn_rewrite.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/mkldnn_rewrite.h
new file mode 100644
index 0000000000000000000000000000000000000000..709eeb2c92b5f441a75de4c114fd953429e0537c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/mkldnn_rewrite.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Config.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+
+#if AT_MKLDNN_ENABLED()
+
+#include <ideep/tensor.hpp>
+
+#endif // AT_MKLDNN_ENABLED()
+
+namespace torch::jit {
+
+#if AT_MKLDNN_ENABLED()
+
+namespace mkldnn {
+
+const static std::map<std::string, std::vector<torch::jit::MatchFilter>>
+    fusion_rewrite_map = {
+        {"none", {}},
+        {"relu", {}},
+};
+
+} // namespace mkldnn
+
+#endif // AT_MKLDNN_ENABLED()
+
+void FuseConvWithEltwise(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/mobile_optimizer_type.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/mobile_optimizer_type.h
new file mode 100644
index 0000000000000000000000000000000000000000..dd88910ecfc21fa505fce9466ede956ef2045b22
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/mobile_optimizer_type.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstdint>
+
+enum class MobileOptimizerType : int8_t {
+  CONV_BN_FUSION,
+  INSERT_FOLD_PREPACK_OPS,
+  REMOVE_DROPOUT,
+  FUSE_ADD_RELU,
+  HOIST_CONV_PACKED_PARAMS,
+  CONV_1D_TO_2D,
+  VULKAN_AUTOMATIC_GPU_TRANSFER,
+};
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/normalize_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/normalize_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..1de085401a1c839de96c1a46a12b9eef5af0ed7d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/normalize_ops.h
@@ -0,0 +1,21 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// This pass converts aten ops to a normalized form. It is
+// run immediately after IR generation in both the tracer and compiler,
+// so downstream consumers of the IR do not need handle ops in their
+// pre-normalized form.
+// Currently only handles normalization of op aliases.
+TORCH_API void NormalizeOps(const std::shared_ptr<Graph>& graph);
+
+const std::unordered_map<Symbol, Symbol>& getOperatorAliasMap();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onednn_graph_fuser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onednn_graph_fuser.h
new file mode 100644
index 0000000000000000000000000000000000000000..79783afd330e4d5df24380d3c1f85e67381c1bb9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onednn_graph_fuser.h
@@ -0,0 +1,66 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/pass_manager.h>
+
+#include <ATen/Config.h>
+
+namespace torch::jit {
+
+namespace fuser::onednn {
+
+static std::atomic<bool> onednn_enabled{true};
+
+static std::atomic<bool>& getLlgaEnabled() {
+  return onednn_enabled;
+}
+
+TORCH_API void fuseGraph(std::shared_ptr<Graph>& g);
+
+} // namespace fuser::onednn
+
+struct C10_EXPORT RegisterLlgaFuseGraph
+    : public PassManager<RegisterLlgaFuseGraph> {
+  static bool setEnabled(bool enabled) {
+    TORCH_CHECK(
+        AT_MKLDNN_ENABLED(),
+        "Running oneDNN Graph fuser is only supported with MKLDNN builds.");
+    bool oldState = fuser::onednn::getLlgaEnabled();
+    fuser::onednn::getLlgaEnabled() = enabled;
+    if (enabled) {
+      registerPass(fuser::onednn::fuseGraph);
+    } else {
+      clearPass();
+    }
+    return oldState;
+  }
+
+  static bool isEnabled() {
+    return fuser::onednn::getLlgaEnabled();
+  }
+
+  // override PassManager::registerPass to register pre-pass
+  static bool registerPass(GraphPass p) {
+    if (!isRegistered()) {
+      passID(registerPrePass(std::move(p)), true);
+      isRegistered(true);
+      return false;
+    }
+    return true;
+  }
+
+  // override PassManager::clearPass to clear pre-pass
+  static void clearPass() {
+    if (isRegistered()) {
+      clearPrePass(passID());
+      isRegistered(true);
+    }
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx.h
new file mode 100644
index 0000000000000000000000000000000000000000..a0bc542b554654a77bda379869958fd2ff9498d6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx.h
@@ -0,0 +1,33 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/onnx/onnx.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::jit {
+
+TORCH_API std::shared_ptr<Graph> ToONNX(
+    std::shared_ptr<Graph>& state,
+    ::torch::onnx::OperatorExportTypes operator_export_type);
+TORCH_API py::dict BlockToONNX(
+    Block* old_block,
+    Block* new_block,
+    ::torch::onnx::OperatorExportTypes operator_export_type,
+    py::dict& env,
+    py::set& values_in_env,
+    bool is_sub_block = false);
+TORCH_API void NodeToONNX(
+    Node* old_node,
+    Block* new_block,
+    ::torch::onnx::OperatorExportTypes operator_export_type,
+    py::dict& env,
+    py::set& values_in_env);
+TORCH_API void RemovePrintOps(std::shared_ptr<Graph>& graph);
+TORCH_API void PreprocessCaffe2Ops(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/cast_all_constant_to_floating.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/cast_all_constant_to_floating.h
new file mode 100644
index 0000000000000000000000000000000000000000..39fbd77a87591372577bf2ce4c6261b4d25e5da0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/cast_all_constant_to_floating.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <memory>
+
+namespace torch::jit {
+// see .cpp for docs
+TORCH_API void CastAllConstantToFloating(const std::shared_ptr<Graph>& graph);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/constant_fold.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/constant_fold.h
new file mode 100644
index 0000000000000000000000000000000000000000..347ed105b0974537e141aceca9457411b545e038
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/constant_fold.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <optional>
+
+namespace torch::jit {
+
+const int ONNX_OPSET_9 = 9;
+const int ONNX_OPSET_10 = 10;
+const int ONNX_OPSET_11 = 11;
+const int ONNX_OPSET_12 = 12;
+const int ONNX_OPSET_13 = 13;
+const int ONNX_OPSET_14 = 14;
+
+namespace onnx_constant_fold {
+
+at::Tensor IntToTensor(int64_t value);
+
+std::optional<at::Tensor> runTorchBackendForOnnx(
+    const Node* node,
+    std::vector<at::Tensor>& inputTensorValues,
+    int opset_version);
+} // namespace onnx_constant_fold
+
+void ConstantFoldONNX(
+    std::shared_ptr<Graph>& g,
+    std::map<std::string, IValue>& paramDict,
+    int opset_version);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/constant_map.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/constant_map.h
new file mode 100644
index 0000000000000000000000000000000000000000..a73ccef68eeea8e56a24b22f034759e3fb2726b3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/constant_map.h
@@ -0,0 +1,117 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Macros.h>
+
+#include <onnx/shape_inference/implementation.h>
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/serialization/export.h>
+#include <unordered_map>
+
+namespace torch::jit {
+
+using ShapeDataMap =
+    std::unordered_map<std::string, ::ONNX_NAMESPACE::TensorShapeProto>;
+
+class ConstantValueMap {
+ public:
+  static ConstantValueMap& getInstance();
+  static void SetRank(const std::string& tensorName, size_t rankValue);
+  static bool HasRank(const std::string& tensorName);
+  static std::optional<size_t> GetRank(const std::string& tensorName);
+
+  static void SetAllGraphInputsStatic(bool all_static);
+  static std::optional<bool> GetAllGraphInputsStatic();
+
+  static void SetAllGraphInputsReliableComputed(bool computed);
+  static bool GetAllGraphInputsReliableComputed();
+
+  static void SetShape(
+      const std::string& tensorName,
+      const c10::SymbolicShape& shapeValue);
+  static bool HasShape(const std::string& tensorName);
+  static std::optional<c10::SymbolicShape> GetShape(
+      const std::string& tensorName);
+
+  static void SetValue(const std::string& tensorName, const at::Tensor& value);
+  static bool HasValue(const std::string& tensorName);
+  static std::optional<at::Tensor> GetValue(const std::string& tensorName);
+  static void EraseValue(const std::string& tensorName);
+
+  static std::vector<int64_t> GetCompleteShapeInto1DInt64Vector(
+      const c10::SymbolicShape& shape);
+  static std::optional<std::vector<int64_t>> GetShapeInto1DInt64Vector(
+      const std::string& value_name);
+  static std::optional<std::vector<int64_t>>
+  GetShapeInto1DInt64VectorWithOneUnknown(const std::string& value_name);
+  static std::vector<int64_t> GetValueInto1DInt64Vector(
+      const std::string& value_name);
+
+  static void SetTypeReliable(const std::string& tensorName, bool reliable);
+  static bool HasTypeReliable(const std::string& tensorName);
+  static std::optional<bool> GetTypeReliable(const std::string& tensorName);
+
+  static void SetUseInferredType(
+      const std::string& tensorName,
+      bool useInferredType);
+  static bool HasUseInferredType(const std::string& tensorName);
+  static std::optional<bool> GetUseInferredType(const std::string& tensorName);
+
+  static void SetShapeValue(
+      const std::string& tensorName,
+      const c10::SymbolicShape& shapeValue);
+  static bool HasShapeValue(const std::string& tensorName);
+  static std::optional<c10::SymbolicShape> GetShapeValue(
+      const std::string& tensorName);
+
+  static ShapeDataMap& GetInferredShapeData();
+
+  static SymbolDimMap& GetSymbolDimMap();
+  static DimSymbolMap& GetDimSymbolMap();
+
+  static void UpdateValueName(
+      const std::string& old_name,
+      const std::string& new_name);
+
+  static void PrintMaps();
+  static void ClearMaps();
+  ~ConstantValueMap() = default;
+
+  ConstantValueMap& operator=(const ConstantValueMap&) = delete;
+
+ private:
+  ConstantValueMap() = default;
+
+  std::unordered_map<std::string, size_t> rankMap;
+  std::unordered_map<std::string, c10::SymbolicShape> shapeMap;
+  std::unordered_map<std::string, at::Tensor> tensorValueMap;
+  // This map indicates whether the current type is reliably estimated or not.
+  std::unordered_map<std::string, bool> typeReliableMap;
+  // This map indicates whether the current type is estimated through inference
+  // or tracer.
+  std::unordered_map<std::string, bool> useInferredTypeMap;
+  // This map indicates a tensor value which represents a shape.
+  // We assume that the rank of the tensor value <= 1, and we ensure this when
+  // we write the processing logic for the operators. When the rank > 1, we
+  // should be able to rewrite the model so that the rank <= 1. The difference
+  // between shapeMap and shapeValueMap: shapeMap stores the shape of the tensor
+  // from a node. shapeValueMap stores the value of the tensor from a node when
+  // this tensor represents a shape.
+  std::unordered_map<std::string, c10::SymbolicShape> shapeValueMap;
+  // Stores earlier data propagation results so that they are accessible
+  // during future node-level shape inference.
+  ShapeDataMap inferredShapeData;
+  SymbolDimMap symbolDimMap;
+  DimSymbolMap dimSymbolMap;
+  // Stores if all graph-level inputs have static shape
+  std::optional<bool> allGraphInputsStatic;
+  // True if reliable has been computed for all graph inputs
+  bool allGraphInputsReliableComputed{};
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/deduplicate_initializers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/deduplicate_initializers.h
new file mode 100644
index 0000000000000000000000000000000000000000..9820f3f7a2345045d849925a5e005e3e2d1de2fd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/deduplicate_initializers.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+void DeduplicateInitializers(
+    std::shared_ptr<Graph>& g,
+    std::map<std::string, IValue>& paramsDict,
+    bool is_train);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/eliminate_unused_items.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/eliminate_unused_items.h
new file mode 100644
index 0000000000000000000000000000000000000000..0b4cd693bd74d270dc6cb440549b1a90362e108b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/eliminate_unused_items.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// EliminateUnusedItemsONNX pass is removing unused
+// initializers and inputs, this is needed because
+// dce pass is only removing unused fork inputs
+void EliminateUnusedItemsONNX(
+    Block* b,
+    std::map<std::string, IValue>& paramDict);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/eval_peephole.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/eval_peephole.h
new file mode 100644
index 0000000000000000000000000000000000000000..4af16a22c7bf84f49d6c727ca6125907e7e81962
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/eval_peephole.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+void EvalPeepholeONNX(
+    std::shared_ptr<Graph>& g,
+    std::map<std::string, IValue>& paramDict);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.h
new file mode 100644
index 0000000000000000000000000000000000000000..9ae93de28bab2d9d41108c76adf0001788e4b246
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+std::vector<Value*> FixupONNXControlflowNode(Node* n, int opset_version);
+void FixupONNXControlflowNodeOutputs(Node* n);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/function_extraction.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/function_extraction.h
new file mode 100644
index 0000000000000000000000000000000000000000..71d441436eff786f0a97b89d40ed54113d6d1ab6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/function_extraction.h
@@ -0,0 +1,69 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+// This api will be used by serialization/export.cpp to extract function
+// information. It should do conversion on graph to
+//    1. Extract subgraph pattern of functions and define as local function
+//    node.
+//    2. Replace subgraph pattern of functions with a single node reflecting
+//    that local function node type.
+// Function attribute map information is also returned, as Torch IR cannot
+// represent these info inside Graph object.
+// export.cpp will serialize the ONNX model with function_proto with
+// above information.
+namespace torch::jit::onnx {
+
+// The following return types are used to track information regarding function
+// attributes, that are unable to be traced through Torch IR.
+// NodeAttrNameMap tracks mapping from attribute name of IR Node inside function
+// subgraph, to function attribute name. Here's an example of exporting CELU and
+// LayerNorm.
+//
+// clang-format off
+// class M(torch.nn.Module):
+//     def __init__(self) -> None:
+//         super().__init__()
+//         self.lns = torch.nn.ModuleList([torch.nn.LayerNorm(3, eps = i) for i in range(2)])
+//         self.celu1 = torch.nn.CELU(1.0)
+//         self.celu2 = torch.nn.CELU(2.0)
+
+//     def forward(self, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+//         res1 = self.celu1(x)
+//         res2 = self.celu2(y)
+//         for ln in self.lns:
+//             z = ln(z)
+//         return res1 + res2 + z
+// clang-format on
+//
+// Returning
+//
+// NodeAttrNameMap:
+// {
+//    %1 : Float(2, 3) = onnx::Celu[alpha=2.](%y) : {
+//      'alpha' : 'Celu_alpha'
+//    }
+// }
+//
+// The info here helps graph._export_onnx to construct function attributes for
+// onnx local FunctionProto.
+using NodeAttrNameMap = std::
+    unordered_map<const Node*, std::unordered_map<std::string, std::string>>;
+
+TORCH_API NodeAttrNameMap ONNXFunctionExtraction(
+    std::shared_ptr<Graph>& graph,
+    const std::unordered_set<std::string>& module_names,
+    const std::vector<std::string>& param_names);
+
+TORCH_API void ONNXClearScopeRecords();
+
+TORCH_API void ONNXTrackScopeAttributes(
+    std::shared_ptr<Graph>& graph,
+    std::map<std::string, IValue>& attributes);
+
+} // namespace torch::jit::onnx
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/function_substitution.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/function_substitution.h
new file mode 100644
index 0000000000000000000000000000000000000000..bfa5fff13e630eeae239c21909b4378ee5258194
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/function_substitution.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void ONNXFunctionCallSubstitution(Graph& graph);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/helper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/helper.h
new file mode 100644
index 0000000000000000000000000000000000000000..9cf0ae45c273fd295124056b89187a1c5b3dd867
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/helper.h
@@ -0,0 +1,77 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Utility functions for PyTorch to ONNX conversion.
+
+static const int OPSET_VERSION_1 = 1;
+static const int OPSET_VERSION_9 = 9;
+static const int OPSET_VERSION_10 = 10;
+static const int OPSET_VERSION_11 = 11;
+static const int OPSET_VERSION_12 = 12;
+static const int OPSET_VERSION_13 = 13;
+static const int OPSET_VERSION_14 = 14;
+static const int OPSET_VERSION_15 = 15;
+static const int OPSET_VERSION_16 = 16;
+
+using ValueToParamPairMap = std::map<Value*, std::pair<std::string, IValue>>;
+
+using ParamMap = std::map<std::string, IValue>;
+
+TORCH_API void buildParamsMapFromValueToParamsMap(
+    const ValueToParamPairMap& valsToParamsMap,
+    ParamMap& paramsDict);
+TORCH_API ValueToParamPairMap
+buildValueToParamsMap(Block* b, const ParamMap& paramsDict);
+TORCH_API void eraseUnusedValuesFromMap(ValueToParamPairMap& valsToParamsMap);
+TORCH_API void eraseUnusedBlockInputs(Block* b);
+
+TORCH_API Node* addNodeToBlock(
+    Block* block,
+    Symbol kind,
+    ArrayRef<Value*> inputs);
+
+TORCH_API Value* addInputToBlock(Block* block);
+
+TORCH_API std::optional<at::ScalarType> ONNXTypeToATenType(int32_t onnx_type);
+
+// Use int return type as no sable way exists to forward declare protobuf enum
+TORCH_API int ATenTypeToOnnxType(at::ScalarType at_type);
+
+TORCH_API void ONNXLintGraph(const std::shared_ptr<Graph>& graph);
+
+Node* createONNXUnsqueeze(
+    Graph* graph,
+    Node* n_to_insert_before,
+    Value* input,
+    int axis,
+    int opset_version);
+Node* createONNXConstant(
+    Graph* graph,
+    Node* n_to_insert_before,
+    at::Tensor value);
+
+bool isValidToTransformToONNXConcatNode(Node* lc_node);
+
+Node* transformToONNXConcatNode(
+    Graph* graph,
+    Node* lc_node,
+    bool need_new_input,
+    int opset_version);
+
+class ScalarTypeHashFunction {
+ public:
+  size_t operator()(const c10::ScalarType& type) const {
+    return static_cast<size_t>(type);
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/list_model_parameters.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/list_model_parameters.h
new file mode 100644
index 0000000000000000000000000000000000000000..9dfad50d210487654dc553a00de925edbdaf1434
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/list_model_parameters.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API std::pair<Module, std::vector<IValue>> list_module_parameters(
+    const Module& module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/naming.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/naming.h
new file mode 100644
index 0000000000000000000000000000000000000000..174dadcf9e3021082fc265f28f76d70e6cf0ce87
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/naming.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::onnx {
+
+namespace ONNXScopeName {
+
+std::string createFullScopeName(
+    const std::string& class_name,
+    const std::string& variable_name);
+std::string variableName(const torch::jit::ScopePtr& scope);
+std::string variableNameFromRoot(
+    const torch::jit::ScopePtr& scope,
+    const std::string& layer_separator);
+std::string className(const torch::jit::ScopePtr& scope);
+std::string classNameFromRoot(
+    const torch::jit::ScopePtr& scope,
+    const std::string& layer_separator);
+bool isCompatibleScope(const torch::jit::ScopePtr& scope);
+
+} // namespace ONNXScopeName
+
+TORCH_API void AssignScopedNamesForNodeAndValue(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit::onnx
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/onnx_log.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/onnx_log.h
new file mode 100644
index 0000000000000000000000000000000000000000..22162bacb8cd05e62a6ef144e6289807cf67180d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/onnx_log.h
@@ -0,0 +1,28 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <memory>
+#include <ostream>
+#include <string>
+
+namespace torch::jit::onnx {
+
+TORCH_API bool is_log_enabled();
+
+TORCH_API void set_log_enabled(bool enabled);
+
+TORCH_API void set_log_output_stream(std::shared_ptr<std::ostream> out_stream);
+
+TORCH_API std::ostream& _get_log_output_stream();
+
+#define ONNX_LOG(...)                            \
+  if (::torch::jit::onnx::is_log_enabled()) {    \
+    ::torch::jit::onnx::_get_log_output_stream() \
+        << ::c10::str(__VA_ARGS__) << std::endl; \
+  }
+
+} // namespace torch::jit::onnx
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/autograd_function_process.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/autograd_function_process.h
new file mode 100644
index 0000000000000000000000000000000000000000..53551bab18164c077c87c3e73802085fa6e8a822
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/autograd_function_process.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void ONNXAutogradFunctionProcess(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/common.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/common.h
new file mode 100644
index 0000000000000000000000000000000000000000..50c1b40cdf180dcdf1684ed1f249db281773014e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/common.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+// Functions used by both encapsulation and conversion.
+
+namespace torch::jit {
+
+struct IndexingPatternFinder {
+ public:
+  static std::vector<Node*> FetchSliceAndSelect(const Node* node);
+
+ private:
+  static bool IsSameSource(const Node* n, const Node* m);
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_conversion.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_conversion.h
new file mode 100644
index 0000000000000000000000000000000000000000..0497bab7457257a94050bab0c73e00a402538d70
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_conversion.h
@@ -0,0 +1,49 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::jit {
+
+// Introduction
+//
+// The conversion part is called inside the onnx pass.
+// In onnx pass, _run_symbolic_function will be called for each node in
+// topological order. When it reaches the placeholder node, this function will
+// be invoked. It will convert the nodes inside the sub-block based on pattern.
+// By that time, it will have shape/type of upstream operators available. After
+// the conversion is complete, the placeholder node will be removed, and nodes
+// inside its sub-block converted. NodeToONNX will be called for these
+// nodes, and they will be converted from ATen operator to ONNX operator.
+//
+// Note: Edit Pattern Conversion
+//
+// Each pattern is differentiated by the name attribute of placeholder node.
+// The placeholder node is part of torch IR graph, After this function, the aten
+// nodes under placeholder node subblock will be converted to ONNX and appended
+// to the new_block, which is under the new ONNX graph. For the pattern
+// conversion code, it can be divided into three parts.
+//      1. Nodes in this pattern should be captured inside the subblock of
+//         Placeholder node after pattern encapsulation[see
+//         pattern_encapsulation.h]. These nodes will be converted based on
+//         pattern. This part of conversion is from aten to aten. It happens on
+//         the torch IR graph inside placeholder node subblock.
+//      2. The second part of conversion is to convert the aten nodes produced
+//         into ONNX. This is done by calling NodeToONNX for each node. The new
+//         ONNX nodes are appended to the new_block, which is under the new ONNX
+//         graph.
+//      3. The last part of conversion is to find and return, in the same order,
+//         the ONNX outputs corresponding to the original output for the
+//         placeholder node.
+TORCH_API std::vector<Value*> ConvertPatternFromSubblock(
+    Block* new_block,
+    Node* old_node,
+    py::dict& env,
+    py::set& values_in_env);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_encapsulation.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_encapsulation.h
new file mode 100644
index 0000000000000000000000000000000000000000..8d5cce99154810fc5537be4e8615d3196d992023
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_encapsulation.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Introduction
+//
+// The encapsulation part will find the nodes of patterns, like how other
+// pre-onnx passes are written. But instead of converting the nodes, it will
+// encapsulate them into a sub-block of a new placeholder node. This part is
+// called before onnx pass, so it runs before calling symbolic functions.
+//
+// Note: Why separate the function into two parts
+//
+// The purpose is to support conversions that depend on shape and type
+// information. Shape and type information is only available after
+// _jit_pass_onnx, which converts aten nodes to onnx nodes. So there is a
+// interdependent issue. _jit_pass_onnx depends on preprocess passes to convert
+// aten nodes into convertible condition, and preprocess passes depend on
+// _jit_pass_onnx to convert upstream nodes and apply onnx shape inference.
+// Separating the pass into two parts breaks the interdependency.
+//
+// Note: Edit Pattern Encapsulation
+//
+// Encapsulation step identifies the pattern, and copies the nodes into
+// the subblock of a new placeholder node. The outputs of the new placeholder
+// node are used in place of the original nodes instead. The category of the
+// pattern is stored as attr::name.
+TORCH_API std::optional<Node*> EncapsulatePatternIntoSubblock(Node* n);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/peephole.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/peephole.h
new file mode 100644
index 0000000000000000000000000000000000000000..e0adf3dda39b0931bde457bc45bf5c65812d6841
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/peephole.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+void PeepholeOptimizeONNX(
+    std::shared_ptr<Graph>& graph,
+    int opset_version,
+    bool fixed_batch_size);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h
new file mode 100644
index 0000000000000000000000000000000000000000..29d86fe75dd7ec7afaa4a6660abfae2a72c74ad9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Prepare division ops for ONNX export. This is necessary for and only used
+// by ONNX export.
+//
+// The pass corrects the following:
+//
+// - aten::div(int, int) -> float is the python truediv operator. This doesn't
+//   exist in ONNX so we cast the ints to FloatTensors
+//
+TORCH_API void PrepareDivisionForONNX(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/preprocess_for_onnx.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/preprocess_for_onnx.h
new file mode 100644
index 0000000000000000000000000000000000000000..b714e438240c5005ad5f3c667ebf1c37bcbedb49
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/preprocess_for_onnx.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+void PreprocessForONNX(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/remove_inplace_ops_for_onnx.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/remove_inplace_ops_for_onnx.h
new file mode 100644
index 0000000000000000000000000000000000000000..37a04f8d6e0c6cb7a5a49326b47493e1499b5b32
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/remove_inplace_ops_for_onnx.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void RemoveInplaceOpsForONNX(
+    const std::shared_ptr<Graph>& graph,
+    Module* model);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/scalar_type_analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/scalar_type_analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..8dde9341f60cfec4e72fc1d3ab408603473ac995
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/scalar_type_analysis.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void ScalarTypeAnalysisForONNX(
+    const std::shared_ptr<Graph>& graph,
+    bool lowprecision_cast,
+    int opset_version);
+void ScalarTypeAnalysisNodeForONNX(Node* n);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/shape_type_inference.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/shape_type_inference.h
new file mode 100644
index 0000000000000000000000000000000000000000..2f684c9cb83f73d78ce85871ad2ca16759470013
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/shape_type_inference.h
@@ -0,0 +1,103 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/onnx/helper.h>
+#include <torch/csrc/jit/python/python_arg_flatten.h>
+
+#include <utility>
+
+namespace torch::jit {
+
+// Merges existing_type and inferred_type.
+// Returns {merged type, whether or not inferred_type was used}.
+//
+// The inferred type will take higher precedence, since it is produced by ONNX
+// shape inference, and is more compatible with ONNX. In cases where ONNX shape
+// inference fails to produce an inferred type, or produces an inferred type
+// that is incomplete, refer to existing type and fill in the gap that is
+// missing. Currently the following cases are supported.
+//  1. existing type: Tensor[], inferred type: Tensor[]
+//    For list of tensors, existing type does not store datatype nor shape for
+//    inner tensor. Thus inferred type always contain more information, and is
+//    returned.
+//  2. existing type: Tensor, inferred type: Tensor
+//    Fill in missing info (shape, data type) for inferred type from existing
+//    type.
+//  3. existing type: Scalar[], inferred type: Tensor
+//    ONNX represents list of scalars by 1-d Tensor. Return inferred type since
+//    it is more compatible with ONNX.
+std::pair<TypePtr, bool> MergeInferredType(
+    const TypePtr& existing_type,
+    const TypePtr& inferred_type);
+
+void MergeInferredTypeAndSetMap(
+    Value* dest_v,
+    const TypePtr& existing_type,
+    const TypePtr& inferred_type);
+
+// Update graph input types with dynamic axes info.
+// Axes that are marked as dynamic will be assigned as dynamic ShapeSymbol.
+// Note it is possible for multiple axes to share the same ShapeSymbol,
+// if they are defined as such in dynamic_axes.
+TORCH_API void ONNXSetDynamicInputShape(
+    std::shared_ptr<Graph>& graph,
+    const std::unordered_map<
+        std::string,
+        std::unordered_map<int64_t, std::string>>& dynamic_axes,
+    const std::vector<std::string>& input_names);
+
+// Update graph output with types of output Tensors.
+// If onnx_shape_inference is true, types of output Tensors will be compared and
+// merged with inferred types. It is possible that inferred types contain
+// dynamic axes, hence it takes precedence over types of output Tensors.
+TORCH_API void ONNXAssignOutputShape(
+    std::shared_ptr<Graph>& graph,
+    at::ArrayRef<at::Tensor> outputs,
+    const python::IODescriptor& desc,
+    bool onnx_shape_inference,
+    bool is_script,
+    int opset_version);
+
+// Replace None in output with Optional node (opset > 15) if it's
+// script model. This helps align the output format in ONNX internal tests
+// when comparing pytorch results with ONNX results, as they have different
+// process for None in output.
+void ReplaceGraphOutputNoneWithOptional(
+    std::shared_ptr<Graph>& graph,
+    size_t outputs_index);
+Node* ONNXOptionalNodeForNone(std::shared_ptr<Graph>& graph);
+
+// Utilize ONNX Shape Inference for node.
+// The node must have ONNX namespace, and is valid ONNX node according to spec.
+// On successful ONNX shape inference runs, the function updates output types of
+// n with inferred shape and type. Otherwise n is unchanged.
+TORCH_API void ONNXShapeTypeInference(
+    Node* n,
+    const ParamMap& params_dict,
+    int opset_version);
+
+// Utilize ONNX Shape Inference for graph.
+// Internally calls ONNXShapeTypeInference for each node, to achieve more
+// coverage that skips only individual nodes if illegal, instead of skipping for
+// the entire graph.
+TORCH_API void ONNXShapeTypeInference(
+    std::shared_ptr<Graph>& g,
+    const ParamMap& params_dict,
+    int opset_version);
+
+bool AllGraphInputsStatic(const Graph* g);
+std::pair<bool, bool> AreInputsReliableOrStatic(Node* n);
+void UpdateReliable(
+    torch::jit::Value* output,
+    const std::pair<bool, bool>& input_reliable,
+    bool no_type_warning = false);
+
+void UpdateReliable(torch::jit::Node* n);
+void UpdateShapeConstantIfReliable(torch::jit::Value* output);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/unpack_quantized_weights.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/unpack_quantized_weights.h
new file mode 100644
index 0000000000000000000000000000000000000000..4461870c64172a68fa33eaf57ba06e8c8776b3d7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/onnx/unpack_quantized_weights.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/onnx/onnx.h>
+
+#include <memory>
+
+namespace torch::jit {
+
+TORCH_API void UnpackQuantizedWeights(
+    std::shared_ptr<Graph>& graph,
+    std::map<std::string, IValue>& paramsDict);
+TORCH_API void insertPermutes(
+    std::shared_ptr<Graph>& graph,
+    std::map<std::string, IValue>& paramsDict);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/pass_manager.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/pass_manager.h
new file mode 100644
index 0000000000000000000000000000000000000000..d3cc4bc3b3417402b0f97fdf15ce658c321ba885
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/pass_manager.h
@@ -0,0 +1,139 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+/* `getCustomPrePasses()` returns a vector of passes that will be executed
+ * after differentiation but before any fusion. This is the de-facto location
+ * for compiler backends to insert passes.
+ *
+ * `getCustomPostPasses()` returns a vector of passes that will be
+ * executed after differentiation and after fusion (if any). This is the
+ * location for fusion cleanup passes if they are needed.
+ *
+ * Static registration of a pass can be done by creating a global
+ * `Register{Pre,Post}Pass r(Pass)` variable in a compilation unit.
+ *
+ * pass_manager.h uses a Meyer's singleton to store a vector of `Pass`es, which
+ * modify the IR graph in place.
+ */
+
+namespace torch::jit {
+
+// A pass modifies a Graph in place.
+using GraphPass = std::function<void(std::shared_ptr<Graph>&)>;
+
+// Since Passes are std::functions, we associate a UUID to each pass, this way
+// if we want to deregister a pass, we have something to reference it by.
+using GraphPassNameType = unsigned int;
+
+// Graph pass entries have a name associated with them
+using GraphPassEntry = std::pair<GraphPass, GraphPassNameType>;
+
+// Return currently registered passes. Passes are stored in a static vector
+TORCH_API std::vector<std::pair<GraphPass, GraphPassNameType>>&
+getCustomPostPasses();
+TORCH_API std::vector<std::pair<GraphPass, GraphPassNameType>>&
+getCustomPrePasses();
+
+TORCH_API GraphPassNameType registerPostPass(GraphPass p);
+TORCH_API GraphPassNameType registerPrePass(GraphPass p);
+
+// Look up pass by name passed in, remove it from registered passes
+TORCH_API void clearPostPass(GraphPassNameType p);
+TORCH_API void clearPrePass(GraphPassNameType p);
+
+// Remove all passes
+TORCH_API void clearAllPostPasses();
+TORCH_API void clearAllPrePasses();
+
+// LEGACY CALL
+struct TORCH_API RegisterPostPass {
+  RegisterPostPass(GraphPass p);
+};
+
+using RegisterPass = RegisterPostPass;
+
+/*
+ * PassManager is a wrapper on the register/clear PostPass functions above. It
+ * will register the pass provided in "registerPass" and will hold on to its
+ * associated name that way clearPass can be later called and will delete the
+ * pass used to register when called.
+ *
+ * PassManager is templated because we want static variables based on a
+ * particular GraphPass. When deriving from PassManager, you should send as the
+ * template parameter your derived class as you would for the curiously
+ * recurring template pattern. This template parameter isn't actually used and
+ * is simply done to prevent static members from being shared across derived
+ * types.
+ */
+template <typename DerivedType>
+struct C10_EXPORT PassManager {
+ private:
+  // We want this class to be abstract because it's
+  virtual void abstract() = 0;
+
+ protected:
+  /*
+   * isRegistered() will return if a pass has been registered
+   * isRegistered(true) will change the value of the internal static bool
+   *
+   * There's an internal static bool to this function to keep track of the
+   * state, this is so when functions are derived from this class, they don't
+   * have to worry about initializing the static members.
+   */
+  static bool isRegistered(bool flip_bit = false) {
+    static bool val = false;
+    if (flip_bit)
+      val = !val;
+    return val;
+  }
+
+  /*
+   * name() will return the name of the registered pass
+   * name(pass_name, true) will set the name of the pass
+   * Similarly to isRegistered we use an internal static variable to hold the
+   * name.
+   */
+  static GraphPassNameType passID(
+      GraphPassNameType PassID = 0,
+      bool set = false) {
+    static GraphPassNameType pass_id = 0;
+    if (set)
+      pass_id = PassID;
+    return pass_id;
+  }
+
+ public:
+  // registerPass(pass) will register the pass provided and set the
+  // name/isRegistered functions appropriately, it returns a bool value
+  // indicating whether the given pass is already registered previously.
+  static bool registerPass(GraphPass p) {
+    if (!isRegistered()) {
+      // If we don't already have a registered pass, register pass
+      // hold on to its name, change isRegistered to true
+      passID(registerPostPass(std::move(p)), true);
+      isRegistered(true);
+      return false;
+    }
+    return true;
+  }
+
+  // Calls ClearPostPass(passID())
+  static void clearPass() {
+    // If the pass is registered, clear it and change isRegistered to false.
+    if (isRegistered()) {
+      clearPostPass(passID());
+      isRegistered(true);
+    }
+  }
+
+  // clang-tidy requires virtual destructor;
+  virtual ~PassManager() = default;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole.h
new file mode 100644
index 0000000000000000000000000000000000000000..ca9437fd7d99358ecfd0cae8287a20d797c5145d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// return true if graph is modified
+TORCH_API bool PeepholeOptimize(
+    const std::shared_ptr<Graph>& graph,
+    bool disable_shape_peepholes = false);
+// return true if graph is modified
+TORCH_API bool PeepholeOptimize(
+    Block* block,
+    bool disable_shape_peepholes = false);
+// return true if graph is modified
+TORCH_API bool FuseAddMM(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_alias_sensitive.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_alias_sensitive.h
new file mode 100644
index 0000000000000000000000000000000000000000..c5ed890ddcd1dda39d85048f1fbd2d9d2c4fe944
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_alias_sensitive.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Peephole Optimizes alias sensitive peepholes
+// Currently this is invoked as part of PeepholeOptimize
+// return true if graph is modified
+// Optimizes on TensorType if shape_peepholes is true
+TORCH_API bool PeepholeOptimizeAliasSensitive(
+    const std::shared_ptr<Graph>& graph,
+    bool shape_peepholes);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_dict_idioms.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_dict_idioms.h
new file mode 100644
index 0000000000000000000000000000000000000000..21bd2f97c159dc7f583b2365f8b3475a4bd0b545
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_dict_idioms.h
@@ -0,0 +1,41 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Peephole Optimizes Dict Ops such as len() and __getitem__
+// 1. getitem optimizations
+// Given a function like this:
+//     def foo():
+//         d = {0 : 1}
+//         x = d[0]
+//         return x
+// This pass produces (after dead code elimination):
+//     def foo(a, b):
+//         return 1
+//
+// This optimization can only happen if the dict is not modified
+// and the dict has constant, non overlapping keys.
+//
+// 2. len optimizations
+// Given a function like this:
+//     def foo():
+//         d = {0 : 1}
+//         return len(d)
+// This pass produces (after dead code elimination):
+//     def foo():
+//         return 1
+//
+// This has the same requirements as the getitem optimizations.
+//
+// Currently this is invoked as part of PeepholeOptimize
+// return true if graph is modified.
+TORCH_API bool PeepholeOptimizeDictIdioms(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_list_idioms.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_list_idioms.h
new file mode 100644
index 0000000000000000000000000000000000000000..e52aaed75d54cb8ad982dc84860372607b9c0fc9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_list_idioms.h
@@ -0,0 +1,75 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Peephole Optimizes List ops such as len(li) and li[1].
+// 1. Construct/Unpack optimizations
+// Given a function like this:
+//    def foo(a, b):
+//        li = [a, b]
+//        x, y = li
+//        return x, y
+// This pass produces (after dead code elimination):
+//    def foo(a, b):
+//        return a, b
+//
+// This is only applied to lists that are not modified.
+//
+// 2. getitem optimizations
+// Given a function like this:
+//     def foo(a, b):
+//         li = [a, b]
+//         x = li[0]
+//         return x
+// This pass produces (after dead code elimination):
+//     def foo(a, b):
+//         return a
+//
+// This optimization can only happen if the list is not modified.
+//
+// 3. len optimizations
+// Given a function like this:
+//     def foo():
+//         li = [1, 2]
+//         return len(li)
+// This pass produces (after dead code elimination):
+//     def foo():
+//         return 2
+//
+// This has the same requirements as the getitem optimizations.
+//
+// 4. ListConstruct + ListConstruct
+// Given a function like this:
+//     def foo():
+//         return [1, 2] + [3, 4]
+// This pass produces (after dead code elimination):
+//     def foo():
+//         return [1, 2, 3, 4]
+//
+// This is only applied to lists that are not modified.
+//
+// 5. Slice
+// Given a function like this:
+//     def foo():
+//         return [1, 2, 3, 4, 5][0:2]
+// This pass produces (after deadcode elimination):
+//     def foo():
+//         return [1, 2]
+//
+// Currently this is invoked as part of PeepholeOptimize
+// return true if graph is modified.
+// If `refine_list_len` is true will attempt to refine the len of lists through
+// len comparisons and assertions. This does not generally optimize pytorch
+// programs so it is not called by default in PeepholeOptimize.
+TORCH_API bool PeepholeOptimizeListIdioms(
+    const std::shared_ptr<Graph>& graph,
+    bool refine_list_len = false);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_non_tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_non_tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..ce4416223c63307ed9e5bac3136f528ea6712d95
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/peephole_non_tensor.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// return true if graph is modified
+// Optimizing General Graph Patterns that
+// are not covered in peephole.cpp and peephole_list_idioms
+TORCH_API bool PeepholeOptimizeNonTensor(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/prepack_folding.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/prepack_folding.h
new file mode 100644
index 0000000000000000000000000000000000000000..d609f83cd5074f084a0af47dbad5bdf171fa7ba8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/prepack_folding.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+using PrePackingOpsFilterFn = std::function<bool(Node*)>;
+
+void PrePackingOpsFolder(
+    script::Module& m,
+    const PrePackingOpsFilterFn& is_foldable_op,
+    const std::string& attr_prefix);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/dedup_module_uses.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/dedup_module_uses.h
new file mode 100644
index 0000000000000000000000000000000000000000..fdaac9df01811de7d826d4f1c816d1f33d5c16cf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/dedup_module_uses.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+
+namespace torch::jit {
+
+/** Recursively deduplicate multiple uses of the same module by
+ *  creating an instance clone for each use of the module, which means
+ *  the type will be the same as before and all the attributes will be
+ *  copied, then we'll change the use of the original module to the use
+ *  of cloned module in the Graph.
+ *
+ *  This is done to ensure that modules can survive destructive passes
+ *  without changing model behavior. For example, here:
+ *
+ *    x = self.conv1(x)
+ *    x = self.relu(x)
+ *    x = self.conv2(x)
+ *    x = self.relu(x)
+ *
+ *  self.relu needs to be deduplicated for potential future destructive passes
+ *  to work properly.
+ */
+TORCH_API void DedupModuleUses(Module& module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/finalize.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/finalize.h
new file mode 100644
index 0000000000000000000000000000000000000000..8fb9403cf4beac11604c4d2a76fc9d17cd8abac6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/finalize.h
@@ -0,0 +1,66 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/quantization/quantization_type.h>
+
+namespace torch::jit {
+
+/** \brief Backend specific pass to fuse dequantize - op - quantize calls
+ * as quantized_op calls.
+ *
+ * Right now this is a fusion for fbgemm backend and only works for quantized
+ * conv op, we'll extend to more ops and more backends in the future.
+ *
+ * Currently supported fusion:
+ * q(conv2d(dq(a), dq(w), dq(b))) --> to_nchw(fbgemm_conv2d(prepack(to_nhwc(a)),
+ *                                                          prepack(to_nhwc(w)),
+ *                                                          prepack(to_nhwc(b))))
+ *
+ * q(linear(dq(a), dq(w), dq(b))) --> to_nchw(fbgemm_linear(prepack(to_nhwc(a)),
+ *                                                          prepack(to_nhwc(w)),
+ *                                                          prepack(to_nhwc(b))))
+ *
+ * \param graph the graph we want to apply fusion
+ */
+TORCH_API void QuantFusion(
+    std::shared_ptr<Graph>& graph,
+    QuantType quant_type = QuantType::STATIC);
+
+/** \brief Insert prepack and unpack function in graph
+ *  We want add pack/unpack functions for quantized weight because later we want
+ * to fold the packed weight as an attribute of the module, in order to reduce
+ * the cost of packing the weight on the fly in quantized models.
+ *
+ *  Each quantized op has it's corresponding prepack/unpack function,
+ *  right now, we only need to do prepack/unpack for quantized::linear
+ * and quantized::conv2d.
+ */
+TORCH_API void InsertPrepackUnpack(std::shared_ptr<Graph>& graph);
+
+/** \brief Insert pack and unpack function in all graphs
+ *   of module
+ *
+ *   Go through graphs of all the methods of all child modules
+ *   and call InsertPrepackUnpack on the graph.
+ */
+TORCH_API void InsertPrepackUnpack(Module& module);
+
+TORCH_API script::Module Finalize(
+    script::Module& module,
+    QuantType quant_type = QuantType::STATIC,
+    const std::vector<std::string>& preserved_attrs =
+        std::vector<std::string>());
+
+TORCH_API void FoldQuantizedPrepackingOps(Module& module);
+
+TORCH_API Module FinalizeOnDevicePTQ(
+    Module& module,
+    QuantType quant_type,
+    const std::string& method_name);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/fusion_passes.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/fusion_passes.h
new file mode 100644
index 0000000000000000000000000000000000000000..063c09a3cdc2aea3215225a315d2ff3c7b48b1f2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/fusion_passes.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+TORCH_API void FuseQuantizedAddRelu(std::shared_ptr<Graph>& graph);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/helper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/helper.h
new file mode 100644
index 0000000000000000000000000000000000000000..5e85cfcae7e732be65edecd8ab886350b53fb2c9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/helper.h
@@ -0,0 +1,219 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/subgraph_matcher.h>
+#include <torch/csrc/jit/passes/graph_rewrite_helper.h>
+#include <torch/csrc/jit/passes/quantization/quantization_type.h>
+
+#include <functional>
+#include <regex>
+
+namespace torch::jit {
+
+using graph_rewrite_helper::getFuncName;
+
+// Vector of a module and the name of its method
+using ModuleMethodVector = std::vector<std::pair<Module, std::string>>;
+// Map of quantization parameter name and value
+// for example _scale, _zero_point,
+// _scalar_type and _axis(for per channel quantization)
+using QParamVector = std::vector<std::pair<std::string, IValue>>;
+
+// =========== helper functions for Value =========
+// Check if a value is weight, since we need to use weight observer
+// for weight
+TORCH_API bool isWeight(Value* v);
+
+// Check if a value is bias for conv and linear, which we do not
+// quantize
+TORCH_API bool isBiasOfConvOrLinear(Value* v);
+
+TORCH_API bool isEmbeddingBagNonInput(Value* v);
+
+// Get the use as scalar input of clamp ops for the input value
+std::optional<Use> getClampScalarInputUse(Value* v);
+
+// For a given value `v`, get the list of values that we need to check
+// if they are observed/quantized or not, if so, we can say the
+// `v` is also observed/quantized, since we can derive
+// the quantization parameters for `v` given the list of values
+TORCH_API std::vector<Value*> getPassThroughInputs(Value* v);
+
+// Clones the method by the name of orig_method_name into new_method_name method
+TORCH_API void cloneMethod(
+    Module& module,
+    const std::string& orig_method_name,
+    const std::string& new_method_name);
+
+// Check if a value in the graph is a Scalar value
+TORCH_API bool isScalar(Value* v);
+
+// Check if value is the input of the graph
+TORCH_API bool hitGraphInput(Value* value);
+
+// Converts a mangled name, such as
+//   __torch__.torch.ao.nn.quantized.modules.conv.___torch_mangle_7.Conv2d
+// into an unmangled name, such as
+//   __torch__.torch.ao.nn.quantized.modules.conv.Conv2d
+TORCH_API std::string removeTorchMangle(const std::string& orig_name);
+
+// Return the module name that corresponds to the value.
+TORCH_API std::optional<std::string> getModuleName(Value* value);
+
+// =========== helper functions for Node =========
+TORCH_API bool isSingleInputGeneralShapeAtenFunction(Node* n);
+
+TORCH_API bool isSingleInputGeneralValueAtenFunction(Node* n);
+
+TORCH_API bool isSingleInputGeneralCallFunction(Node* n);
+
+TORCH_API bool isSingleInputGeneralAtenFunction(Node* n);
+
+TORCH_API bool isClamp(Node* n);
+
+// Check if the node will produce the same result regardless of whether
+// the input tensor is quantized or not, example: aten::size
+TORCH_API bool isTensorInfoNode(Node* n);
+
+// Check if this the propagate op that has single input, e.g. aten::cat
+TORCH_API bool isPropagateQuantSingleInputOp(Node* n);
+
+// Check if this is the propagate op that has two inputs, e.g. aten::add
+TORCH_API bool isPropagateQuantBinaryOp(Node* n);
+
+// Check if this is the node that we'll quantize or not quantize depending on
+// whether the input of the node is quantized, example: aten::cat
+TORCH_API bool isPropagateQuantOp(Node* n);
+
+// Check if the node is a binary op like aten::add and aten::mul and
+// if the input 1 is a scalar, these ops will be quantized to
+// quantized::{op}_scalar
+TORCH_API bool isBinaryOpWithScalarInput(Node* n);
+
+TORCH_API std::optional<std::tuple<c10::QScheme, QParamVector>> getFixedQParams(
+    Node* n);
+
+// We don't want to analyze the graph for some `builtin` CallFunctions
+// like `linear` because we want to preserve the op boundary
+TORCH_API bool userDefinedCallFunction(Node* n);
+
+// Check if the node has scalar input
+TORCH_API bool hasScalarInput(Node* n);
+
+// Check if a node is quantizable
+TORCH_API bool nodeQuantizable(
+    Node* n,
+    QuantType quant_type = QuantType::STATIC);
+
+// Nodes which only require quantization of weight value, eg. embedding_bag
+bool isWeightOnlyStaticQuantOp(Node* n);
+
+// Check if a use of the value is quantizable, this depends on
+// both the use node and the offset
+TORCH_API bool useQuantizable(const Use& use, QuantType quant_type);
+
+// Given a CallFunction node, extract the graph of the called function
+TORCH_API std::shared_ptr<Graph> getCallFunctionGraph(Node* n);
+
+// Check if `use` is a CallFunction of name `func_name` and if value
+// `v` is the nth argument (if provided) of the function
+bool matchCallFuncToUse(
+    const Use& use,
+    const std::string& func_name,
+    std::optional<int> nth_arg);
+
+// Check if `use` is a AtenFunction of name `func_name` and if value
+// `v` is the nth argument (if provided) of the function
+bool matchAtenFuncToUse(
+    const Use& use,
+    const std::string& func_name,
+    std::optional<int> nth_arg);
+
+// =========== helper functions for Block =========
+// checks if a block will always raise an Exception
+TORCH_API bool alwaysRaisesException(Block* block);
+
+// =========== helper functions for Module  ==========
+// TODO: remove
+TORCH_API std::vector<std::string> getModuleAccessPath(
+    Value* instance,
+    Value* self);
+// TODO: remove
+TORCH_API Module
+findChildModule(const Module& module, const std::vector<std::string>& path);
+
+// Given an CallMethod node, get the module instance corresponding
+// to the instance Value
+// TODO: refactor all current uses of this function to the Opt one
+TORCH_API Module getInvokedModule(Module& module, Node* n, Value* self);
+
+// Given an CallMethod node, get the module instance corresponding
+// to the instance Value if the instance is a module, otherwise return
+// std::nullopt
+std::optional<Module> getInvokedModuleOpt(
+    const Module& module,
+    Node* n,
+    Value* self);
+
+// ==================== filter functions for matches ==============
+// filter to check Value `vname` is a constant of int value `value`
+bool is_int_constant(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap,
+    const std::string& vname,
+    int value);
+
+// filter to check if the %alpha argument of aten::add is constant 1
+bool aten_add_alpha_is_one(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+// filter to check if the functional in CallFunction is relu
+bool is_functional_relu(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+// filter to check if the module is torch.nn.ReLU
+bool is_relu_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+bool is_linear_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+// TODO: add a macro to declare the filters
+bool is_conv1d_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+bool is_conv2d_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+bool is_conv3d_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+bool is_conv_transpose1d_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+bool is_conv_transpose2d_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+bool is_batchnorm2d_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+bool is_batchnorm3d_module(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/insert_observers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/insert_observers.h
new file mode 100644
index 0000000000000000000000000000000000000000..319f216fa0851a0bcd8a5f07ca9d081533684e0a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/insert_observers.h
@@ -0,0 +1,71 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/passes/quantization/quantization_type.h>
+
+namespace std {
+
+template <>
+struct hash<torch::jit::Module> {
+  inline size_t operator()(const torch::jit::Module& arg) const {
+    return std::hash<c10::intrusive_ptr<c10::ivalue::Object>>()(arg._ivalue());
+  }
+};
+
+} // namespace std
+
+namespace torch::jit {
+
+using QConfig = std::tuple<Module, Module>;
+using QConfigDict = std::unordered_map<std::string, std::optional<QConfig>>;
+
+/** \brief Insert observer module and observer function call for
+ *  the Tensors that needs to be observed.
+ *
+ * For each Tensor that needs to be observed in the method, insert observer
+ * module to the input module and add forward calls of observer to the specified
+ * method.
+ *
+ * \param module the input module
+ * \param method_name the method we want to insert observers for
+ * \param qconfig_dict the qconfig dictionary that specifies how
+ * each module is going to be quantized
+ * \param inplace whether we want to do inplace modification to the input module
+ * or clone the module
+ * \param is_dynamic whether the dynamic quantization script is being used.
+ */
+TORCH_API Module InsertObservers(
+    Module& module,
+    const std::string& method_name,
+    const QConfigDict& qconfig_dict,
+    bool inplace,
+    QuantType quant_type = QuantType::STATIC);
+
+/** \brief Insert observer module and observer method for
+ *  the Tensors that needs to be observed.
+ *
+ * For each Tensor that needs to be observed in the method, insert observer
+ * module to the input module and observe_<method-name> methods to the module.
+ * This method is clone of method_name with forward calls of observer added.
+ *
+ * \param module the input module
+ * \param method_name the method we want to insert observers for
+ * \param qconfig_dict the qconfig dictionary that specifies how
+ * each module is going to be quantized
+ * \param inplace whether we want to do inplace modification to the input module
+ * or clone the module
+ * \param is_dynamic whether the dynamic quantization script is being used.
+ */
+TORCH_API Module InsertObserversForOnDevicePTQ(
+    Module& module,
+    const std::string& method_name,
+    const QConfigDict& qconfig_dict,
+    bool inplace,
+    QuantType quant_type = QuantType::STATIC);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/insert_quant_dequant.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/insert_quant_dequant.h
new file mode 100644
index 0000000000000000000000000000000000000000..2410ddf652b6641ce11520a49764b1b356bc220a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/insert_quant_dequant.h
@@ -0,0 +1,49 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/quantization/quantization_type.h>
+
+namespace torch::jit {
+
+/** Replicate quantize node for prim::If blocks, so that we can match
+ *  quantization patterns in prim::If blocks
+ */
+TORCH_API void ReplicateQuant(std::shared_ptr<Graph>& graph);
+
+/** Replicate dequantize node for each use, so that we can match
+ *  quantization patterns
+ */
+TORCH_API void ReplicateDeQuant(std::shared_ptr<Graph>& graph);
+
+/** \brief Insert quantize - dequantize calls to the Tensors
+ *  that are observed in insert_observers pass
+ *
+ * For each Tensor that is observed, get the observer module and call
+ * calculate_qparam on the observer module to get quantization parameters
+ * and add quantize - int_repr - dequantize function calls using these
+ * parameters we also have special handling for quantizing "bias" right now.
+ *
+ * \param module the input module
+ * \param method_name the method we want to insert quantization calls for
+ */
+TORCH_API Module InsertQuantDeQuant(
+    Module& module,
+    const std::string& method_name,
+    bool inplace,
+    bool debug,
+    QuantType quant_type = QuantType::STATIC);
+
+TORCH_API Module InsertQuantDeQuantOnDevicePTQ(
+    Module& module,
+    const std::string& method_name,
+    bool inplace,
+    bool debug,
+    QuantType quant_type = QuantType::STATIC);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/quantization_patterns.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/quantization_patterns.h
new file mode 100644
index 0000000000000000000000000000000000000000..8ce511b1e375368baabd643634940c39742eb9ca
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/quantization_patterns.h
@@ -0,0 +1,1269 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/irange.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/subgraph_matcher.h>
+#include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/passes/quantization/helper.h>
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+namespace torch::jit {
+
+struct QuantFusionInfo {
+  std::string quantized_op_name;
+  std::string pattern;
+  std::string replacement;
+  std::vector<MatchFilter> filters;
+};
+
+namespace {
+std::string getExtraArgList(std::vector<std::string> extra_args) {
+  return std::accumulate(
+      extra_args.begin(),
+      extra_args.end(),
+      std::string(),
+      [](const std::string& acc, const std::string& arg) {
+        return acc + ", " + arg;
+      });
+}
+
+// Get the pattern we want to replace the match with
+std::string getAtenOpPattern(
+    const std::string& graph_header,
+    const std::string& op_name,
+    const std::vector<std::string>& extra_op_args,
+    bool scalar_args = false) {
+  std::vector<std::string> _extra_op_args = extra_op_args;
+  std::string aten_op_pattern = graph_header;
+  if (scalar_args) {
+    for (const auto& extra_arg : _extra_op_args) {
+      aten_op_pattern
+          .append(R"(
+          )")
+          .append(extra_arg)
+          .append("_scalar = aten::item(")
+          .append(extra_arg)
+          .append(")");
+    }
+
+    for (auto& _extra_op_arg : _extra_op_args) {
+      _extra_op_arg.append("_scalar");
+    }
+  }
+  const auto& extra_op_arg_list = getExtraArgList(std::move(_extra_op_args));
+  aten_op_pattern += R"(
+          %r = )";
+  aten_op_pattern += op_name + "(" + "%a_quant" + extra_op_arg_list + ")";
+  aten_op_pattern += R"(
+          return (%r) )";
+  return aten_op_pattern;
+}
+
+// generate ops for quantize pattern for a scalar value
+std::string getQuantizeForScalar(const std::string& value) {
+  // 6 is `torch.float` ScalarType, we are creating a float scalar
+  // tensor from a scalar value
+  std::string quantize_pattern = R"(
+          )" +
+      value + "_float_scalar_type : int = prim::Constant[value=6]()";
+  quantize_pattern += R"(
+          )" +
+      value + "_none : None = prim::Constant()";
+  quantize_pattern += R"(
+          )" +
+      value + "_tensor : Tensor = aten::scalar_tensor(" + value + ", " + value +
+      "_float_scalar_type";
+  for ([[maybe_unused]] const auto i : c10::irange(3)) {
+    quantize_pattern += ", " + value + "_none";
+  }
+  quantize_pattern += ")";
+  quantize_pattern +=
+      R"(
+          )" +
+      value + "_quant = aten::quantize_per_tensor(" + value + "_tensor" +
+      getExtraArgList(
+          {value + "_scale", value + "_zero_point", value + "_dtype"}) +
+      ")";
+  return quantize_pattern;
+}
+
+std::string getDequantize(const std::string& value) {
+  return R"(
+          )" +
+      value + "_dequant = aten::dequantize(" + value + "_quant)";
+}
+
+std::string getItem(const std::string& value) {
+  return R"(
+          )" +
+      value + "_scalar : float = aten::item(" + value + "_dequant)";
+}
+
+// Patterns for the ops that inherit parameters from input
+std::string getInputTensorQParamOpPattern(
+    const std::string& op_name,
+    const std::vector<std::string>& extra_op_args) {
+  const auto& extra_op_arg_list = getExtraArgList(extra_op_args);
+  std::string op_pattern = "graph(%a_quant" + extra_op_arg_list + "):" + R"(
+          %a_dequant = aten::dequantize(%a_quant)
+          %r = )" +
+      op_name + "(" + "%a_dequant" + extra_op_arg_list + ")" + R"(
+          %r_scale : float = aten::q_scale(%a_quant)
+          %r_zero_point : int = aten::q_zero_point(%a_quant)
+          %r_dtype : int = prim::dtype(%a_quant)
+          %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+          return (%r_quant) )";
+  return op_pattern;
+}
+
+// QuantFusionInfo for the ops that inherit parameters from input
+QuantFusionInfo getInputTensorQParamOpFusionInfo(
+    const std::string& op_name,
+    const std::vector<std::string>& extra_op_args) {
+  std::string op_pattern =
+      getInputTensorQParamOpPattern(op_name, extra_op_args);
+  const auto& extra_op_arg_list = getExtraArgList(extra_op_args);
+  std::string graph_header = "graph(%a_quant" + extra_op_arg_list + "):";
+  std::string op_replacement =
+      getAtenOpPattern(graph_header, op_name, extra_op_args);
+
+  return {op_name, std::move(op_pattern), std::move(op_replacement)};
+}
+
+// quant fusion for ops like `quantized::add_scalar`, `quantized::mul_scalar`
+QuantFusionInfo getBinaryOpScalarFusionInfo(
+    const std::string& op_name,
+    const std::vector<std::string>& extra_op_args,
+    const std::string& quantized_op_name,
+    const std::vector<std::string>& extra_quantized_op_args,
+    const std::vector<MatchFilter>& filters = {}) {
+  std::string op_pattern =
+      getInputTensorQParamOpPattern(op_name, extra_op_args);
+
+  const auto& extra_op_arg_list = getExtraArgList(extra_op_args);
+  std::string graph_header = "graph(%a_quant" + extra_op_arg_list + "):";
+  std::string op_replacement = getAtenOpPattern(
+      graph_header, quantized_op_name, extra_quantized_op_args);
+
+  return {op_name, std::move(op_pattern), std::move(op_replacement), filters};
+}
+
+QuantFusionInfo getClampOpFusionInfo(
+    const std::string& op_name,
+    const std::vector<std::string>& extra_op_args) {
+  std::vector<std::string> header_args = extra_op_args;
+  std::vector<std::string> input_qparams = {"_scale", "_zero_point", "_dtype"};
+  for (const auto& arg : extra_op_args) {
+    for (const auto& qparam : input_qparams) {
+      header_args.push_back(arg + qparam);
+    }
+  }
+  for (const auto& qparam : input_qparams) {
+    header_args.push_back("%r" + qparam);
+  }
+  const auto& extra_header_arg_list = getExtraArgList(std::move(header_args));
+  std::string graph_header = "graph(%a_quant" + extra_header_arg_list + "):";
+  std::string op_pattern = graph_header;
+  for (const auto& arg : extra_op_args) {
+    op_pattern += getQuantizeForScalar(arg);
+    op_pattern += getDequantize(arg);
+    op_pattern += getItem(arg);
+  }
+  op_pattern += getDequantize("%a");
+  op_pattern += R"(
+          %r = )";
+  std::vector<std::string> scalar_extra_args;
+  scalar_extra_args.reserve(extra_op_args.size());
+  for (const auto& arg : extra_op_args) {
+    scalar_extra_args.push_back(arg + "_scalar");
+  }
+  op_pattern += op_name + "(" + "%a_dequant" +
+      getExtraArgList(std::move(scalar_extra_args)) + ")";
+  // IR pattern common to all ops that inherit qparam from input
+  op_pattern += R"(
+          %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+          return (%r_quant) )";
+
+  std::string aten_op_pattern =
+      getAtenOpPattern(graph_header, op_name, extra_op_args);
+
+  return {op_name, std::move(op_pattern), std::move(aten_op_pattern)};
+}
+
+// Patterns for the ops that has fixed quantization parameters
+QuantFusionInfo getFixedQParamOpFusionInfo(
+    const std::string& op_name,
+    const std::vector<std::string>& extra_op_args,
+    bool is_symmetric) {
+  const auto& extra_op_arg_list = getExtraArgList(extra_op_args);
+  std::string graph_header = "graph(%a_quant" + extra_op_arg_list + "):";
+  std::string op_pattern = graph_header;
+  op_pattern += R"(
+          %a_dequant = aten::dequantize(%a_quant)
+          %r = )";
+  op_pattern += op_name + "(" + "%a_dequant" + extra_op_arg_list + ")";
+  // IR pattern common to all ops with fixed quantization parameters for
+  // asymmetric quantization
+  std::string asym_fixed_qparam_op_suffix = R"(
+          %r_scale : float = prim::Constant[value=0.00390625]()
+          %r_zero_point : int = prim::Constant[value=0]()
+          %r_dtype : int = prim::Constant[value=13]()
+          %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+          return (%r_quant) )";
+
+  std::string sym_fixed_qparam_op_suffix = R"(
+          %r_scale : float = prim::Constant[value=0.0078125]()
+          %r_zero_point : int = prim::Constant[value=128]()
+          %r_dtype : int = prim::Constant[value=13]()
+          %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+          return (%r_quant) )";
+  op_pattern +=
+      is_symmetric ? sym_fixed_qparam_op_suffix : asym_fixed_qparam_op_suffix;
+
+  std::string aten_op_pattern =
+      getAtenOpPattern(graph_header, op_name, extra_op_args);
+
+  return {op_name, std::move(op_pattern), std::move(aten_op_pattern)};
+}
+
+// filter that checks %b_scalar is a scalar
+bool input_b_is_scalar(
+    const Match& match,
+    const std::unordered_map<std::string, Value*>& vmap) {
+  const auto& match_vmap = match.values_map;
+  auto b_scalar = match_vmap.at(vmap.at("b_scalar"));
+  return isScalar(b_scalar);
+}
+
+// Patterns for ops that require observation for output quantization parameters
+// Example:
+//
+// before fusion:
+//
+// graph(%a_quant, %r_scale, %r_zero_point, %r_dtype):
+//     %a_dequant = aten::dequantize(%a_quant)
+//     %r = {op_name}(%a_dequant, {extra_args})
+//     %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point,
+//     %r_dtype) return (%r_quant)
+//
+// after fusion:
+//
+// graph(%a_quant, %r_scale, %r_zero_point, %r_dtype):
+//     %r_quant = {quantized_op_name}(%a_quant, {extra_args}, %r_scale,
+//     %r_zero_point) return (%r_quant)
+QuantFusionInfo getObservedQParamOpFusionInfo(
+    const std::string& fp_op_name,
+    const std::string& q_op_name,
+    const std::vector<std::string>& fp_extra_args,
+    const std::vector<std::string>& q_extra_args) {
+  const auto& fp_extra_arg_list = getExtraArgList(fp_extra_args);
+  const auto& q_extra_arg_list = getExtraArgList(q_extra_args);
+
+  std::string op_pattern = "graph(%a_quant" + fp_extra_arg_list +
+      ", %r_scale, %r_zero_point, %r_dtype):" + R"(
+          %a_dequant = aten::dequantize(%a_quant)
+          %r = )" +
+      fp_op_name + "(" + "%a_dequant" + fp_extra_arg_list + ")" + R"(
+          %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+          return (%r_quant) )";
+
+  std::string aten_op_pattern = "graph(%a_quant" + fp_extra_arg_list +
+      ", %r_scale, %r_zero_point, %r_dtype):" + R"(
+          %r_quant = )" +
+      q_op_name + "(%a_quant" + q_extra_arg_list +
+      ", %r_scale, %r_zero_point)" + R"(
+          return (%r_quant) )";
+
+  return {q_op_name, std::move(op_pattern), std::move(aten_op_pattern)};
+}
+
+} // namespace
+
+static std::vector<QuantFusionInfo> quant_fusion_pattern_and_replacements() {
+  // aten::conv1d
+  std::string conv1d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv1d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv1d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // aten::conv1d - aten::relu
+  std::string conv1d_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv1d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %conv_out = aten::conv1d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r = aten::relu(%conv_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // aten::conv1d - aten::relu_
+  std::string conv1d_inplace_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv1d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %conv_out = aten::conv1d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r = aten::relu_(%conv_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // quantized::conv1d
+  std::string quantized_conv1d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %r_quant = quantized::conv1d(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  // quantized::conv1d_relu
+  std::string quantized_conv1d_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %r_quant = quantized::conv1d_relu(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  // aten::conv2d
+  std::string conv2d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv2d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv2d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // aten::conv2d - aten::relu
+  std::string conv2d_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv2d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %conv_out = aten::conv2d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r = aten::relu(%conv_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // aten::conv2d - aten::relu_
+  std::string conv2d_inplace_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv2d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %conv_out = aten::conv2d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r = aten::relu_(%conv_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // quantized::conv2d
+  std::string quantized_conv2d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %r_quant = quantized::conv2d(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  // quantized::conv2d_relu
+  std::string quantized_conv2d_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %r_quant = quantized::conv2d_relu(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  // aten::conv3d
+  std::string conv3d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv3d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv3d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // aten::conv3d - aten::relu
+  std::string conv3d_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv3d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %conv_out = aten::conv3d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r = aten::relu(%conv_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // aten::conv3d - aten::relu_
+  std::string conv3d_inplace_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv3d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %conv_out = aten::conv3d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        %r = aten::relu_(%conv_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // quantized::conv3d
+  std::string quantized_conv3d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %r_quant = quantized::conv3d(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  // quantized::conv3d_relu
+  std::string quantized_conv3d_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %dilation, %groups):
+        %r_quant = quantized::conv3d_relu(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  // aten::conv_transpose1d
+  std::string conv_transpose1d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %output_padding, %groups, %dilation):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv_transpose1d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv_transpose1d(%a_dequant, %w_dequant, %b, %stride, %padding, %output_padding, %groups, %dilation)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // quantized::conv_transpose1d
+  std::string quantized_conv_transpose1d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %output_padding, %groups, %dilation):
+        %r_quant = quantized::conv_transpose1d(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  // aten::conv_transpose2d
+  std::string conv_transpose2d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %output_padding, %groups, %dilation):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::conv_transpose2d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv_transpose2d(%a_dequant, %w_dequant, %b, %stride, %padding, %output_padding, %groups, %dilation)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // quantized::conv_transpose1d
+  std::string quantized_conv_transpose2d = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype, %stride, %padding, %output_padding, %groups, %dilation):
+        %r_quant = quantized::conv_transpose2d(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r_quant) )";
+
+  std::string add_relu = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_add = aten::add(%a_dequant, %b_dequant, %alpha)
+         %r_relu = aten::relu(%r_add)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string add_inplace_relu = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_add = aten::add(%a_dequant, %b_dequant, %alpha)
+         %r_relu = aten::relu_(%r_add)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string inplace_add_relu = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_add = aten::add_(%a_dequant, %b_dequant, %alpha)
+         %r_relu = aten::relu(%r_add)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string inplace_add_inplace_relu = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_add = aten::add_(%a_dequant, %b_dequant, %alpha)
+         %r_relu = aten::relu_(%r_add)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string quantized_add_relu = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %r = quantized::add_relu(%a_quant, %b_quant, %scale, %zero_point)
+         return (%r) )";
+
+  // aten::linear
+  std::string linear = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::linear_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::linear(%a_dequant, %w_dequant, %b)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  std::string linear_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::linear_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %linear_out = aten::linear(%a_dequant, %w_dequant, %b)
+        %r = aten::relu(%linear_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  std::string linear_inplace_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype):
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::linear_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %linear_out = aten::linear(%a_dequant, %w_dequant, %b)
+        %r = aten::relu_(%linear_out)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  // quantized::linear
+  std::string quantized_linear = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype):
+        %r = quantized::linear(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r) )";
+
+  std::string quantized_linear_relu = R"(
+graph(%a_quant, %packed_params, %r_scale, %r_zero_point, %r_dtype):
+        %r = quantized::linear_relu(%a_quant, %packed_params, %r_scale, %r_zero_point)
+        return (%r) )";
+
+  std::string cat = R"(
+graph(%input_quant, %dim, %r_scale, %r_zero_point, %r_dtype):
+        %input_dequant = aten::dequantize(%input_quant)
+        %r = aten::cat(%input_dequant, %dim)
+        %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+        return (%r_quant) )";
+
+  std::string quantized_cat = R"(
+graph(%input_quant, %dim, %r_scale, %r_zero_point, %r_dtype):
+         %r_quant = quantized::cat(%input_quant, %dim, %r_scale, %r_zero_point)
+         return (%r_quant) )";
+
+  // aten::add
+  std::string add = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_add = aten::add(%a_dequant, %b_dequant, %alpha)
+         %r = aten::quantize_per_tensor(%r_add, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  // TODO: add %dtype after when https://github.com/pytorch/pytorch/issues/34351
+  // is fixed
+  // quantized::add
+  std::string quantized_add = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %r = quantized::add(%a_quant, %b_quant, %scale, %zero_point)
+         return (%r) )";
+
+  // aten::add_
+  std::string inplace_add = R"(
+graph(%a_quant, %b_quant, %alpha, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_add = aten::add_(%a_dequant, %b_dequant, %alpha)
+         %r = aten::quantize_per_tensor(%r_add, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  auto add_scalar = getBinaryOpScalarFusionInfo(
+      "aten::add",
+      {"%b_scalar", "%alpha"},
+      "quantized::add_scalar",
+      {"%b_scalar"},
+      {aten_add_alpha_is_one, input_b_is_scalar});
+
+  auto add_scalar_out = getBinaryOpScalarFusionInfo(
+      "aten::add_",
+      {"%b_scalar", "%alpha"},
+      "quantized::add_scalar_out",
+      {"%b_scalar", "%a_quant"},
+      {aten_add_alpha_is_one, input_b_is_scalar});
+
+  // quantized::add_scalar_relu -- fusing quantized::add_scalar
+  // and aten::relu
+  auto quantized_add_scalar_relu_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_add = quantized::add_scalar(%a_quant, %b_scalar)
+         %r = aten::relu(%r_add)
+         return (%r) )";
+
+  auto quantized_add_scalar_inplace_relu_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_add = quantized::add_scalar(%a_quant, %b_scalar)
+         %r = aten::relu_(%r_add)
+         return (%r) )";
+
+  auto quantized_add_scalar_relu_replacement = R"(
+graph(%a_quant, %b_scalar):
+         %r = quantized::add_scalar_relu(%a_quant, %b_scalar)
+         return (%r) )";
+
+  // quantized::add_scalar_relu_out -- fusing quantized::add_scalarOut
+  // and aten::relu
+  auto quantized_add_scalar_relu_out_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_add = quantized::add_scalar_out(%a_quant, %b_scalar, %a_quant)
+         %r = aten::relu(%r_add)
+         return (%r) )";
+
+  auto quantized_add_scalar_inplace_relu_out_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_add = quantized::add_scalar_out(%a_quant, %b_scalar, %a_quant)
+         %r = aten::relu_(%r_add)
+         return (%r) )";
+
+  auto quantized_add_scalar_relu_out_replacement = R"(
+graph(%a_quant, %b_scalar):
+         %r = quantized::add_scalar_relu_out(%a_quant, %b_scalar, %a_quant)
+         return (%r) )";
+
+  // quantized::batch_norm
+  std::string batch_norm = R"(
+graph(%a_quant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7, %scale, %zero_point, %scalar_type):
+         %a_dequant = aten::dequantize(%a_quant)
+         %r_bn = aten::batch_norm(%a_dequant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7)
+         %r = aten::quantize_per_tensor(%r_bn, %scale, %zero_point, %scalar_type)
+         return (%r) )";
+  std::string quantized_batch_norm = R"(
+graph(%a_quant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7, %scale, %zero_point, %scalar_type):
+         %r = quantized::batch_norm(%a_quant, %weight, %bias, %mean, %var, %eps, %scale, %zero_point)
+         return (%r) )";
+
+  std::string batch_norm_relu = R"(
+graph(%a_quant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7, %scale, %zero_point, %scalar_type):
+         %a_dequant = aten::dequantize(%a_quant)
+         %bn_out = aten::batch_norm(%a_dequant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7)
+         %relu = aten::relu(%bn_out)
+         %r = aten::quantize_per_tensor(%relu, %scale, %zero_point, %scalar_type)
+         return (%r) )";
+  std::string batch_norm_inplace_relu = R"(
+graph(%a_quant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7, %scale, %zero_point, %scalar_type):
+         %a_dequant = aten::dequantize(%a_quant)
+         %bn_out = aten::batch_norm(%a_dequant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7)
+         %relu = aten::relu_(%bn_out)
+         %r = aten::quantize_per_tensor(%relu, %scale, %zero_point, %scalar_type)
+         return (%r) )";
+
+  std::string quantized_batch_norm_relu = R"(
+graph(%a_quant, %weight, %bias, %mean, %var, %training, %eaf, %eps, %7, %scale, %zero_point, %scalar_type):
+         %r = quantized::batch_norm_relu(%a_quant, %weight, %bias, %mean, %var, %eps, %scale, %zero_point)
+         return (%r) )";
+
+  // aten::mul
+  std::string mul = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_mul = aten::mul(%a_dequant, %b_dequant)
+         %r = aten::quantize_per_tensor(%r_mul, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  // aten::mul_
+  std::string inplace_mul = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_mul = aten::mul_(%a_dequant, %b_dequant)
+         %r = aten::quantize_per_tensor(%r_mul, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  // quantized::mul
+  std::string quantized_mul = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %r = quantized::mul(%a_quant, %b_quant, %scale, %zero_point)
+         return (%r) )";
+
+  auto mul_scalar = getBinaryOpScalarFusionInfo(
+      "aten::mul",
+      {"%b_scalar"},
+      "quantized::mul_scalar",
+      {"%b_scalar"},
+      {input_b_is_scalar});
+
+  auto mul_scalar_out = getBinaryOpScalarFusionInfo(
+      "aten::mul_",
+      {"%b_scalar"},
+      "quantized::mul_scalar_out",
+      {"%b_scalar", "%a_quant"},
+      {input_b_is_scalar});
+
+  // quantized::mul_relu
+  std::string mul_relu = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_mul = aten::mul(%a_dequant, %b_dequant)
+         %r_relu = aten::relu(%r_mul)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string mul_inplace_relu = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_mul = aten::mul(%a_dequant, %b_dequant)
+         %r_relu = aten::relu_(%r_mul)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string inplace_mul_relu = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_mul = aten::mul_(%a_dequant, %b_dequant)
+         %r_relu = aten::relu(%r_mul)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string inplace_mul_inplace_relu = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %b_dequant = aten::dequantize(%b_quant)
+         %r_mul = aten::mul_(%a_dequant, %b_dequant)
+         %r_relu = aten::relu_(%r_mul)
+         %r = aten::quantize_per_tensor(%r_relu, %scale, %zero_point, %dtype)
+         return (%r) )";
+
+  std::string quantized_mul_relu = R"(
+graph(%a_quant, %b_quant, %scale, %zero_point, %dtype):
+         %r = quantized::mul_relu(%a_quant, %b_quant, %scale, %zero_point)
+         return (%r) )";
+
+  // quantized::mul_scalar_relu -- fusing quantized::mul_scalar
+  // and aten::relu
+  auto quantized_mul_scalar_relu_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_mul = quantized::mul_scalar(%a_quant, %b_scalar)
+         %r = aten::relu(%r_mul)
+         return (%r) )";
+
+  auto quantized_mul_scalar_inplace_relu_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_mul = quantized::mul_scalar(%a_quant, %b_scalar)
+         %r = aten::relu_(%r_mul)
+         return (%r) )";
+
+  auto quantized_mul_scalar_relu_replacement = R"(
+graph(%a_quant, %b_scalar):
+         %r = quantized::mul_scalar_relu(%a_quant, %b_scalar)
+         return (%r) )";
+
+  // quantized::mul_scalar_relu_out -- fusing quantized::mul_scalarOut
+  // and aten::relu
+  auto quantized_mul_scalar_relu_out_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_mul = quantized::mul_scalar_out(%a_quant, %b_scalar, %a_quant)
+         %r = aten::relu(%r_mul)
+         return (%r) )";
+
+  auto quantized_mul_scalar_inplace_relu_out_pattern = R"(
+graph(%a_quant, %b_scalar):
+         %r_mul = quantized::mul_scalar_out(%a_quant, %b_scalar, %a_quant)
+         %r = aten::relu_(%r_mul)
+         return (%r) )";
+
+  auto quantized_mul_scalar_relu_out_replacement = R"(
+graph(%a_quant, %b_scalar):
+         %r = quantized::mul_scalar_relu_out(%a_quant, %b_scalar, %a_quant)
+         return (%r) )";
+
+  // quantized::elu
+  std::string elu = R"(
+graph(%a_quant, %alpha, %scale, %input_scale, %r_scale, %r_zero_point, %r_dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %r = aten::elu(%a_dequant, %alpha, %scale, %input_scale)
+         %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+         return (%r_quant) )";
+
+  std::string quantized_elu = R"(
+graph(%a_quant, %alpha, %scale, %input_scale, %r_scale, %r_zero_point, %r_dtype):
+         %r_quant = quantized::elu(%a_quant, %r_scale, %r_zero_point, %alpha, %scale, %input_scale)
+         return (%r_quant) )";
+
+  std::string elu_ = R"(
+graph(%a_quant, %alpha, %scale, %input_scale, %r_scale, %r_zero_point, %r_dtype):
+         %a_dequant = aten::dequantize(%a_quant)
+         %r = aten::elu_(%a_dequant, %alpha, %scale, %input_scale)
+         %r_quant = aten::quantize_per_tensor(%r, %r_scale, %r_zero_point, %r_dtype)
+         return (%r_quant) )";
+
+  // ============= General Ops that inherit quantization parameters from input
+  // tensor =============
+  auto avg_pool1d = getInputTensorQParamOpFusionInfo(
+      "aten::avg_pool1d",
+      {"%kernel_size",
+       "%stride",
+       "%padding",
+       "%ceil_mode",
+       "%count_include_pad"});
+
+  auto avg_pool2d = getInputTensorQParamOpFusionInfo(
+      "aten::avg_pool2d",
+      {"%kernel_size",
+       "%stride",
+       "%padding",
+       "%ceil_mode",
+       "%count_include_pad",
+       "%divisor_override"});
+
+  auto avg_pool3d = getInputTensorQParamOpFusionInfo(
+      "aten::avg_pool3d",
+      {"%kernel_size",
+       "%stride",
+       "%padding",
+       "%ceil_mode",
+       "%count_include_pad",
+       "%divisor_override"});
+
+  auto adaptive_avg_pool1d = getInputTensorQParamOpFusionInfo(
+      "aten::adaptive_avg_pool1d", {"%output_size"});
+
+  auto adaptive_avg_pool2d = getInputTensorQParamOpFusionInfo(
+      "aten::adaptive_avg_pool2d", {"%output_size"});
+
+  auto adaptive_avg_pool3d = getInputTensorQParamOpFusionInfo(
+      "aten::adaptive_avg_pool3d", {"%output_size"});
+
+  auto mean1 = getInputTensorQParamOpFusionInfo("aten::mean", {"%dim"});
+
+  auto mean2 = getInputTensorQParamOpFusionInfo(
+      "aten::mean", {"%dim", "%keepdim", "%out"});
+
+  auto upsample_nearest1d_vec = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_nearest1d", {"%output_size", "%scale_factors"});
+
+  auto upsample_nearest2d_vec = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_nearest2d", {"%output_size", "%scale_factors"});
+
+  auto upsample_nearest3d_vec = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_nearest3d", {"%output_size", "%scale_factors"});
+
+  auto upsample_linear1d_vec = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_linear1d",
+      {"%output_size", "%align_corners", "%scale_factors"});
+
+  auto upsample_bilinear2d_vec = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_bilinear2d",
+      {"%output_size", "%align_corners", "%scale_factors"});
+
+  auto upsample_trilinear3d_vec = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_trilinear3d",
+      {"%output_size", "%align_corners", "%scale_factors"});
+
+  auto upsample_nearest1d = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_nearest1d", {"%output_size", "%scales"});
+
+  auto upsample_nearest2d = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_nearest2d", {"%output_size", "%scale_h", "%scale_w"});
+
+  auto upsample_nearest3d = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_nearest3d",
+      {"%output_size", "%scale_d", "%scale_h", "%scale_w"});
+
+  auto upsample_linear1d = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_linear1d", {"%output_size", "%align_corners", "%scales"});
+
+  auto upsample_bilinear2d = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_bilinear2d",
+      {"%output_size", "%align_corners", "%scale_h", "%scale_w"});
+
+  auto upsample_trilinear3d = getInputTensorQParamOpFusionInfo(
+      "aten::upsample_trilinear3d",
+      {"%output_size", "%align_corners", "%scale_d", "%scale_h", "%scale_w"});
+
+  auto clamp = getClampOpFusionInfo("aten::clamp", {"%min", "%max"});
+
+  auto hardtanh = getClampOpFusionInfo("aten::hardtanh", {"%min", "%max"});
+
+  auto hardtanh_ = getClampOpFusionInfo("aten::hardtanh_", {"%min", "%max"});
+
+  auto leaky_relu =
+      getInputTensorQParamOpFusionInfo("aten::leaky_relu", {"%negative_slope"});
+
+  auto leaky_relu_ = getInputTensorQParamOpFusionInfo(
+      "aten::leaky_relu_", {"%negative_slope"});
+
+  // Ops with fixed quantization parameters
+  auto hardsigmoid = getFixedQParamOpFusionInfo("aten::hardsigmoid", {}, false);
+
+  auto hardsigmoid_ =
+      getFixedQParamOpFusionInfo("aten::hardsigmoid_", {}, false);
+
+  auto sigmoid = getFixedQParamOpFusionInfo("aten::sigmoid", {}, false);
+
+  auto sigmoid_ = getFixedQParamOpFusionInfo("aten::sigmoid_", {}, false);
+
+  auto tanh = getFixedQParamOpFusionInfo("aten::tanh", {}, true);
+
+  auto tanh_ = getFixedQParamOpFusionInfo("aten::tanh_", {}, true);
+
+  auto hardswish = getObservedQParamOpFusionInfo(
+      "aten::hardswish", "quantized::hardswish", {}, {});
+
+  auto hardswish_ = getObservedQParamOpFusionInfo(
+      "aten::hardswish_", "quantized::hardswish", {}, {});
+
+  auto layer_norm = getObservedQParamOpFusionInfo(
+      "aten::layer_norm",
+      "quantized::layer_norm",
+      {"%normalized_shape", "%weight", "%bias", "%eps", "%cudnn_enabled"},
+      {"%normalized_shape", "%weight", "%bias", "%eps"});
+
+  auto group_norm = getObservedQParamOpFusionInfo(
+      "aten::group_norm",
+      "quantized::group_norm",
+      {"%num_groups", "%weight", "%bias", "%eps", "%cudnn_enabled"},
+      {"%num_groups", "%weight", "%bias", "%eps"});
+
+  auto instance_norm = getObservedQParamOpFusionInfo(
+      "aten::instance_norm",
+      "quantized::instance_norm",
+      {"%weight",
+       "%bias",
+       "%running_mean",
+       "%running_var",
+       "%use_input_stats",
+       "%momentum",
+       "%eps",
+       "%cudnn_enabled"},
+      {"%weight", "%bias", "%eps"});
+
+  return {
+      {"quantized::conv1d", std::move(conv1d), std::move(quantized_conv1d)},
+      {"quantized::conv1d_relu", std::move(conv1d_relu), quantized_conv1d_relu},
+      {"quantized::conv1d_relu",
+       std::move(conv1d_inplace_relu),
+       std::move(quantized_conv1d_relu)},
+      {"quantized::conv2d", std::move(conv2d), std::move(quantized_conv2d)},
+      {"quantized::conv2d_relu", std::move(conv2d_relu), quantized_conv2d_relu},
+      {"quantized::conv2d_relu",
+       std::move(conv2d_inplace_relu),
+       std::move(quantized_conv2d_relu)},
+      {"quantized::conv3d", std::move(conv3d), std::move(quantized_conv3d)},
+      {"quantized::conv3d_relu", std::move(conv3d_relu), quantized_conv3d_relu},
+      {"quantized::conv3d_relu",
+       std::move(conv3d_inplace_relu),
+       std::move(quantized_conv3d_relu)},
+      {"quantized::conv_transpose1d",
+       std::move(conv_transpose1d),
+       std::move(quantized_conv_transpose1d)},
+      {"quantized::conv_transpose2d",
+       std::move(conv_transpose2d),
+       std::move(quantized_conv_transpose2d)},
+      {"quantized::linear", std::move(linear), std::move(quantized_linear)},
+      {"quantized::linear_relu", std::move(linear_relu), quantized_linear_relu},
+      {"quantized::linear_relu",
+       std::move(linear_inplace_relu),
+       std::move(quantized_linear_relu)},
+      {"quantized::add_relu",
+       std::move(add_relu),
+       quantized_add_relu,
+       {aten_add_alpha_is_one}},
+      {"quantized::add_relu",
+       std::move(add_inplace_relu),
+       quantized_add_relu,
+       {aten_add_alpha_is_one}},
+      {"quantized::add_relu",
+       std::move(inplace_add_relu),
+       quantized_add_relu,
+       {aten_add_alpha_is_one}},
+      {"quantized::add_relu",
+       std::move(inplace_add_inplace_relu),
+       std::move(quantized_add_relu),
+       {aten_add_alpha_is_one}},
+      std::move(add_scalar),
+      std::move(add_scalar_out),
+      // note that these must come after quantized::add_scalar and
+      // quantized::add_scalar_out patterns
+      {"quantized::add_scalar_relu",
+       quantized_add_scalar_relu_pattern,
+       quantized_add_scalar_relu_replacement},
+      {"quantized::add_scalar_relu",
+       quantized_add_scalar_inplace_relu_pattern,
+       quantized_add_scalar_relu_replacement},
+      {"quantized::add_scalar_relu_out",
+       quantized_add_scalar_relu_out_pattern,
+       quantized_add_scalar_relu_out_replacement},
+      {"quantized::add_scalar_relu_out",
+       quantized_add_scalar_inplace_relu_out_pattern,
+       quantized_add_scalar_relu_out_replacement},
+      {"quantized::add",
+       std::move(add),
+       quantized_add,
+       {aten_add_alpha_is_one}},
+      {"quantized::add",
+       std::move(inplace_add),
+       std::move(quantized_add),
+       {aten_add_alpha_is_one}},
+      {"quantized::cat", std::move(cat), std::move(quantized_cat)},
+      {"quantized::batch_norm",
+       std::move(batch_norm),
+       std::move(quantized_batch_norm)},
+      {"quantized::batch_norm_relu",
+       std::move(batch_norm_relu),
+       quantized_batch_norm_relu},
+      {"quantized::batch_norm_relu",
+       std::move(batch_norm_inplace_relu),
+       std::move(quantized_batch_norm_relu)},
+      std::move(mul_scalar),
+      std::move(mul_scalar_out),
+      // note that these must come after quantized::mul_scalar and
+      // quantized::mul_scalar_out patterns
+      {"quantized::mul_scalar_relu",
+       quantized_mul_scalar_relu_pattern,
+       quantized_mul_scalar_relu_replacement},
+      {"quantized::mul_scalar_relu",
+       quantized_mul_scalar_inplace_relu_pattern,
+       quantized_mul_scalar_relu_replacement},
+      {"quantized::mul_scalar_relu_out",
+       quantized_mul_scalar_relu_out_pattern,
+       quantized_mul_scalar_relu_out_replacement},
+      {"quantized::mul_scalar_relu_out",
+       quantized_mul_scalar_inplace_relu_out_pattern,
+       quantized_mul_scalar_relu_out_replacement},
+      {"quantized::mul_relu", std::move(mul_relu), quantized_mul_relu},
+      {"quantized::mul_relu", std::move(mul_inplace_relu), quantized_mul_relu},
+      {"quantized::mul_relu", std::move(inplace_mul_relu), quantized_mul_relu},
+      {"quantized::mul_relu",
+       std::move(inplace_mul_inplace_relu),
+       std::move(quantized_mul_relu)},
+      {"quantized::mul", std::move(mul), quantized_mul},
+      {"quantized::mul", std::move(inplace_mul), std::move(quantized_mul)},
+      std::move(hardswish),
+      std::move(hardswish_),
+      std::move(layer_norm),
+      std::move(group_norm),
+      std::move(instance_norm),
+      {"quantized::elu", std::move(elu), quantized_elu},
+      {"quantized::elu_", std::move(elu_), std::move(quantized_elu)},
+      std::move(avg_pool1d),
+      std::move(avg_pool2d),
+      std::move(avg_pool3d),
+      std::move(adaptive_avg_pool1d),
+      std::move(adaptive_avg_pool2d),
+      std::move(adaptive_avg_pool3d),
+      std::move(mean1),
+      std::move(mean2),
+      std::move(upsample_nearest1d),
+      std::move(upsample_nearest2d),
+      std::move(upsample_nearest3d),
+      std::move(upsample_linear1d),
+      std::move(upsample_bilinear2d),
+      std::move(upsample_trilinear3d),
+      std::move(upsample_nearest1d_vec),
+      std::move(upsample_nearest2d_vec),
+      std::move(upsample_nearest3d_vec),
+      std::move(upsample_linear1d_vec),
+      std::move(upsample_bilinear2d_vec),
+      std::move(upsample_trilinear3d_vec),
+      std::move(clamp),
+      std::move(hardtanh),
+      std::move(hardtanh_),
+      std::move(leaky_relu),
+      std::move(leaky_relu_),
+      // fixed qparam ops
+      std::move(hardsigmoid),
+      std::move(hardsigmoid_),
+      std::move(sigmoid),
+      std::move(sigmoid_),
+      std::move(tanh),
+      std::move(tanh_),
+  };
+}
+
+inline std::vector<QuantFusionInfo>
+dynamic_quantized_linear_pattern_and_replacements() {
+  std::string linear_dynamic = R"(
+graph(%packed_params, %a):
+        %w_quant : Tensor, %b : Tensor? = quantized::linear_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::linear(%a, %w_dequant, %b)
+        return (%r) )";
+
+  // This pattern ignores reduce range
+  // Set the reduce range to default to true, since qnnpack backend ignores this
+  // argument.
+  std::string quantized_linear_dynamic = R"(
+graph(%packed_params, %a):
+        %reduce_range : bool = prim::Constant[value=1]()
+        %r = quantized::linear_dynamic(%a, %packed_params, %reduce_range)
+        return (%r) )";
+
+  return {
+      {"quantized::linear_dynamic",
+       std::move(linear_dynamic),
+       std::move(quantized_linear_dynamic)},
+  };
+}
+
+static std::vector<QuantFusionInfo>
+dynamic_quant_fusion_pattern_and_replacements() {
+  std::string linear_dynamic = R"(
+graph(%packed_params, %a, %reduce_range, %a_dtype):
+        %a_scale : float, %a_zero_point : int = aten::_choose_qparams_per_tensor(%a, %reduce_range)
+        %a_quant = aten::quantize_per_tensor(%a, %a_scale, %a_zero_point, %a_dtype)
+        %a_dequant = aten::dequantize(%a_quant)
+        %w_quant : Tensor, %b : Tensor? = quantized::linear_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::linear(%a_dequant, %w_dequant, %b)
+        return (%r) )";
+
+  std::string quantized_linear_dynamic = R"(
+graph(%packed_params, %a, %reduce_range, %a_dtype):
+        %r = quantized::linear_dynamic(%a, %packed_params, %reduce_range)
+        return (%r) )";
+
+  std::string linear_dynamic_fp16 = R"(
+graph(%packed_params, %a):
+        %w_unpacked : Tensor, %b : Tensor? = quantized::linear_unpack_fp16(%packed_params)
+        %r = aten::linear(%a, %w_unpacked, %b)
+        return (%r) )";
+
+  std::string quantized_linear_dynamic_fp16 = R"(
+graph(%packed_params, %a):
+        %r = quantized::linear_dynamic_fp16(%a, %packed_params)
+        return (%r) )";
+
+  return {
+      {"quantized::linear_dynamic",
+       std::move(linear_dynamic),
+       std::move(quantized_linear_dynamic)},
+      {"quantized::linear_dynamic_fp16",
+       std::move(linear_dynamic_fp16),
+       std::move(quantized_linear_dynamic_fp16)},
+  };
+}
+
+static std::vector<QuantFusionInfo> linear_prepack_unpack_patterns() {
+  std::string linear_with_quant = R"(
+graph(%a_dequant, %w_quant, %b):
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::linear(%a_dequant, %w_dequant, %b)
+        return (%r) )";
+
+  std::string linear_with_quant_prepack = R"(
+graph(%a_dequant, %w_quant, %b):
+        %packed_params = quantized::linear_prepack(%w_quant, %b)
+        %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::linear_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant_unpacked)
+        %r = aten::linear(%a_dequant, %w_dequant, %b_unpacked)
+        return (%r) )";
+  std::string linear_fp16_with_cast = R"(
+graph(%w, %a_dq, %b):
+        %fp16_tensor = aten::_saturate_weight_to_fp16(%w)
+        %r = aten::linear(%a_dq, %fp16_tensor, %b)
+        return (%r) )";
+  std::string linear_fp16_with_prepack = R"(
+graph(%w, %a_dq, %b):
+        %packed_params = quantized::linear_prepack_fp16(%w, %b)
+        %w_unpacked : Tensor, %b_unpacked : Tensor? = quantized::linear_unpack_fp16(%packed_params)
+        %r = aten::linear(%a_dq, %w_unpacked, %b_unpacked)
+        return (%r) )";
+
+  return {
+      {"linear_prepack_unpack",
+       std::move(linear_with_quant),
+       std::move(linear_with_quant_prepack)},
+      {"linear_fp16_prepack_unpack",
+       std::move(linear_fp16_with_cast),
+       std::move(linear_fp16_with_prepack)},
+  };
+}
+
+static std::vector<QuantFusionInfo> conv_prepack_unpack_patterns() {
+  std::string conv1d_with_quant = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %dilation, %groups):
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv1d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        return (%r) )";
+
+  std::string conv1d_with_quant_prepack = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %dilation, %groups):
+        %packed_params : __torch__.torch.classes.quantized.Conv2dPackedParamsBase = quantized::conv1d_prepack(%w_quant, %b, %stride, %padding, %dilation, %groups)
+        %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::conv1d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant_unpacked)
+        %r = aten::conv1d(%a_dequant, %w_dequant, %b_unpacked, %stride, %padding, %dilation, %groups)
+        return (%r) )";
+
+  std::string conv2d_with_quant = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %dilation, %groups):
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv2d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        return (%r) )";
+
+  std::string conv2d_with_quant_prepack = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %dilation, %groups):
+        %packed_params : __torch__.torch.classes.quantized.Conv2dPackedParamsBase = quantized::conv2d_prepack(%w_quant, %b, %stride, %padding, %dilation, %groups)
+        %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::conv2d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant_unpacked)
+        %r = aten::conv2d(%a_dequant, %w_dequant, %b_unpacked, %stride, %padding, %dilation, %groups)
+        return (%r) )";
+
+  std::string conv3d_with_quant = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %dilation, %groups):
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv3d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups)
+        return (%r) )";
+
+  std::string conv3d_with_quant_prepack = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %dilation, %groups):
+        %packed_params : __torch__.torch.classes.quantized.Conv3dPackedParamsBase = quantized::conv3d_prepack(%w_quant, %b, %stride, %padding, %dilation, %groups)
+        %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::conv3d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant_unpacked)
+        %r = aten::conv3d(%a_dequant, %w_dequant, %b_unpacked, %stride, %padding, %dilation, %groups)
+        return (%r) )";
+
+  std::string conv_transpose1d_with_quant = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %output_padding, %groups, %dilation):
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv_transpose1d(%a_dequant, %w_dequant, %b, %stride, %padding, %output_padding, %groups, %dilation)
+        return (%r) )";
+
+  std::string conv_transpose1d_with_quant_prepack = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %output_padding, %groups, %dilation):
+        %packed_params : __torch__.torch.classes.quantized.Conv2dPackedParamsBase = quantized::conv_transpose1d_prepack(%w_quant, %b, %stride, %padding, %output_padding, %dilation, %groups)
+        %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::conv_transpose1d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant_unpacked)
+        %r = aten::conv_transpose1d(%a_dequant, %w_dequant, %b_unpacked, %stride, %padding, %output_padding, %groups, %dilation)
+        return (%r) )";
+
+  std::string conv_transpose2d_with_quant = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %output_padding, %groups, %dilation):
+        %w_dequant = aten::dequantize(%w_quant)
+        %r = aten::conv_transpose2d(%a_dequant, %w_dequant, %b, %stride, %padding, %output_padding, %groups, %dilation)
+        return (%r) )";
+
+  std::string conv_transpose2d_with_quant_prepack = R"(
+graph(%a_dequant, %w_quant, %b, %stride, %padding, %output_padding, %groups, %dilation):
+        %packed_params : __torch__.torch.classes.quantized.Conv2dPackedParamsBase = quantized::conv_transpose2d_prepack(%w_quant, %b, %stride, %padding, %output_padding, %dilation, %groups)
+        %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::conv_transpose2d_unpack(%packed_params)
+        %w_dequant = aten::dequantize(%w_quant_unpacked)
+        %r = aten::conv_transpose2d(%a_dequant, %w_dequant, %b_unpacked, %stride, %padding, %output_padding, %groups, %dilation)
+        return (%r) )";
+
+  return {
+      {"conv1d_prepack_unpack",
+       std::move(conv1d_with_quant),
+       std::move(conv1d_with_quant_prepack)},
+      {"conv2d_prepack_unpack",
+       std::move(conv2d_with_quant),
+       std::move(conv2d_with_quant_prepack)},
+      {"conv3d_prepack_unpack",
+       std::move(conv3d_with_quant),
+       std::move(conv3d_with_quant_prepack)},
+      {"conv_transpose1d_prepack_unpack",
+       std::move(conv_transpose1d_with_quant),
+       std::move(conv_transpose1d_with_quant_prepack)},
+      {"conv_transpose2d_prepack_unpack",
+       std::move(conv_transpose2d_with_quant),
+       std::move(conv_transpose2d_with_quant_prepack)}};
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/quantization_type.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/quantization_type.h
new file mode 100644
index 0000000000000000000000000000000000000000..7f9f123ba89e777e6144097dcd10f5b81cb25402
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/quantization_type.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <cstdint>
+#include <ostream>
+
+namespace torch::jit {
+
+// Quantization type (dynamic quantization, static quantization).
+// Should match the Python enum in quantize_jit.py
+enum QuantType : std::uint8_t { DYNAMIC = 0, STATIC };
+
+std::ostream& operator<<(std::ostream& os, QuantType t);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/register_packed_params.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/register_packed_params.h
new file mode 100644
index 0000000000000000000000000000000000000000..fbc1a48c9d791fdf3e7b714a118a373a9660d708
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/quantization/register_packed_params.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <memory>
+
+namespace torch::jit {
+
+using PrePackParamFilterFn = std::function<bool(Node*)>;
+
+TORCH_API std::unordered_set<std::string> RegisterPrePackParams(
+    Module& m,
+    const std::string& method_name,
+    const PrePackParamFilterFn& is_packed_param,
+    const std::string& attr_prefix);
+
+TORCH_API std::string joinPaths(const std::vector<std::string>& paths);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/refine_tuple_types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/refine_tuple_types.h
new file mode 100644
index 0000000000000000000000000000000000000000..bff1d48ba87e957602c00304d24856291f952610
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/refine_tuple_types.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// updates the types of tuples according to the type of their current inputs.
+TORCH_API void RefineTupleTypes(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_dropout.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_dropout.h
new file mode 100644
index 0000000000000000000000000000000000000000..e97600c3288e4fa14f7d812048e17cf891471906
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_dropout.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void removeDropout(std::shared_ptr<Graph>& graph);
+
+TORCH_API void removeDropout(script::Module& module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_exceptions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_exceptions.h
new file mode 100644
index 0000000000000000000000000000000000000000..00a824427644dca888f307bcd48de11d6f05db95
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_exceptions.h
@@ -0,0 +1,26 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Considering prim::RaiseException nodes unreachable, simplify prim::If nodes
+// when one of the branches contains prim::RaiseException.
+//
+// This pass is illegal in general case as the modified graph might not throw
+// an exception that the original graph would throw. The purpose of the pass is
+// to cleanup the graph in a "risky" way by removing pathways leading to
+// RaiseExceptions nodes. In some sense, this pass could be considered as a
+// "Release" mode, while the original graph was in a "Debug" mode.
+// The pass should only be used when such transformation is guaranteed to be
+// safe by some other mechanisms. For instance, when we know exact shapes of
+// tensors flowing through the graph and tensors with such shapes never cause
+// exceptions.
+TORCH_API void EliminateExceptions(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_expands.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_expands.h
new file mode 100644
index 0000000000000000000000000000000000000000..472ed1e84f458042c024273287235719cd23fac2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_expands.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void RemoveExpands(const std::shared_ptr<Graph>& graph);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_inplace_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_inplace_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..e39a5ffe5a9b59de5107529b382dd27e02d0ad58
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_inplace_ops.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <memory>
+
+namespace torch::jit {
+// see .cpp for docs
+TORCH_API void RemoveInplaceOps(const std::shared_ptr<Graph>& graph);
+
+TORCH_API void ImplicitCastForBinaryInplaceOps(Block* block);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_mutation.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_mutation.h
new file mode 100644
index 0000000000000000000000000000000000000000..ed541091465c766e000e9b449050f495dac34d25
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_mutation.h
@@ -0,0 +1,86 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/Exception.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/alias_analysis.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <utility>
+
+namespace torch::jit {
+
+struct TORCH_API MutationRemover {
+  MutationRemover(
+      std::shared_ptr<Graph> graph,
+      std::optional<std::function<bool(Node*)>> mutation_filter = std::nullopt)
+      : mutation_filter_(std::move(mutation_filter)),
+        aliasDb_(nullptr),
+        graph_(std::move(graph)) {}
+
+  // return true if graph is modified
+  bool removeListMutation();
+
+  // return true if graph is modified
+  bool removeTensorMutation();
+
+  bool isSpecialMappedOp(Node* n) {
+    return n->matches("aten::zero_(Tensor(a!) self) -> Tensor(a!)") ||
+        n->matches(
+            "aten::fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!)") ||
+        n->matches(
+            "aten::normal_(Tensor(a!) self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor(a!)");
+  }
+
+  bool inplaceOpVariant(Node* n);
+
+  static bool hasSideEffectOrAlias(Value* v, AliasDb* aliasDb);
+
+ private:
+  Node* createSpecialMappedOp(Node* n);
+  bool listMutationFollowingListConstruct(Node* n);
+  bool tryMakeCreationAndMutationAtomic(
+      Value* mutated_value,
+      Node* mutating_op);
+  bool tryMakeUnaliasedIfOutputAndMutationAtomic(
+      Value* mutated_value,
+      Node* mutating_op);
+  // return true if graph is modified
+  bool RemoveListMutation(Block* block);
+  // return true if graph is modified
+  bool RemoveTensorMutation(Block* block);
+
+  AliasDb* getOrCreateAliasDb() {
+    if (!aliasDb_) {
+      aliasDb_ = std::make_unique<AliasDb>(graph_);
+    }
+    return aliasDb_.get();
+  }
+
+  std::optional<std::function<bool(Node*)>> mutation_filter_;
+  std::unique_ptr<AliasDb> aliasDb_ = nullptr;
+  std::shared_ptr<Graph> graph_;
+};
+
+// Removes list mutation with functional equivalents
+// return true if graph is modified
+TORCH_API bool RemoveListMutation(const std::shared_ptr<Graph>& graph);
+
+// Replaces in-place aten ops with their functional equivalents
+// when it can be proven that this does not change graph semantics
+// if `mutation_filter` is present, the pass will only attempt to
+// remove mutation on nodes which return true for the filter
+// return true if graph is modified
+TORCH_API bool RemoveTensorMutation(
+    const std::shared_ptr<Graph>& graph,
+    std::optional<std::function<bool(Node*)>> mutation_filter = std::nullopt);
+
+// Replaces in-place aten activation ops with their functional equivalence
+TORCH_API bool InplaceToFunctionalActivation(
+    const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_redundant_profiles.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_redundant_profiles.h
new file mode 100644
index 0000000000000000000000000000000000000000..913a9a4fe5019bd78ac600d3fae9e45b8dc4b3ae
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/remove_redundant_profiles.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void RemoveRedundantProfiles(std::shared_ptr<Graph>& graph);
+TORCH_API void RemoveRedundantProfiles(Block* block, AliasDb& db);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/replacement_of_old_operators.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/replacement_of_old_operators.h
new file mode 100644
index 0000000000000000000000000000000000000000..d3cb6abcabb0b94d186894caefd2ace2b68e791b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/replacement_of_old_operators.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Find the valid upgrader graph for the upgrader and cache the result
+// for later lookups. Will error out if there is no valid upgrader graph
+// provided for the upgrader name.
+std::shared_ptr<Graph> getUpgraderGraph(const std::string& upgrader_name);
+
+TORCH_API void ReplaceOldOperatorsWithUpgraders(std::shared_ptr<Graph> graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/requires_grad_analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/requires_grad_analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..d143d219bb19b862261ee96ea44e75881ff319ce
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/requires_grad_analysis.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+
+#include <memory>
+
+namespace torch::jit {
+
+struct Graph;
+struct ArgumentSpec;
+
+TORCH_API void PropagateRequiresGrad(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/restore_mutation.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/restore_mutation.h
new file mode 100644
index 0000000000000000000000000000000000000000..eb92d07c161326426bb07066ecb8478303cb8e44
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/restore_mutation.h
@@ -0,0 +1,66 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/symbol.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/alias_analysis.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// A map which stores if an activation operator can perform type promotion
+const std::unordered_map<Symbol, bool> activation_type_promotion_mapping = {
+    {aten::sigmoid, true},
+    {aten::tanh, true},
+    {aten::celu, false},
+    {aten::elu, false},
+    {aten::gelu, false},
+    {aten::glu, false},
+    {aten::hardshrink, false},
+    {aten::hardsigmoid, false},
+    {aten::hardswish, false},
+    {aten::hardtanh, false},
+    {aten::leaky_relu, false},
+    {aten::prelu, false},
+    {aten::relu6, false},
+    {aten::relu, false},
+    {aten::rrelu, false},
+    {aten::selu, false},
+    {aten::silu, false}};
+
+class FunctionalToInplaceRewriter {
+ public:
+  FunctionalToInplaceRewriter(std::shared_ptr<Graph> graph);
+
+  bool FunctionalToInplace(Block* block);
+
+ private:
+  AliasDb* getOrCreateAliasDb() {
+    if (!aliasDb_) {
+      aliasDb_ = std::make_unique<AliasDb>(graph_);
+    }
+    return aliasDb_.get();
+  }
+
+  bool CanBeInplace(Node* node);
+
+  std::unique_ptr<AliasDb> aliasDb_ = nullptr;
+  std::shared_ptr<Graph> graph_;
+};
+
+// A common application scenario is to apply InplaceToFunctionalActivation
+// before some JIT optimization passes, so that those passes are less
+// constrained by in-place ops. After those passes are done, we can call
+// FunctionalToInplaceActivation to recover in-place activation ops,
+// so that we won't lose the performance benefit coming from memory reduction.
+
+// Replaces functional aten activation ops with their in-place equivalents
+TORCH_API bool FunctionalToInplaceActivation(
+    const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/shape_analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/shape_analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..224ec887d0b63b4d4e47660bed4b94825e823309
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/shape_analysis.h
@@ -0,0 +1,46 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <memory>
+
+namespace torch::jit {
+
+struct Graph;
+
+struct propagation_error : std::exception {};
+
+class PropertyPropBase {
+  // Used for both Shape Propagation and Dtype/Device Propagation
+ public:
+  explicit PropertyPropBase(std::shared_ptr<Graph> graph)
+      : graph_(std::move(graph)) {}
+  virtual ~PropertyPropBase() = default;
+
+  void propagateBlock(Block* block, bool insert_expands = true);
+  // insert_expands is used for shape inference
+
+  void processIf(Node* node);
+  void processLoop(Node* node);
+
+ protected:
+  virtual void propagateNode(Node* node, bool insert_expands = true) = 0;
+  void setUnshapedType(Value* o);
+  void setUnshapedType(Node* node);
+  std::shared_ptr<Graph> graph_;
+};
+
+TORCH_API void EraseShapeInformation(const std::shared_ptr<Graph>& graph);
+TORCH_API void PropagateInputShapes(const std::shared_ptr<Graph>& graph);
+
+TORCH_API bool mergeTypes(
+    ArrayRef<Value*> lhs,
+    ArrayRef<Value*> rhs,
+    ArrayRef<Value*> outputs);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/specialize_autogradzero.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/specialize_autogradzero.h
new file mode 100644
index 0000000000000000000000000000000000000000..37e8ce67eb967be9800bbf9eb16d8265de80adc2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/specialize_autogradzero.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// propagate autograd zero information through a gradient graph and
+// remove grad_of blocks if present.
+// Note: this is a very limited pass. It only propagates autograd zeros for
+// operations generated by the symbolic autodiff code and cleans up
+// AutogradAdds when possible. Outputs of other nodes are conservatively
+// marked Unknown and not optimized.
+TORCH_API void specializeAutogradZero(std::shared_ptr<Graph> g);
+
+struct ProfilingRecord;
+
+TORCH_API void InsertProfileNodesForSpecializeAutogradZero(ProfilingRecord* pr);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/subgraph_rewrite.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/subgraph_rewrite.h
new file mode 100644
index 0000000000000000000000000000000000000000..596089b2ec3051f1e4f741cf32e6001e94cb2e59
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/subgraph_rewrite.h
@@ -0,0 +1,119 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/** This file defines API for pattern-based subgraph rewrites.
+ *
+ * The API can be used for finding concrete patterns in the model and replacing
+ * the corresponding subgraphs with another subgraph. A special case of such
+ * rewrites is fusion, where the new subgraph consists of just a single node.
+ *
+ * There is a default set of the most common patterns that everyone could use.
+ * Alternatively, an arbitrary pattern can be registered.
+ */
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <functional>
+#include <unordered_set>
+#include <vector>
+
+namespace torch::jit {
+
+/** Rewrite pattern descriptor.
+ *
+ * This structure is used in the implementation of `SubgraphRewriter` and
+ * is not supposed to be used externally.
+ */
+struct RewritePatternDescr {
+  std::string pattern;
+  std::string replacement;
+  std::unordered_map<std::string, std::string> value_name_map;
+};
+
+// Forward declarations.
+struct Match;
+
+using MatchFilter = std::function<
+    bool(const Match&, const std::unordered_map<std::string, Value*>&)>;
+
+/** Run pattern-based subgraph rewrites on all methods in the module.
+ *
+ * This pass will go through all methods in the module and try to replace all
+ * recognized patterns (see SubgraphRewriter::RegisterDefaultPatterns for the
+ * list of these patterns).
+ */
+TORCH_API Module PatternBasedRewrite(const Module& module);
+
+/** A class implementing API for pattern-based subgraph rewrites.
+ *
+ * To perform pattern-based subgraph rewrites on a module using this API, one
+ * needs to create an object of such class, register rewrite patterns and run
+ * the transformation pass (`runOnModule`).
+ *
+ * To use standard patterns, one could use `RegisterDefaultPatterns`.
+ *
+ * To enable rewrites of custom patterns, the custom patterns must be registered
+ * with `RegisterRewritePattern`.
+ */
+class TORCH_API SubgraphRewriter {
+ public:
+  // Run pattern-based subgraph rewrite pass on the module.
+  Module runOnModule(const Module& module);
+
+  // Run pattern-based subgraph rewrite pass on the graph (used in testing).
+  // `filter` is a function that does extra filtering on the match. If it
+  // returns false for a given Match, we'll skip the Match. The filter
+  // function's arguments consist of a Match and a value map from parsing the
+  // pattern graph. Both the Match and the value map are necessary because we
+  // need to 1) do extra filtering on the matched result as well as 2) refer to
+  // the values in the matched result through the values in the pattern graph.
+  void runOnGraph(
+      std::shared_ptr<Graph>& graph,
+      const std::vector<MatchFilter>& filters);
+
+  void runOnGraph(
+      std::shared_ptr<Graph>& graph,
+      const MatchFilter& filter =
+          [](const Match&, const std::unordered_map<std::string, Value*>&) {
+            return true;
+          }) {
+    runOnGraph(graph, std::vector<MatchFilter>({filter}));
+  }
+
+  // Register standard rewrite patterns.
+  void RegisterDefaultPatterns();
+
+  /** Register a custom rewrite pattern.
+   *
+   * The method takes two parameters specifying the pattern:
+   * \p PATTERN - IR string representing the pattern subgraph.
+   * \p REPLACEMENT - IR string representing the replacement subgraph.
+   * \p value name map - vector of pairs mapping values in the replacement graph
+   * to the values in the pattern graph. Used for preserving source range info
+   * across graph rewrite.
+   *
+   * See examples of pattern registering in `RegisterDefaultPatterns`.
+   */
+  void RegisterRewritePattern(
+      const std::string& pattern,
+      const std::string& replacement,
+      const std::vector<std::pair<std::string, std::string>>& value_name_pair =
+          {});
+
+ private:
+  std::vector<RewritePatternDescr> patterns_;
+  std::unordered_set<Node*> nodes_to_delete_;
+
+  void rewriteSinglePatternOnGraph(
+      std::shared_ptr<Graph>& graph,
+      const RewritePatternDescr& pattern,
+      const std::vector<MatchFilter>& filters);
+
+  bool overlapsWithPreviousMatches(const Match* match);
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..9afac592de18fb5350a34f97c8bda30a16293671
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_analysis.h
@@ -0,0 +1,61 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <unordered_map>
+#include <utility>
+#include <variant>
+
+namespace torch::jit {
+
+// CAUTION NOT TO BE USED, STILL A WIP, NOT STABLE
+
+TORCH_API void PropagateShapesOnGraph(std::shared_ptr<Graph>& graph);
+
+// CAUTION NOT TO BE USED, STILL A WIP, NOT STABLE
+// From [beg, end) attempt to propagate shapes and
+// build up a graph that will compute all remaining symbolic
+// shapes in [beg, end) that can be executed before beg
+
+struct ShapeComputeGraphMapping {
+  ShapeComputeGraphMapping(
+      std::shared_ptr<Graph> partial_eval_shape_graph,
+      std::unordered_map<Value*, Value*>
+          enclosing_graph_value_to_shape_graph_input,
+      std::unordered_map<Value*, int64_t> graph_output_to_symbolic_shape_dim)
+      : partial_eval_shape_graph(std::move(partial_eval_shape_graph)),
+        enclosing_graph_value_to_shape_graph_input_(
+            std::move(enclosing_graph_value_to_shape_graph_input)),
+        graph_output_to_symbolic_shape_dim_(
+            std::move(graph_output_to_symbolic_shape_dim)) {}
+
+  std::shared_ptr<Graph> partial_eval_shape_graph;
+  std::unordered_map<Value*, Value*>
+      enclosing_graph_value_to_shape_graph_input_;
+  std::unordered_map<Value*, int64_t> graph_output_to_symbolic_shape_dim_;
+};
+
+TORCH_API std::optional<ShapeComputeGraphMapping>
+PropagateShapesAndBuildLargeShapeComputeGraph(
+    std::shared_ptr<Graph>& graph,
+    Node* beg,
+    Node* end);
+
+// don't insert complete tensor shapes in shape compute graphs and instead
+// rely on our partial evaluation pipeline to propagate information.
+// this is a good proxy for our ability to propagate non-complete shape
+// information.
+TORCH_API bool setSymbolicShapeAnalysisTestMode(bool value);
+TORCH_API bool symbolicShapeAnalysisTestModeEnabled();
+
+using SSAInput = std::variant<IValue, c10::SymbolicShape>;
+TORCH_API std::optional<std::vector<c10::SymbolicShape>>
+calculateSymbolicShapesOnOp(
+    const FunctionSchema* schema,
+    const std::vector<SSAInput>& inputs);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_cache.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_cache.h
new file mode 100644
index 0000000000000000000000000000000000000000..df37fe446397e6e8419b5b3d9eb21fa1e5b58bf8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_cache.h
@@ -0,0 +1,60 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/symbolic_shape_analysis.h>
+
+namespace torch::jit {
+
+struct TORCH_API CanonicalizedSymbolicShape {
+  // TODO: Consider in the future if it is reasonable to
+  // merge code with SymbolicShape or VaryingShape while keeping
+  // the two not implicitly convertible (and cause bugs).
+  CanonicalizedSymbolicShape(
+      const c10::SymbolicShape& orig_shape,
+      std::unordered_map<int64_t, int64_t>& ss_map) {
+    init(orig_shape, ss_map);
+  }
+
+  CanonicalizedSymbolicShape(c10::SymbolicShape& orig_shape) {
+    std::unordered_map<int64_t, int64_t> new_ssmap;
+    init(orig_shape, new_ssmap);
+  }
+
+  size_t hash() const;
+
+  c10::SymbolicShape toSymbolicShape(
+      std::unordered_map<int64_t, int64_t>& inverse_ss_map) const;
+
+  TORCH_API friend bool operator==(
+      const CanonicalizedSymbolicShape& a,
+      const CanonicalizedSymbolicShape& b);
+
+ private:
+  std::optional<std::vector<int64_t>> values_;
+
+  void init(
+      const c10::SymbolicShape& orig_shape,
+      std::unordered_map<int64_t, int64_t>& ss_map);
+};
+
+// SHAPE CACHE API
+TORCH_API std::optional<std::vector<at::SymbolicShape>>
+get_cached_shape_function(
+    const FunctionSchema* schema,
+    const std::vector<SSAInput>& arg_vec);
+
+TORCH_API void cache_shape_function(
+    const FunctionSchema* schema,
+    const std::vector<SSAInput>& arg_vec,
+    const std::vector<at::SymbolicShape>& ret_vec);
+
+// For use in test code
+TORCH_API void clear_shape_cache();
+TORCH_API size_t get_shape_cache_size();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.h
new file mode 100644
index 0000000000000000000000000000000000000000..3c356073c2d8013000cdc4858c4b54fd1b46d4a5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.h
@@ -0,0 +1,56 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/symbolic_shape_analysis.h>
+
+namespace torch::jit {
+
+// Takes in a TensorExprGraph of static shapes and generalizes the input shapes
+// to symbolic dimensions. Dimensions of value 1 will be preserved, otherwise
+// dimensions with the same value will be bucketed to the same symbolic shape.
+// E.g. Tensor(5, 3), Tensor(3, 1) -> Tensor(SS(-1), SS(-2)), Tensor(SS(-2), 1)
+// From there, runs symbolic shape inference on the graph, and creates a
+// versioning if in the graph with prim::TensorExprDynamicGuard checking if
+// the inputs at runtime match the Generalized Symbolic Shapes that are inputs
+// to the TE Kernel. The computate to calculate all symbolic dimensions is
+// inlined in to the if block with the TE Kernel. All Sym Dim Value* are
+// appended to the end of the TE Kernel Graph/Node inputs, and the Node is
+// augmented with a integer list attr `symbolic_shape_inputs` that gives the
+// mapping from Value * -> Symbolic Shape int64_t value. For more lengthy IR
+// examples and walkthrough look at ShapeAnalysisTest.DynamicShapesFusion in
+// `test_shape_analysis` Returns True on Success, False on Failure, can fail if
+// shape propagation fails to propagate # of dims or if complete shapes on
+// inputs not set
+
+TORCH_API bool GenerateGuard(
+    Node* tensorexpr_graph_node,
+    bool add_composed_op = false);
+
+TORCH_API void runTensorExprDynamicGroup(const Code& code, Stack& stack);
+
+enum class StrideInput {
+  // Tensors natively store whether they are contiguous or not as a property
+  // this makes it faster to query `is_contiguous` or
+  // `is_contiguous(memory_format=channels_last)`
+  // than looping through the sizes/strides yourself
+  // For tensors with these properties, we only store one value:
+  TENSOR_CONT,
+  TENSOR_CONT_CHANNELS_LAST,
+  // now, we describe other cases, where there is one stride enum
+  // per dimension
+  S_ONE, // STRIDE_ONE: packed
+  S_CONT, // STRIDE_CONTIGUOUS: stride[i + 1] * sizes[i + 1]
+  S_TRAN_CONT, // STRIDE_TRANSPOSED_CONTIGUOUS: stride[i-1] * sizes[i-1]
+  S_AS_ARG, // STRIDE_AS_ARG: stride passed in as runtime value
+};
+
+TORCH_API std::string toString(StrideInput si);
+TORCH_API StrideInput strideInputFromString(const std::string& si);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/tensorexpr_fuser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/tensorexpr_fuser.h
new file mode 100644
index 0000000000000000000000000000000000000000..a54e6c6e88c5610516892c39a9bd793666590cd8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/tensorexpr_fuser.h
@@ -0,0 +1,82 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <memory>
+
+namespace torch::jit {
+
+// Run TensorExpressions-based fuser.
+// If add_composed_op is true, creates a single operation that
+// performs both the runtime check that types align
+// and then the dispatch to the kernel/unoptimized graph
+TORCH_API void FuseTensorExprs(
+    std::shared_ptr<Graph>& graph,
+    size_t min_group_size = 2,
+    bool add_composed_op = false,
+    bool fuse_to_dynamic_shapes = false);
+
+TORCH_API void setTensorExprFuserEnabled(bool val);
+TORCH_API bool tensorExprFuserEnabled();
+TORCH_API void setTensorExprDynamicShapeFusionEnabled(bool val);
+TORCH_API bool tensorExprDynamicShapeFusionEnabled();
+TORCH_API bool setTexprReductionsEnabled(bool value);
+TORCH_API bool texprReductionsEnabled();
+
+TORCH_API void RemoveProfileNodesAndSpecializeTypes(
+    std::shared_ptr<Graph>& graph);
+TORCH_API bool hasTensorTypeSpecialization(Value* v);
+TORCH_API void RemoveTensorTypeSpecializations(std::shared_ptr<Graph>& graph);
+TORCH_API void removeTensorTypeSpecializations(Block* block);
+
+using tensor_type_converter_t =
+    c10::function_ref<TensorTypePtr(const TensorTypePtr& t)>;
+
+// inserts a TypeCheck pattern
+//
+// around the guarded node that has a Subgraph attribute, this inserts a pattern
+//
+//   if TypeCheck(...):
+//     guarded_node
+//   else:
+//     FallbackGraph(...)
+//
+// The TypeCheck includes the types of all Tensor inputs to the guarded_node,
+// as processed by the type_converter, a lambda
+// TensorTypePtr(const TensorTypePtr& t). This allows to erase irrelevant
+// aspects of the type.
+//
+// The Fallback graph will have the same subgraph as the guarded node (with the
+// expectation that the guarded_node's subgraph will then be optimized.
+TORCH_API void insertTypeGuard(
+    Node* guarded_node,
+    tensor_type_converter_t type_converter,
+    c10::Symbol kind);
+
+TORCH_API bool usedOnlyInSize(Value* v);
+TORCH_API Value* broadcastSizes(at::ArrayRef<Value*> sizes, AliasDb* db);
+
+namespace tensorexpr {
+TORCH_API bool isSupported(Node* node);
+
+/// Get the modifiable custom operator set object.
+///
+/// For static shapes, if a custom operator has been added to the custom
+/// operator set, it will be pulled into the NNC fusion group. But it doesn't
+/// work with dynamic shapes unless explicitly register the shape function via
+/// `torch::jit::RegisterShapeComputeGraphForSchema` for the custom operator.
+///
+/// @return Reference of the custom operator set
+///
+TORCH_API OperatorSet& getCustomOperatorSet();
+
+} // namespace tensorexpr
+} // namespace torch::jit
+
+C10_DECLARE_bool(torch_jit_disable_cat);
+C10_DECLARE_bool(torch_jit_enable_dynamic_shape_fusion);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/update_differentiable_graph_requires_grad.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/update_differentiable_graph_requires_grad.h
new file mode 100644
index 0000000000000000000000000000000000000000..b1721a4714ba2243e104dca50346f1ba79493930
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/update_differentiable_graph_requires_grad.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Because differentiable graphs detach the gradients of input Tensors,
+// creating and inlining differentiable graphs changes the requires_grad
+// property of tensors in the graph. This pass updates prim::profiles
+// requires_grad to keep profiled properties up to date, it does not update
+// grad properties of other nodes like graph inputs bc the only downstream
+// user of the grad property is the profiling executor, which just uses
+// the types of prim::profiles
+TORCH_API void UpdateDifferentiableGraphRequiresGrad(
+    std::shared_ptr<Graph>& diff_forward_graph,
+    std::optional<bool> new_requires_grad);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/check_alias_annotation.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/check_alias_annotation.h
new file mode 100644
index 0000000000000000000000000000000000000000..9f038180f3b85dbba36ac10b811edc76733f8ea8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/check_alias_annotation.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace torch::jit {
+
+// Verify that alias annotations are correct. See impl for definition of
+// "correct".
+//
+// This function expects a graph with a single op with `unqualifiedOpName`, plus
+// the inputs that you would otherwise have passed to the graph executor.
+TORCH_API void checkAliasAnnotation(
+    const std::shared_ptr<Graph>& graph,
+    std::vector<IValue> pythonInputs,
+    const std::string& unqualifiedOpName);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/memory_dag.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/memory_dag.h
new file mode 100644
index 0000000000000000000000000000000000000000..755aead9375fbc1a563bb555871973755dfa8c5c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/memory_dag.h
@@ -0,0 +1,179 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/jit_type.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/flat_hash_map.h>
+#include <c10/util/sparse_bitset.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/ir/type_hashing.h>
+#include <memory>
+#include <optional>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include <torch/csrc/Export.h>
+
+// Uses a compressed index representation for faster comparisons
+typedef c10::SparseBitVector<256> MemoryLocations;
+namespace torch::jit {
+
+struct Value;
+
+using AliasTypeSet = std::vector<TypePtr>;
+
+// `Element` represents a vertex in the points-to graph. It represents
+// anything that could have an aliasing relationship--mostly IR
+// `Value`s, but also wildcards or the type inside a container (e.g. `T`
+// in `List[T]`)
+struct Element {
+  Element(const Value* value_, unsigned index_);
+  // wildcard constructor
+  explicit Element(unsigned index_);
+
+  // Index into the owning DAG's bit vector that represents this element.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  unsigned index;
+
+  // All elements that this element *may* point to. It's possible to have
+  // multiple elements that you might point to due to control flow/complex ops
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  MemoryLocations pointsTo;
+  // Backreference for points-to.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  MemoryLocations pointedFrom;
+
+  // Elements can contain other elements (e.g. List[Tensor])
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  MemoryLocations containedElements;
+
+  // The values that this element corresponds to. May be empty if this element
+  // doesn't represent a first-class value.
+  // This is for debug information only.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::unordered_set<const Value*> values;
+
+ private:
+  // Make `from` point at `to`.
+  void makePointerTo(Element* from, Element* to);
+
+  friend class MemoryDAG;
+  // We memoize the results of `getMemoryLocations` to speed up queries.
+  // A nullopt means that this cache is not yet populated. Since `MemoryDAG` is
+  // immutable, this cache should never need to be invalidated.
+  mutable std::optional<MemoryLocations> cachedMemoryLocations_;
+
+  mutable std::optional<MemoryLocations> cachedAllContainedMemoryLocations_;
+};
+
+// class MemoryDAG
+//
+// This class tracks the "A points to B" graph for all values. It is used by
+// AliasDb to provide a higher-level API.
+//
+// We maintain a DAG where:
+//   - Vertices (called "Elements") represent Values and
+//     other aliasing entities (e.g. the stuff inside a list)
+//   - Edges represent a "points-to" relationship.
+//
+// Leaves in this DAG are entities that don't point to anything, and thus
+// correspond to unique "memory locations".
+//
+// So, by traversing the "points-to" graph to the leaves, you can determine
+// which memory locations an element may point to.
+class TORCH_API MemoryDAG {
+ public:
+  explicit MemoryDAG(std::vector<std::unique_ptr<Element>> indexToElementMap)
+      : indexToElementMap_(std::move(indexToElementMap)) {}
+  // explicitly delete copy constructor because otherwise windows build is
+  // confused for an exported class see
+  // https://stackoverflow.com/a/51033485/105137
+  MemoryDAG(const MemoryDAG&) = delete;
+  MemoryDAG& operator=(const MemoryDAG&) = delete;
+
+  // Return the unique memory locations that `Element` might represent.
+  const MemoryLocations& getMemoryLocations(const Element* e) const;
+
+  // Do `a` and `b` potentially share a memory location?
+  bool mayAlias(const Element* a, const Element* b) const;
+
+  // Does `a` hold reference to any memory that is stored in `b`, or vice versa?
+  bool mayContainAlias(const Element* a, const Element* b) const;
+
+  bool mayContainAlias(const Element* a, const at::ArrayRef<Element*> b) const;
+
+  bool mayContainAlias(
+      const at::ArrayRef<Element*> a,
+      const at::ArrayRef<Element*> b) const;
+
+  // Converts from the compressed index representation
+  const Element* fromIndex(unsigned x) const;
+  Element* fromIndex(unsigned x);
+  void collectAllContainedMemoryLocations(
+      const Element* elem,
+      MemoryLocations& cont) const;
+
+  /**
+   * The following methods are special cases where we need to mutate the
+   * internals of MemoryDAG for efficiency reasons. Don't call them unless you
+   * know what you're doing! In particular, don't add new mutating methods
+   * without ensuring that you are maintaining cache consistency for memory
+   * locations.
+   */
+
+  // Adding wildcards can trigger extremely expensive cache invalidations. This
+  // method adds them in a more efficient cache-aware way.
+  void setWildcards(
+      const std::unordered_set<const Value*>& wildcards,
+      const ska::flat_hash_map<const Value*, Element*>& elementMap,
+      const std::function<Element*(const Value*)>& getWildcardElement);
+  Element* unsafeMakeFreshValue(const Value* v);
+
+ private:
+  const MemoryLocations& getAllContainedMemoryLocations(
+      const Element* elem) const;
+  void collectAllContainedMemoryLocationsImpl(
+      const Element* elem,
+      MemoryLocations& cont) const;
+  std::vector<std::unique_ptr<Element>> indexToElementMap_;
+};
+
+/**
+ * Helper to build up the points-to graph.
+ *
+ * We separate the "building" into a different class because it allows us to
+ * cache internally to MemoryDAG without worrying about how the DAG structure
+ * is mutated.
+ */
+class TORCH_API MemoryDAGBuilder {
+ public:
+  MemoryDAGBuilder() = default;
+  MemoryDAGBuilder(const MemoryDAGBuilder&) = delete;
+  MemoryDAGBuilder& operator=(const MemoryDAGBuilder&) = delete;
+
+  // Make `from` point at `to`.
+  void makePointerTo(Element* from, Element* to);
+
+  void addToContainedElements(Element* contained, Element* container);
+
+  std::unique_ptr<MemoryDAG> createMemoryDAG() && {
+    return std::make_unique<MemoryDAG>(std::move(indexToElementMap_));
+  }
+
+  // Make a fresh Element (i.e. an Element that doesn't point to anything) and
+  // return it.
+  Element* makeFreshValue(const Value* v);
+
+  friend MemoryDAG;
+
+ private:
+  // `MemoryDAGBuilder` builds up `indexToElementMap_`, then uses
+  // the map to construct the `MemoryDAG`
+  std::vector<std::unique_ptr<Element>> indexToElementMap_;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/op_registry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/op_registry.h
new file mode 100644
index 0000000000000000000000000000000000000000..d6a566d1ef06843ff0d27b6b07b42970f8c89825
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/op_registry.h
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <memory>
+
+namespace torch::jit {
+// Moved from shape_analysis.cpp
+
+// Requirements:
+//   dims           : preserved from the first argument
+//   scalar type    : preserved from the first argument (doesn't have to
+//                    match other arguments)
+//   device         : always matching and preserved
+//   tensor inputs  : *
+//   tensor outputs : 1
+// NB: those ops (with slight adjustments) are good candidates for restarts.
+//     Knowing the type and device of weights or biases is usually enough to
+//     infer the output type.
+std::shared_ptr<OperatorSet> nn_ops_first_input_preserving();
+
+// Requirements:
+//   dims           : Changed from first argument
+//   scalar type    : preserved from the first argument
+//   device         : always matching and preserved
+//   tensor inputs  : 1
+//   tensor outputs : 1
+std::shared_ptr<OperatorSet> ops_one_tensor_in_shape_transform();
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/optimization_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/optimization_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..189d8128354445927ca03f457ca51951e81453e3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/optimization_utils.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Checks if the parameters, not including the
+// first param are all constants.
+bool nonConstantParameters(Node* n);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/subgraph_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/subgraph_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..eb84520636361f0353ddafbb0d27d12c6de0df58
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/utils/subgraph_utils.h
@@ -0,0 +1,75 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/alias_analysis.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+// Utilities for dealing with nodes that contain subgraphs.
+//
+// They handle the complexity of editing inputs/outputs as you merge nodes in
+// and out of subgraphs.
+namespace torch::jit::SubgraphUtils {
+
+// Create a new subgraph node that contains only `n`. The new subgraph will have
+// `subgraphKind` as its type.
+//
+// `n` is destroyed.
+//
+// Returns the new subgraph node.
+TORCH_API Node* createSingletonSubgraph(Node* n, Symbol subgraphKind);
+
+// Creates a new subgraph that only contains `n`, amd updates the new outputs
+// of the subgraph to have the aliasing properties of the original `n` outputs
+TORCH_API Node* createSingletonSubgraphAndUpdateAliasing(
+    Node* to_merge,
+    Symbol subgraphKind,
+    AliasDb& db);
+
+// Merge a node into a subgraph node. If `toMerge` is also a subgraph, the
+// subgraphs are merged.
+// If `destroyNode` is true `toMerge` is destroyed.
+// An optional argument 'vmap' could be used to retrieve value mappings.
+// Values will be mapped to their new subgraph values
+TORCH_API void mergeNodeIntoSubgraph(
+    Node* toMerge,
+    Node* subgraphNode,
+    bool destroyNode = true);
+
+// Merges a node into a subgraph node, and updates the new outputs of the
+// subgraph to have the aliasing properties of the corresponding `to_merge`
+// outputs
+TORCH_API void mergeNodeIntoSubgraphAndUpdateAliasing(
+    Node* to_merge,
+    Node* subgraphNode,
+    AliasDb& db);
+
+TORCH_API std::vector<Node*> unmergeAliasedOutputs(
+    Node* subgraphNode,
+    AliasDb& db);
+
+// Move nodes from a subgraph node to the outer graph.
+// `subgraphNode` is destroyed.
+TORCH_API void unmergeSubgraph(Node* subgraphNode);
+
+// Move `node_to_unmerge` and its descendants after `subgraphNode`
+// promotes any dependencies of `node_to_unmerge` to subgraphNode outputs
+TORCH_API void unmergeNode(Node* node_to_unmerge, Node* subgraphNode);
+
+TORCH_API bool unmergeOutputsAlisingInputs(Node* subgraphNode);
+
+TORCH_API bool unmergeAliasedOutputs(Node* subgraphNode);
+
+// Convenience function
+std::shared_ptr<Graph> getSubgraph(Node* n);
+
+TORCH_API std::string generateNameForGraph(
+    const std::shared_ptr<Graph>& graph,
+    size_t maxlen = 40,
+    const std::string& prefix = "fused");
+
+} // namespace torch::jit::SubgraphUtils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/value_refinement_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/value_refinement_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..87922b88f957eed96baaad5b6443435d96af0157
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/value_refinement_utils.h
@@ -0,0 +1,84 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/jit/ir/alias_analysis.h>
+#include <torch/csrc/jit/ir/ir_views.h>
+#include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/passes/dead_code_elimination.h>
+#include <torch/csrc/jit/passes/peephole.h>
+#include <torch/csrc/jit/passes/peephole_list_idioms.h>
+#include <torch/csrc/jit/runtime/graph_executor.h>
+
+namespace torch::jit {
+
+// Refine from Value of type List -> len of list
+// If a refinement mapping of List Value * -> len is present in a block
+// the list is guaranteed to be that length
+// TODO: vector may be faster
+using ListRefinement = std::unordered_map<Value*, int64_t>;
+
+TORCH_API ListRefinement
+intersectRefinements(const ListRefinement& ref1, const ListRefinement& ref2);
+
+TORCH_API ListRefinement
+unionRefinements(const ListRefinement& ref1, const ListRefinement& ref2);
+
+// Represents the refinement information that can be carried on a boolean
+struct BooleanRefinementMapping {
+  BooleanRefinementMapping(
+      ListRefinement true_refine,
+      ListRefinement false_refine)
+      : true_refine_(std::move(true_refine)),
+        false_refine_(std::move(false_refine)) {}
+  BooleanRefinementMapping() = default; // empty
+
+  static BooleanRefinementMapping FalseRefinements(
+      ListRefinement false_refine) {
+    return BooleanRefinementMapping({}, std::move(false_refine));
+  }
+
+  static BooleanRefinementMapping TrueRefinements(ListRefinement true_refine) {
+    return BooleanRefinementMapping(std::move(true_refine), {});
+  }
+
+  BooleanRefinementMapping intersectBooleanRefinementMapping(
+      BooleanRefinementMapping& other) {
+    return BooleanRefinementMapping(
+        intersectRefinements(true_refine_, other.true_refine()),
+        intersectRefinements(false_refine_, other.false_refine()));
+  }
+
+  ListRefinement& true_refine() {
+    return true_refine_;
+  }
+
+  ListRefinement& false_refine() {
+    return false_refine_;
+  }
+
+ private:
+  ListRefinement true_refine_;
+  ListRefinement false_refine_;
+};
+
+TORCH_API void joinIfRefinements(
+    Node* if_node,
+    std::unordered_set<Block*>& throwing_blocks,
+    ListRefinement& curr_block_refinements,
+    ListRefinement& true_block_refinements,
+    ListRefinement& false_block_refinements,
+    std::unordered_map<Value*, BooleanRefinementMapping>& info);
+
+// handles adding blocks to throwing blocks and propagating refinements via
+// boolean comparisons
+TORCH_API bool handleCommonRefinentOperators(
+    Node* n,
+    std::unordered_set<Block*>& throwing_blocks,
+    std::unordered_map<Value*, BooleanRefinementMapping>& info);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/variadic_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/variadic_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..a017a91293e330a0b9afc7029acdd0463b55193f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/variadic_ops.h
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// Try to replace an op that takes a list input with another op that takes a
+// variadic number of arguments.
+TORCH_API bool UseVariadicOp(
+    const std::shared_ptr<Graph>& graph,
+    NodeKind op,
+    NodeKind variadic_op);
+
+TORCH_API bool RemoveListMutationAndUseVariadicOp(
+    const std::shared_ptr<Graph>& graph,
+    NodeKind op,
+    NodeKind variadic_op);
+
+// Convenient functions for replacing aten::stack/aten::cat with their
+// variadic versions.
+TORCH_API bool UseVariadicCat(const std::shared_ptr<Graph>& graph);
+TORCH_API bool RemoveListMutationAndUseVariadicCat(
+    const std::shared_ptr<Graph>& graph);
+
+TORCH_API bool UseVariadicStack(const std::shared_ptr<Graph>& graph);
+TORCH_API bool RemoveListMutationAndUseVariadicStack(
+    const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/vulkan_rewrite.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/vulkan_rewrite.h
new file mode 100644
index 0000000000000000000000000000000000000000..c1dbb8d23e58a416dea4c95d5e45df787b51025e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/vulkan_rewrite.h
@@ -0,0 +1,21 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/mobile_optimizer_type.h>
+
+namespace torch::jit {
+TORCH_API void vulkanInsertPrePackedOps(std::shared_ptr<Graph>& graph);
+TORCH_API void vulkanInsertPrePackedOps(script::Module& module);
+TORCH_API void vulkanFusePrePackedConvWithClamp(script::Module& module);
+TORCH_API void vulkanFoldPrePackingOps(script::Module& module);
+TORCH_API script::Module vulkanOptimizeForMobile(
+    const script::Module& module,
+    const std::set<MobileOptimizerType>& optimization_blocklist,
+    const std::vector<std::string>& preserved_methods);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/xnnpack_rewrite.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/xnnpack_rewrite.h
new file mode 100644
index 0000000000000000000000000000000000000000..39e08ec04f9ae75eba189a644f028d4e05fe56a4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/passes/xnnpack_rewrite.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/mobile_optimizer_type.h>
+
+namespace torch::jit {
+
+TORCH_API void transformConv1dToConv2d(std::shared_ptr<Graph>& graph);
+TORCH_API void transformConv1dToConv2d(script::Module& module);
+TORCH_API void insertPrePackedOps(std::shared_ptr<Graph>& graph);
+TORCH_API void insertPrePackedOps(script::Module& module);
+TORCH_API void fusePrePackedLinearConvWithClamp(script::Module& module);
+TORCH_API void FoldPrePackingOps(script::Module& module);
+TORCH_API script::Module optimizeForMobile(
+    const script::Module& module,
+    const std::set<MobileOptimizerType>& optimization_blocklist = {},
+    const std::vector<std::string>& preserved_methods = {});
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..29180349f5aade9591e36fe892eabf91cb92d064
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/init.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::jit {
+
+void initJITBindings(PyObject* module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/module_python.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/module_python.h
new file mode 100644
index 0000000000000000000000000000000000000000..c127ff8b5a6ac87672a7a1e77158351b9d5dc591
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/module_python.h
@@ -0,0 +1,69 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/utils/pybind.h>
+#include <tuple>
+
+namespace py = pybind11;
+
+namespace torch::jit {
+
+inline std::optional<Module> as_module(py::handle obj) {
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+      storage;
+  auto& ScriptModule =
+      storage
+          .call_once_and_store_result([]() -> py::object {
+            return py::module_::import("torch.jit").attr("ScriptModule");
+          })
+          .get_stored();
+#else
+  static py::handle ScriptModule =
+      py::module::import("torch.jit").attr("ScriptModule");
+#endif
+  if (py::isinstance(obj, ScriptModule)) {
+    return py::cast<Module>(obj.attr("_c"));
+  }
+  return std::nullopt;
+}
+
+inline std::optional<Object> as_object(py::handle obj) {
+#if IS_PYBIND_2_13_PLUS
+  PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<
+      std::tuple<py::object, py::object>>
+      storage;
+  auto& [ScriptObject, RecursiveScriptClass] =
+      storage
+          .call_once_and_store_result(
+              []() -> std::tuple<py::object, py::object> {
+                return {
+                    py::module_::import("torch").attr("ScriptObject"),
+                    py::module_::import("torch.jit")
+                        .attr("RecursiveScriptClass")};
+              })
+          .get_stored();
+#else
+  static py::handle ScriptObject =
+      py::module::import("torch").attr("ScriptObject");
+
+  static py::handle RecursiveScriptClass =
+      py::module::import("torch.jit").attr("RecursiveScriptClass");
+#endif
+
+  if (py::isinstance(obj, ScriptObject)) {
+    return py::cast<Object>(obj);
+  }
+  if (py::isinstance(obj, RecursiveScriptClass)) {
+    return py::cast<Object>(obj.attr("_c"));
+  }
+  return std::nullopt;
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/opaque_obj.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/opaque_obj.h
new file mode 100644
index 0000000000000000000000000000000000000000..8f6b6c324b0ac2642fc38ef05a323a544cc2c1b3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/opaque_obj.h
@@ -0,0 +1,84 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <utility>
+
+#include <c10/macros/Macros.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/utils/pybind.h>
+#include <torch/custom_class.h>
+
+namespace torch::jit {
+struct OpaqueObject : public CustomClassHolder {
+  OpaqueObject(py::object payload) : payload_(std::move(payload)) {}
+
+  void setPayload(py::object payload) {
+    payload_ = std::move(payload);
+  }
+
+  py::object getPayload() {
+    return payload_;
+  }
+
+  py::object payload_;
+};
+
+static auto register_opaque_obj_class =
+    torch::class_<OpaqueObject>("aten", "OpaqueObject")
+        .def(
+            "__eq__",
+            [](const c10::intrusive_ptr<OpaqueObject>& self,
+               const c10::intrusive_ptr<OpaqueObject>& other) {
+              auto self_payload = self->getPayload();
+              auto other_payload = other->getPayload();
+
+              if (!self_payload.ptr() || !other_payload.ptr()) {
+                return false;
+              }
+
+              py::gil_scoped_acquire gil;
+              auto res = PyObject_RichCompareBool(
+                  self_payload.ptr(), other_payload.ptr(), Py_EQ);
+              if (res == -1) {
+                throw py::error_already_set();
+              }
+              return res > 0;
+            })
+        .def_pickle(
+            [](const c10::intrusive_ptr<OpaqueObject>& self) { // __getstate__
+              // Since we cannot directly return the py::object due to
+              // CustomClassHolder's signature limitations, we will have to
+              // serialize it directly here. We also can't return py::bytes so
+              // need to encode it into a string.
+              py::module_ pickle = py::module_::import("pickle");
+              py::module_ base64 = py::module_::import("base64");
+              py::bytes pickled_payload =
+                  pickle.attr("dumps")(self->getPayload());
+              py::bytes encoded_payload =
+                  base64.attr("b64encode")(pickled_payload);
+              return std::string(encoded_payload);
+            },
+            [](const std::string& state) { // __setstate__
+              py::module_ pickle = py::module_::import("pickle");
+              py::module_ base64 = py::module_::import("base64");
+              py::bytes state_bytes(state);
+              py::bytes decoded_payload = base64.attr("b64decode")(state_bytes);
+              py::object restored_payload =
+                  pickle.attr("loads")(decoded_payload);
+              return c10::make_intrusive<OpaqueObject>(restored_payload);
+            })
+        .def(
+            "__obj_flatten__",
+            [](const c10::intrusive_ptr<OpaqueObject>& self) {
+              throw std::runtime_error(
+                  "Unable to implement __obj_flatten__ for opaque objects.");
+            });
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/pybind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/pybind.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ded51cf15914dea299a1e76b7abb534153a08b5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/pybind.h
@@ -0,0 +1,218 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+#include <ATen/core/ivalue.h>
+#include <ATen/core/symbol.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/DynamicTypes.h>
+#include <torch/csrc/THP.h>
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/jit/frontend/tracer.h>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include <pybind11/functional.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+namespace torch::jit {
+
+// This is a variant of shared_ptr that "sees through" a wrapper.
+// We use it to convert Value, Node, Block and node to "wrapped" Python
+// values. When we destruct the C++ object, the wrapper's pointer will
+// be set to 0 and any future dereferencing will throw. We need this
+// because the Python objects may hang around after the C++ object
+// has already been destroyed.
+// This also needs the magic type_caster below, which is from the
+// workaround offered in https://github.com/pybind/pybind11/issues/2751
+template <typename T>
+class unwrapping_shared_ptr {
+  static_assert(
+      std::is_same_v<T, torch::jit::Value> ||
+          std::is_same_v<T, torch::jit::Node> ||
+          std::is_same_v<T, torch::jit::Block>,
+      "unwrapping type only defined for Graph object types");
+
+ private:
+  std::shared_ptr<torch::jit::Wrap<T>> impl;
+
+ public:
+  unwrapping_shared_ptr() : impl({}) {}
+  explicit unwrapping_shared_ptr(T* p) : impl(p->wrap()) {
+    impl->clear_cb = &clear_registered_instances;
+  }
+  T* get() const {
+    if (!impl->elem) {
+      throw std::logic_error("has been invalidated");
+    }
+    return impl->elem;
+  }
+  // we need to disable the overloaded & for PyBind11 < 2.3 due.
+  // see https://github.com/pybind/pybind11/pull/1435
+#if (PYBIND11_VERSION_MAJOR > 2) || \
+    ((PYBIND11_VERSION_MAJOR == 2) && (PYBIND11_VERSION_MINOR >= 3))
+  T** operator&() {
+    if (!impl->elem) {
+      throw std::logic_error("has been invalidated");
+    }
+    return &(impl->elem);
+  }
+#endif
+};
+
+} // namespace torch::jit
+
+PYBIND11_DECLARE_HOLDER_TYPE(T, torch::jit::unwrapping_shared_ptr<T>, true)
+
+namespace pybind11::detail {
+
+#define CREATE_UNWRAPPING_CASTER(Class)                                                   \
+  template <>                                                                             \
+  struct type_caster<Class> : public type_caster_base<Class> {                            \
+   public:                                                                                \
+    using type = Class;                                                                   \
+    using holder_type = torch::jit::unwrapping_shared_ptr<Class>;                         \
+                                                                                          \
+    bool load(handle src, bool convert) {                                                 \
+      return load_impl<type_caster<Class>>(src, convert);                                 \
+    }                                                                                     \
+                                                                                          \
+    explicit operator type*() {                                                           \
+      return static_cast<type*>(value);                                                   \
+    }                                                                                     \
+    explicit operator type&() {                                                           \
+      return *static_cast<type*>(value);                                                  \
+    }                                                                                     \
+                                                                                          \
+   protected:                                                                             \
+    friend class type_caster_generic;                                                     \
+                                                                                          \
+    bool load_value(const value_and_holder& v_h) {                                        \
+      if (v_h.holder_constructed()) {                                                     \
+        value = v_h.template holder<holder_type>().get();                                 \
+        return true;                                                                      \
+      } else {                                                                            \
+        throw cast_error(                                                                 \
+            "Unable to cast from non-held to held instance (#Class& to Holder<#Class>)"); \
+      }                                                                                   \
+    }                                                                                     \
+  }
+
+CREATE_UNWRAPPING_CASTER(torch::jit::Node);
+CREATE_UNWRAPPING_CASTER(torch::jit::Value);
+CREATE_UNWRAPPING_CASTER(torch::jit::Block);
+
+#undef CREATE_UNWRAPPING_CASTER
+
+template <>
+struct type_caster<torch::jit::IValue> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(torch::jit::IValue, _("IValue"));
+
+  bool load(handle src, bool /*unused*/) {
+    try {
+      value = torch::jit::toTypeInferredIValue(src);
+      return true;
+    } catch (std::exception&) {
+      return false;
+    }
+  }
+
+  static handle cast(
+      torch::jit::IValue src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return torch::jit::toPyObject(std::move(src)).release();
+  }
+};
+
+template <>
+struct type_caster<torch::jit::Symbol> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(torch::jit::Symbol, _("Symbol"));
+
+  bool load(handle src, bool /*unused*/) {
+    // TODO: Is there a way to py::cast that doesn't raise an exception on
+    // failure?  Can we catch pybind11::cast_error here instead?
+    std::string src_str;
+    try {
+      src_str = py::cast<std::string>(src);
+    } catch (std::exception&) {
+      return false;
+    }
+    value = torch::jit::Symbol::fromQualString(src_str);
+    return true;
+  }
+
+  static handle cast(
+      torch::jit::Symbol src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return py::cast(std::string(src.toQualString()), return_value_policy::copy)
+        .release();
+  }
+};
+
+template <>
+struct type_caster<torch::jit::AttributeKind> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(torch::jit::AttributeKind, _("AttributeKind"));
+
+  bool load(handle src, bool /*unused*/) {
+    return false;
+  }
+
+  static handle cast(
+      torch::jit::AttributeKind src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return py::cast(
+               std::string(torch::jit::toString(src)),
+               return_value_policy::copy)
+        .release();
+  }
+};
+
+// See https://github.com/pybind/pybind11/issues/637
+using ListCasterBase = pybind11::detail::
+    list_caster<std::vector<torch::jit::Node*>, torch::jit::Node*>;
+template <>
+struct type_caster<std::vector<torch::jit::Node*>> : ListCasterBase {
+  static handle cast(
+      const std::vector<torch::jit::Node*>& src,
+      return_value_policy /*unused*/,
+      handle parent) {
+    return ListCasterBase::cast(src, return_value_policy::reference, parent);
+  }
+  static handle cast(
+      const std::vector<torch::jit::Node*>* src,
+      return_value_policy pol,
+      handle parent) {
+    return cast(*src, pol, parent);
+  }
+};
+
+} // namespace pybind11::detail
+
+namespace torch::jit {
+
+static inline py::tuple tuple_tail(const py::tuple& tup) {
+  py::tuple r(tup.size() - 1);
+  for (const auto i : c10::irange(1, tup.size())) {
+    r[i - 1] = tup[i];
+  }
+  return r;
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/pybind_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/pybind_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..d7f28e7013ce99d1b8fd75b0fd7521078bfd7c30
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/pybind_utils.h
@@ -0,0 +1,1351 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/qualified_name.h>
+#include <ATen/core/stack.h>
+#include <pybind11/complex.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <torch/csrc/Device.h>
+#include <torch/csrc/Dtype.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/Layout.h>
+#include <torch/csrc/QScheme.h>
+#include <torch/csrc/Stream.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/frontend/schema_matching.h>
+#include <torch/csrc/jit/frontend/tracer.h>
+#include <torch/csrc/jit/python/module_python.h>
+#include <torch/csrc/jit/python/python_custom_class.h>
+#include <torch/csrc/jit/python/python_tracer.h>
+#include <torch/csrc/jit/resource_guard.h>
+#include <torch/csrc/jit/runtime/operator.h>
+#include <torch/csrc/utils/pybind.h>
+#include <torch/csrc/utils/python_arg_parser.h>
+#ifdef USE_DISTRIBUTED
+#include <torch/csrc/distributed/rpc/py_rref.h>
+#include <torch/csrc/distributed/rpc/rref_impl.h>
+#endif
+
+#include <ATen/core/function_schema.h>
+#include <c10/core/Stream.h>
+#include <c10/util/Exception.h>
+#include <c10/util/irange.h>
+#include <optional>
+
+#include <algorithm>
+#include <cstddef>
+#include <string>
+#include <utility>
+#include <vector>
+
+// The visibility attribute is to avoid a warning about storing a field in the
+// struct that has a different visibility (from pybind) than the struct.
+#ifdef _WIN32
+#define VISIBILITY_HIDDEN
+#else
+#define VISIBILITY_HIDDEN __attribute__((visibility("hidden")))
+#endif
+
+namespace torch::jit {
+
+using ResolutionCallback = std::function<py::object(std::string)>;
+
+void clear_registered_instances(void* ptr);
+
+TORCH_PYTHON_API IValue toIValue(
+    py::handle obj,
+    const TypePtr& type,
+    std::optional<int32_t> N = std::nullopt);
+
+TORCH_PYTHON_API py::object toPyObject(IValue ivalue);
+
+// Hack to overload the behavior of toIValue to accept Python
+// numbers in places where a Tensor is expected
+// See also torch::should_allow_numbers_as_tensors
+class TORCH_PYTHON_API ToIValueAllowNumbersAsTensors {
+  bool old_;
+
+ public:
+  ToIValueAllowNumbersAsTensors(bool enable);
+  ~ToIValueAllowNumbersAsTensors();
+};
+
+// Wrap Python function to guard deref
+// NB: Need VISIBILITY_HIDDEN for silencing compiler error,
+// 'torch::jit::PythonFunctionGuard' declared with greater visibility than the
+// type of its field 'torch::jit::PythonFunctionGuard::func_'
+struct VISIBILITY_HIDDEN PythonFunctionGuard {
+  explicit PythonFunctionGuard(py::function func) : func_(std::move(func)) {}
+  PythonFunctionGuard(const PythonFunctionGuard&) = delete;
+  PythonFunctionGuard(PythonFunctionGuard&&) = delete;
+  PythonFunctionGuard& operator=(const PythonFunctionGuard&) = delete;
+  PythonFunctionGuard& operator=(PythonFunctionGuard&&) = delete;
+
+  ~PythonFunctionGuard() {
+    pybind11::gil_scoped_acquire ag;
+    func_.dec_ref();
+    // explicitly setting PyObject* to nullptr to prevent py::object's dtor to
+    // decref on the PyObject again.
+    // See Note [Destructing py::object] in python_ivalue.h
+    func_.ptr() = nullptr;
+  }
+
+  py::function func_;
+};
+
+// The PythonFutureWrapper for ivalue::Future
+//
+// NB: VISIBILITY_HIDDEN is for silencing compiling error,
+// "error: 'torch::jit::PythonFutureWrapper' declared with greater visibility
+// than the type of its field 'torch::jit::PythonFutureWrapper::unwrap_func'
+// [-Werror=attributes]"
+//
+// NB: inherit from enable_shared_from_this because then(py::function) needs to
+//     get a shared_ptr from this pointer.
+struct VISIBILITY_HIDDEN PythonFutureWrapper
+    : std::enable_shared_from_this<PythonFutureWrapper> {
+  using UnwrapFunc = std::function<void(py::object)>;
+
+  explicit PythonFutureWrapper(
+      c10::intrusive_ptr<c10::ivalue::Future> fut,
+      std::optional<UnwrapFunc> unwrap_func = std::nullopt)
+      : fut(std::move(fut)), unwrap_func(std::move(unwrap_func)) {}
+
+  explicit PythonFutureWrapper(const PythonFutureWrapper&) = delete;
+  PythonFutureWrapper& operator=(const PythonFutureWrapper&) = delete;
+  PythonFutureWrapper(PythonFutureWrapper&&) = default;
+  PythonFutureWrapper& operator=(PythonFutureWrapper&&) = default;
+  ~PythonFutureWrapper() = default;
+
+  bool done() {
+    return fut->completed();
+  }
+
+  py::object value() {
+    // acquiring GIL as toPyObject creates new py::object
+    // without grabbing the GIL.
+    py::gil_scoped_acquire acquire;
+    py::object py_obj = toPyObject(fut->value());
+    // unwrap_func is a general compositional function that takes in a
+    // py::object and executes some python function. It is currently mostly used
+    // to throw python exceptions.
+    if (unwrap_func) {
+      (*unwrap_func)(py_obj);
+    }
+    return py_obj;
+  }
+
+  py::object wait() {
+    fut->wait();
+    if (jit::tracer::isTracing()) {
+      auto graph = jit::tracer::getTracingState()->graph;
+
+      Value* fut_val = jit::tracer::getValueTrace(fut);
+      auto output = graph->insert(aten::wait, {fut_val});
+      jit::tracer::setValueTrace(fut->value(), output);
+    }
+    return value();
+  }
+
+  // The py::function cb arg must take a std::shared_ptr<PythonFutureWrapper>
+  // (i.e., torch._C.Future) as the only argument. If the type mismatches, an
+  // error will be thrown when waiting for the value of this returned Future.
+  std::shared_ptr<PythonFutureWrapper> then(py::function cb) {
+    // We need this an additional layer of wrapper here to guard the
+    // destruction of the py::function object. Because, the
+    // Future owns a reference to the py::function in its callback
+    // vector, but Future does not acquire GIL on destruction.
+    auto pf = std::make_shared<PythonFunctionGuard>(std::move(cb));
+
+    return std::make_shared<jit::PythonFutureWrapper>(fut->then(
+        // Capture a copy of the ivalue::Future instead of the `this` pointer
+        // because the PythonFutureWrapper object could have been deleted
+        // when the callbacks are fired. For example, RPC only captures the
+        // ivalue::Future instead of PythonFutureWrapper in JitFuture's
+        // callback functions. Hence, if user code does not hold a reference to
+        // this PythonFutureWrapper object, there is no guarantee that the
+        // PythonFutureWrapper is still valid when running the callback.
+        [pyFut(this->getPtr()),
+         pf(std::move(pf))](c10::ivalue::Future& /* unused */) -> IValue {
+          try {
+            pybind11::gil_scoped_acquire ag;
+            return toIValue(pf->func_(pyFut), PyObjectType::get());
+          } catch (py::error_already_set& e) {
+            auto err = std::runtime_error(c10::str(
+                "Got the following error when running the callback: ",
+                e.what()));
+            {
+              pybind11::gil_scoped_acquire ag;
+              // Release ownership on py::objects and also restore Python
+              // Error Indicator.
+              e.restore();
+              // Clear the Python Error Indicator as we has recorded the
+              // exception in the response message.
+              PyErr_Clear();
+            }
+
+            throw std::runtime_error(err);
+          }
+        },
+        PyObjectType::get()));
+  }
+
+  void add_done_callback(py::function cb) {
+    auto pf = std::make_shared<PythonFunctionGuard>(std::move(cb));
+    // NOLINTNEXTLINE(modernize-avoid-bind)
+    fut->addCallback(std::bind(
+        [pyFut(this->getPtr())](
+            const std::shared_ptr<PythonFunctionGuard>& pf) {
+          try {
+            pybind11::gil_scoped_acquire ag;
+            pf->func_(pyFut);
+          } catch (py::error_already_set& e) {
+            {
+              pybind11::gil_scoped_acquire ag;
+              // Release ownership on py::objects and also restore Python
+              // Error Indicator.
+              e.restore();
+              // Clear the Python Error Indicator as we has recorded the
+              // exception in the response message.
+              PyErr_Clear();
+            }
+            // Log and ignore exceptions raised through the callback
+            LOG(ERROR) << "Got the following error when running the callback: "
+                       << e.what();
+
+          } catch (const std::exception& e) {
+            // Log and ignore exceptions raised through the callback
+            LOG(ERROR) << "Got the following error when running the callback: "
+                       << e.what();
+          }
+        },
+        std::move(pf)));
+  }
+
+  void markCompleted(const py::object& pyValue) {
+    DCHECK(PyGILState_Check());
+    IValue value = toIValue(pyValue, PyObjectType::get());
+
+    py::gil_scoped_release release;
+    fut->markCompleted(std::move(value));
+  }
+
+  c10::intrusive_ptr<c10::ivalue::Future> fut;
+  // unwrap_func works like a callback for the value returned by
+  // PythonFutureWrapper::wait().
+  std::optional<UnwrapFunc> unwrap_func;
+
+ private:
+  std::shared_ptr<PythonFutureWrapper> getPtr() {
+    return shared_from_this();
+  }
+};
+
+// The PythonAwaitWrapper for ivalue::Await
+//
+// Expresses delayed function execution with Lazy semantic.
+// i.e. Await[W] in eager mode can be used as W.
+// When the attribute of W type is requested, Await[W] will return the
+// attribute of W, transparently calling wait() beforehand.
+// No Lazy semantic for script, explicit wait(Await[W]) -> W must be called to
+// convert to type W.
+//
+// The Await object takes shared ownership of specified function and the
+// arguments. After first call for wait() it owns the result. Deliberately no
+// type inference for eager mode.
+struct VISIBILITY_HIDDEN PythonAwaitWrapper
+    : std::enable_shared_from_this<PythonAwaitWrapper> {
+  explicit PythonAwaitWrapper(c10::intrusive_ptr<c10::ivalue::Await> aw)
+      : aw_(std::move(aw)) {}
+  explicit PythonAwaitWrapper(py::handle input) {
+    args_ = py::tuple(1u);
+    args_[0] = input;
+    auto type = PyObjectType::get();
+    aw_ = c10::make_intrusive<c10::ivalue::Await>(type);
+    aw_->markCompleted(toIValue(input, type));
+  }
+
+  explicit PythonAwaitWrapper(py::function pf, py::tuple args)
+      : args_(std::move(args)) {
+    pyfg_ = std::make_shared<torch::jit::PythonFunctionGuard>(std::move(pf));
+
+    std::function<IValue()> f = [fg(pyfg_), &args(args_)]() {
+      pybind11::gil_scoped_acquire ag;
+      return toIValue(fg->func_(*args), PyObjectType::get());
+    };
+    aw_ = c10::make_intrusive<c10::ivalue::Await>(
+        PyObjectType::get(), std::move(f));
+  }
+
+  explicit PythonAwaitWrapper(const PythonAwaitWrapper&) = delete;
+  PythonAwaitWrapper& operator=(const PythonAwaitWrapper&) = delete;
+  PythonAwaitWrapper(PythonAwaitWrapper&&) = default;
+  PythonAwaitWrapper& operator=(PythonAwaitWrapper&&) = default;
+  ~PythonAwaitWrapper() = default;
+
+  py::object wait() {
+    py::gil_scoped_acquire acquire;
+    return toPyObject(aw_->wait());
+  }
+
+  // Nowait semantic means trivial case when Await is constructed from the
+  // result
+  bool is_nowait() {
+    return pyfg_ == nullptr;
+  }
+
+  const py::function fn() {
+    TORCH_CHECK(
+        pyfg_, "Await constructed as awaitable_nowait does not have fn");
+    return pyfg_->func_;
+  }
+
+  const py::tuple args() {
+    return args_;
+  }
+
+  TypePtr type() {
+    return aw_->type();
+  }
+
+  c10::intrusive_ptr<c10::ivalue::Await> aw_;
+  std::shared_ptr<torch::jit::PythonFunctionGuard> pyfg_;
+  py::tuple args_;
+
+ private:
+  std::shared_ptr<PythonAwaitWrapper> getPtr() {
+    return shared_from_this();
+  }
+};
+
+// error reporting: when reporting user-caused errors, these functions should
+// not use AT_ERROR macros, since these macros add stack trace information
+// that is confusing to display to the end user since it always reports
+// locations in libtorch code rather than user code.
+
+inline std::shared_ptr<CompilationUnit> get_python_cu() {
+  return py::module::import("torch.jit._state")
+      .attr("_python_cu")
+      .cast<std::shared_ptr<CompilationUnit>>();
+}
+
+struct TypedIValue : public std::pair<IValue, TypePtr> {
+  using pair::pair;
+
+  IValue& ivalue() {
+    return this->first;
+  }
+  TypePtr& type() {
+    return this->second;
+  }
+};
+
+inline TypedIValue toDictKeyIValue(py::handle key) {
+  if (py::isinstance<py::str>(key)) {
+    return TypedIValue(
+        ConstantString::create(py::cast<std::string>(key)), StringType::get());
+  } else if (py::isinstance<py::int_>(key)) {
+    return TypedIValue(py::cast<int64_t>(key), IntType::get());
+  } else if (py::isinstance<py::float_>(key)) {
+    return TypedIValue(py::cast<double>(key), FloatType::get());
+  } else {
+    TORCH_CHECK(
+        false, "Dictionary inputs may only have string, int, or float keys");
+  }
+}
+
+inline std::optional<TypePtr> unifyOrInitializeType(
+    const TypePtr& accum,
+    const TypePtr& unify) {
+  if (!accum) {
+    return unify;
+  }
+  return unifyTypes(accum, unify);
+}
+
+using InferredType = c10::InferredType;
+
+InferredType tryToInferContainerType(py::handle input, bool primitiveTypeOnly);
+
+namespace detail {
+
+// Additional implementations for tryToInferType().
+std::optional<InferredType> _tryToInferTypeImpl(py::handle input);
+
+} // namespace detail
+
+// Try to infer the type of a Python object
+// The type cannot be inferred if:
+//   input is an empty container (list, dict)
+//   input is an list with element types that cannot be unified
+//   input is an dict with key or value types that cannot be unified
+inline InferredType tryToInferType(py::handle input) {
+  // Try tensor types
+  if (THPVariable_Check(input.ptr())) {
+    return InferredType(TensorType::get());
+  }
+
+  if (input.is_none()) {
+    return InferredType(NoneType::get());
+  }
+
+  if (py::isinstance<StrongFunctionPtr>(input)) {
+    auto fn = py::cast<StrongFunctionPtr>(input).function_;
+    return InferredType(FunctionType::create(fn));
+  }
+
+  // Try basic types first
+  if (py::isinstance<py::bool_>(input)) {
+    return InferredType(BoolType::get());
+    // NOLINTNEXTLINE(bugprone-branch-clone)
+  } else if (py::isinstance<py::int_>(input)) {
+    return InferredType(IntType::get());
+  } else if (py::isinstance<py::float_>(input)) {
+    return InferredType(FloatType::get());
+  } else if (PyComplex_CheckExact(input.ptr())) {
+    return InferredType(ComplexType::get());
+    // NOLINTNEXTLINE(bugprone-branch-clone)
+  } else if (py::isinstance<py::bytes>(input)) {
+    // NOTE: We may need a ByteType in the future
+    return InferredType(StringType::get());
+  } else if (py::isinstance<py::str>(input)) {
+    return InferredType(StringType::get());
+  } else if (THPLayout_Check(input.ptr())) {
+    return InferredType(IntType::get());
+  } else if (THPDevice_Check(input.ptr())) {
+    return InferredType(DeviceObjType::get());
+  } else if (THPGenerator_Check(input.ptr())) {
+    return InferredType(GeneratorType::get());
+  } else if (THPStream_Check(input.ptr())) {
+    return InferredType(StreamObjType::get());
+  } else if (THPDtype_Check(input.ptr())) {
+    return InferredType(IntType::get());
+  } else if (THPQScheme_Check(input.ptr())) {
+    return InferredType(IntType::get());
+  } else if (THPLayout_Check(input.ptr())) {
+    return InferredType(IntType::get());
+  }
+
+  // Check for types registered in _tryToInferTypeImpl (e.g. ProcessGroup)
+  // before falling through to the expensive inspect.isclass / JIT compilation
+  // path below.
+  auto ty = detail::_tryToInferTypeImpl(input);
+  if (ty.has_value()) {
+    return ty.value();
+  }
+
+  auto enum_type = py::module::import("enum").attr("Enum");
+  py::bool_ isEnumValue = py::isinstance(input, enum_type);
+  if (py::cast<bool>(isEnumValue)) {
+    auto enum_class = input.attr("__class__");
+    auto enum_type = py::cast<TypePtr>(
+        py::module::import("torch.jit.annotations")
+            .attr("try_ann_to_type")(enum_class, SourceRange()));
+    return InferredType(std::move(enum_type));
+  }
+
+  py::bool_ isClass =
+      py::module::import("inspect").attr("isclass")(py::type::handle_of(input));
+  if (py::cast<bool>(isClass)) {
+    // Assume that the class is compiled already or will compile. Invalidate
+    // this later if needed.
+    bool class_compiled = true;
+
+    // Check if the type is already compiled.
+    py::object existing_ty =
+        py::module::import("torch.jit._state")
+            .attr("_get_script_class")(py::type::handle_of(input));
+
+    if (existing_ty.is_none()) {
+      // If not, try to compile it.
+      py::bool_ can_compile =
+          py::module::import("torch._jit_internal")
+              .attr("can_compile_class")(py::type::handle_of(input));
+
+      if (py::cast<bool>(can_compile)) {
+        // Try to compile the class. This is wrapped in a try-catch because
+        // compilation of class types can raise an Exception and in that case,
+        // we want to defer to other attempts at type inference below rather
+        // than fail compilation altogether.
+        try {
+          py::module::import("torch.jit._script")
+              .attr("_recursive_compile_class")(
+                  py::type::handle_of(input), SourceRange());
+        } catch (...) {
+          // Invalidate the assumption that the class compiled so that we don't
+          // look up and return its JIT type as the type for the input.
+          class_compiled = false;
+        }
+      }
+    }
+
+    // If the class compiled successfully, look up the existing JIT type by
+    // qualified name and return it.
+    if (class_compiled) {
+      auto script_class =
+          py::module::import("torch.jit._state")
+              .attr("_get_script_class")(py::type::handle_of(input));
+
+      if (!script_class.is_none()) {
+        auto class_type = py::cast<ClassTypePtr>(script_class);
+
+        if (class_type && !class_type->is_module()) {
+          return InferredType(std::move(class_type));
+        }
+      }
+    }
+  }
+
+  if (py::isinstance<Object>(input)) {
+    auto object = py::cast<Object>(input);
+    return InferredType(object.type());
+#ifdef USE_RPC
+  } else if (py::isinstance<torch::distributed::rpc::PyRRef>(input)) {
+    auto rref_ivalue = input.cast<torch::distributed::rpc::PyRRef>().toIValue();
+    return InferredType(rref_ivalue.type());
+#endif
+  }
+
+  auto await_type = py::module::import("torch._awaits").attr("_Await");
+  py::bool_ is_await = py::isinstance(input, await_type);
+  if (py::cast<bool>(is_await)) {
+    auto awptr = input.cast<std::shared_ptr<PythonAwaitWrapper>>();
+    return InferredType(AwaitType::create(awptr->aw_->elementType()));
+  }
+
+  if (as_module(py::cast<py::object>(input))) {
+    return InferredType("Cannot infer type of ScriptModule");
+  }
+
+  auto module_type = py::module::import("torch.nn").attr("Module");
+  py::bool_ is_module = py::isinstance(input, module_type);
+  if (py::cast<bool>(is_module)) {
+    return InferredType("Cannot infer concrete type of torch.nn.Module");
+  }
+
+  // Try container types
+  return tryToInferContainerType(input, false);
+}
+
+// This function is similar to tryToInferType, but it only tries to infer
+// primitive types (int, float, bool, complex) or nested container of primitive
+// types.
+inline InferredType tryToInferPrimitiveType(py::handle input) {
+  if (input.is_none()) {
+    return InferredType(NoneType::get());
+  }
+
+  // Only primitive data type
+  if (py::isinstance<py::bool_>(input)) {
+    return InferredType(BoolType::get());
+    // NOLINTNEXTLINE(bugprone-branch-clone)
+  } else if (py::isinstance<py::int_>(input)) {
+    return InferredType(IntType::get());
+  } else if (py::isinstance<py::float_>(input)) {
+    return InferredType(FloatType::get());
+  } else if (PyComplex_CheckExact(input.ptr())) {
+    return InferredType(ComplexType::get());
+  }
+
+  // Try container types
+  return tryToInferContainerType(input, true);
+}
+
+inline InferredType tryToInferContainerType(
+    py::handle input,
+    bool primitiveTypeOnly = false) {
+  if (PyTuple_Check(input.ptr())) {
+    py::tuple tuple = py::cast<py::tuple>(input);
+    std::vector<TypePtr> element_types;
+    element_types.reserve(tuple.size());
+
+    for (py::handle elem : tuple) {
+      auto type_match = primitiveTypeOnly ? tryToInferPrimitiveType(elem)
+                                          : tryToInferType(elem);
+      if (type_match.success()) {
+        element_types.push_back(type_match.type());
+      } else {
+        // Forward error message along
+        return type_match.reason();
+      }
+    }
+    return InferredType(TupleType::create(std::move(element_types)));
+  } else if (PyDict_Check(input.ptr())) {
+    // Check to make sure we can generate useful input/output types
+    auto dict = py::cast<py::dict>(input);
+    size_t len = py::len(dict);
+    if (!len) {
+      return InferredType("Dictionary inputs must have entries");
+    }
+
+    TypePtr key_type = nullptr;
+    TypePtr value_type = nullptr;
+
+    for (auto entry : dict) {
+      // Try to infer the key type and unify it with the existing one
+      auto entry_key_type_match = primitiveTypeOnly
+          ? tryToInferPrimitiveType(entry.first)
+          : tryToInferType(entry.first);
+      if (!entry_key_type_match.success()) {
+        return entry_key_type_match.reason();
+      }
+      auto unified_key =
+          unifyOrInitializeType(key_type, entry_key_type_match.type());
+      if (!unified_key) {
+        return InferredType(c10::str(
+            "Dictionary inputs to traced functions must have consistent type. Found ",
+            key_type->repr_str(),
+            " and ",
+            (entry_key_type_match.type())->repr_str()));
+      }
+
+      // Try to infer the value type and unify it with the existing one
+      auto entry_value_type_match = primitiveTypeOnly
+          ? tryToInferPrimitiveType(entry.second)
+          : tryToInferType(entry.second);
+      if (!entry_value_type_match.success()) {
+        return entry_value_type_match.reason();
+      }
+      auto unified_value =
+          unifyOrInitializeType(value_type, entry_value_type_match.type());
+      if (!unified_value) {
+        return InferredType(c10::str(
+            "Dictionary inputs to traced functions must have consistent type. Found ",
+            value_type->repr_str(),
+            " and ",
+            (entry_value_type_match.type())->repr_str()));
+      }
+
+      key_type = *unified_key;
+      value_type = *unified_value;
+    }
+    return InferredType(
+        DictType::create(std::move(key_type), std::move(value_type)));
+  } else if (PyList_Check(input.ptr())) {
+    auto list = py::cast<py::list>(input);
+    size_t len = py::len(list);
+    if (!len) {
+      return InferredType("List trace inputs must have elements");
+    }
+
+    TypePtr element_type = nullptr;
+    for (auto elem : list) {
+      auto element_type_match = primitiveTypeOnly
+          ? tryToInferPrimitiveType(elem)
+          : tryToInferType(elem);
+      if (!element_type_match.success()) {
+        return InferredType(c10::str(
+            "Could not infer type of list element: ",
+            element_type_match.reason()));
+      }
+      auto unified_type =
+          unifyOrInitializeType(element_type, element_type_match.type());
+      if (!unified_type) {
+        return InferredType(c10::str(
+            "List inputs to traced functions must have consistent element type. Found ",
+            element_type->repr_str(),
+            " and ",
+            (element_type_match.type())->repr_str()));
+      }
+      element_type = *unified_type;
+    }
+    return InferredType(ListType::create(element_type));
+  } else {
+    if (primitiveTypeOnly) {
+      return InferredType(c10::str(
+          "Only tuple, list, or dict (possibly nested) of primitive types (bool, float, int, complex)",
+          "are supported ",
+          "as inputs or outputs of traced functions",
+          ", but instead got value of type ",
+          py::str(py::type::handle_of(input).attr("__name__")),
+          "."));
+    } else {
+      // TODO: this message is not correct anymore, since this InferredType is
+      // used from a bunch of circumstances unrelated to tracing. We can reuse
+      // this instead of the attribute_failure stuff in concreteType
+      return InferredType(c10::str(
+          "Only tensors and (possibly nested) tuples of tensors, lists, or dicts ",
+          "are supported ",
+          "as inputs or outputs of traced functions",
+          ", but instead got value of type ",
+          py::str(py::type::handle_of(input).attr("__name__")),
+          "."));
+    }
+  }
+}
+
+inline bool isTraceableType(const TypePtr& type) {
+  if (type->isSubtypeOf(*TensorType::get())) {
+    return true;
+  }
+
+  if (auto list_type = type->cast<ListType>()) {
+    return isTraceableType(list_type->getElementType());
+  }
+
+  if (auto tuple_type = type->cast<TupleType>()) {
+    return std::all_of(
+        tuple_type->elements().begin(),
+        tuple_type->elements().end(),
+        [](const TypePtr& element_type) {
+          return isTraceableType(element_type);
+        });
+  }
+
+  if (auto dict_type = type->cast<DictType>()) {
+    return isTraceableType(dict_type->getValueType());
+  }
+
+  return false;
+}
+
+inline IValue toTypeInferredIValue(py::handle input) {
+  auto match = tryToInferType(input);
+  if (!match.success()) {
+    auto object = py::cast<py::object>(input);
+    if (auto mod = as_module(object)) {
+      // if obj is already a ScriptModule, just return its ivalue
+      auto ptr = mod.value()._ivalue();
+      // explicit copy semantics for strong ownership of the resource.
+      return c10::intrusive_ptr<c10::ivalue::Object>::reclaim_copy(
+          ptr.release());
+    }
+
+    // Check if the obj is a ScriptObject.
+    if (auto script_obj = as_object(object)) {
+      auto ptr = script_obj.value()._ivalue();
+      return c10::intrusive_ptr<c10::ivalue::Object>::reclaim_copy(
+          ptr.release());
+    }
+    TORCH_CHECK(
+        false,
+        "Tracer cannot infer type of ",
+        py::str(input),
+        "\n:",
+        match.reason());
+  }
+  return toIValue(input, match.type());
+}
+
+inline Stack toTraceableStack(const py::tuple& inputs) {
+  auto info = toTypeInferredIValue(inputs);
+  TORCH_CHECK(
+      isTraceableType(info.type()),
+      "Type '",
+      info.type()->repr_str(),
+      "' cannot be traced. Only Tensors and (possibly nested) Lists, Dicts, and"
+      " Tuples of Tensors can be traced");
+  return info.toTupleRef().elements().vec();
+}
+
+// Serialize the python dictionary into a traceable stack.
+inline Stack toTraceableStack(const py::dict& inputs) {
+  Stack res;
+  for (auto it = inputs.begin(); it != inputs.end(); it++) {
+    if (THPVariable_Check(it->second.ptr())) {
+      res.push_back(toIValue(it->second, tryToInferType(it->second).type()));
+    }
+  }
+  return res;
+}
+
+inline IValue createGenericList(py::handle obj, const TypePtr& elem_type) {
+  auto elems = c10::impl::GenericList(elem_type);
+  for (auto elem : obj) {
+    elems.push_back(toIValue(elem, elem_type));
+  }
+  return IValue(elems);
+}
+
+inline IValue createGenericDict(
+    const py::dict& obj,
+    const TypePtr& key_type,
+    const TypePtr& value_type) {
+  c10::impl::GenericDict elems(key_type, value_type);
+  elems.reserve(py::len(obj));
+  for (auto& entry : obj) {
+    elems.insert(
+        toIValue(entry.first, key_type), toIValue(entry.second, value_type));
+  }
+  return IValue(elems);
+}
+
+template <class T>
+inline void guardAgainstNamedTensor(const T& var) {
+  TORCH_CHECK(
+      !var.has_names(),
+      "NYI: Named tensors are currently unsupported in TorchScript. As a  "
+      "workaround please drop names via `tensor = tensor.rename(None)`.");
+}
+
+// Extract custom class registered with torchbind
+template <typename T>
+c10::intrusive_ptr<T> toCustomClass(py::handle obj) {
+  static_assert(
+      std::is_base_of_v<CustomClassHolder, T>, "T is not a CustomClass");
+  const auto& type = c10::getCustomClassType<c10::intrusive_ptr<T>>();
+  c10::IValue ivalue = toIValue(obj, type);
+  return std::move(ivalue).toCustomClass<T>();
+}
+
+// Small wrapper around getting the type name string from Python to make
+// types easier to interpret, e.g. give the structural type for a NamedTuple
+inline std::string friendlyTypeName(py::handle obj) {
+  if (py::isinstance<py::tuple>(obj) && py::hasattr(obj, "_fields")) {
+    auto field_names =
+        py::cast<std::vector<std::string>>(py::getattr(obj, "_fields"));
+    std::stringstream ss;
+    ss << py::str(py::type::handle_of(obj).attr("__name__"));
+    ss << " (aka NamedTuple(";
+    bool first = true;
+    for (auto& field_name : field_names) {
+      if (!first) {
+        ss << ", ";
+      }
+      ss << field_name;
+      first = false;
+    }
+    ss << "))";
+    return ss.str();
+  } else {
+    return py::str(py::type::handle_of(obj).attr("__name__"));
+  }
+}
+
+// Thrown when trying to create a schema for a list of python
+// arguments that cannot be converted.
+// Can be caught by the caller to attempt to use other schema
+// when there is an overloaded operator.
+struct schema_match_error : public std::runtime_error {
+  using std::runtime_error::runtime_error;
+};
+
+inline IValue argumentToIValue(
+    const FunctionSchema& schema,
+    size_t argumentPosition,
+    py::handle object) {
+  const auto& argument = schema.arguments().at(argumentPosition);
+  try {
+    return toIValue(object, argument.real_type(), argument.N());
+  } catch (const py::cast_error& error) {
+    throw schema_match_error(c10::str(
+        schema.formatTypeMismatchMsg(
+            argument,
+            friendlyTypeName(object),
+            argumentPosition,
+            py::repr(object)),
+        "\nCast error details: ",
+        error.what()));
+  } catch (const py::error_already_set& error) {
+    throw schema_match_error(c10::str(
+        schema.formatTypeMismatchMsg(
+            argument,
+            friendlyTypeName(object),
+            argumentPosition,
+            py::repr(object)),
+        "\n Python error details: ",
+        error.what()));
+  }
+}
+
+inline IValue returnToIValue(const TypePtr& type, py::handle object) {
+  try {
+    return toIValue(object, type);
+  } catch (const py::cast_error& error) {
+    throw std::runtime_error(c10::str(
+        " expected value of type ",
+        type->str(),
+        " for return value but instead got value of type ",
+        py::str(py::type::handle_of(object).attr("__name__")),
+        ".",
+        "\nValue: ",
+        py::repr(object),
+        "\nCast error details: ",
+        error.what()));
+  }
+}
+
+inline py::object getScriptedClassOrError(const c10::NamedTypePtr& classType) {
+  auto py_class =
+      py::module::import("torch.jit._state")
+          .attr("_get_python_class")(classType->name()->qualifiedName());
+  if (py_class.is_none()) {
+    std::stringstream err;
+    err << "Unknown reference to ScriptClass ";
+    err << classType->name()->qualifiedName();
+    err << ". (Did you forget to import it?)";
+    throw std::runtime_error(err.str());
+  }
+  return py_class;
+}
+
+struct VISIBILITY_HIDDEN tuple_slice {
+  /*implicit*/ tuple_slice(py::tuple tup_)
+      : tup(std::move(tup_)), b(0), e(static_cast<int64_t>(tup.size())) {}
+  tuple_slice(py::tuple tup_, int64_t b_)
+      : tup(std::move(tup_)), b(b_), e(static_cast<int64_t>(tup.size())) {}
+  tuple_slice(py::tuple tup_, int64_t b_, int64_t e_)
+      : tup(std::move(tup_)), b(b_), e(e_) {}
+  py::detail::tuple_iterator begin() const {
+    return {tup, static_cast<pybind11::ssize_t>(b)};
+  }
+  py::detail::tuple_iterator end() const {
+    return {tup, static_cast<pybind11::ssize_t>(e)};
+  }
+  size_t size() const {
+    return e - b;
+  }
+  py::detail::tuple_accessor operator[](size_t index) const {
+    return {tup, static_cast<size_t>(b + index)};
+  }
+
+ private:
+  py::tuple tup;
+  int64_t b;
+  int64_t e;
+};
+
+inline bool validateFakeScriptObjectSchema(
+    const c10::FunctionSchema& schema,
+    size_t argumentPosition,
+    py::handle object) {
+  auto argument = schema.arguments().at(argumentPosition);
+  auto class_type = argument.real_type()->expect<c10::ClassType>();
+  auto fake_class_registry =
+      py::module::import("torch._library.fake_class_registry");
+  auto fake_class = fake_class_registry.attr("find_fake_class")(
+      class_type->name().value().qualifiedName());
+  if (!py::isinstance(object.attr("wrapped_obj"), fake_class)) {
+    throw schema_match_error(c10::str(
+        schema.formatTypeMismatchMsg(
+            argument,
+            friendlyTypeName(object),
+            argumentPosition,
+            py::repr(object.attr("wrapped_obj"))),
+        "\nCast error details: ",
+        argument.name(),
+        " is expected to be a FakeScriptObject of ",
+        class_type->name().value().qualifiedName()));
+  }
+  return true;
+}
+
+inline bool matchSchemaAllowFakeScriptObject(
+    const FunctionSchema& schema,
+    const tuple_slice& args,
+    const py::kwargs& kwargs) {
+  size_t all_arguments = args.size() + kwargs.size();
+  if (all_arguments > schema.arguments().size()) {
+    throw schema_match_error(c10::str(
+        schema.name(),
+        "() expected at most ",
+        schema.arguments().size(),
+        " argument(s) but received ",
+        all_arguments,
+        " argument(s). Declaration: ",
+        schema));
+  }
+
+  int64_t arg_idx = 0;
+  auto fake_class_registry =
+      py::module::import("torch._library.fake_class_registry");
+
+  // First push all positional args.
+  for (const auto& arg : args) {
+    // ...but refuse to do it if the schema says that this was supposed
+    // to be keyword only
+    if (schema.arguments()[arg_idx].kwarg_only()) {
+      throw schema_match_error(c10::str(
+          schema.name(),
+          "() takes ",
+          arg_idx,
+          " positional argument(s) but ",
+          args.size(),
+          " was/were given.  Declaration: ",
+          schema));
+    }
+    // Use the type information from the schema to convert the PyObject.
+    const auto& argument = schema.arguments().at(arg_idx);
+    if (argument.real_type()->kind() == TypeKind::ClassType &&
+        py::isinstance(arg, fake_class_registry.attr("FakeScriptObject"))) {
+      validateFakeScriptObjectSchema(schema, arg_idx, arg);
+    } else {
+      argumentToIValue(schema, arg_idx, arg);
+    }
+
+    arg_idx++;
+  }
+
+  // Now for every remaining non-positional argument in the schema, look for it
+  // in the kwargs dict and push it if found, or use its default value if it
+  // has one.
+  size_t consumed_kwargs = 0;
+  for (size_t i = arg_idx; i < schema.arguments().size(); ++i) {
+    const auto& arg = schema.arguments()[i];
+    if (kwargs.contains(arg.name().c_str())) {
+      auto cur_kwarg = kwargs[arg.name().c_str()];
+      if (arg.real_type()->kind() == TypeKind::ClassType &&
+          py::isinstance(
+              cur_kwarg, fake_class_registry.attr("FakeScriptObject"))) {
+        validateFakeScriptObjectSchema(schema, i, cur_kwarg);
+      } else {
+        argumentToIValue(schema, i, cur_kwarg);
+      }
+      consumed_kwargs += 1;
+    } else if (arg.default_value()) {
+      continue;
+    } else {
+      throw schema_match_error(c10::str(
+          schema.name(),
+          "() is missing value for argument '",
+          arg.name(),
+          "'. Declaration: ",
+          schema));
+    }
+  }
+
+  if (consumed_kwargs != kwargs.size()) {
+    std::vector<std::string> names;
+    for (const auto& kwarg : kwargs) {
+      names.emplace_back(py::cast<std::string>(kwarg.first));
+    }
+    throw schema_match_error(schema.findErrorInKwargs(names));
+  }
+
+  return true;
+}
+
+inline Stack createStackForSchema(
+    const FunctionSchema& schema,
+    const tuple_slice& args,
+    const py::kwargs& kwargs,
+    std::optional<IValue> self) {
+  size_t all_arguments = (self ? 1 : 0) + args.size() + kwargs.size();
+  if (all_arguments > schema.arguments().size()) {
+    throw schema_match_error(c10::str(
+        schema.name(),
+        "() expected at most ",
+        schema.arguments().size(),
+        " argument(s) but received ",
+        all_arguments,
+        " argument(s). Declaration: ",
+        schema));
+  }
+  Stack stack;
+  stack.reserve(schema.arguments().size());
+
+  int64_t arg_idx = 0;
+  if (self) {
+    push(stack, std::move(*self));
+    arg_idx++;
+  }
+  // First push all positional args.
+  for (const auto& arg : args) {
+    // ...but refuse to do it if the schema says that this was supposed
+    // to be keyword only
+    if (schema.arguments()[arg_idx].kwarg_only()) {
+      throw schema_match_error(c10::str(
+          schema.name(),
+          "() takes ",
+          arg_idx,
+          " positional argument(s) but ",
+          self ? 1 + args.size() : args.size(),
+          " was/were given.  Declaration: ",
+          schema));
+    }
+    // Use the type information from the schema to convert the PyObject.
+    push(stack, argumentToIValue(schema, stack.size(), arg));
+    arg_idx++;
+  }
+
+  // Now for every remaining non-positional argument in the schema, look for it
+  // in the kwargs dict and push it if found, or use its default value if it
+  // has one.
+  size_t consumed_kwargs = 0;
+  for (size_t i = stack.size(); i < schema.arguments().size(); ++i) {
+    const auto& arg = schema.arguments()[i];
+    if (kwargs.contains(arg.name().c_str())) {
+      push(stack, argumentToIValue(schema, i, kwargs[arg.name().c_str()]));
+      consumed_kwargs += 1;
+    } else if (arg.default_value()) {
+      push(stack, *arg.default_value());
+    } else {
+      throw schema_match_error(c10::str(
+          schema.name(),
+          "() is missing value for argument '",
+          arg.name(),
+          "'. Declaration: ",
+          schema));
+    }
+  }
+
+  if (consumed_kwargs != kwargs.size()) {
+    std::vector<std::string> names;
+    for (const auto& kwarg : kwargs) {
+      names.emplace_back(py::cast<std::string>(kwarg.first));
+    }
+    throw schema_match_error(schema.findErrorInKwargs(names));
+  }
+
+  return stack;
+}
+
+// NOLINTNEXTLINE(cppcoreguidelines-rvalue-reference-param-not-moved)
+inline py::object createPyObjectForStack(Stack&& stack) {
+  if (stack.empty()) {
+    return py::none();
+  }
+
+  // Return a simple value and not a single-element tuple if there is only one
+  // return value.
+  if (stack.size() == 1) {
+    return toPyObject(std::move(stack[0]));
+  }
+
+  // If there is more than one return value, pop them into a py::tuple.
+  py::tuple return_values(stack.size());
+  for (const auto ret : c10::irange(return_values.size())) {
+    return_values[ret] = toPyObject(std::move(stack[ret]));
+  }
+
+#if defined(__clang__)
+  return std::move(return_values);
+#else
+  return return_values;
+#endif
+}
+
+// TODO: Remove once we clean up the GraphExecutor usage.
+inline Stack evilDeprecatedBadCreateStackDoNotUse(
+    const py::tuple& tuple,
+    at::ArrayRef<Value*> inputs,
+    size_t reserve_extra_space = 0) {
+  if (tuple.size() != inputs.size()) {
+    TORCH_CHECK(
+        false,
+        "expected " + std::to_string(inputs.size()) + " inputs, but got " +
+            std::to_string(tuple.size()));
+  }
+  Stack result;
+  result.reserve(tuple.size() + reserve_extra_space);
+  for (const auto i : c10::irange(inputs.size())) {
+    result.push_back(toIValue(std::move(tuple[i]), inputs[i]->type()));
+  }
+  return result;
+}
+
+// Run `callee`, potentially inserting a CallFunction/CallMethod node into the
+// tracing graph.
+inline py::object runAndInsertCall(
+    Function& callee,
+    const tuple_slice& args,
+    const py::kwargs& kwargs,
+    std::optional<IValue> self,
+    // Lambda that tells this function how to insert `callee` into the graph if
+    // we're tracing.
+    const std::function<Value*(Graph&, const MatchedSchema& match)>&
+        callInserter) {
+  auto stack =
+      createStackForSchema(callee.getSchema(), args, kwargs, std::move(self));
+  const auto& tracing_state = tracer::getTracingState();
+  if (!tracing_state) {
+    pybind11::gil_scoped_release no_gil_guard;
+    // If we're not tracing, just run the callee as normal.
+    callee.run(stack);
+  } else {
+    // If we are tracing, insert the appropriate CallFunction or CallMethod node
+    // and then run the callee with tracing disabled.
+
+    // Get the graph `Value`s that represent the input IValues
+    auto inputs = last(stack, callee.num_inputs());
+    auto input_values =
+        fmap(inputs, [](const IValue& v) { return tracer::getValueTrace(v); });
+    TORCH_INTERNAL_ASSERT(callee.getSchema().returns().size() == 1)
+    auto return_type = callee.getSchema().returns().at(0).type();
+    auto graph = tracing_state->graph;
+    std::vector<NamedValue> named_values;
+    named_values.reserve(input_values.size());
+    for (Value* v : input_values) {
+      named_values.emplace_back(v);
+    }
+
+    // Add a call node.
+    MatchedSchema match = matchSchema(
+        callee.getSchema(),
+        tracer::getPythonInterpreterSourceRange(),
+        *graph,
+        named_values,
+        {});
+    auto output_value = callInserter(*graph, match);
+
+    // Actually run the callee. Pause the tracer so that we don't double-add the
+    // callee nodes.
+    {
+      pybind11::gil_scoped_release no_gil_guard;
+      ResourceGuard guard(tracer::pauseTracing());
+      callee.run(stack);
+    }
+
+    // Associate the output IValues with the output `Value`s in the graph
+    tracer::setValueTrace(stack.back(), output_value);
+  }
+
+  TORCH_CHECK(
+      !stack.empty(),
+      "Expected values in the stack after execution but found none");
+  return toPyObject(std::move(stack.back()));
+}
+
+inline std::optional<py::object> maybeTorchFunctionDispatch(
+    const py::object& callee,
+    const tuple_slice& args_no_self,
+    const py::kwargs& kwargs,
+    const c10::QualifiedName& qualname) {
+  std::vector<py::handle> args_vec;
+  for (const auto& arg : args_no_self) {
+    args_vec.push_back(arg);
+  }
+  py::tuple args = py::cast(args_vec);
+
+  // Handle __torch_function__ dispatch
+  std::vector<PyObject*> overloaded_args;
+  size_t total_arg_num = args.size() + kwargs.size();
+  for (const auto& arg : args) {
+    is_tensor_and_append_overloaded(arg.ptr(), &overloaded_args);
+    is_tensor_list_and_append_overloaded(
+        arg.ptr(),
+        &overloaded_args,
+        static_cast<int>(total_arg_num),
+        false /* throw_error */);
+  }
+  // NB: for kwargs, we cannot guarantee the order of appending
+  // is the same as the argument order in operator's schema.
+  // This is suboptimal, but should be fine. Later when we have
+  // better schema matching and argument parsing, we could
+  // match the operator in `operations` first, then the order will
+  // be guaranteed.
+  for (auto item : kwargs) {
+    is_tensor_and_append_overloaded(item.second.ptr(), &overloaded_args);
+    is_tensor_list_and_append_overloaded(
+        item.second.ptr(),
+        &overloaded_args,
+        total_arg_num,
+        false /* throw_error */);
+  }
+  if (!overloaded_args.empty()) {
+    return pybind11::reinterpret_steal<py::object>(
+        handle_torch_function_no_python_arg_parser(
+            /*overloaded_args=*/overloaded_args,
+            /*args=*/args.ptr(),
+            /*kwargs=*/kwargs.ptr(),
+            /*func_name=*/qualname.name().c_str(),
+            /*torch_api_function=*/callee.ptr(),
+            /*module_name=*/qualname.prefix().c_str()));
+  }
+
+  return std::nullopt;
+}
+
+inline py::object invokeScriptFunctionFromPython(
+    Function& callee,
+    const tuple_slice& args,
+    const py::kwargs& kwargs) {
+  // TODO: we could add __torch_function__ dispatch here but I don't know
+  // the implications of doing so
+
+  return runAndInsertCall(
+      callee,
+      args,
+      kwargs,
+      /*self=*/std::nullopt,
+      [&](Graph& graph, const MatchedSchema& match) {
+        return graph.insertFunctionCall(&callee, match);
+      });
+}
+
+inline py::object invokeScriptMethodFromPython(
+    Method& callee,
+    const tuple_slice& args,
+    const py::kwargs& kwargs) {
+  auto self = callee.owner()._ivalue();
+
+  if (auto torch_fn_result = maybeTorchFunctionDispatch(
+          py::cast(callee), args, kwargs, callee.name())) {
+    return *torch_fn_result;
+  }
+
+  return runAndInsertCall(
+      callee.function(),
+      args,
+      kwargs,
+      self,
+      [&](Graph& graph, const MatchedSchema& match) {
+        return graph.insertMethodCall(callee.name(), match);
+      });
+}
+
+TORCH_PYTHON_API std::pair<std::shared_ptr<Operator>, Stack> getOpWithStack(
+    const std::vector<std::shared_ptr<Operator>>& operations,
+    const py::args& args,
+    const py::kwargs& kwargs);
+
+// Efficient overload (does not require vector allocation) of the
+// above for use from C++ code.
+std::pair<std::shared_ptr<Operator>, Stack> getOpWithStack(
+    c10::ArrayRef<std::shared_ptr<Operator>> operations,
+    const py::args& args,
+    const py::kwargs& kwargs);
+
+TORCH_PYTHON_API py::object invokeOperatorFromPython(
+    const std::vector<std::shared_ptr<Operator>>& operations,
+    const py::args& args,
+    const py::kwargs& kwargs,
+    std::optional<c10::DispatchKey> dk = std::nullopt);
+
+// Efficient overload (does not require vector allocation) of the
+// above for use from C++ code.
+py::object invokeOperatorFromPython(
+    c10::ArrayRef<std::shared_ptr<Operator>> operations,
+    const py::args& args,
+    const py::kwargs& kwargs,
+    std::optional<c10::DispatchKey> dk = std::nullopt);
+
+TORCH_PYTHON_API std::optional<py::object> _maybe_handle_torch_function(
+    const std::string& ns,
+    const std::string& method_name,
+    const std::string& overload_name,
+    bool is_overload,
+    const py::args& args,
+    const py::kwargs& kwargs);
+
+TORCH_PYTHON_API bool checkSchemaAllowFakeScriptObject(
+    const FunctionSchema& schema,
+    const py::args& args,
+    const py::kwargs& kwargs);
+
+TORCH_PYTHON_API py::object _get_operation_for_overload_or_packet(
+    const std::vector<std::shared_ptr<Operator>>& operations,
+    Symbol symbol,
+    const py::args& args,
+    const py::kwargs& kwargs,
+    bool is_overload,
+    std::optional<c10::DispatchKey> dk = std::nullopt);
+
+// Efficient overload (does not require vector allocation) of the
+// above for use from C++ code.
+py::object _get_operation_for_overload_or_packet(
+    c10::ArrayRef<std::shared_ptr<Operator>> operations,
+    Symbol symbol,
+    const py::args& args,
+    const py::kwargs& kwargs,
+    bool is_overload,
+    std::optional<c10::DispatchKey> dk = std::nullopt);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_arg_flatten.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_arg_flatten.h
new file mode 100644
index 0000000000000000000000000000000000000000..873901a15199313fc3358fa0013b7333ce570d46
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_arg_flatten.h
@@ -0,0 +1,124 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/hash.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/jit/python/pybind.h>
+
+#include <ATen/ATen.h>
+#include <functional>
+#include <tuple>
+#include <vector>
+
+namespace torch::jit::python {
+
+struct IODescriptor {
+  struct VariableMetadata {
+    VariableMetadata(const autograd::Variable& var)
+        : sizes(var.sizes().vec()),
+          type(var.scalar_type()),
+          device(var.device()),
+          requires_grad(var.requires_grad()) {}
+
+    bool operator==(const VariableMetadata& o) const {
+      return std::tie(device, requires_grad, type, sizes) ==
+          std::tie(o.device, o.requires_grad, o.type, o.sizes);
+    }
+
+    static size_t hash(const VariableMetadata& m) {
+      return c10::get_hash(m.sizes, m.device, m.requires_grad, m.type);
+    }
+
+    std::vector<int64_t> sizes;
+    at::ScalarType type;
+    at::Device device;
+    bool requires_grad;
+  };
+
+  bool operator==(const IODescriptor& o) const {
+    return std::tie(structure, metadata, grad_enabled) ==
+        std::tie(o.structure, o.metadata, o.grad_enabled);
+  }
+
+  static size_t hash(const IODescriptor& o) {
+    return c10::get_hash(o.structure, o.metadata, o.grad_enabled);
+  }
+
+  void extend(const autograd::variable_list& list) {
+    metadata.reserve(metadata.size() + list.size());
+    for (auto& var : list)
+      metadata.emplace_back(var);
+  }
+
+  // Description of argument structure. Variables are replaced with
+  // different characters, depending on their flags, beginnings and
+  // ends of tuples and lists are denoted by a pair of parenthesis
+  // of their corresponding kind. They should always be paired.
+  // Example desc: (vv[v(v)v])
+  // NOTE: if extend() was ever called then metadata.size() can be
+  // different than the number of 'v's in structure.
+  std::string structure;
+  std::vector<std::string> strings;
+  std::vector<VariableMetadata> metadata;
+  bool grad_enabled = false;
+};
+
+static inline std::ostream& operator<<(
+    std::ostream& out,
+    const IODescriptor::VariableMetadata& meta) {
+  at::Device meta_device = meta.device;
+  auto& t = at::getDeprecatedTypeProperties(
+      meta_device.is_cpu() ? at::Backend::CPU : at::Backend::CUDA, meta.type);
+  out << t << "(requires_grad=" << meta.requires_grad;
+  if (meta_device.is_cuda()) {
+    out << ", device=" << meta_device.index();
+  }
+  out << ") {";
+  for (const auto i : c10::irange(meta.sizes.size())) {
+    if (i > 0)
+      out << ", ";
+    out << meta.sizes[i];
+  }
+  out << '}';
+  return out;
+}
+
+static inline std::ostream& operator<<(
+    std::ostream& out,
+    const IODescriptor& desc) {
+  out << desc.structure << '\n';
+  out << "  with grad_enabled=" << desc.grad_enabled << '\n';
+  for (const auto i : c10::irange(desc.metadata.size())) {
+    out << "  with v" << i << " having type " << desc.metadata[i] << '\n';
+  }
+  return out;
+}
+
+struct ParsedArgs {
+  // Flat vector of Variables found in arguments
+  autograd::variable_list vars;
+  // Metadata describing nesting of objects received from Python and
+  // metadata of vars and whether grad is enabled.
+  IODescriptor desc;
+
+  void extend(const autograd::variable_list& list) {
+    if (list.empty())
+      return;
+    vars.reserve(vars.size() + list.size());
+    for (auto& var : list)
+      vars.emplace_back(var);
+    desc.extend(list);
+  }
+};
+
+ParsedArgs flatten(py::handle obj);
+PyObject* unflatten(
+    at::ArrayRef<autograd::Variable> vars,
+    const IODescriptor& structure);
+
+} // namespace torch::jit::python
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_custom_class.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_custom_class.h
new file mode 100644
index 0000000000000000000000000000000000000000..50af856c35ee8bf3e0c92676e2b71724631d5a82
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_custom_class.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/pybind.h>
+#include <torch/custom_class.h>
+
+namespace torch::jit {
+
+void initPythonCustomClassBindings(PyObject* module);
+
+struct ScriptClass {
+  ScriptClass(c10::StrongTypePtr class_type)
+      : class_type_(std::move(class_type)) {}
+
+  py::object __call__(const py::args& args, const py::kwargs& kwargs);
+
+  c10::StrongTypePtr class_type_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_dict.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_dict.h
new file mode 100644
index 0000000000000000000000000000000000000000..e698d5ad0ebfa6b500e22be43c8a787e9689883a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_dict.h
@@ -0,0 +1,132 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/Dict.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::jit {
+
+void initScriptDictBindings(PyObject* module);
+
+/// An iterator over the keys of ScriptDict. This is used to support
+/// .keys() and iteration.
+class ScriptDictKeyIterator final {
+ public:
+  ScriptDictKeyIterator(
+      c10::impl::GenericDict::iterator iter,
+      c10::impl::GenericDict::iterator end)
+      : iter_(std::move(iter)), end_(std::move(end)) {}
+  at::IValue next();
+
+ private:
+  c10::impl::GenericDict::iterator iter_;
+  c10::impl::GenericDict::iterator end_;
+};
+
+/// An iterator over the key-value pairs of ScriptDict. This is used to support
+/// .items().
+class ScriptDictIterator final {
+ public:
+  ScriptDictIterator(
+      c10::impl::GenericDict::iterator iter,
+      c10::impl::GenericDict::iterator end)
+      : iter_(std::move(iter)), end_(std::move(end)) {}
+  at::IValue next();
+
+ private:
+  c10::impl::GenericDict::iterator iter_;
+  c10::impl::GenericDict::iterator end_;
+};
+
+/// A wrapper around c10::Dict that can be exposed in Python via pybind
+/// with an API identical to the Python dictionary class. This allows
+/// dictionaries to have reference semantics across the Python/TorchScript
+/// boundary.
+class ScriptDict final {
+ public:
+  // Constructor.
+  ScriptDict(const at::IValue& data)
+      : dict_(at::AnyType::get(), at::AnyType::get()) {
+    TORCH_INTERNAL_ASSERT(data.isGenericDict());
+    dict_ = data.toGenericDict();
+  }
+
+  // Get the type of the dictionary.
+  at::DictTypePtr type() const {
+    return at::DictType::create(dict_.keyType(), dict_.valueType());
+  }
+
+  // Return a string representation that can be used
+  // to reconstruct the instance.
+  std::string repr() const {
+    std::ostringstream s;
+    s << '{';
+    bool f = false;
+    for (auto const& kv : dict_) {
+      if (f) {
+        s << ", ";
+      }
+      s << kv.key() << ": " << kv.value();
+      f = true;
+    }
+    s << '}';
+    return s.str();
+  }
+
+  // Return an iterator over the keys of the dictionary.
+  ScriptDictKeyIterator iter() const {
+    auto begin = dict_.begin();
+    auto end = dict_.end();
+    return ScriptDictKeyIterator(begin, end);
+  }
+
+  // Return an iterator over the key-value pairs of the dictionary.
+  ScriptDictIterator items() const {
+    auto begin = dict_.begin();
+    auto end = dict_.end();
+    return ScriptDictIterator(begin, end);
+  }
+
+  // Interpret the dictionary as a boolean; empty means false, non-empty means
+  // true.
+  bool toBool() const {
+    return !(dict_.empty());
+  }
+
+  // Get the value for the given key. Throws std::out_of_range if the key does
+  // not exist.
+  at::IValue getItem(const at::IValue& key) {
+    return dict_.at(key);
+  }
+
+  // Set the value for the given key.
+  void setItem(const at::IValue& key, const at::IValue& value) {
+    dict_.insert_or_assign(key, value);
+  }
+
+  // Check whether the dictionary contains the given key.
+  bool contains(const at::IValue& key) {
+    return dict_.contains(key);
+  }
+
+  // Delete the given key from the dictionary.
+  bool delItem(const at::IValue& key) {
+    return dict_.erase(key);
+  }
+
+  // Get the size of the dictionary.
+  int64_t len() const {
+    return dict_.size();
+  }
+
+  // A c10::Dict instance that holds the actual data.
+  c10::impl::GenericDict dict_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_ir.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_ir.h
new file mode 100644
index 0000000000000000000000000000000000000000..28c757a3d95dc33283441ea815486b67f1ef8eb9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_ir.h
@@ -0,0 +1,55 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/utils/object_ptr.h>
+
+namespace torch::jit {
+
+void initPythonIRBindings(PyObject* module);
+
+// execute a Python function, used for Ops we can't optimize but that we want to
+// optimize around
+struct ConcretePythonOp : public PythonOp {
+  static Symbol Kind;
+
+  ConcretePythonOp(Graph* graph) : PythonOp(graph, ::c10::prim::PythonOp) {}
+  ConcretePythonOp* init(
+      THPObjectPtr&& pyobj,
+      const std::string& cconv,
+      pyobj_list&& scalar_args) {
+    this->pyobj = std::move(pyobj);
+    this->scalar_args = std::move(scalar_args);
+    this->cconv = cconv;
+    return this;
+  }
+  // The Python object which contains the implementation of this function.
+  // This is either a class (non-legacy) or an object (legacy).  See
+  // TraceInterpreterState for execution semantics.
+  THPObjectPtr pyobj;
+  // The calling convention for the Python function.
+  // 'c' -- constant argument
+  // 'd' -- dynamic argument
+  std::string cconv;
+  // Scalar arguments to the Python function.  Not necessarily passed to
+  // the function in this order; see cconv for the correct order.
+  std::vector<THPObjectPtr> scalar_args;
+
+  std::string name() const override;
+  void cloneFrom(Node* other_) override;
+  Node* allocNewInstance(Graph* g) override {
+    return new ConcretePythonOp(g);
+  }
+  // recover the autograd.Function instance, if this PythonOp's function
+  // was originally SomeFunction.apply
+  // used in ONNX for discovering symbolics
+  std::optional<THPObjectPtr> autogradFunction() const override;
+  void writeScalars(std::ostream& out) const override;
+  void lint_python() const override;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_ivalue.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_ivalue.h
new file mode 100644
index 0000000000000000000000000000000000000000..fdb65d222f85a8551ae73f3532226a2d82e87804
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_ivalue.h
@@ -0,0 +1,116 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/ivalue.h>
+#include <pybind11/pybind11.h>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace py = pybind11;
+
+namespace c10::ivalue {
+
+// concrete ivalue Holder that hold a py::object
+struct C10_EXPORT ConcretePyObjectHolder final : PyObjectHolder {
+ public:
+  static c10::intrusive_ptr<PyObjectHolder> create(py::object py_obj) {
+    return c10::make_intrusive<ConcretePyObjectHolder>(std::move(py_obj));
+  }
+
+  static c10::intrusive_ptr<PyObjectHolder> create(const py::handle& handle) {
+    py::gil_scoped_acquire ag;
+    return c10::make_intrusive<ConcretePyObjectHolder>(
+        handle.cast<py::object>());
+  }
+
+  PyObject* getPyObject() override {
+    return py_obj_.ptr();
+  }
+
+  InferredType tryToInferType() override {
+    pybind11::gil_scoped_acquire ag;
+    return torch::jit::tryToInferType(py_obj_);
+  }
+
+  IValue toIValue(const TypePtr& type, std::optional<int32_t> N = std::nullopt)
+      override {
+    pybind11::gil_scoped_acquire ag;
+    return torch::jit::toIValue(py_obj_, type, N);
+  }
+
+  std::string toStr() override {
+    pybind11::gil_scoped_acquire ag;
+    return py::str(py_obj_);
+  }
+
+  std::vector<at::Tensor> extractTensors() override {
+    // We could implement this entirely in C++ via pybind11 but it turns out to
+    // be substantially slower. Namely, the total time taken by markCompleted on
+    // a CUDAFuture is 21.5us with this implementation, but goes up to 58.7us
+    // when using C++. The reason is unclear.
+    try {
+      pybind11::gil_scoped_acquire ag;
+
+#if IS_PYBIND_2_13_PLUS
+      PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object>
+          storage;
+      auto& extractorFn =
+          storage
+              .call_once_and_store_result([]() -> py::object {
+                return py::module_::import("torch._jit_internal")
+                    .attr("_extract_tensors");
+              })
+              .get_stored();
+#else
+      static py::object& extractorFn = *new py::object(
+          py::module::import("torch._jit_internal").attr("_extract_tensors"));
+#endif
+
+      return extractorFn(py_obj_).cast<std::vector<at::Tensor>>();
+    } catch (py::error_already_set& e) {
+      auto err = std::runtime_error(
+          c10::str("Cannot extract tensors from value: ", e.what()));
+      {
+        pybind11::gil_scoped_acquire ag;
+        e.restore();
+        PyErr_Clear();
+      }
+      throw std::runtime_error(err);
+    }
+  }
+
+  // Note [Destructing py::object]
+  // ~~~~~~~~~~~~~~~~~~~~~~~~~~
+  //
+  // (1) Why py_obj_ = py::none(); does not work. Because we also need to
+  // acquire GIL when destructing py::object of None that de-references None.
+  // https://docs.python.org/3/c-api/none.html#c.Py_RETURN_NONE
+  //
+  // https://stackoverflow.com/questions/15287590/why-should-py-increfpy-none-be-required-before-returning-py-none-in-c
+  //
+  // (2) Why we need to call dec_ref() explicitly. Because py::object of
+  // nullptr, on destruction, effectively does nothing because of it calls
+  // Py_XDECREF(NULL) underlying.
+  // https://docs.python.org/3/c-api/refcounting.html#c.Py_XDECREF
+  ~ConcretePyObjectHolder() override {
+    pybind11::gil_scoped_acquire ag;
+    py_obj_.dec_ref();
+    // explicitly setting PyObject* to nullptr to prevent py::object's dtor to
+    // decref on the PyObject again.
+    py_obj_.ptr() = nullptr;
+  }
+
+  // explicit construction to avoid erroneous implicit conversion and
+  // copy-initialization
+  explicit ConcretePyObjectHolder(py::object py_obj)
+      : py_obj_(std::move(py_obj)) {}
+
+ private:
+  py::object py_obj_;
+};
+
+} // namespace c10::ivalue
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_list.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_list.h
new file mode 100644
index 0000000000000000000000000000000000000000..a7e53f537833efdf013b13f67ba1b49a6762a10b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_list.h
@@ -0,0 +1,233 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/Dict.h>
+#include <ATen/core/List.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <pybind11/detail/common.h>
+#include <torch/csrc/utils/pybind.h>
+#include <cstddef>
+#include <optional>
+#include <stdexcept>
+
+namespace torch::jit {
+
+void initScriptListBindings(PyObject* module);
+
+/// An iterator over the elements of ScriptList. This is used to support
+/// __iter__(), .
+class ScriptListIterator final {
+ public:
+  ScriptListIterator(
+      c10::impl::GenericList::iterator iter,
+      c10::impl::GenericList::iterator end)
+      : iter_(iter), end_(end) {}
+  at::IValue next();
+  bool done() const;
+
+ private:
+  c10::impl::GenericList::iterator iter_;
+  c10::impl::GenericList::iterator end_;
+};
+
+/// A wrapper around c10::List that can be exposed in Python via pybind
+/// with an API identical to the Python list class. This allows
+/// lists to have reference semantics across the Python/TorchScript
+/// boundary.
+class ScriptList final {
+ public:
+  // TODO: Do these make sense?
+  using size_type = size_t;
+  using diff_type = ptrdiff_t;
+  using ssize_t = Py_ssize_t;
+
+  // Constructor for empty lists created during slicing, extending, etc.
+  ScriptList(const at::TypePtr& type) : list_(at::AnyType::get()) {
+    auto list_type = type->expect<at::ListType>();
+    list_ = c10::impl::GenericList(list_type);
+  }
+
+  // Constructor for instances based on existing lists (e.g. a
+  // Python instance or a list nested inside another).
+  ScriptList(const at::IValue& data) : list_(at::AnyType::get()) {
+    TORCH_INTERNAL_ASSERT(data.isList());
+    list_ = data.toList();
+  }
+
+  at::ListTypePtr type() const {
+    return at::ListType::create(list_.elementType());
+  }
+
+  // Return a string representation that can be used
+  // to reconstruct the instance.
+  std::string repr() const {
+    std::ostringstream s;
+    s << '[';
+    bool f = false;
+    for (auto const& elem : list_) {
+      if (f) {
+        s << ", ";
+      }
+      s << at::IValue(elem);
+      f = true;
+    }
+    s << ']';
+    return s.str();
+  }
+
+  // Return an iterator over the elements of the list.
+  ScriptListIterator iter() const {
+    auto begin = list_.begin();
+    auto end = list_.end();
+    return ScriptListIterator(begin, end);
+  }
+
+  // Interpret the list as a boolean; empty means false, non-empty means
+  // true.
+  bool toBool() const {
+    return !(list_.empty());
+  }
+
+  // Get the value for the given index.
+  at::IValue getItem(diff_type idx) {
+    idx = wrap_index(idx);
+    return list_.get(idx);
+  }
+
+  // Set the value corresponding to the given index.
+  void setItem(diff_type idx, const at::IValue& value) {
+    idx = wrap_index(idx);
+    return list_.set(idx, value);
+  }
+
+  // Check whether the list contains the given value.
+  bool contains(const at::IValue& value) {
+    for (const auto& elem : list_) {
+      if (elem == value) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  // Delete the item at the given index from the list.
+  void delItem(diff_type idx) {
+    idx = wrap_index(idx);
+    auto iter = list_.begin() + idx;
+    list_.erase(iter);
+  }
+
+  // Get the size of the list.
+  ssize_t len() const {
+    return list_.size();
+  }
+
+  // Count the number of times a value appears in the list.
+  ssize_t count(const at::IValue& value) const {
+    ssize_t total = 0;
+
+    for (const auto& elem : list_) {
+      if (elem == value) {
+        ++total;
+      }
+    }
+
+    return total;
+  }
+
+  // Remove the first occurrence of a value from the list.
+  void remove(const at::IValue& value) {
+    auto list = list_;
+
+    int64_t idx = -1, i = 0;
+
+    for (const auto& elem : list) {
+      if (elem == value) {
+        idx = i;
+        break;
+      }
+
+      ++i;
+    }
+
+    if (idx == -1) {
+      throw py::value_error();
+    }
+
+    list.erase(list.begin() + idx);
+  }
+
+  // Append a value to the end of the list.
+  void append(const at::IValue& value) {
+    list_.emplace_back(value);
+  }
+
+  // Clear the contents of the list.
+  void clear() {
+    list_.clear();
+  }
+
+  // Append the contents of an iterable to the list.
+  void extend(const at::IValue& iterable) {
+    list_.append(iterable.toList());
+  }
+
+  // Remove and return the element at the specified index from the list. If no
+  // index is passed, the last element is removed and returned.
+  at::IValue pop(std::optional<size_type> idx = std::nullopt) {
+    at::IValue ret;
+
+    if (idx) {
+      idx = wrap_index(*idx);
+      ret = list_.get(*idx);
+      list_.erase(list_.begin() + *idx);
+    } else {
+      ret = list_.get(list_.size() - 1);
+      list_.pop_back();
+    }
+
+    return ret;
+  }
+
+  // Insert a value before the given index.
+  void insert(const at::IValue& value, diff_type idx) {
+    // wrap_index cannot be used; idx == len() is allowed
+    if (idx < 0) {
+      idx += len();
+    }
+
+    if (idx < 0 || idx > len()) {
+      throw std::out_of_range("list index out of range");
+    }
+
+    list_.insert(list_.begin() + idx, value);
+  }
+
+  // A c10::List instance that holds the actual data.
+  c10::impl::GenericList list_;
+
+ private:
+  // Wrap an index so that it can safely be used to access
+  // the list. For list of size sz, this function can successfully
+  // wrap indices in the range [-sz, sz-1]
+  diff_type wrap_index(diff_type idx) {
+    auto sz = len();
+    if (idx < 0) {
+      idx += sz;
+    }
+
+    if (idx < 0 || idx >= sz) {
+      throw std::out_of_range("list index out of range");
+    }
+
+    return idx;
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_sugared_value.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_sugared_value.h
new file mode 100644
index 0000000000000000000000000000000000000000..8d907dfd8166b6c8270bee466b87627732652654
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_sugared_value.h
@@ -0,0 +1,383 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/frontend/concrete_module_type.h>
+#include <torch/csrc/jit/frontend/sugared_value.h>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace torch::jit {
+
+std::string typeString(py::handle h);
+
+inline std::shared_ptr<SugaredValue> toSimple(Value* v) {
+  return std::make_shared<SimpleValue>(v);
+}
+
+// NB: This should be the single entry-point for instantiating a SugaredValue
+// from a Python object. If you are adding support for converting a new Python
+// type, *add it in this function's implementation*.
+std::shared_ptr<SugaredValue> toSugaredValue(
+    py::object obj,
+    GraphFunction& m,
+    const SourceRange& loc,
+    bool is_constant = false);
+
+std::optional<StrongFunctionPtr> as_function(const py::object& obj);
+
+struct VISIBILITY_HIDDEN PythonValue : public SugaredValue {
+  PythonValue(
+      py::object the_self,
+      std::optional<py::object> rcb = std::nullopt,
+      Value* module_self = nullptr)
+      : self(std::move(the_self)),
+        rcb(std::move(rcb)),
+        moduleSelf_(module_self) {}
+
+  FunctionSchema getSchema(
+      const size_t n_args,
+      const size_t n_binders,
+      const SourceRange& loc);
+
+  // call it like a function, e.g. `outputs = this(inputs)`
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& m,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+  std::string kind() const override;
+
+  std::vector<std::shared_ptr<SugaredValue>> asTuple(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::optional<size_t>& size_hint = {}) override;
+
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  Value* asValue(const SourceRange& loc, GraphFunction& m) override {
+    throw(
+        ErrorReport(loc)
+        << kind() << " cannot be used as a value. "
+        << "Perhaps it is a closed over global variable? If so, please "
+        << "consider passing it in as an argument or use a local variable "
+        << "instead.");
+  }
+
+ protected:
+  py::object getattr(const SourceRange& loc, const std::string& name);
+
+  void checkForAddToConstantsError(std::stringstream& ss);
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  py::object self;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::optional<py::object> rcb;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  Value* moduleSelf_ = nullptr;
+};
+
+struct VISIBILITY_HIDDEN PythonModuleValue : public PythonValue {
+  explicit PythonModuleValue(py::object mod) : PythonValue(std::move(mod)) {}
+
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+};
+
+// Used for desugaring uses of the torch.cuda module. All the CUDA APIs with
+// torch.cuda.* are resolved using CUDAPythonModuleValue.
+struct VISIBILITY_HIDDEN CUDAPythonModuleValue : public PythonValue {
+  explicit CUDAPythonModuleValue(py::object mod)
+      : PythonValue(std::move(mod)) {}
+
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+};
+
+// Represents all the parameters of a module as a List[Tensor]
+struct VISIBILITY_HIDDEN ConstantParameterList : public SugaredValue {
+  ConstantParameterList(Value* the_list) : the_list_(the_list) {}
+  std::string kind() const override {
+    return "constant parameter list";
+  }
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& caller,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override {
+    return toSimple(the_list_);
+  }
+
+ private:
+  Value* the_list_;
+};
+
+struct VISIBILITY_HIDDEN ModuleDictMethod : public SugaredValue {
+  explicit ModuleDictMethod(SugaredValuePtr iterable, std::string name)
+      : iterable_(std::move(iterable)), name_(std::move(name)) {}
+
+  std::string kind() const override {
+    return name_;
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& f,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override {
+    if (!args.empty() || !kwargs.empty()) {
+      throw(
+          ErrorReport(loc) << name_ << " method does not accept any arguments");
+    }
+    return iterable_;
+  }
+
+  SugaredValuePtr iterable_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const std::string name_;
+};
+
+struct SugaredDict;
+
+// defines how modules/methods behave inside the script subset.
+// for now this does not have any interaction with python.
+// in the future, we will add the ability to resolve `self.foo` to python
+// {functions, modules, constants} so this SugaredValue is defined here
+// anticipating we will eventually need to replace Module with a py::object
+// holding the actual nn.Module class.
+
+struct VISIBILITY_HIDDEN ModuleValue : public SugaredValue {
+  ModuleValue(Value* self, std::shared_ptr<ConcreteModuleType> concreteType)
+      : self_(self), concreteType_(std::move(concreteType)) {}
+
+  std::string kind() const override {
+    return "module";
+  }
+
+  Value* asValue(const SourceRange& loc, GraphFunction& m) override;
+
+  SugaredValuePtr asTupleValue(const SourceRange& loc, GraphFunction& m)
+      override;
+
+  // select an attribute on it, e.g. `this.field`
+  std::shared_ptr<SugaredValue> tryGetAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field);
+
+  // select an attribute on it, e.g. `this.field`
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  // select an attribute on it, e.g. `this.field`
+  bool hasAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  // call module.forward with pre_hooks and hooks
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& caller,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+  std::shared_ptr<SugaredDict> getSugaredDict(
+      const SourceRange& loc,
+      GraphFunction& m);
+
+  std::shared_ptr<SugaredDict> getSugaredNamedBufferDict(
+      const SourceRange& loc,
+      GraphFunction& m);
+
+  std::shared_ptr<SugaredDict> getSugaredNamedParameterList(
+      const SourceRange& loc,
+      GraphFunction& m);
+
+  std::shared_ptr<SugaredDict> getSugaredNamedParameterDict(
+      const SourceRange& loc,
+      GraphFunction& m);
+
+  void setAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field,
+      Value* newValue) override;
+
+  SugaredValuePtr iter(const SourceRange& loc, GraphFunction& m) override;
+
+  std::shared_ptr<SugaredValue> getitem(
+      const SourceRange& loc,
+      GraphFunction& m,
+      Value* idx,
+      TypePtr type_hint) override;
+
+ private:
+  // Check that the type of all submodules is a subtype of ty. If the function
+  // returns false, more information about why it returns false (e.g. which
+  // submodule's type is not a subtype of ty) is printed it why_not if it is not
+  // null.
+  bool areAllSubmodulesSubtypeOf(
+      const TypePtr& ty,
+      std::ostream* why_not = nullptr) const;
+
+  Value* self_;
+  std::shared_ptr<ConcreteModuleType> concreteType_;
+};
+
+bool isNamedTupleClass(const py::object& obj);
+TypePtr registerNamedTuple(
+    const py::object& obj,
+    const SourceRange& loc,
+    const ResolutionCallback& rcb);
+
+void recurseThroughNestedModules(
+    const SourceRange& loc,
+    GraphFunction& m,
+    std::vector<SugaredValuePtr>& keys,
+    std::vector<SugaredValuePtr>& values,
+    std::shared_ptr<ModuleValue>& self,
+    const std::string& prefix,
+    const std::string& field);
+
+// Used to support named_modules()
+struct VISIBILITY_HIDDEN SugaredDict : public SugaredValue {
+  explicit SugaredDict(
+      std::shared_ptr<ModuleValue> self,
+      std::shared_ptr<SugaredTupleValue> keys,
+      std::shared_ptr<SugaredTupleValue> modules)
+      : self_(std::move(self)),
+        keys_(std::move(keys)),
+        modules_(std::move(modules)) {}
+
+  std::string kind() const override {
+    return "ModuleDict";
+  }
+
+  std::shared_ptr<SugaredTupleValue> getKeys() {
+    return keys_;
+  }
+
+  std::shared_ptr<SugaredTupleValue> getModules() {
+    return modules_;
+  }
+
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  SugaredValuePtr iter(const SourceRange& loc, GraphFunction& m) override {
+    return keys_;
+  }
+
+  std::shared_ptr<ModuleValue> self_;
+  std::shared_ptr<SugaredTupleValue> keys_;
+  std::shared_ptr<SugaredTupleValue> modules_;
+};
+
+struct VISIBILITY_HIDDEN BooleanDispatchValue : public SugaredValue {
+  BooleanDispatchValue(py::dict dispatched_fn)
+      : dispatched_fn_(std::move(dispatched_fn)) {}
+
+  std::string kind() const override {
+    return "boolean dispatch";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& caller,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+ private:
+  py::dict dispatched_fn_;
+};
+
+struct VISIBILITY_HIDDEN PythonClassValue : public ClassValue {
+  PythonClassValue(ClassTypePtr type, py::object py_type)
+      : ClassValue(std::move(type)), py_type_(std::move(py_type)) {}
+
+  std::string kind() const override {
+    return "Python type";
+  }
+
+  std::shared_ptr<SugaredValue> attr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+  bool hasAttr(
+      const SourceRange& loc,
+      GraphFunction& m,
+      const std::string& field) override;
+
+ private:
+  py::object py_type_;
+};
+
+struct VISIBILITY_HIDDEN PythonExceptionValue : public ExceptionValue {
+  explicit PythonExceptionValue(const py::object& exception_class)
+      : ExceptionValue(
+            py::str(py::getattr(exception_class, "__name__", py::str("")))),
+        exception_class_qualified_name_(
+            py::str(py::module::import("torch._jit_internal")
+                        .attr("_qualified_name")(
+                            exception_class,
+                            /*mangle_name=*/false))) {}
+
+  std::string kind() const override {
+    return "Python exception";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& caller,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+
+ private:
+  std::string exception_class_qualified_name_;
+};
+
+// Python Slice class.
+struct VISIBILITY_HIDDEN PythonSliceClass : public SugaredValue {
+  explicit PythonSliceClass() = default;
+
+  std::string kind() const override {
+    return "Python slice class";
+  }
+
+  std::shared_ptr<SugaredValue> call(
+      const SourceRange& loc,
+      GraphFunction& caller,
+      at::ArrayRef<NamedValue> args,
+      at::ArrayRef<NamedValue> kwargs,
+      size_t n_binders) override;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_tracer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_tracer.h
new file mode 100644
index 0000000000000000000000000000000000000000..d6fc5f502fb4820b8d45343d1fa104e4eaabb32e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_tracer.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/frontend/tracer.h>
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include <memory>
+#include <string>
+
+namespace torch::jit {
+
+struct Module;
+
+namespace tracer {
+void initPythonTracerBindings(PyObject* module);
+
+SourceRange getPythonInterpreterSourceRange();
+
+Node* preRecordPythonTrace(
+    THPObjectPtr pyobj,
+    const std::string& arg_types,
+    at::ArrayRef<autograd::Variable> inputs,
+    std::vector<THPObjectPtr> scalar_args);
+
+std::pair<std::shared_ptr<Graph>, Stack> createGraphByTracingWithDict(
+    const py::function& func,
+    const py::dict& inputs_dict,
+    const Stack& inputs,
+    const py::function& var_name_lookup_fn,
+    bool strict,
+    bool force_outplace,
+    Module* self = nullptr,
+    const std::vector<std::string>& argument_names = {});
+
+std::pair<std::shared_ptr<Graph>, Stack> createGraphByTracing(
+    const py::function& func,
+    Stack inputs,
+    const py::function& var_name_lookup_fn,
+    bool strict,
+    bool force_outplace,
+    Module* self = nullptr,
+    const std::vector<std::string>& argument_names = {});
+} // namespace tracer
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_tree_views.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_tree_views.h
new file mode 100644
index 0000000000000000000000000000000000000000..922babfff9ddd475c6f6bbf1da32080f71c39a32
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/python_tree_views.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::jit {
+
+void initTreeViewBindings(PyObject* module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/script_init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/script_init.h
new file mode 100644
index 0000000000000000000000000000000000000000..29aea54b291cb43ce5c921dd30ebaf98e209d27c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/script_init.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/python/pybind.h>
+
+namespace torch::jit {
+void initJitScriptBindings(PyObject* module);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/update_graph_executor_opt.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/update_graph_executor_opt.h
new file mode 100644
index 0000000000000000000000000000000000000000..0ee12ba1cfabddba9b497b189fe756d7cdf3d221
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/update_graph_executor_opt.h
@@ -0,0 +1,11 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+namespace torch::jit {
+TORCH_API void setGraphExecutorOptimize(bool o);
+TORCH_API bool getGraphExecutorOptimize();
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/utf8_decoding_ignore.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/utf8_decoding_ignore.h
new file mode 100644
index 0000000000000000000000000000000000000000..8e0f57b2114203af5e68b519cf3b2d4ef2baaaf2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/python/utf8_decoding_ignore.h
@@ -0,0 +1,11 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+namespace torch::jit {
+TORCH_API void setUTF8DecodingIgnore(bool o);
+TORCH_API bool getUTF8DecodingIgnore();
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/resource_guard.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/resource_guard.h
new file mode 100644
index 0000000000000000000000000000000000000000..fb3655628ebcc7af02ad8b2c8d8a0ec32a74d6c5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/resource_guard.h
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <functional>
+
+namespace torch::jit {
+
+class ResourceGuard {
+  std::function<void()> _destructor;
+  bool _released{false};
+
+ public:
+  ResourceGuard(std::function<void()> destructor)
+      : _destructor(std::move(destructor)) {}
+
+  // NOLINTNEXTLINE(bugprone-exception-escape)
+  ~ResourceGuard() {
+    if (!_released)
+      _destructor();
+  }
+
+  void release() {
+    _released = true;
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/argument_spec.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/argument_spec.h
new file mode 100644
index 0000000000000000000000000000000000000000..72ec0fb8db841cae4d3558a1d5b9d170b1549477
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/argument_spec.h
@@ -0,0 +1,509 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+#include <c10/util/hash.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <ostream>
+#include <vector>
+
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wshorten-64-to-32")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wshorten-64-to-32")
+#endif
+
+namespace torch::jit {
+
+// GraphExecutor creates specializations of Graphs for different
+// dimensionalitities and types of inputs.
+
+struct ArgumentInfo {
+  friend struct ArgumentSpec;
+  using plain_data_type = uint64_t;
+
+  bool defined() const {
+    return defined_;
+  }
+  at::Device device() const {
+    return at::Device(DeviceType(dev_type_), device_);
+  }
+  // XXX: It is guaranteed that this will return false when called on non-tensor
+  // arguments
+  bool requires_grad() const {
+    return requires_grad_;
+  }
+  int dim() const {
+    return dim_;
+  }
+  at::ScalarType type() const {
+    return at::ScalarType(type_);
+  }
+  TypePtr toType() const {
+    if (!defined())
+      return TensorType::get();
+
+    return TensorType::create(
+        type(), device(), std::optional<size_t>(dim()), requires_grad());
+  }
+  operator TypePtr() const {
+    return toType();
+  }
+
+ private:
+  unsigned defined_ : 1;
+  unsigned requires_grad_ : 1;
+  unsigned : 5;
+  unsigned dim_ : 8;
+  unsigned device_ : 8;
+  unsigned type_ : 8;
+  unsigned dev_type_ : 16;
+  unsigned : 16;
+};
+
+static_assert(
+    std::is_standard_layout_v<ArgumentInfo>,
+    "ArgumentInfo is to be a POD struct");
+static_assert(
+    sizeof(ArgumentInfo) == sizeof(ArgumentInfo::plain_data_type),
+    "ArgumentInfo is expected to be a 32-bit struct");
+
+struct ArgumentSpec {
+  ArgumentSpec(size_t num_flat_tensor_inputs, size_t num_flat_optional_inputs)
+      : hash_code(c10::hash_combine(
+            num_flat_tensor_inputs,
+            num_flat_optional_inputs)) {
+    tensor_args.reserve(num_flat_tensor_inputs);
+    optional_presence.reserve(num_flat_optional_inputs);
+  }
+
+  void addOptional(const IValue& input) {
+    bool is_present = !input.isNone();
+    optional_presence.push_back(is_present);
+    hash_code = c10::hash_combine(hash_code, is_present);
+  }
+
+  void addTensor(const IValue& input, bool with_grad) {
+    AT_ASSERT(input.isTensor(), "Expected Tensor but found ", input.tagKind());
+    tensor_args.emplace_back();
+    auto& arg = tensor_args.back();
+    // Initialize all fields to 0. This is convenient, because e.g.
+    // requires_grad() can be checked even on tensors AND will make
+    // padding bits all 0s.
+    std::memset(&arg, 0, sizeof(ArgumentInfo));
+
+    // [argspec refcounting] reinterpret the IValue to avoid having to refcount
+    // the Tensor microbenchmarks
+    // https://github.com/zdevito/pytorch/commit/21e7200a0a0fc456bea2f10e95b1781f83933d10
+    // show overhead in extra refcounting along this path
+    const at::Tensor* t = reinterpret_cast<const at::Tensor*>(&input);
+    arg.defined_ = t->defined();
+    if (arg.defined_) {
+      arg.requires_grad_ = with_grad && t->requires_grad();
+      arg.dim_ = t->dim();
+      at::Device device = t->device();
+      arg.dev_type_ =
+          // NOLINTNEXTLINE(bugprone-signed-char-misuse)
+          static_cast<std::underlying_type_t<DeviceType>>(device.type());
+      // NOLINTNEXTLINE(bugprone-signed-char-misuse)
+      arg.device_ = device.index();
+      arg.type_ = static_cast<unsigned>(t->scalar_type());
+    }
+    combineHash(arg);
+  }
+
+  void combineHash(const ArgumentInfo& arg) {
+    ArgumentInfo::plain_data_type arg_data = 0;
+    std::memcpy(&arg_data, &arg, sizeof(ArgumentInfo));
+    hash_code = c10::hash_combine(hash_code, arg_data);
+  }
+
+  // equality is fast: check ninputs, and then check the raw array data,
+  // there are no size/stride indirections
+  // hopefully std::vector<bool> has fast equality
+  bool operator==(const ArgumentSpec& spec) const {
+    if (optional_presence != spec.optional_presence) {
+      return false;
+    }
+    if (tensor_args.size() != spec.tensor_args.size())
+      return false;
+    // NB: we need to break out early when there are no elements, because
+    // passing a nullptr to memcmp is UB.
+    if (tensor_args.empty())
+      return true;
+    return std::memcmp(
+               tensor_args.data(),
+               spec.tensor_args.data(),
+               tensor_args.size() * sizeof(ArgumentInfo)) == 0;
+  }
+  bool operator!=(const ArgumentSpec& spec) const {
+    return !(*this == spec);
+  }
+  size_t numTensors() const {
+    return tensor_args.size();
+  }
+  const ArgumentInfo& tensorAt(size_t i) const {
+    return tensor_args[i];
+  }
+  size_t numOptionals() const {
+    return optional_presence.size();
+  }
+  bool isPresent(size_t i) const {
+    return optional_presence[i];
+  }
+  size_t hashCode() const noexcept {
+    return hash_code;
+  }
+
+ private:
+  size_t hash_code; // precomputed on construction
+  std::vector<ArgumentInfo> tensor_args;
+  std::vector<bool> optional_presence;
+};
+
+namespace {
+static constexpr size_t ARG_SPEC_DEPTH_LIMIT = 128;
+}
+
+// ArgumentSpecCreator takes an initial graph and comes up with a set
+// of simple instructions to compute the ArgumentSpec given a set of
+// input tensors.
+struct TORCH_API ArgumentSpecCreator {
+  // instructs acts on a stack of a list of input IValues
+  // at the beginning the stack contains a single list of the inputs to the
+  // function the ENTER_ instructs descend into subobjects and push new lists
+  // onto the stack
+  enum Inst : char {
+    ENTER_TUPLE, // consume a tuple ivalue from the top-most list, and push the
+                 // list of its elements onto the stack as a new list
+    ENTER_OBJECT, // same as ENTER_TUPLE, but the input is a class
+    LEAVE, // pop the top-most list from the stack
+    SKIP, // consume an element from the top-most list, and discard
+    SPECIALIZE_OPTIONAL_TENSOR, // consume a optional tensor for the top-most
+                                // list, and add it to the ArgSpec key being
+                                // created
+    SPECIALIZE_TENSOR, // consume a tensor for the top-most
+                       // list, and add it to the ArgSpec key being created
+    SPECIALIZE_OPTIONAL,
+    // consume a nontensor optional from the top-most list,
+    // and add it to the ArgSpec key being created
+  };
+  ArgumentSpecCreator(Graph& graph);
+  ArgumentSpec create(bool with_grad, const Stack& stack) const;
+  void specializeTypes(Graph& g, const ArgumentSpec& spec) const;
+  void dump() const;
+  using WrittenSlots = std::unordered_set<std::string>;
+
+ private:
+  void scan(
+      const TypePtr& typ,
+      size_t depth,
+      const WrittenSlots& written_slots);
+  size_t num_inputs_;
+  size_t num_tensors_ = 0;
+  size_t num_optionals_ = 0;
+  std::vector<Inst> instructions_;
+};
+
+// CompleteArgumentSpec represents one particular specialization.
+// It is designed so that it can be created, hashed, and compared quickly
+// since it is used along the hot-path of the JIT to check if the code
+// we have created is valid for the given inputs.
+
+// COmpleteArgumentInfoPOD is only used internally in CompleteArgumentSpec
+// API users should use ArgumentInfo
+struct CompleteArgumentInfoPOD {
+  // total size is 64-bit
+  unsigned is_tensor : 8; // all other fields are invalid if this is false
+  unsigned type : 8; // scalar type
+  unsigned defined : 1;
+  unsigned requires_grad : 1;
+  signed device : 14;
+  unsigned dev_type : 16;
+  unsigned
+      total_dims : 16; // all TensorInfoPODs are in CompleteArgumentSpec's
+                       // tensor_info() array. total_dims is the total number of
+                       // dimensions seen so far in all previous members of
+                       // tensor_info(), including this tensor 2*total_dims
+                       // becomes the offset into the sizes_strides list for the
+                       // _next_ tensor in the tensor_info array for tensor 0,
+                       // the offset is always 0
+};
+
+static_assert(
+    sizeof(CompleteArgumentInfoPOD) == sizeof(int64_t),
+    "CompleteArgumentInfoPOD must be 64-bit struct for CompleteArgumentSpec encoding to work");
+
+struct CompleteArgumentInfo;
+
+struct CompleteArgumentSpec {
+  CompleteArgumentSpec(bool with_grad, at::ArrayRef<IValue> inputs)
+      : ninputs(inputs.size()) {
+    int64_t all_dims = 0;
+    const auto num_inputs = inputs.size();
+    for (const auto i : c10::irange(num_inputs)) {
+      if (!inputs[i].isTensor())
+        continue;
+      auto& tensor = inputs[i].toTensor();
+      all_dims += tensor.defined() ? tensor.ndimension() : 0;
+    }
+    // allocate enough room for all TensorPODs and dimensions
+    data.resize(ninputs + all_dims * 2);
+
+    // and reinterpret our data array as these structs
+    auto* pods = reinterpret_cast<CompleteArgumentInfoPOD*>(data.data());
+    int64_t* next_dim = sizes_strides();
+    int32_t total_dims = 0;
+    for (const auto i : c10::irange(num_inputs)) {
+      auto& pod = pods[i];
+      pod.is_tensor = static_cast<uint32_t>(inputs[i].isTensor());
+      if (pod.is_tensor) {
+        at::Tensor t = inputs[i].toTensor();
+        pod.defined = t.defined();
+        if (pod.defined) {
+          pod.type = static_cast<int>(t.scalar_type());
+          at::Device device = t.device();
+          // NOLINTNEXTLINE(bugprone-signed-char-misuse)
+          pod.dev_type =
+              static_cast<std::underlying_type_t<DeviceType>>(device.type());
+          // NOLINTNEXTLINE(bugprone-signed-char-misuse)
+          pod.device = device.index();
+          pod.requires_grad = with_grad && t.requires_grad();
+          total_dims += t.ndimension();
+          auto sizes = t.sizes();
+          std::copy(sizes.begin(), sizes.end(), next_dim);
+          next_dim += sizes.size();
+          auto strides = t.strides();
+          std::copy(strides.begin(), strides.end(), next_dim);
+          next_dim += strides.size();
+        }
+      }
+      // each POD has a running tally of all dimensions including its own
+      TORCH_CHECK(
+          total_dims < std::numeric_limits<uint16_t>::max(),
+          "The number of dims cannot be packed into CompleteArgumentSpec:",
+          total_dims);
+      pod.total_dims = total_dims;
+    }
+    // we precompute the hash_code to minimize the time inside of hash
+    // table operations where we may need to hold a compiler cache lock.
+    hash_code = c10::hash_combine(0, ninputs);
+    for (auto d : data) {
+      hash_code = c10::hash_combine(hash_code, d);
+    }
+  }
+
+  // equality is fast: check ninputs, and then check the raw array data,
+  // there are no size/stride indirections
+  bool operator==(const CompleteArgumentSpec& spec) const {
+    return ninputs == spec.ninputs && data == spec.data;
+  }
+  bool operator!=(const CompleteArgumentSpec& spec) const {
+    return !(*this == spec);
+  }
+  friend struct CompleteArgumentInfo;
+  CompleteArgumentInfo at(size_t i) const;
+  size_t size() const {
+    return ninputs;
+  }
+  size_t hashCode() const noexcept {
+    return hash_code;
+  }
+
+ private:
+  ArrayRef<CompleteArgumentInfoPOD> tensor_info() const {
+    return ArrayRef<CompleteArgumentInfoPOD>(
+        reinterpret_cast<const CompleteArgumentInfoPOD*>(data.data()), ninputs);
+  }
+  // the start of the sizes_strides information, which comes after the
+  // CompleteArgumentInfoPOD list.
+  const int64_t* sizes_strides() const {
+    return data.data() + ninputs;
+  }
+  int64_t* sizes_strides() {
+    return data.data() + ninputs;
+  }
+  size_t hash_code{0}; // precomputed on construction
+  size_t ninputs;
+  // layout is ninputs of TensorPOD (each 64-bit) followed by their size and
+  // stride info for 3 tensors:
+  // [t0POD][t1POD][t2POD]...
+  // [t0 sizes][t0 strides][t1 sizes][t1 strides][t2 sizes][t2 strides]
+  std::vector<int64_t> data;
+};
+
+// public view of compressed CompleteArgumentInfo
+struct CompleteArgumentInfo {
+  CompleteArgumentInfo(const CompleteArgumentSpec& spec, const int i)
+      : spec(spec), i(i) {}
+  bool isTensor() const {
+    return pod(i).is_tensor;
+  }
+  at::ScalarType type() const {
+    return at::ScalarType(pod(i).type);
+  }
+  bool defined() const {
+    return pod(i).defined;
+  }
+  bool requires_grad() const {
+    return pod(i).requires_grad;
+  }
+  at::Device device() const {
+    return at::Device(
+        DeviceType(pod(i).dev_type),
+        static_cast<c10::DeviceIndex>(pod(i).device));
+  }
+  int ndimension() const {
+    // See [valid range], it is always valid to ask for offset for (i + 1)
+    return (sizes_strides_offset(i + 1) - sizes_strides_offset(i)) / 2;
+  }
+  at::IntArrayRef sizes() const {
+    return at::IntArrayRef(
+        spec.sizes_strides() + sizes_strides_offset(i), ndimension());
+  }
+  at::IntArrayRef strides() const {
+    int ndim = ndimension();
+    return at::IntArrayRef(
+        spec.sizes_strides() + sizes_strides_offset(i) + ndim, ndim);
+  }
+  operator TypePtr() const {
+    if (!defined())
+      return TensorType::get();
+    return TensorType::create(
+        type(),
+        device(),
+        c10::VaryingShape<int64_t>{sizes()},
+        c10::VaryingShape<int64_t>{strides()},
+        requires_grad());
+  }
+
+ private:
+  // offsetinto sizes_strides() array where the sizes start for tensor j
+  // [valid range] valid range is [0, ninputs]
+  // (i.e. you can ask for the offset at ninputs, which would be the offset of
+  // the next tensor if it existed)
+  int sizes_strides_offset(int j) const {
+    if (j == 0)
+      return 0;
+    return 2 * pod(j - 1).total_dims;
+  }
+  const CompleteArgumentInfoPOD& pod(int j) const {
+    return spec.tensor_info().at(j);
+  }
+  const CompleteArgumentSpec& spec;
+  const int i;
+};
+
+inline std::ostream& operator<<(std::ostream& out, const ArgumentInfo& info) {
+  if (!info.defined()) {
+    return out << "<undefined>";
+  }
+  out << "Tensor(device=" << info.device() << ", type=" << toString(info.type())
+      << ", requires_grad=" << info.requires_grad() << ", dims=" << info.dim()
+      << ')';
+  return out;
+}
+
+inline std::ostream& operator<<(std::ostream& out, const ArgumentSpec& spec) {
+  out << '{';
+  for (const auto i : c10::irange(spec.numTensors())) {
+    if (i > 0)
+      out << ", ";
+    out << spec.tensorAt(i);
+  }
+  out << "; ";
+  for (const auto i : c10::irange(spec.numOptionals())) {
+    if (i > 0)
+      out << ", ";
+    out << spec.isPresent(i);
+  }
+  out << '}';
+  return out;
+}
+
+inline std::ostream& operator<<(
+    std::ostream& out,
+    const CompleteArgumentInfo& info) {
+  if (!info.defined()) {
+    return out << "<undefined>";
+  }
+  out << "Tensor(device=" << info.device() << ", type=" << toString(info.type())
+      << ", requires_grad=" << info.requires_grad()
+      << ", sizes=" << info.sizes() << ", strides=" << info.strides() << ')';
+  return out;
+}
+
+inline std::ostream& operator<<(
+    std::ostream& out,
+    const CompleteArgumentSpec& spec) {
+  out << '{';
+  for (const auto i : c10::irange(spec.size())) {
+    if (i > 0)
+      out << ", ";
+    out << spec.at(i);
+  }
+  out << '}';
+  return out;
+}
+
+inline CompleteArgumentInfo CompleteArgumentSpec::at(size_t i) const {
+  return CompleteArgumentInfo(*this, i);
+}
+
+inline std::optional<int8_t> convertOptional(
+    std::optional<c10::ScalarType> const& from) {
+  return from ? std::optional<int8_t>(static_cast<int8_t>(*from))
+              : std::optional<int8_t>{};
+}
+
+} // namespace torch::jit
+
+namespace std {
+
+template <typename T>
+struct hash<c10::VaryingShape<T>> {
+  size_t operator()(const c10::VaryingShape<T>& vs) const {
+    return c10::get_hash(
+        vs.size(),
+        vs.size() ? vs.sizes().value() : std::vector<std::optional<T>>());
+  }
+};
+
+template <>
+struct hash<c10::TensorType> {
+  size_t operator()(const c10::TensorType& ptt) const {
+    return c10::get_hash<
+        std::optional<int8_t>,
+        c10::VaryingShape<int64_t>,
+        c10::VaryingShape<int64_t>,
+        std::optional<bool>>(
+        torch::jit::convertOptional(ptt.scalarType()),
+        ptt.sizes(),
+        ptt.strides(),
+        ptt.requiresGrad());
+  }
+};
+
+template <>
+struct hash<torch::jit::ArgumentSpec> {
+  size_t operator()(const torch::jit::ArgumentSpec& spec) const noexcept {
+    return spec.hashCode();
+  }
+};
+template <>
+struct hash<torch::jit::CompleteArgumentSpec> {
+  size_t operator()(
+      const torch::jit::CompleteArgumentSpec& spec) const noexcept {
+    return spec.hashCode();
+  }
+};
+} // namespace std
+
+C10_CLANG_DIAGNOSTIC_POP()
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/autodiff.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/autodiff.h
new file mode 100644
index 0000000000000000000000000000000000000000..9c47f7a9f08042a997b52e1d579042ad133c02f5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/autodiff.h
@@ -0,0 +1,99 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <memory>
+#include <vector>
+
+namespace torch::jit {
+
+using value_list = std::vector<Value*>;
+// clang-format off
+// Example showcasing how Gradient is constructed:
+//
+// Let's assume we have a function f, `m` and `n` do not require grad
+// (`n` can depend only on `m`):
+//   y, n = f(x, m)
+//
+// Now, let's assume that the reverse of f (called f') needs to use values of `x`, `t` and `y`.
+// `t` is an intermediate value produced in the body of f, and let's assume that it requires
+// grad too.
+//
+// In this case differentiate(f) will return this:
+//   y, n, t = f(x, m)        // `t` is appended to the output list
+//   dx = f'(dy, dt, x, t, y) // No `dm` or `dn` because they do not require gradient
+//                            // All needed values from f are prepended to the input list
+//
+//   f_real_outputs = 2       // Only first two outputs were present in f originally
+//   df_input_vjps = {0, 2}   // i.e. connect grad_fn of y and t variables produced by f,
+//                    y  t    // with y's output_nr = 0 and t's output_nr = 1
+//   df_input_captures = {I0, O2, O0} // Order matches the prefix of inputs to df
+//                        x   t   y
+//   df_output_vjps = {0}     // i.e. connect next_edge[0] of grad_fn to x's (grad_fn, output_nr).
+//
+// Terminology: vjp = vector-jacobian product
+// clang-format on
+
+struct Gradient {
+  explicit operator bool() const {
+    return df != nullptr;
+  }
+  std::shared_ptr<Graph> f;
+  std::shared_ptr<Graph> df;
+
+  // Describes how to construct outputs of f from what its graph will return.
+  // This is necessary because some trailing outputs are intermediates produced
+  // only to be saved for df (and should be ignored).
+  size_t f_real_outputs = 0; // initialized for safety.
+
+  // df inputs are split into two sections: vjps (aka grad_outputs) and
+  // captures. VJPs are "seeds" for the gradient computation given for each
+  // input capture of an Output kind. Captures are values the need to be saved
+  // when f is run. We handle inputs specially, because this allows us to avoid
+  // adding extra vjps as df inputs.
+
+  std::vector<size_t> df_input_vjps; // Offsets into f's outputs.
+  // capture can come from inputs or outputs
+  std::vector<size_t> df_input_captured_inputs; // Offsets into f's inputs
+  std::vector<size_t> df_input_captured_outputs; // Offsets into f's outputs
+
+  // df will produce vjps for a subset of inputs of f that required grad.
+  // df_output_vjps[idx] == inp_idx means that idx-th output of df produces a
+  // vjp for inp_idx-th input of f.
+  std::vector<size_t> df_output_vjps; // Offsets into f's inputs.
+
+  // How to use gradient to implement a differentiable autograd function:
+  // When running f:
+  //   - Unwrap input Variables
+  //   - Run f's graph
+  //   - Create grad_fn
+  //   - Wrap outputs in Variables (assume we have a tensor_outputs array):
+  //       outputs = map(Variable, tensor_output)
+  //       for i, offset in enumerate(df_input_vjps):
+  //         outputs[offset].set_grad_fn(grad_fn, output_nr=i)
+  //   - Use df_output_vjps to connect next_edges of grad_fn:
+  //       for idx in df_output_vjps:
+  //         grad_fn.add_next_edge(inputs[idx].gradient_edge())
+  //   - Save captures for df (care needs to be taken to use SavedVariables for
+  //                           inputs and outputs that we will actually return)
+  //   - Return outputs[:f_real_outputs]
+  //
+  // When running df:
+  //   - Concatenate received vjps and captured Variables
+  //   - Interpret df
+  //   - Wrap outputs of df into Variables (that don't require grad)
+};
+TORCH_API Gradient differentiate(std::shared_ptr<Graph>& graph);
+
+// can we take a derivative of this node symbolically?
+TORCH_API bool isDifferentiable(const Node* n);
+TORCH_API bool isDifferentiable(Graph& g);
+TORCH_API bool isZero(Value* v);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/calculate_necessary_args.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/calculate_necessary_args.h
new file mode 100644
index 0000000000000000000000000000000000000000..4951616fbded8b83d0cf9d75459a29849c65822e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/calculate_necessary_args.h
@@ -0,0 +1,74 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/schema_matching.h>
+#include <cstddef>
+
+namespace torch::jit {
+
+// Calculates the number of args that need to be passed in.
+// Less args may be needed if defaults are provided.
+// Returns: {number args needed, number of out args}
+inline std::pair<int64_t, int64_t> CalculateNecessaryArgs(
+    const std::vector<Argument>& schema_args,
+    at::ArrayRef<Value*> actual_inputs,
+    bool allow_trailing_out_args) {
+  if (schema_args.empty()) {
+    return std::make_pair(0, 0);
+  }
+
+  // count number of out arguments
+  int64_t schema_idx = static_cast<int64_t>(schema_args.size()) - 1;
+  if (allow_trailing_out_args) {
+    // skip over out arguments in the end.
+    while (schema_idx >= 0) {
+      const auto& current_arg = schema_args.at(schema_idx);
+      if (!current_arg.is_out()) {
+        break;
+      }
+      schema_idx--;
+    }
+  }
+
+  int64_t num_out = static_cast<int64_t>(schema_args.size()) - schema_idx - 1;
+
+  if (schema_args.size() < actual_inputs.size()) {
+    return std::make_pair(actual_inputs.size(), num_out);
+  }
+
+  // if it is the default args, we reset the index to the last element
+  if (!allow_trailing_out_args) {
+    schema_idx = schema_args.size() - 1;
+  }
+  // keeps track of trailing unnecessary args
+  while (schema_idx >= 0) {
+    // this means it is not default argument, so it is necessary
+    if (!schema_args.at(schema_idx).default_value().has_value()) {
+      return std::make_pair(schema_idx + 1, num_out);
+    } else {
+      auto schema_value =
+          schema_args.at(schema_idx).default_value().value().toIValue();
+      // non-const value will become nullptr here, so will be marked necessary
+      // non-const would include prim::ListConstruct, prim::DictConstruct as
+      // well.
+      auto actual_value = toIValue(actual_inputs[schema_idx]);
+      if (!actual_value.has_value()) {
+        return std::make_pair(schema_idx + 1, num_out);
+      }
+      // if the IR has same value as default value of the schema,
+      // it is not necessary argument.
+      if (schema_value != actual_value.value()) {
+        return std::make_pair(schema_idx + 1, num_out);
+      }
+    }
+    schema_idx--;
+  }
+  return std::make_pair(0, num_out);
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/custom_operator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/custom_operator.h
new file mode 100644
index 0000000000000000000000000000000000000000..07b2ca1264bb178b526eb4d2bb6c0daf0a8f7c78
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/custom_operator.h
@@ -0,0 +1,35 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/op_registration/op_registration.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/jit/runtime/operator.h>
+
+namespace torch::jit {
+
+/// Registration class for new operators. Effectively calls
+/// `torch::jit::registerOperator` for every supplied operator, but allows doing
+/// so in the global scope when a `RegisterOperators` object is assigned to a
+/// static variable.
+/// Note: This is *not* the custom operator API. If you want to register custom
+/// operators, take a look at torch::RegisterOperators.
+struct TORCH_API RegisterOperators {
+  RegisterOperators() = default;
+
+  /// Registers a vector of already created `Operator`s.
+  /// The operator element is now optional to filter null ops. It's backward
+  /// compatible and works for selective operator registration.
+  explicit RegisterOperators(std::vector<std::optional<Operator>> operators) {
+    for (std::optional<Operator>& o : operators) {
+      if (o) {
+        registerOperator(std::move(o.value()));
+      }
+    }
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/decomposition_registry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/decomposition_registry.h
new file mode 100644
index 0000000000000000000000000000000000000000..1761821aaeb846a379310913ca9bed4e408e6f72
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/decomposition_registry.h
@@ -0,0 +1,38 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+// This file is temporary until native_functions.yaml and derivatives.yaml are
+// merged. Ideally this should all go into native_functions.yaml
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API std::optional<std::shared_ptr<Graph>> GetDecomposition(
+    const FunctionSchema& schema);
+
+TORCH_API void RegisterDecomposition(
+    const FunctionSchema& schema,
+    std::shared_ptr<Graph> g);
+
+TORCH_API void RunDecompositions(std::shared_ptr<Graph> g);
+
+TORCH_API std::optional<GraphFunction*> GetDecompositionFunction(
+    const FunctionSchema& schema);
+
+// For invocation in C++, recommended is to assign to static local variable
+TORCH_API Function* GetDecompositionExecutor(const char* schema_literal);
+
+TORCH_API Function* GetDecompositionExecutor(const FunctionSchema& schema);
+
+TORCH_API void run_jit_decomposition(
+    const c10::OperatorHandle& op,
+    torch::jit::Stack* stack);
+
+TORCH_API bool has_jit_decomposition(const FunctionSchema& schema);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/decomposition_registry_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/decomposition_registry_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..87da540f6da8ab34bcff01c70817dcc1411d6862
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/decomposition_registry_util.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API const std::string& GetSerializedDecompositions();
+
+TORCH_API const OperatorMap<std::string>& GetDecompositionMapping();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/exception_message.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/exception_message.h
new file mode 100644
index 0000000000000000000000000000000000000000..2c88b716dd9d0917b6f07dc2f2ab135d4f459035
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/exception_message.h
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Exception.h>
+#include <stdexcept>
+
+namespace torch::jit {
+
+struct ExceptionMessage {
+  ExceptionMessage(const std::exception& e) : e_(e) {}
+
+ private:
+  const std::exception& e_;
+  friend std::ostream& operator<<(
+      std::ostream& out,
+      const ExceptionMessage& msg);
+};
+
+inline std::ostream& operator<<(
+    std::ostream& out,
+    const ExceptionMessage& msg) {
+  auto c10_error = dynamic_cast<const c10::Error*>(&msg.e_);
+  if (c10_error) {
+    out << c10_error->what_without_backtrace();
+  } else {
+    out << msg.e_.what();
+  }
+  return out;
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_executor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_executor.h
new file mode 100644
index 0000000000000000000000000000000000000000..4d73a96f44cd8ac2063776c402ccf5b2595854a2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_executor.h
@@ -0,0 +1,157 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <atomic>
+#include <memory>
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/python/update_graph_executor_opt.h>
+#include <torch/csrc/jit/runtime/argument_spec.h>
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <torch/csrc/jit/runtime/variable_tensor_list.h>
+
+TORCH_DECLARE_bool(torch_jit_enable_new_executor);
+
+TORCH_DECLARE_bool(torch_jit_execution_plan_reuse_code_graph);
+
+namespace torch::jit {
+struct GraphExecutorState;
+struct Code;
+
+enum ExecutorExecutionMode {
+  SIMPLE,
+  PROFILING,
+};
+
+struct ExecutionPlan {
+  ExecutionPlan() = default;
+  ExecutionPlan(std::shared_ptr<Graph> graph, std::string function_name)
+      : code(graph, std::move(function_name)),
+        graph(
+            FLAGS_torch_jit_execution_plan_reuse_code_graph
+                ? code.graph()
+                : std::move(graph)) {}
+
+  operator bool() const {
+    return static_cast<bool>(graph);
+  }
+
+  Code code;
+  std::shared_ptr<Graph> graph;
+};
+
+// Notice that those structs don't manage lifetime of their members.
+// They are only valid only right after you call getDebugState() and should
+// never be used again once another GraphExecutor function is called.
+
+struct GraphExecutorState {
+  const Graph* graph = nullptr;
+  ExecutionPlan fallback; // XXX: members of this field are optional
+  std::unordered_map<ArgumentSpec, ExecutionPlan> execution_plans;
+};
+
+struct TORCH_API EnableProfilingGuard {
+  EnableProfilingGuard();
+  ~EnableProfilingGuard();
+
+ private:
+  bool old_executor_mode = false;
+  bool old_get_optimize = false;
+};
+
+struct GraphExecutorImplBase;
+struct TORCH_API GraphExecutor {
+  GraphExecutor() = default;
+  GraphExecutor(const std::shared_ptr<Graph>& graph, std::string function_name);
+
+  GraphExecutor(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name,
+      ExecutorExecutionMode executor_mode);
+
+  void run(Stack& inputs);
+  c10::intrusive_ptr<Future> runAsync(
+      Stack& stack,
+      TaskLauncher taskLauncher = at::launch);
+
+  // `remaining_bailout_depth` stands for the maximum number of profiled and
+  // specialized recompilations allowed for the current `GraphExecutor`. if
+  // remaining_bailout_depth is equal to 0, `GraphExecutor` won't perform any
+  // profiling and specialization. This is also equivalent to the
+  // SIMPLE_EXECUTOR mode. if remaining_bailout_depth is greater than 0,
+  // `GraphExecutor` will profile and specialize its input graph based on the
+  // profiled information whenever a bailout check is failed/triggered, a new
+  // `GraphExecutor` will be created. This new `GraphExecutor`'s
+  // remaining_bailout_depth will be reduced by 1.
+  // If no bailout depth is passed, the depth will be initialized from the
+  // current global fusion strategy settings.
+  const ExecutionPlan& getPlanFor(
+      Stack& inputs,
+      std::optional<size_t> remaining_bailout_depth = std::nullopt);
+  // Returns an optimized execution plan without requiring input arguments.
+  // Runs input-independent optimization passes (e.g. inlining, constant
+  // propagation, peephole, CSE) but skips profiling-based specializations
+  // that require runtime type/shape information.
+  const ExecutionPlan& getInputIndependentPlan();
+  GraphExecutorState getDebugState();
+
+  void debugFlushCompilationCache();
+
+  bool isOptimized() const;
+
+ private:
+  std::shared_ptr<GraphExecutorImplBase> pImpl;
+};
+
+TORCH_API Node* replaceBlockWithFallbackGraph(
+    Block* b,
+    ArrayRef<Value*> inputs);
+
+// These passes need to run before it is valid to pass to the interpreter
+// regardless of whether sizes have been specialized or not.
+TORCH_API void runRequiredPasses(const std::shared_ptr<Graph>& g);
+
+TORCH_API void debugSetFusionGroupInlining(bool state);
+TORCH_API bool getFusionGroupInlining();
+
+TORCH_API void debugSetAutodiffSubgraphInlining(bool state);
+TORCH_API std::shared_ptr<Graph> lastExecutedOptimizedGraph();
+
+TORCH_API std::atomic<bool>& getProfilingMode();
+TORCH_API std::atomic<bool>& getExecutorMode();
+TORCH_API std::atomic<size_t>& getNumProfiledRuns();
+TORCH_API size_t getBailoutDepth();
+TORCH_API bool IsNewExecutorEnabled();
+
+struct TORCH_API GraphOptimizerEnabledGuard {
+  GraphOptimizerEnabledGuard(bool state)
+      : old_state_(getGraphExecutorOptimize()) {
+    setGraphExecutorOptimize(state);
+  }
+
+  ~GraphOptimizerEnabledGuard() {
+    setGraphExecutorOptimize(old_state_);
+  }
+
+  bool old_state_;
+};
+
+namespace detail {
+
+GraphExecutor* getGradExecutor(Operation& op);
+
+GraphExecutor* getDifferentiableGraphOpExecutor(Operation& op);
+
+// for debugging information we expose a way to get the last actually
+// run graph. Previous approaches allowed querying the GraphExecutor
+// for what graph it would run in certain circumstances (graphFor), but
+// this is fragile because we sometimes change how these decisions are made.
+// This interface still allows our tests to look at optimized graphs, but
+// with less plumbing.
+} // namespace detail
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_executor_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_executor_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..c40ac7ac15cbd3ab0255070e13a97a9471014902
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_executor_impl.h
@@ -0,0 +1,123 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/jit/runtime/graph_executor.h>
+
+#include <ATen/core/ivalue.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/autograd/grad_mode.h>
+#include <torch/csrc/jit/frontend/tracer.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/shape_analysis.h>
+#include <torch/csrc/jit/resource_guard.h>
+#include <torch/csrc/jit/runtime/argument_spec.h>
+#include <torch/csrc/jit/runtime/autodiff.h>
+#include <torch/csrc/jit/runtime/custom_operator.h>
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <torch/csrc/jit/runtime/profiling_record.h>
+
+#include <torch/csrc/autograd/edge.h>
+#include <torch/csrc/autograd/function.h>
+#include <torch/csrc/jit/frontend/ir_emitter.h>
+#include <torch/csrc/jit/runtime/logging.h>
+
+#include <cstdint>
+#include <iterator>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace torch::jit {
+
+void packGradient(const Gradient& gradient, Node* dnode);
+bool needsGradient(const std::shared_ptr<const Graph>& graph);
+void runOptimization(
+    std::shared_ptr<Graph>& graph,
+    bool unroll_non_constant_loops = true,
+    bool const_prop_user_classes = true);
+void runNondiffOptimization(
+    std::shared_ptr<Graph>& graph,
+    bool strict_fuser_check = false);
+void debugSetAutodiffSubgraphInlining(bool state);
+bool TORCH_API getAutodiffSubgraphInlining();
+
+void debugSetFusionGroupInlining(bool state);
+bool getFusionGroupInlining();
+
+// Tunable parameters for deciding when to create/keep subgraphs of
+// differentiable code
+const size_t autodiffSubgraphNodeThreshold = 2;
+const size_t autodiffSubgraphInlineThreshold = 5;
+
+// a Graph can be created via tracing, or via a language-based frontend
+// GraphExecutor runs it. It can run the same graph on many different sizes
+// and different requires_grad states, and handles specializations for each
+// situation. GraphExecutor is completely unaware of tracing or module
+// parameters to keep the tracing concerns separated.
+struct GraphExecutorImplBase {
+  static std::shared_ptr<Graph> prepareGraph(
+      const std::shared_ptr<Graph>& graph) {
+    auto copy = graph->copy();
+    EraseShapeInformation(copy);
+    return copy;
+  }
+
+  GraphExecutorImplBase(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name)
+      : graph(prepareGraph(graph)),
+        function_name_(std::move(function_name)),
+        num_inputs(this->graph->inputs().size()),
+        num_outputs(this->graph->outputs().size()) {}
+
+  // entry point where execution begins
+  void run(Stack& stack);
+  c10::intrusive_ptr<Future> runAsync(
+      Stack& stack,
+      TaskLauncher taskLauncher = at::launch);
+
+  virtual const ExecutionPlan& getPlanFor(
+      Stack& stack,
+      std::optional<size_t> remaining_bailout_depth = std::nullopt) = 0;
+  // Returns an optimized execution plan without requiring input arguments.
+  // Runs input-independent optimization passes (e.g. inlining, constant
+  // propagation, peephole, CSE) but skips profiling-based specializations
+  // that require runtime type/shape information.
+  virtual const ExecutionPlan& getInputIndependentPlan() = 0;
+  virtual GraphExecutorState getDebugState() = 0;
+  virtual ~GraphExecutorImplBase() = default;
+
+  virtual bool isOptimized() const {
+    return false;
+  }
+
+ protected:
+  friend struct GraphExecutor;
+
+  // The unoptimized starting graph. This field is effectively const, but we
+  // can't make it so because Graph::copy() is not const (and making it const is
+  // not that easy at this point).
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::shared_ptr<Graph> graph;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::string function_name_;
+
+  // If false, we'll run the graph as we get it, without any optimizations.
+  // Useful for debugging.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const size_t num_inputs;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  const size_t num_outputs;
+
+  // GraphExecutors can be accessed from multiple threads, so this thread needs
+  // to be held every time we access the fallback or plan_cache.
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::mutex compile_mutex;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_iterator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_iterator.h
new file mode 100644
index 0000000000000000000000000000000000000000..4d0db56a19b489224168b58ebdc854f4d3ed574c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/graph_iterator.h
@@ -0,0 +1,152 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+// This class facilitates depth-first iteration over all nodes in a graph.
+class DepthFirstGraphNodeIterator {
+  Node* current_;
+
+ public:
+  // Constructor.
+  explicit DepthFirstGraphNodeIterator(std::shared_ptr<Graph>& graph)
+      : current_(*(graph->block()->nodes().begin())) {}
+
+  // Moves up and to the next node (may move up recursively).
+  void move_up() {
+    if (current_ == nullptr) {
+      return;
+    }
+    // Basically we start from the child block (which is current_)
+    // and we try to find the block that owns it. Now we need to check
+    // if that block is the graph root block, or if it is an If/Loop/etc
+    // block.
+    //
+    // If it's the graph root block we can stop because there is no "up"
+    // but if it is a node (e.g. If/Loop/etc) we need to apply logic
+    // based on where we are coming from to move to the next block.
+    // This might mean that we need to traverse up again (e.g. if we've
+    // reached the end of the else clause in an if block we need to go)
+    // up to the parent block that contains the if.
+    //
+    // Similarly if we've reached the end of the parent block containing
+    // the else clause we might need to go up again so this is a recursive
+    // function.
+    //
+    //              BlockNode (if/loop/with)
+    //                       |
+    //            [Block1]  ... [Block2]
+    //                |
+    //   [ Node1, Node2, Node3, FromNode]
+    //
+    auto parent_block = current_->owningBlock();
+    TORCH_INTERNAL_ASSERT(parent_block, "Every node must be owned by a block");
+
+    // Get the node that owns the parent block. This node has to be an if,
+    // loop, or with.
+    auto parent_node = parent_block->owningNode();
+    if (parent_node == nullptr) {
+      // If there's no node that owns this current block then we're at the
+      // top of the graph and since we're trying to move up we have reached
+      // the end of the traversal.
+      current_ = nullptr;
+      return;
+    }
+
+    // Check the type of node this root is.
+    if (parent_node->kind() == prim::If) {
+      // Need to check if we came from the `then` branch or the `else` branch.
+      auto* then_block = parent_node->blocks().at(0);
+      auto* else_block = parent_node->blocks().at(1);
+
+      if (parent_block == else_block) {
+        // If else block then we move to the next node in the parent block.
+        current_ = parent_node->next();
+        if (current_->kind() == prim::Return) {
+          move_up();
+        }
+      } else {
+        // If then block then move to the else block if it is not empty.
+        TORCH_INTERNAL_ASSERT(parent_block == then_block);
+        bool else_block_empty =
+            else_block->nodes().begin() == else_block->nodes().end();
+
+        if (!else_block_empty) {
+          current_ = *(else_block->nodes().begin());
+        } else {
+          // Since it's empty we move to the next node.
+          current_ = parent_node->next();
+          if (current_->kind() == prim::Return) {
+            move_up();
+          }
+        }
+      }
+    } else if (
+        parent_node->kind() == prim::Loop ||
+        parent_node->kind() == prim::With) {
+      current_ = parent_node->next();
+      if (current_->kind() == prim::Return) {
+        move_up();
+      }
+    } else {
+      TORCH_INTERNAL_ASSERT(
+          false, "Only if/loop/with nodes should have child blocks");
+    }
+  }
+
+  // Moves to the next adjacent node or up in to the parent if that is not
+  // possible.
+  void move_next() {
+    if (current_ == nullptr) {
+      return;
+    }
+
+    // Increment to the next node in the current block.
+    current_ = current_->next();
+
+    // Check if we're at the end of the block. If so we need
+    // to move upwards (if it makes sense to).
+    if (current_->kind() == prim::Return) {
+      move_up();
+    }
+  }
+
+  // Moves to the next node in the graph into children if it can.
+  void move_into() {
+    if (current_ == nullptr) {
+      return;
+    }
+
+    // Check if we're currently on a node that contains sub-nodes.
+    if (current_->kind() == prim::If || current_->kind() == prim::Loop ||
+        current_->kind() == prim::With) {
+      auto* first_block = current_->blocks().at(0);
+      current_ = first_block->param_node();
+      // Move next will move up and out of the current node if the block is
+      // empty. `move_up` which is called by `move_next` will handle the
+      // difference between If, Loop, and With blocks appropriately.
+      move_next();
+    } else {
+      move_next();
+    }
+  }
+
+  // Get the next Node in the graph. \returns nullptr if there are no nodes
+  // left.
+  Node* next() {
+    auto result = current_;
+
+    // Try move into the existing node to set the next node to be returned.
+    // This will move to the next node if not possible, or move upwards and
+    // to the next.
+    move_into();
+
+    return result;
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/instruction.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/instruction.h
new file mode 100644
index 0000000000000000000000000000000000000000..56eda3d01057cfc19ef2363dc5a7d34d7117a354
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/instruction.h
@@ -0,0 +1,103 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstdint>
+#include <typeinfo>
+#include <unordered_set>
+
+namespace torch::jit {
+// instruction look like:
+// op_code X, N
+// meaning of X, N depend on the op:
+// O - index into operator table
+// R - index into register table
+// I - literal integer
+// C - index into constant table
+// P - jump offset relative to beginning of current instruction
+// F - index into function table
+// T - index into the type table, used for guard instructions
+// S - index into object slots
+// C - index into code table
+
+#define FORALL_OPCODES(_)                                                      \
+  _(OP, "O") /* invoke operator X */                                           \
+  _(OPN, "OI") /* invoke vararg operator X with N arguments */                 \
+  _(LOAD, "R") /* push a value from a register X */                            \
+  _(MOVE, "R") /* push a value from register X, clearing the register */       \
+  _(STOREN, "RI") /* store N values to registers [X, X+N) */                   \
+  _(STORE, "R") /* store 1 value to registers X */                             \
+  _(DROP, "") /* drop 1 value from the top of the stack */                     \
+  _(DROPR, "R") /* clear register X */                                         \
+  _(LOADC, "C") /* push the constant X */                                      \
+  _(JF, "P") /* pop the top of the stack, if false, branch to P */             \
+  _(JMP, "P") /* unconditional branch to X */                                  \
+  _(LOOP, "PI") /* perform a loop, X is where to branch if cond is false */    \
+  _(RET, "") /* exit execution */                                              \
+  _(WAIT, "") /* wait for a future to be complete */                           \
+  _(CALL, "F") /* call function X */                                           \
+  _(GUARD, "T") /* check a guard against type_table, true if passes */         \
+  _(TYPECHECK, "TN") /* check each type of input[i] against type_table[X+N] */ \
+  _(FAIL_GUARD, "T") /* fail a guard, patch back to GUARD */                   \
+  _(PROFILE_OP, "F") /* get a callback from profile_function_table at X */     \
+  _(TAIL_CALL, "F") /* replace current frame with function F */                \
+  _(INTERFACE_CALL, "CI") /* call method X on the first argument (of N) */     \
+  _(GET_ATTR, "S") /* get attribute from slot X in an Object */                \
+  _(SET_ATTR, "S") /* set attribute to slot X in an Object */                  \
+  _(LIST_UNPACK, "I") /* unpack list expecting length I */                     \
+  _(TUPLE_CONSTRUCT, "I") /* construct a tuple using X inputs */               \
+  _(NAMED_TUPLE_CONSTRUCT,                                                     \
+    "TI") /* construct a tuple of type X, using N inputs */                    \
+  _(LIST_CONSTRUCT, "TI") /* construct a list of type X, using N inputs */     \
+  _(DICT_CONSTRUCT, "TI") /* construct a dict of type X, using N inputs */     \
+  _(CREATE_OBJECT, "T") /* create an object of type X */                       \
+  _(ISINSTANCE, "TI") /* check object is one of  types[X:X+N]  */              \
+  _(TUPLE_SLICE, "II") /* slice tup[X:(X+N)] */                                \
+  _(TUPLE_INDEX, "") /* get the value from a tuple at that index */            \
+  _(RAISE_EXCEPTION, "") /* throws the exception from Python */                \
+  _(DICT_INDEX, "") /* gets the value from the dict for given key */           \
+  _(UNCHECKED_CAST, "") /* perform an unchecked cast operation */              \
+  _(__IS__, "") /* performs `is` operator from Python */                       \
+  _(UN_INITIALIZED,                                                            \
+    "") /* sets default values to variables that are uninitialized */          \
+  _(__ISNOT__, "") /* performs `is not` operator from Python  */               \
+  _(FORMAT, "I") /* performs string format function `f strings` or `{}.format` \
+                     the number of inputs in stored in X */                    \
+  _(DEVICE, "") /* invokes aten::device for a Tensor */                        \
+  _(DTYPE, "") /* invokes aten::dtype for a Tensor */                          \
+  _(DIM, "") /* invokes aten::dim for a Tensor */                              \
+  _(__NOT__, "") /* performs `not` operator from Python  */                    \
+  _(TO_LIST, "") /* convert the input to a list */                             \
+  _(NUM_TO_TENSOR,                                                             \
+    "") /* performs the conversion of a number/scalar to Tensor */             \
+  _(IS_CUDA, "") /* invokes aten::is_cuda for a Tensor */                      \
+  _(FORK, "CN") /* launch a thread to run code entry x with N inputs  */       \
+  _(WARN, "I") /* emit a warning with line information */                      \
+  _(ENTER, "EN") /* enter scope of a contextmanager */                         \
+  _(EXIT, "EX") /* exit the last entered contextmanager */                     \
+  _(AWAITABLE, "CN") /* initialize await for code entry x with N inputs  */
+
+enum OpCode : uint8_t {
+#define DEFINE_OP(op, _) op,
+  FORALL_OPCODES(DEFINE_OP)
+#undef DEFINE_OP
+};
+
+struct Instruction {
+  OpCode op;
+  uint8_t unused{0};
+  uint16_t N;
+  int32_t X;
+  // TODO: check for overflow
+  Instruction(OpCode op, int32_t X, uint16_t N) : op(op), N(N), X(X) {}
+};
+std::ostream& operator<<(std::ostream& out, Instruction inst);
+
+bool isOpSupportedInMobile(OpCode op);
+char const* toString(OpCode op);
+OpCode parseOpCode(const char* str);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter.h
new file mode 100644
index 0000000000000000000000000000000000000000..b5b5b3da8f19a94d8cf666148595ac0736826769
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter.h
@@ -0,0 +1,165 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include <ATen/ThreadLocalState.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+
+TORCH_DECLARE_bool(torch_jit_disable_warning_prints);
+TORCH_DECLARE_bool(torch_jit_enable_rethrow_caught_exception);
+
+namespace at {
+class Tensor;
+TORCH_API void launch(std::function<void()> func);
+} // namespace at
+namespace c10 {
+struct IValue;
+struct OperatorName;
+} // namespace c10
+
+namespace torch::jit {
+
+// The interpreter run Graphs with Tensor inputs and Tensor outputs
+// a separate component in the autograd handles unwrapping and wrapping
+// variable objects for use in the interpreter.
+namespace interpreter {
+struct CodeImpl;
+}
+
+struct Node;
+struct GraphExecutor;
+struct InterpreterStateImpl;
+struct Graph;
+struct Node;
+struct Instruction;
+using Stack = std::vector<c10::IValue>;
+using c10::ivalue::Future;
+using TaskLauncher = std::function<void(std::function<void()>)>;
+
+bool TORCH_API in_torchscript_runtime();
+
+struct TORCH_API Code {
+  Code() = default;
+  explicit Code(interpreter::CodeImpl* pImpl);
+  // remaining_bailout_depth is irrelevant in a `Code` object unless the `Code`
+  // is directly created by `GraphExecutor` in which case it's likely to contain
+  // `prim::BailOut`s to control the maximum depth of bailout chains
+  explicit Code(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name,
+      size_t remaining_bailout_depth = 0);
+
+  const std::vector<GraphExecutor*>& grad_executors();
+  const std::vector<GraphExecutor*>& diff_graph_op_executors();
+
+  explicit operator bool() const {
+    return pImpl != nullptr;
+  }
+  size_t num_inputs() const;
+  size_t num_outputs() const;
+  size_t num_bailouts() const;
+  const std::vector<c10::IValue>& constant_table() const;
+  const std::vector<c10::TypePtr>& type_table() const;
+  const std::vector<Instruction>& instructions() const;
+  const std::unordered_map<std::string, size_t>& op_to_num_specified_args()
+      const;
+  const std::vector<Node*>& instructions_source() const;
+  void request_bailout(size_t index);
+  size_t register_size() const;
+  std::shared_ptr<Graph> graph() const;
+
+ private:
+  std::shared_ptr<interpreter::CodeImpl> pImpl;
+  friend struct InterpreterStateImpl;
+  friend std::ostream& operator<<(std::ostream& out, const Code& code);
+};
+
+struct TORCH_API MobileCode : Code {
+  explicit MobileCode(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name,
+      bool emit_default_input_instructions = true,
+      bool support_default_args_before_out = true,
+      bool emit_promoted_ops = true,
+      size_t remaining_bailout_depth = 0);
+};
+
+struct InterpreterState {
+  TORCH_API InterpreterState(
+      const Code& code,
+      TaskLauncher taskLauncher = at::launch);
+  TORCH_API void run(Stack& stack);
+  TORCH_API c10::intrusive_ptr<Future> runAsync(Stack& stack);
+  c10::intrusive_ptr<Future> getFuture();
+
+ private:
+  InterpreterState(c10::intrusive_ptr<c10::intrusive_ptr_target> pImpl);
+  // Ideally we should use c10::intrusive_ptr<InterpreterStateImpl> for pImpl;
+  // but intrusive_ptr requires full definition of InterpreterStateImpl,
+  // which we need to hide in the header.
+  c10::intrusive_ptr<c10::intrusive_ptr_target> pImpl;
+  friend struct InterpreterStateImpl;
+};
+
+// Created by wait()
+struct Suspend : public std::exception {
+  const char* what() const noexcept override {
+    return "Suspend";
+  }
+
+  explicit Suspend(c10::intrusive_ptr<Future> future_)
+      : future(std::move(future_)) {}
+
+  c10::intrusive_ptr<Future> future;
+};
+
+// InterpreterContinuation propagates dist_autograd_context_id
+// through (and only through) the forward pass manually, other
+// thread local settings are propagated with ThreadLocalState
+struct InterpreterContinuation {
+  InterpreterContinuation(
+      InterpreterState state_,
+      Stack stack_,
+      int64_t dist_autograd_context_id = 0,
+      std::optional<at::ThreadLocalState> tls_state = std::nullopt)
+      : state(std::move(state_)),
+        stack(std::move(stack_)),
+        tls_state_(std::move(tls_state))
+#ifdef USE_DISTRIBUTED
+        ,
+        dist_autograd_context_id_(dist_autograd_context_id)
+#endif
+  {
+  }
+
+  void operator()();
+
+ private:
+  InterpreterState state;
+  Stack stack;
+  std::optional<at::ThreadLocalState> tls_state_ = std::nullopt;
+#ifdef USE_DISTRIBUTED
+  int64_t dist_autograd_context_id_;
+#endif
+};
+
+// what is the tensors type, including state from the current execution context
+// that modifies how the tensor behaves. For instance if no_grad is enabled
+// this will cause the TensorType to have requires_grad=False.
+TORCH_API at::TensorTypePtr tensorTypeInCurrentExecutionContext(
+    const at::Tensor& t);
+
+// current (TLS) TorchScript interpreter callstack
+TORCH_API std::vector<StackEntry> currentCallstack();
+TORCH_API std::vector<std::string> currentModuleHierarchy();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/can_emit_inline.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/can_emit_inline.h
new file mode 100644
index 0000000000000000000000000000000000000000..65479623bac2d18d83a213a1c95163844462da08
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/can_emit_inline.h
@@ -0,0 +1,111 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::interpreter {
+/*
+This is an optimization that reduces the number of store/load/move nodes needed
+by recognizing that parts of the graph are simple trees like a*x + b*y. When
+this happens it is possible to work directly off of the stack by emitting the
+tree in a depth-first left-to-right manner:
+  load a
+  load x
+  mul # stack now is a*x
+  load b
+  load y
+  mul # stack now is a*x, b*y
+  add
+
+can_emit_inline_[node] == true means that this node participates as a non-root
+member of one of these trees. The code emitter will not emit this node when
+it is encountered in the node. Instead the node is emitted in a depth first
+traversal from where it is used in a tree.
+
+To participate in a tree a node must have a single use (otherwise it is not
+tree-like) and output a single value (for simplicity.) If our IR was functional,
+these would be the only constraints. However, many nodes have side effects, so
+we must ensure that emitting the nodes in depth first order from the tree's root
+_does not reorder the emission of the nodes_. To ensure this, we work backward
+from the root of a potential tree, visiting its inputs in reverse depth first
+order, while scanning the node list backward (with the block_point node). When
+these traversal line up we know it is safe to emit the tree in this way. We
+ignore constant nodes, which do not have side effects.
+*/
+struct CanEmitInline {
+  explicit CanEmitInline(Graph& graph) {
+    scanBlock(graph.block());
+  }
+  bool canInline(Value* v) {
+    return v->node()->kind() != prim::Param &&
+        // without this a BailOut may float downstream past some later
+        // BailOut
+        // and receive a higher jf_index. Then a GUARD instruction
+        // we generated for the floated BailOut will get popped up from the
+        // instruction stack
+        // by the later BailOut in createBailoutBlock and its jf_index
+        // will become invalid.
+        v->node()->kind() != prim::TensorExprGroup &&
+        v->node()->kind() != prim::TensorExprDynamicGroup &&
+        v->node()->kind() != prim::StaticSubgraph &&
+        v->node()->kind() != prim::CudaFusionGroup &&
+        v->node()->kind() != prim::FusionGroup &&
+        v->node()->kind() != prim::BailOut && v->uses().size() == 1 &&
+        v->node()->outputs().size() == 1;
+  }
+
+  Node* previousNonConstant(Node* n) {
+    do {
+      n = n->prev();
+    } while (n->kind() == prim::Constant);
+    return n;
+  }
+
+  Node* scanValue(Node* block_point, Value* v) {
+    // this node is a candidate for inline, if our reverse scan of the
+    // node list lines up with the use of v, we know it will be emitted in
+    // tree order, and we can inlining. Scan continues for further nodes.
+    if (v->node() == block_point && canInline(v)) {
+      // since we inlined this node, we may be able to recursively inline
+      // its inputs, so we continue scanning it
+      block_point = scanNode(v->node());
+      can_emit_inline_[v->node()] = true;
+    }
+    // if it does not line up, we can't inline 'v', and will just generate
+    // a load/move for it. However, other inputs may still appear in tree
+    // order so we continue the scan of the inputs.
+    return block_point;
+  }
+
+  Node* scanNode(Node* n) {
+    // don't bother to scan nodes we have already determined to be inline
+    if (can_emit_inline_.count(n)) {
+      return nullptr;
+    }
+    for (auto b : n->blocks()) {
+      scanBlock(b);
+    }
+    Node* block_point = previousNonConstant(n);
+    for (auto it = n->inputs().rbegin(), end = n->inputs().rend(); it != end;
+         ++it) {
+      block_point = scanValue(block_point, *it);
+    }
+    return block_point;
+  }
+
+  void scanBlock(Block* b) {
+    scanNode(b->return_node());
+    for (auto node : b->nodes().reverse()) {
+      scanNode(node);
+    }
+  }
+  std::unordered_map<Node*, bool> can_emit_inline_;
+};
+
+} // namespace torch::jit::interpreter
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/code_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/code_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..8b3563f42b919259fd3e682f25f6470ebac6325d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/code_impl.h
@@ -0,0 +1,1066 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <c10/util/irange.h>
+#include <torch/csrc/jit/api/function_impl.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/passes/bailout_graph.h>
+#include <torch/csrc/jit/runtime/calculate_necessary_args.h>
+#include <torch/csrc/jit/runtime/graph_iterator.h>
+#include <torch/csrc/jit/runtime/instruction.h>
+#include <torch/csrc/jit/runtime/interpreter/preprocess_graph.h>
+
+TORCH_DECLARE_bool(torch_jit_enable_expanded_stacks);
+TORCH_DECLARE_bool(torch_jit_expanded_stacks_mangled);
+
+namespace torch::jit::interpreter {
+
+template <class Ttarget, class Tsource>
+Ttarget safe_narrow_cast(Tsource v) {
+  Ttarget res = static_cast<Ttarget>(v);
+  // Casting it back to check whether it overflew.
+  if (static_cast<Tsource>(res) != v) {
+    TORCH_WARN(
+        "ATTENTION: your model computation is overflowing, safe_narrow_cast<>() failed");
+    return v;
+  }
+  return res;
+}
+
+// BailoutBlocks are used to temporarily store
+// instructions (typically, argument LOADs and TAIL_CALL)
+// generated for prim::BailOut nodes
+// before they are merged back into
+// CodeImpl._instructions_ by insertBailoutBlocks
+struct BailoutBlock {
+  size_t jf_instruction_index; // this node gets patched to jump here on failure
+  std::vector<Instruction> instructions; // ends in a TAIL_CALL
+
+  explicit BailoutBlock(size_t jf_index) : jf_instruction_index(jf_index) {}
+};
+
+// for keeping track of the current node
+struct WithCurrentNode {
+  WithCurrentNode(Node** loc, Node* new_value) : loc_(loc), old_value_(*loc_) {
+    *loc = new_value;
+  }
+  ~WithCurrentNode() {
+    *loc_ = old_value_;
+  }
+
+ private:
+  Node** loc_;
+  Node* old_value_;
+};
+
+struct NodeSourceInfo {
+  const char* func_name_{nullptr};
+  const char* file_name_{nullptr};
+  size_t line_{0};
+  NodeSourceInfo() = default;
+};
+
+struct CodeImpl {
+  friend struct InterpreterState;
+  std::vector<Instruction> instructions_;
+
+  const c10::unique_t node_stack_attr_symbol_ =
+      static_cast<c10::unique_t>(attr::node_stack_idx);
+  // Expanded inlined stacks as pointers to values in inlined call stack.
+  std::vector<std::vector<NodeSourceInfo>> expanded_node_stacks_;
+
+  // same length as instructions.
+  // what node in the graph cause this
+  // instruction to be emitted?
+  std::vector<Node*> instructions_source_;
+  std::vector<IValue> constant_table_;
+  std::vector<Operation> operator_table_;
+#ifndef NDEBUG
+  std::vector<Operator> full_operator_table_;
+#endif
+  // map<(op name, num inputs), index in operator table>, to avoid duplicates,
+  // not including vararg operators
+  std::unordered_map<
+      std::pair<std::string, int>,
+      int,
+      std::function<size_t(const std::pair<std::string, int>& p)>>
+      operator_table_inv_;
+  std::vector<Function*> function_table_;
+  std::vector<std::unique_ptr<GraphFunction>> forked_functions_;
+  std::vector<std::unique_ptr<GraphFunction>> awaited_functions_;
+  std::vector<TypePtr> type_table_;
+  std::vector<std::function<void(std::vector<IValue>&)>>
+      profile_function_table_;
+
+  int register_size_ = 0;
+  size_t n_outputs;
+  size_t n_inputs;
+  TypePtr return_type_;
+  std::string function_name_;
+
+  // We MUST hold onto graph here because some Operators stored in the
+  // instruction lists have dependencies on meta-data stored in the graph
+  // that would be dead otherwise.
+  // It is also very useful for debugging interpreter problems to
+  // keep this around.
+  std::shared_ptr<Graph> graph_;
+  std::optional<std::vector<GraphExecutor*>> grad_executors_;
+  std::optional<std::vector<GraphExecutor*>> forward_executors_;
+  PreprocessGraph preprocess_;
+
+  // map from unique of nodes to register in register table
+  std::unordered_map<Value*, int> value_to_reg_;
+
+  // map from operator name to specified arguments
+  // Example: for a schema of aten::foo.str
+  // aten::foo.str(arg0: str="default", arg1: int=0,
+  //               arg2: bool=False, arg3: float=0.0)
+  // If the usages in a graph is:
+  //    aten::foo("somestr", arg1=0, arg2=True, arg3=0.0)
+  //    aten::foo("somestr", arg1=1, arg2=False, arg3=0.0)
+  // op_to_num_specified_args_["aten::foo.str"] = 3
+  // This is because for all usages, at most 3 args are used.
+  std::unordered_map<std::string, size_t> op_to_num_specified_args_;
+
+  std::unordered_map<std::string, size_t> op_to_num_out_args_;
+
+  // running count of uses as we emit. When we reach use_count_[v] =
+  // v.uses().size() we know it is the final use and we can move rather than
+  // load.
+  std::unordered_map<Value*, size_t> use_count_;
+
+  Node* current_node_; // used in creation of code to keep track
+                       // of node being emitted
+  Node* last_inserted_op_ = nullptr;
+
+  // out-of-line jumps for bailouts that are patched in at the end
+  std::vector<BailoutBlock> bailout_blocks_;
+  std::vector<std::unique_ptr<Function>> bailout_functions_;
+  size_t remaining_bailout_depth_;
+
+  CodeImpl(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name,
+      size_t remaining_bailout_depth,
+      bool emit_instructions = true)
+      : operator_table_inv_(
+            0,
+            [](const std::pair<std::string, int>& p) {
+              return std::hash<std::string>()(p.first) ^
+                  std::hash<int>()(p.second);
+            }),
+        function_name_(std::move(function_name)),
+        preprocess_(*graph),
+        current_node_(preprocess_.graph->return_node()),
+        remaining_bailout_depth_(remaining_bailout_depth) {
+    graph_ = preprocess_.graph;
+    n_outputs = graph_->outputs().size();
+    if (n_outputs == 1) {
+      return_type_ = graph->outputs().at(0)->type();
+    } else {
+      return_type_ = TupleType::create(
+          fmap(graph->outputs(), [](const Value* v) { return v->type(); }));
+    }
+    n_inputs = graph_->inputs().size();
+    if (emit_instructions) {
+      // NOLINTNEXTLINE(clang-analyzer-optin.cplusplus.VirtualCall)
+      run();
+    }
+  }
+
+  virtual ~CodeImpl() = default;
+
+  // since subclass of CodeImpl needs to populate
+  // op_to_num_specified_args, we separate the calls
+  // that changes internals of CodeImpl into a separate
+  // function.
+  virtual void run() {
+    emitCodeForBlock(graph_->block());
+    insertInstruction(RET);
+    // we deferred the emission of bailout blocks so they appear at the end
+    // emit them now and patch up the jumps
+    insertBailoutBlocks();
+  }
+
+  const std::vector<c10::IValue>& constant_table() const {
+    return constant_table_;
+  }
+
+  void request_bailout(size_t index) {
+    auto count = index;
+    for (const auto instr_index : c10::irange(instructions_.size())) {
+      if (instructions_[instr_index].op == GUARD ||
+          instructions_[instr_index].op == FAIL_GUARD) {
+        if (count-- == 0) {
+          // patching GUARD to FAIL_GUARD
+          instructions_[instr_index].op = FAIL_GUARD;
+          GRAPH_DEBUG(
+              "Added a bailout request for ",
+              index,
+              " at instruction ",
+              instr_index);
+          break;
+        }
+      }
+    }
+  }
+
+  const std::vector<Instruction>& instructions() const {
+    return instructions_;
+  }
+
+  const std::unordered_map<std::string, size_t>& op_to_num_specified_args()
+      const {
+    return op_to_num_specified_args_;
+  }
+
+  const std::vector<Node*>& instructions_source() const {
+    return instructions_source_;
+  }
+
+  NodeSourceInfo getSourceInfoFromSourceRange(const SourceRange& range) {
+    NodeSourceInfo nodeSource;
+    SourceRange r = range;
+    if (!FLAGS_torch_jit_expanded_stacks_mangled && range.source()) {
+      if (auto orig = range.source()->findSourceRangeThatGenerated(r)) {
+        r = *orig;
+      }
+    }
+    if (r.source()) {
+      auto lineno = r.source()->lineno_for_offset(r.start());
+      nodeSource.line_ = r.source()->lineno_to_source_lineno(lineno);
+      if (r.source()->filename()) {
+        nodeSource.file_name_ = r.source()->filename().value().c_str();
+      }
+    }
+    return nodeSource;
+  }
+
+  void expandInlinedNodeStack(
+      const InlinedCallStackPtr& cs,
+      std::vector<NodeSourceInfo>* expandedstack) {
+    auto nodeSourceInfo = getSourceInfoFromSourceRange(cs->source_range());
+    nodeSourceInfo.func_name_ = cs->function_name().c_str();
+    expandedstack->emplace_back(nodeSourceInfo);
+
+    if (cs->callee()) {
+      expandInlinedNodeStack(cs->callee().value(), expandedstack);
+    }
+  }
+
+  void getNodeStack(
+      const Node* node,
+      std::vector<NodeSourceInfo>* expandedstack) {
+    if (current_node_->callstack()) {
+      expandInlinedNodeStack(current_node_->callstack().value(), expandedstack);
+    }
+    auto nodeSourceInfo = getSourceInfoFromSourceRange(node->sourceRange());
+    expandedstack->emplace_back(nodeSourceInfo);
+  }
+
+  void insertInstruction(OpCode op, int64_t X = 0, uint64_t N = 0) {
+    instructions_.emplace_back(
+        op,
+        safe_narrow_cast<int32_t, int64_t>(X),
+        safe_narrow_cast<uint16_t, uint64_t>(N));
+    instructions_source_.emplace_back(current_node_);
+
+    if (FLAGS_torch_jit_enable_expanded_stacks &&
+        !current_node_->hasAttribute(attr::node_stack_idx)) {
+      std::vector<NodeSourceInfo> expandedStack;
+      getNodeStack(current_node_, &expandedStack);
+      auto insertIdx = expanded_node_stacks_.size();
+      expanded_node_stacks_.emplace_back(expandedStack);
+      current_node_->i_(attr::node_stack_idx, insertIdx);
+    }
+
+    // check that we didn't accidentally emit nodes out of topological order
+    if (op == OP) {
+      if (last_inserted_op_ != nullptr && current_node_ != last_inserted_op_ &&
+          current_node_->owningBlock() == last_inserted_op_->owningBlock()) {
+        TORCH_INTERNAL_ASSERT(
+            current_node_->isAfter(last_inserted_op_),
+            *current_node_,
+            " is not after ",
+            *last_inserted_op_);
+      }
+      last_inserted_op_ = current_node_;
+    }
+  }
+
+  void truncateInstructions(size_t size) {
+    while (instructions_.size() > size) {
+      instructions_.pop_back();
+      instructions_source_.pop_back();
+    }
+  }
+
+  void createBailoutBlock(size_t jf_index) {
+    bailout_blocks_.emplace_back(jf_index);
+    auto& bailout_instructions = bailout_blocks_.back().instructions;
+
+    bailout_instructions.insert(
+        bailout_instructions.end(),
+        instructions_.begin() + jf_index + 1,
+        instructions_.end());
+    truncateInstructions(jf_index + 1);
+  }
+
+  int allocRegs(at::ArrayRef<Value*> vs) {
+    int result = register_size_ + 1;
+    for (Value* v : vs) {
+      AT_ASSERT(value_to_reg_.count(v) == 0);
+      value_to_reg_[v] = ++register_size_;
+    }
+    return result;
+  }
+
+  int registerFor(Value* v) {
+    return value_to_reg_.at(v);
+  }
+
+  void emitUse(Value* input, bool drop) {
+    // drop - if true, we are not actually going to use this thing
+    // and we can short circuit doing many instructions here
+    // by either clearing the register (DROPR) or just popping the stack
+    // (DROP)
+    if (preprocess_.can_emit_inline[input->node()]) {
+      emitNode(input->node());
+      if (drop) {
+        insertInstruction(DROP);
+      }
+    } else {
+      int reg = registerFor(input);
+      bool moved = input->uses().size() == ++use_count_[input];
+
+      OpCode op{};
+      if (input->node()->kind() == prim::Constant) {
+        op = LOADC;
+      } else if (moved) {
+        op = MOVE;
+      } else {
+        op = LOAD;
+      }
+
+      if (drop) {
+        op = DROPR;
+      }
+      insertInstruction(op, reg);
+    }
+  }
+
+  void emitLoadInputs(at::ArrayRef<Value*> inputs) {
+    for (Value* input : inputs) {
+      emitUse(input, false);
+    }
+  }
+
+  void emitLoadInputs(at::ArrayRef<Value*> inputs, int num_include) {
+    int count = 0;
+    for (Value* input : inputs) {
+      if (count < num_include) {
+        emitUse(input, false);
+        count++;
+      }
+    }
+  }
+
+  void emitLoadInputs(at::ArrayRef<Value*> inputs, size_t start, size_t end) {
+    for (size_t i = start; i < end; i++) {
+      emitUse(inputs[i], false);
+    }
+  }
+
+  virtual void emitOperator(Node* node) {
+    emitLoadInputs(node->inputs());
+    const Operator& op = node->getOperator();
+    int num_inputs = node->inputs().size();
+    bool is_vararg = op.schema().is_vararg();
+
+    int operation_index = add_to_operator_table(
+        op,
+        node,
+        c10::toString(op.schema().operator_name()),
+        num_inputs,
+        is_vararg);
+
+    if (op.hasOperation() && is_vararg) {
+      insertInstruction(OPN, operation_index, num_inputs);
+    } else {
+      insertInstruction(OP, operation_index);
+    }
+  }
+
+  void emitWait(Node* node) {
+    emitLoadInputs(node->inputs());
+    insertInstruction(WAIT);
+  }
+
+  void emitDrop(at::ArrayRef<Value*> to_drop) {
+    for (Value* input : to_drop) {
+      emitUse(input, true);
+    }
+  }
+
+  void emitStoreOutputs(Node* node) {
+    size_t N = node->outputs().size();
+    if (N == 0) {
+      return;
+    }
+    int regs = allocRegs(node->outputs());
+    if (N == 1) {
+      insertInstruction(STORE, regs);
+    } else {
+      insertInstruction(STOREN, regs, node->outputs().size());
+    }
+  }
+
+  int insertConstant(IValue value) {
+    int result = constant_table_.size();
+    constant_table_.emplace_back(std::move(value));
+    return result;
+  }
+
+  virtual void emitOperatorOrInstruction(
+      Node* node,
+      OpCode op,
+      int64_t X = 0,
+      uint64_t N = 0,
+      bool emit_inputs = true) {
+    if (emit_inputs) {
+      emitLoadInputs(node->inputs());
+    }
+    insertInstruction(op, X, N);
+  }
+
+  void emitFormat(Node* node) {
+    emitOperatorOrInstruction(node, FORMAT, node->inputs().size(), 0);
+  }
+
+  void checkNodeAndEmit(Node* node) {
+    // check if the node should be emitted as instruction or operator
+    const Operator& op = node->getOperator();
+    std::string unique_op_name = c10::toString(op.schema().operator_name());
+    if (unique_op_name.find("aten::__getitem__.Dict") == 0) {
+      // __get_item__ overloaded operator for Dict
+      // needs to be emitted an instruction
+      emitOperatorOrInstruction(node, DICT_INDEX);
+    } else {
+      emitOperator(node);
+    }
+  }
+
+  void emitConstant(Node* node) {
+    if (node->output()->type()->kind() == FunctionType::Kind) {
+      return;
+    }
+    // constants are just put in the constant table
+    value_to_reg_[node->output()] =
+        insertConstant(toIValue(node->output()).value());
+  }
+
+  void emitIf(Node* node) {
+    emitLoadInputs(node->inputs());
+    size_t start_if = instructions_.size();
+    insertInstruction(JF, 0); // dummy offset to be filled in
+    emitCodeForBlock(node->blocks().at(0));
+    insertInstruction(JMP, 0); // dummy offset
+    size_t start_else = instructions_.size();
+    instructions_[start_if].X = start_else - start_if;
+    emitCodeForBlock(node->blocks().at(1));
+    instructions_[start_else - 1].X = instructions_.size() - (start_else - 1);
+  }
+
+  void emitLoop(Node* loop) {
+    insertInstruction(LOADC, insertConstant(0));
+    emitLoadInputs(loop->inputs());
+    size_t start = instructions_.size();
+    insertInstruction(LOOP, 0, loop->inputs().size()); // dummy offset
+    emitCodeForBlock(loop->blocks().at(0));
+    insertInstruction(JMP, start - instructions_.size());
+    instructions_[start].X = instructions_.size() - start;
+  }
+
+  void emitCall(Function* func, at::ArrayRef<Value*> inputs) {
+    emitLoadInputs(inputs);
+    insertInstruction(CALL, function_table_.size());
+    function_table_.emplace_back(func);
+  }
+
+  void emitNodeAtBlockLevel(Node* node) {
+    WithCurrentNode guard(&current_node_, node);
+    switch (node->kind()) {
+      case prim::Constant:
+        emitConstant(node);
+        break;
+      case prim::Return:
+        emitLoadInputs(node->inputs());
+        break;
+      default:
+        if (!preprocess_.can_emit_inline[node]) {
+          emitNode(node);
+          emitStoreOutputs(node);
+        }
+        break;
+    }
+  }
+
+  size_t emitType(TypePtr t) {
+    size_t r = type_table_.size();
+    type_table_.emplace_back(std::move(t));
+    return r;
+  }
+
+  void emitTypeCheck(Node* node) {
+    auto num_inputs = node->inputs().size();
+
+    // Check that TypeCheck has at least one input.
+    TORCH_INTERNAL_ASSERT(
+        num_inputs && num_inputs + 1 == node->outputs().size());
+    emitLoadInputs(node->inputs());
+
+    // Emit the expected type.
+    size_t types_start = type_table_.size();
+    auto types = node->tys(attr::types);
+    for (const auto i : c10::irange(num_inputs)) {
+      emitType(types[i]);
+    }
+    insertInstruction(TYPECHECK, types_start, num_inputs);
+  }
+
+  size_t emitGuard(Node* node) {
+    // unoptimized graph is at index 0
+    // guarded input is at index 1
+    // the rest of args follow
+    emitLoadInputs(node->inputs().slice(1, 1));
+    insertInstruction(GUARD, emitType(node->outputs().at(0)->type()));
+    insertInstruction(JF, 0 /* to be patched */);
+    return instructions_.size() - 1;
+  }
+
+  void emitBailOut(Node* node) {
+    auto jf_index = emitGuard(node);
+    auto unoptimized_graph = node->inputs().at(0)->node()->g(attr::Subgraph);
+    // note, guaded input is already loaded onto the stack
+    // for GUARD instruction
+    emitLoadInputs(node->inputs().slice(2));
+    insertInstruction(TAIL_CALL, function_table_.size());
+    TORCH_INTERNAL_ASSERT(node->kind() == prim::BailOut);
+    auto bailout_index = node->i(attr::index);
+    TORCH_INTERNAL_ASSERT(bailout_index >= 0);
+
+    auto build_bailout_graph = [bailout_index,
+                                unoptimized_graph](GraphFunction& func) {
+      BuildBailOutGraphFrom(bailout_index, unoptimized_graph, func.graph());
+    };
+
+    auto empty_graph = std::make_shared<Graph>();
+    auto func = std::make_unique<GraphFunction>(
+        "bailout", empty_graph, build_bailout_graph);
+    function_table_.emplace_back(func.get());
+    bailout_functions_.emplace_back(std::move(func));
+    createBailoutBlock(jf_index);
+  }
+
+  void emitProfile(Node* node) {
+    emitLoadInputs(node->inputs());
+    insertInstruction(PROFILE_OP, profile_function_table_.size());
+    if (node->cast<ProfileOp>()) {
+      profile_function_table_.push_back(node->cast<ProfileOp>()->getCallback());
+    } else if (node->cast<ProfileIValueOp>()) {
+      profile_function_table_.push_back(
+          node->cast<ProfileIValueOp>()->getCallback());
+    } else {
+      TORCH_INTERNAL_ASSERT(false);
+    }
+  }
+
+  void emitGetAttr(Node* node) {
+    emitLoadInputs(node->inputs());
+    const auto type = node->input()->type()->expect<ClassType>();
+    const auto& field = node->s(attr::name);
+    const auto slot = type->getAttributeSlot(field);
+    insertInstruction(GET_ATTR, slot);
+  }
+
+  void emitSetAttr(Node* node) {
+    emitLoadInputs(node->inputs());
+    const auto type = node->inputs().at(0)->type()->expect<ClassType>();
+    const auto& field = node->s(attr::name);
+    const auto slot = type->getAttributeSlot(field);
+    insertInstruction(SET_ATTR, slot);
+  }
+
+  void insertBailoutBlocks() {
+    for (const BailoutBlock& block : bailout_blocks_) {
+      TORCH_INTERNAL_ASSERT(instructions_[block.jf_instruction_index].op == JF)
+      instructions_[block.jf_instruction_index].X =
+          instructions_.size() - block.jf_instruction_index;
+      instructions_.insert(
+          instructions_.end(),
+          block.instructions.begin(),
+          block.instructions.end());
+      instructions_source_.insert(
+          instructions_source_.end(),
+          block.instructions.size(),
+          instructions_source_[block.jf_instruction_index]);
+    }
+  }
+  void emitInterfaceCall(
+      std::string method_name_str,
+      c10::ArrayRef<Value*> inputs) {
+    emitLoadInputs(inputs);
+    auto method_name = insertConstant(std::move(method_name_str));
+    insertInstruction(INTERFACE_CALL, method_name, inputs.size());
+  }
+
+  void emitListUnpack(Node* node) {
+    emitLoadInputs(node->inputs());
+    insertInstruction(LIST_UNPACK, node->outputs().size());
+  }
+
+  void emitTupleConstruct(Node* node) {
+    bool named =
+        node->output()->type()->expectRef<TupleType>().name().has_value();
+    if (named) {
+      emitContainerConstruct(NAMED_TUPLE_CONSTRUCT, node);
+    } else {
+      emitLoadInputs(node->inputs());
+      insertInstruction(TUPLE_CONSTRUCT, node->inputs().size());
+    }
+  }
+
+  void emitContainerConstruct(OpCode op, Node* node) {
+    emitLoadInputs(node->inputs());
+    insertInstruction(
+        op, emitType(node->output()->type()), node->inputs().size());
+  }
+
+  void emitCreateObject(Node* node) {
+    insertInstruction(CREATE_OBJECT, emitType(node->output()->type()));
+  }
+  void emitIsinstance(Node* node) {
+    emitLoadInputs(node->inputs());
+    std::vector<TypePtr> types = node->tys(attr::types);
+    size_t types_start = type_table_.size();
+    for (const auto& typ : types) {
+      emitType(typ);
+    }
+    insertInstruction(ISINSTANCE, types_start, types.size());
+  }
+
+  void emitTupleSlice(Node* node) {
+    emitLoadInputs(node->inputs());
+    int64_t beg_ind = node->i(attr::beg);
+    int64_t end_ind = node->i(attr::end);
+    insertInstruction(TUPLE_SLICE, beg_ind, end_ind - beg_ind);
+  }
+
+  void emitFork(Node* node) {
+    emitLoadInputs(node->inputs());
+    auto forked_fn = std::make_unique<GraphFunction>(
+        "<forked function>", node->g(attr::Subgraph), nullptr);
+    forked_functions_.emplace_back(std::move(forked_fn));
+    function_table_.emplace_back(forked_functions_.back().get());
+    insertInstruction(FORK, function_table_.size() - 1, node->inputs().size());
+  }
+
+  void emitAwaitable(Node* node) {
+    emitLoadInputs(node->inputs());
+    auto await_fn = std::make_unique<GraphFunction>(
+        "<awaitable function>", node->g(attr::Subgraph), nullptr);
+    awaited_functions_.emplace_back(std::move(await_fn));
+    function_table_.emplace_back(awaited_functions_.back().get());
+    insertInstruction(
+        AWAITABLE, function_table_.size() - 1, node->inputs().size());
+  }
+
+  void emitWarn(Node* node) {
+    if (FLAGS_torch_jit_disable_warning_prints) {
+      return;
+    }
+
+    emitLoadInputs(node->inputs());
+    int32_t idx = -1;
+    if (node->hasAttribute(attr::warn_id)) {
+      idx = static_cast<int32_t>(node->i(attr::warn_id));
+    }
+    insertInstruction(WARN, idx);
+  }
+
+  void emitEnter(Node* node) {
+    emitLoadInputs(node->inputs());
+    insertInstruction(ENTER);
+  }
+
+  void emitExit(Node* /* node */) {
+    insertInstruction(EXIT);
+  }
+
+  void emitNode(Node* node) {
+    WithCurrentNode guard(&current_node_, node);
+    switch (node->kind()) {
+      default:
+        // NOLINTNEXTLINE(clang-analyzer-optin.cplusplus.VirtualCall)
+        checkNodeAndEmit(node);
+        // emitOperator(node);
+        break;
+      case prim::RaiseException:
+        emitOperatorOrInstruction(node, RAISE_EXCEPTION);
+        break;
+      case prim::TupleIndex:
+        emitOperatorOrInstruction(node, TUPLE_INDEX);
+        break;
+      case prim::Drop:
+        emitDrop(node->inputs());
+        break;
+      case prim::Constant:
+        emitConstant(node);
+        break;
+      case prim::If:
+        emitIf(node);
+        break;
+      case prim::Loop:
+        emitLoop(node);
+        break;
+      case aten::wait:
+        emitWait(node);
+        break;
+      case prim::Param:
+        break;
+      case prim::CallFunction:
+        emitCall(
+            node->inputs().at(0)->type()->expectRef<FunctionType>().function(),
+            node->inputs().slice(1));
+        break;
+      case prim::CallMethod:
+        if (auto class_type = node->inputs().at(0)->type()->cast<ClassType>()) {
+          emitCall(&class_type->getMethod(node->s(attr::name)), node->inputs());
+        } else {
+          emitInterfaceCall(node->s(attr::name), node->inputs());
+        }
+        break;
+      case prim::TypeCheck:
+        emitTypeCheck(node);
+        break;
+      case prim::BailOut:
+        emitBailOut(node);
+        break;
+      case prim::profile_ivalue:
+      case prim::profile:
+        emitProfile(node);
+        break;
+      case prim::GetAttr:
+        emitGetAttr(node);
+        break;
+      case prim::SetAttr:
+        emitSetAttr(node);
+        break;
+      case prim::ListUnpack:
+        emitListUnpack(node);
+        break;
+      case prim::TupleConstruct:
+        emitTupleConstruct(node);
+        break;
+      case prim::ListConstruct:
+        emitContainerConstruct(LIST_CONSTRUCT, node);
+        break;
+      case prim::DictConstruct:
+        emitContainerConstruct(DICT_CONSTRUCT, node);
+        break;
+      case prim::CreateObject:
+        emitCreateObject(node);
+        break;
+      case prim::isinstance:
+        emitIsinstance(node);
+        break;
+      case prim::TupleSlice:
+        emitTupleSlice(node);
+        break;
+      case prim::fork:
+        emitFork(node);
+        break;
+      case prim::awaitable:
+        emitAwaitable(node);
+        break;
+      case aten::warn:
+        emitWarn(node);
+        break;
+      case prim::Enter:
+        emitEnter(node);
+        break;
+      case prim::Exit:
+        emitExit(node);
+        break;
+      case prim::Uninitialized:
+        emitOperatorOrInstruction(node, UN_INITIALIZED, 0, 0, false);
+        break;
+      case prim::dtype:
+        emitOperatorOrInstruction(node, DTYPE);
+        break;
+      case prim::device:
+        emitOperatorOrInstruction(node, DEVICE);
+        break;
+      case aten::dim:
+        emitOperatorOrInstruction(node, DIM);
+        break;
+      case prim::is_cuda:
+        emitOperatorOrInstruction(node, IS_CUDA);
+        break;
+      case aten::__not__:
+        emitOperatorOrInstruction(node, __NOT__);
+        break;
+      case aten::format:
+        emitFormat(node);
+        break;
+      case aten::__is__:
+        emitOperatorOrInstruction(node, __IS__);
+        break;
+      case aten::__isnot__:
+        emitOperatorOrInstruction(node, __ISNOT__);
+        break;
+      case prim::NumToTensor:
+        emitOperatorOrInstruction(node, NUM_TO_TENSOR);
+        break;
+      case prim::tolist:
+        emitOperatorOrInstruction(node, TO_LIST);
+        break;
+    }
+  }
+
+  void emitCodeForBlock(Block* block) {
+    emitNodeAtBlockLevel(block->param_node());
+    for (auto node : block->nodes()) {
+      emitNodeAtBlockLevel(node);
+    }
+    emitNodeAtBlockLevel(block->return_node());
+  }
+
+  const std::vector<GraphExecutor*>& grad_executors() {
+    if (!grad_executors_) {
+      grad_executors_.emplace();
+      for (Operation& op : operator_table_) {
+        if (auto executor = detail::getGradExecutor(op)) {
+          grad_executors_->push_back(executor);
+        }
+      }
+    }
+    return *grad_executors_;
+  }
+
+  const std::vector<GraphExecutor*>& diff_graph_op_executors() {
+    if (!forward_executors_) {
+      forward_executors_.emplace();
+      for (Operation& op : operator_table_) {
+        if (auto executor = detail::getDifferentiableGraphOpExecutor(op)) {
+          forward_executors_->push_back(executor);
+        }
+      }
+    }
+    return *forward_executors_;
+  }
+
+  void dump(std::ostream& out, size_t i) const {
+    out << i << ' ' << instructions_[i];
+    if (instructions_[i].op == OP || instructions_[i].op == CALL ||
+        instructions_[i].op == OPN) {
+      out << " # " << *instructions_source_[i];
+    } else {
+      out << '\n';
+    }
+  }
+
+  void dump(std::ostream& out) const {
+    out << *graph_ << '\n';
+    for (const auto i : c10::irange(instructions_.size())) {
+      dump(out, i);
+    }
+  }
+
+  /**
+   * Add an operation to operator_table_ if not a duplicate and return its index
+   */
+  int add_to_operator_table(
+      const Operator& op,
+      const Node* node,
+      const std::string& op_name,
+      const int num_inputs,
+      const bool is_vararg) {
+    int size = operator_table_.size();
+
+    const Operation& oper = op.getOperation(node);
+
+    if (!is_vararg) {
+      std::pair<std::string, int> key(op_name, num_inputs);
+      auto found = operator_table_inv_.find(key);
+
+      if (found != operator_table_inv_.end()) {
+        return found->second;
+      }
+
+      operator_table_inv_.emplace(key, size);
+    }
+
+    operator_table_.emplace_back(oper);
+#ifndef NDEBUG
+    full_operator_table_.emplace_back(op);
+#endif
+    return size;
+  }
+
+  inline void assert_stack_size(
+      int32_t instruction_index,
+      size_t init_size,
+      size_t actual_size) const {
+#ifndef NDEBUG
+    const auto& schema = full_operator_table_[instruction_index].schema();
+    int64_t expected_size = static_cast<int64_t>(init_size) -
+        static_cast<int64_t>(schema.arguments().size()) +
+        static_cast<int64_t>(schema.returns().size());
+    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
+        static_cast<size_t>(expected_size) == actual_size ||
+            schema.is_varret() || schema.is_vararg(),
+        "Expected to find ",
+        expected_size,
+        " values on the stack, but found ",
+        actual_size,
+        " on the stack after ",
+        toString(full_operator_table_[instruction_index].schema()));
+#endif
+  }
+};
+
+struct MobileCodeImpl : CodeImpl {
+  MobileCodeImpl(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name,
+      bool emit_default_input_instructions,
+      bool support_default_args_before_out,
+      bool emit_promoted_ops,
+      size_t remaining_bailout_depth)
+      : CodeImpl(
+            graph,
+            std::move(function_name),
+            remaining_bailout_depth,
+            false),
+        emit_default_input_instructions_(emit_default_input_instructions),
+        support_default_args_before_out_(support_default_args_before_out),
+        emit_promoted_ops_(emit_promoted_ops) {
+    // NOLINTNEXTLINE(clang-analyzer-optin.cplusplus.VirtualCall)
+    run();
+  }
+
+  void run() override {
+    process_ops_for_mobile();
+    emitCodeForBlock(graph_->block());
+    insertInstruction(RET);
+    // we deferred the emission of bailout blocks so they appear at the end
+    // emit them now and patch up the jumps
+    insertBailoutBlocks();
+  }
+
+  void process_ops_for_mobile() {
+    DepthFirstGraphNodeIterator graph_it(graph_);
+    Node* node = graph_it.next();
+    while (node) {
+      if (node->maybeOperator()) {
+        auto op_schema = node->getOperator().schema();
+        // skip if schema has vararg
+        if (!op_schema.is_vararg()) {
+          auto specifiedArgs = CalculateNecessaryArgs(
+              op_schema.arguments(),
+              node->inputs(),
+              support_default_args_before_out_);
+
+          size_t numInclude = specifiedArgs.first +
+              (support_default_args_before_out_ ? specifiedArgs.second : 0);
+          auto unique_name = !op_schema.overload_name().empty()
+              ? op_schema.name() + "." + op_schema.overload_name()
+              : op_schema.name();
+          auto it = op_to_num_specified_args_.insert(
+              std::pair<std::string, size_t>(unique_name, 0));
+          op_to_num_out_args_.insert(std::pair<std::string, size_t>(
+              unique_name, specifiedArgs.second));
+          auto prev_value = it.first->second;
+          it.first->second = std::max(numInclude, prev_value);
+        }
+      }
+      node = graph_it.next();
+    }
+  }
+
+ private:
+  void emitOperator(Node* node) override {
+    if (emit_default_input_instructions_) {
+      CodeImpl::emitOperator(node);
+    } else {
+      const Operator& op = node->getOperator();
+      std::string unique_op_name = c10::toString(op.schema().operator_name());
+      int num_inputs = node->inputs().size();
+      bool is_vararg = op.schema().is_vararg();
+
+      if (op.hasOperation() && is_vararg) {
+        emitLoadInputs(node->inputs());
+        int operation_index = add_to_operator_table(
+            op,
+            node,
+            unique_op_name,
+            num_inputs,
+            /* is_vararg */ true);
+        insertInstruction(OPN, operation_index, num_inputs);
+      } else {
+        auto num_include = num_inputs;
+        auto it = op_to_num_specified_args_.find(unique_op_name);
+        if (it != op_to_num_specified_args_.end()) {
+          num_include = it->second;
+        }
+        if (support_default_args_before_out_) {
+          auto num_out = op_to_num_out_args_.find(unique_op_name)->second;
+          auto num_specified_before_out = num_include - num_out;
+          emitLoadInputs(node->inputs(), 0, num_specified_before_out);
+          emitLoadInputs(
+              node->inputs(),
+              node->inputs().size() - num_out,
+              node->inputs().size());
+        } else {
+          emitLoadInputs(node->inputs(), num_include);
+        }
+        int operation_index = add_to_operator_table(
+            op, node, unique_op_name, num_inputs, is_vararg);
+        insertInstruction(OP, operation_index);
+      }
+    }
+  }
+
+  void emitOperatorOrInstruction(
+      Node* node,
+      OpCode op,
+      int64_t X = 0,
+      uint64_t N = 0,
+      bool emit_inputs = true) override {
+    if (emit_promoted_ops_) {
+      CodeImpl::emitOperatorOrInstruction(node, op, X, N, emit_inputs);
+    } else {
+      CodeImpl::emitOperator(node);
+    }
+  }
+
+  // To support forward compatibility for bytecode version bump from v5 to v6
+  bool emit_default_input_instructions_;
+  // To support forward compatibility for bytecode version bump from v6 to v7
+  bool support_default_args_before_out_;
+  // To support forward compatibility for bytecode version bump from v7 to v8
+  bool emit_promoted_ops_;
+};
+
+} // namespace torch::jit::interpreter
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/frame.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/frame.h
new file mode 100644
index 0000000000000000000000000000000000000000..16b4efc0497fdf3b943e52075ae8c2129ddc1caf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/frame.h
@@ -0,0 +1,45 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <atomic>
+#include <memory>
+
+#include <torch/csrc/jit/runtime/interpreter/code_impl.h>
+#include <torch/csrc/jit/runtime/profiling_record.h>
+
+namespace torch::jit::interpreter {
+
+// A Frame captures function's state
+// (e.g. `pc` and `base_pointer`)
+// Each Frame corresponds to a call to a `Frame::function`
+// which has not yet returned
+// The arguments for `Frame::function`
+// are located at [base_pointer + arg_number]
+struct Frame {
+  std::shared_ptr<CodeImpl> function;
+  // program counter corresponds to the index
+  // of the currently executed instruction
+  size_t pc;
+  // marks the start index of the frame
+  // base_pointer is used by TAIL_CALL
+  // to replace the current frame
+  // with a frame of a bailout graph
+  size_t base_pointer;
+
+  // unique to every frame with prim::profile across all threads
+  std::optional<size_t> id;
+
+  // RecordFunction object associated with this frame
+  std::unique_ptr<at::RecordFunction> record_function;
+
+  // symbol table for a frame
+  ShapeSymbolTable symbols2dims;
+
+  static size_t genId();
+};
+
+} // namespace torch::jit::interpreter
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/preprocess_graph.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/preprocess_graph.h
new file mode 100644
index 0000000000000000000000000000000000000000..62cb9e5367ba77f9c9e364f0cd775cca48039952
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/interpreter/preprocess_graph.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <unordered_map>
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::interpreter {
+
+// pre-processing that happens once per graph
+struct PreprocessGraph {
+  explicit PreprocessGraph(Graph& g);
+
+  // Outputs of the preprocessing:
+  std::shared_ptr<Graph> graph;
+  std::unordered_map<Node*, bool> can_emit_inline;
+};
+
+} // namespace torch::jit::interpreter
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/jit_exception.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/jit_exception.h
new file mode 100644
index 0000000000000000000000000000000000000000..1ce95c6271f6f2ec91155d36793e2769f1da0fa4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/jit_exception.h
@@ -0,0 +1,43 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <stdexcept>
+
+#include <torch/csrc/Export.h>
+#include <optional>
+#include <string>
+
+namespace torch::jit {
+
+struct TORCH_API JITException : public std::runtime_error {
+  explicit JITException(
+      const std::string& msg,
+      std::optional<std::string> python_class_name = std::nullopt,
+      std::optional<std::string> original_msg = std::nullopt);
+
+  std::optional<std::string> getPythonClassName() const {
+    return python_class_name_;
+  }
+
+  // the original msg if this is from a python exception. The interpreter has
+  // changed the original message by adding "The following operation failed in
+  // the TorchScript interpreter." in front of it in the handleError function.
+  std::optional<std::string> getOriginalMsg() const {
+    return original_msg_;
+  }
+
+  static const std::string& getCaughtOriginalMsg();
+  static const std::string& getCaughtPythonClassName();
+  static void setCaughtOriginalMsg(const std::string& msg);
+  static void setCaughtPythonClassName(const std::string& pythonClassName);
+
+ private:
+  std::optional<std::string> python_class_name_;
+  std::optional<std::string> original_msg_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/jit_trace.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/jit_trace.h
new file mode 100644
index 0000000000000000000000000000000000000000..d6c7d15e94c0785223a7e96b5e02981c5cc4c3c2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/jit_trace.h
@@ -0,0 +1,13 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/jit/ir/ir.h>
+#include <memory>
+
+namespace torch::jit {
+TORCH_API std::shared_ptr<Graph> TraceGraph(
+    const std::shared_ptr<Graph>& graph,
+    Stack& stack);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/logging.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/logging.h
new file mode 100644
index 0000000000000000000000000000000000000000..b56c20b3ba24fdedd04c9e31b3700b8d98057680
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/logging.h
@@ -0,0 +1,94 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <chrono>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <torch/csrc/Export.h>
+
+namespace torch::jit::logging {
+
+class LoggerBase {
+ public:
+  TORCH_API virtual void addStatValue(
+      const std::string& stat_name,
+      int64_t val) = 0;
+  virtual ~LoggerBase() = default;
+};
+
+TORCH_API LoggerBase* getLogger();
+TORCH_API LoggerBase* setLogger(LoggerBase* logger);
+
+// No-op logger. This is the default and is meant to incur almost no runtime
+// overhead.
+
+class NoopLogger : public LoggerBase {
+ public:
+  void addStatValue(
+      const std::string& stat_name [[maybe_unused]],
+      int64_t val [[maybe_unused]]) override {}
+  ~NoopLogger() override = default;
+};
+
+// Trivial locking logger. Pass in an instance of this to setLogger() to use it.
+// This keeps track of the sum of all statistics.
+//
+// NOTE: this is not written in a scalable way and should probably only be used
+// in the single-threaded case or for testing.
+class TORCH_API LockingLogger : public LoggerBase {
+ public:
+  void addStatValue(const std::string& stat_name, int64_t val) override;
+  virtual int64_t getCounterValue(const std::string& name) const;
+  enum class AggregationType { SUM = 0, AVG = 1 };
+  void setAggregationType(const std::string& stat_name, AggregationType type);
+  ~LockingLogger() override = default;
+
+ private:
+  mutable std::mutex m;
+  struct RawCounter {
+    RawCounter() = default;
+    int64_t sum{0};
+    size_t count{0};
+  };
+  std::unordered_map<std::string, RawCounter> raw_counters;
+  std::unordered_map<std::string, AggregationType> agg_types;
+};
+
+// Make this struct so the timer internals are opaque to the user.
+struct JITTimePoint {
+  std::chrono::time_point<std::chrono::high_resolution_clock> point;
+};
+
+TORCH_API JITTimePoint timePoint();
+TORCH_API void recordDurationSince(
+    const std::string& name,
+    const JITTimePoint& tp);
+
+namespace runtime_counters {
+constexpr const char* GRAPH_EXECUTORS_CONSTRUCTED =
+    "pytorch_runtime.graph_executors_constructed";
+constexpr const char* GRAPH_EXECUTOR_INVOCATIONS =
+    "pytorch_runtime.graph_executor_invocations";
+constexpr const char* EXECUTION_PLAN_CACHE_HIT =
+    "pytorch_runtime.execution_plan_cache_hit";
+constexpr const char* EXECUTION_PLAN_CACHE_MISS =
+    "pytorch_runtime.execution_plan_cache_miss";
+
+inline std::vector<const char*> allRuntimeCounters() {
+  return {
+      GRAPH_EXECUTORS_CONSTRUCTED,
+      GRAPH_EXECUTOR_INVOCATIONS,
+      EXECUTION_PLAN_CACHE_HIT,
+      EXECUTION_PLAN_CACHE_MISS};
+}
+
+} // namespace runtime_counters
+
+} // namespace torch::jit::logging
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/operator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/operator.h
new file mode 100644
index 0000000000000000000000000000000000000000..619965550581c0ebe99a436aea2de536b0c2868c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/operator.h
@@ -0,0 +1,348 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// in memory description of all ATen Ops similar to Caffe2 schema
+// once C10 exists this can be removed, or stubbed out, but we need
+// it now to implement correct semantic checking for script
+#pragma once
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/core/dispatch/OperatorOptions.h>
+#include <ATen/core/op_registration/op_allowlist.h>
+#include <ATen/core/stack.h>
+#include <c10/util/Exception.h>
+#include <c10/util/overloaded.h>
+#include <torch/csrc/jit/frontend/function_schema_parser.h>
+#include <torch/csrc/jit/runtime/operator_options.h>
+#include <torch/library.h>
+
+#include <ATen/core/function_schema.h>
+#include <ATen/core/symbol.h>
+
+#include <functional>
+#include <initializer_list>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <variant>
+#include <vector>
+
+namespace torch::jit {
+
+struct Node;
+using ::c10::Argument;
+using ::c10::FunctionSchema;
+using ::c10::Symbol;
+
+using OperationCreator = Operation (*)(const Node*);
+
+namespace {
+const std::array<at::Tag, 1> kJitOnlyOperatorTags = {
+    at::Tag::pt2_compliant_tag};
+}
+
+/*
+ * Note: JIT relies on Operator instances having static lifetime, because
+ * it for example stores a non-owning FunctionSchema* pointer in the Node class,
+ * which points to the function schema stored in the Operator instance.
+ * Also, jit::Operator is meant to store more operator related information like
+ * symbolic derivatives, which also requires them to have static lifetime
+ * so that changes to symbolic derivatives are remembered.
+ *
+ * Currently, the JIT operator library contains a jit::Operator instance
+ * with a wrapper for each c10 operator. The c10 operator library registers
+ * those wrappers using listeners in register_c10_ops.cpp.
+ * TODO Instead of doing it this way, we should only have pure-jit ops in
+ * the jit library but have the JIT operator lookup look into the c10 library
+ * too.
+ */
+
+// An Operator is a thin wrapper around either a pure JIT operator (e.g. prim
+// ops) or a c10 operator, allowing some common operations and abstracting away
+// the concrete operator nature.
+struct TORCH_API Operator {
+ private:
+  struct C10Operator final {
+    c10::OperatorHandle handle_;
+    Operation op_;
+  };
+  struct UnparsedFunctionSchema final {
+    std::string schema_string_;
+    mutable std::optional<c10::AliasAnalysisKind> alias_analysis_;
+  };
+  struct JitOnlyOperator final {
+    // The only valid transition for schema_ is from right->left, i.e.
+    // when the schema gets parsed.
+    mutable std::variant<FunctionSchema, UnparsedFunctionSchema> schema_;
+
+    std::variant<Operation, OperationCreator> op_;
+  };
+
+ public:
+  Operator(c10::OperatorHandle opHandle, Operation operation)
+      : op_(C10Operator{std::move(opHandle), std::move(operation)}) {}
+
+  Operator(
+      std::string schema,
+      Operation op,
+      c10::AliasAnalysisKind alias_analysis)
+      : op_(JitOnlyOperator{
+            UnparsedFunctionSchema{std::move(schema), alias_analysis},
+            Operation(std::move(op))}) {}
+
+  Operator(
+      std::string name,
+      std::string overload_name,
+      std::vector<Argument> arguments,
+      std::vector<Argument> returns,
+      Operation op,
+      c10::AliasAnalysisKind alias_analysis)
+      : op_(JitOnlyOperator{
+            FunctionSchema(varArgSchemaWithName(
+                std::move(name),
+                std::move(overload_name),
+                std::move(arguments),
+                std::move(returns),
+                alias_analysis)),
+            std::move(op)}) {}
+
+  Operator(
+      std::string schema,
+      OperationCreator op_creator,
+      c10::AliasAnalysisKind alias_analysis)
+      : op_(JitOnlyOperator{
+            UnparsedFunctionSchema{std::move(schema), alias_analysis},
+            op_creator}) {}
+
+  // Helper constructor to register `op` to run
+  // run for _every_ IR Node where n.kind() == name, regardless of arguments.
+  // This is accomplished by marking the schema varargs and having no required
+  // arguments.
+  Operator(
+      Symbol name,
+      OperationCreator op_creator,
+      c10::AliasAnalysisKind alias_analysis)
+      : op_(JitOnlyOperator{
+            FunctionSchema(varArgSchemaWithName(name, alias_analysis)),
+            op_creator}) {}
+
+  Operation getOperation(const Node* node = nullptr) const {
+    return std::visit(
+        c10::overloaded(
+            [](const C10Operator& op) { return op.op_; },
+            [node](const JitOnlyOperator& op) {
+              return std::visit(
+                  c10::overloaded(
+                      [](const Operation& op) { return op; },
+                      [node](const OperationCreator& op_creator) {
+                        return op_creator(node);
+                      }),
+                  op.op_);
+            }),
+        op_);
+  }
+
+  Operation getOperationForDispatchKey(c10::DispatchKey dk) const {
+    // TODO: some sort of caching mechanism?
+    return std::visit(
+        c10::overloaded(
+            [dk](const C10Operator& op) {
+              return Operation([op, dk](Stack& stack) {
+                op.handle_.callBoxedForDispatchKey(dk, stack);
+              });
+            },
+            [](const JitOnlyOperator& op) {
+              TORCH_CHECK(
+                  false,
+                  "calling a JIT operator for dispatch key is not supported");
+              return Operation(nullptr);
+            }),
+        op_);
+  }
+
+  const FunctionSchema& schema() const {
+    return std::visit(
+        c10::overloaded(
+            [](const C10Operator& op) -> const FunctionSchema& {
+              return op.handle_.schema();
+            },
+            [](const JitOnlyOperator& op) -> const FunctionSchema& {
+              // we lazily parse schema initialized from strings so that
+              // we do less work during static operator registration
+              if (op.schema_.index() == 1) {
+                auto& unmaterializedSchema =
+                    std::get<UnparsedFunctionSchema>(op.schema_);
+                FunctionSchema schema =
+                    parseSchema(unmaterializedSchema.schema_string_);
+                if (unmaterializedSchema.alias_analysis_.has_value()) {
+                  // TODO What if it gets set later?
+                  schema.setAliasAnalysis(
+                      *unmaterializedSchema.alias_analysis_);
+                }
+                op.schema_ = std::move(schema);
+              }
+              return std::get<FunctionSchema>(op.schema_);
+            }),
+        op_);
+  }
+
+  c10::ArrayRef<at::Tag> getTags() const {
+    return std::visit(
+        c10::overloaded(
+            [](const C10Operator& op) { return op.handle_.getTags(); },
+            [](const JitOnlyOperator& op) {
+              // JitOnlyOperators don't have an c10::OperatorHandle or a way to
+              // specify tags. We're grandfathering them all into
+              // pt2_compliant_tag, but for anything else, please just stop
+              // using JitOnlyOperator.
+              return c10::ArrayRef<at::Tag>(kJitOnlyOperatorTags);
+            }),
+        op_);
+  }
+
+  bool isC10Op() const {
+    return op_.index() == 0;
+  }
+
+  c10::AliasAnalysisKind aliasAnalysisKind() const {
+    const FunctionSchema& schemaRef = schema();
+    c10::AliasAnalysisKind alias_analysis = schemaRef.aliasAnalysis();
+
+    TORCH_CHECK(
+        alias_analysis == AliasAnalysisKind::FROM_SCHEMA ||
+            !schemaRef.hasAnyAliasInfo(),
+        "In operator registration: Tried to register operator ",
+        schemaRef,
+        " with aliasing information in the schema but without AliasAnalysisKind::FROM_SCHEMA.");
+    return alias_analysis;
+  }
+
+  bool hasOperation() const {
+    return std::visit(
+        c10::overloaded(
+            [](const C10Operator&) { return true; },
+            [](const JitOnlyOperator& op) { return op.op_.index() == 0; }),
+        op_);
+  }
+
+ private:
+  static FunctionSchema varArgSchemaWithName(
+      Symbol name,
+      AliasAnalysisKind alias_analysis) {
+    auto result = FunctionSchema(
+        name,
+        "",
+        {},
+        {},
+        /*is_vararg*/ true,
+        /*is_varret*/ true);
+    result.setAliasAnalysis(alias_analysis);
+    return result;
+  }
+
+  static FunctionSchema varArgSchemaWithName(
+      std::string name,
+      std::string overload_name,
+      std::vector<Argument> arguments,
+      std::vector<Argument> returns,
+      AliasAnalysisKind alias_analysis) {
+    auto result = FunctionSchema(
+        std::move(name),
+        std::move(overload_name),
+        std::move(arguments),
+        std::move(returns),
+        /*is_vararg*/ false,
+        /*is_varret*/ false);
+    result.setAliasAnalysis(alias_analysis);
+    return result;
+  }
+
+  std::variant<C10Operator, JitOnlyOperator> op_;
+};
+
+TORCH_API std::string canonicalSchemaString(const FunctionSchema& schema);
+
+TORCH_API const std::vector<std::shared_ptr<Operator>> getAllOperators();
+TORCH_API const std::vector<std::shared_ptr<Operator>>& getAllOperatorsFor(
+    Symbol name);
+// Returns operators in the order which OpOverloadPacket resolves them.
+TORCH_API std::vector<std::shared_ptr<Operator>> getAllSortedOperatorsFor(
+    Symbol name);
+
+// given a operator with an overload name, find the specific operator related to
+// it, may return nullptr if no operator exists.
+TORCH_API std::shared_ptr<Operator> findOperatorFor(
+    const c10::OperatorName& full_name);
+
+TORCH_API std::vector<Symbol> findSimilarOperators(Symbol input_op);
+
+TORCH_API void registerOperator(Operator&& op);
+TORCH_API void deregisterOperator(const FunctionSchema& schema);
+
+// XXX: this function is meant to be used with string literals only!
+TORCH_API std::shared_ptr<Operator> getOperatorForLiteral(
+    const char* signature);
+
+// Ensure the thing that registers c10 ops is defined.
+// Otherwise, our registry will not have c10 ops. You can run into this
+// scenario if you're querying registered ops during static init.
+//
+// This fn is defined in register_c10_ops.cpp
+TORCH_API void ensure_c10_registerer_defined();
+
+// Used to assert that unschematized operators have an analysis method written
+TORCH_API bool aliasAnalysisHasSpecialCaseFor(c10::Symbol sym);
+
+// A factory function to generate an optional operator. It has two
+// instantiations depending on the template bool arg value. The arg can be a
+// compile-time function for the selective op registration based on schema
+// string.
+template <typename Func>
+std::optional<Operator> OperatorGenerator(
+    const char* schema_str,
+    Func&& op,
+    AliasAnalysisKind alias_analysis) {
+  return std::optional<Operator>(Operator(
+      std::string(schema_str), std::forward<Func>(op), alias_analysis));
+}
+
+template <typename Func>
+std::optional<Operator> OperatorGenerator(
+    torch::detail::SelectiveStr<true> schema_str,
+    Func&& op,
+    AliasAnalysisKind alias_analysis) {
+  return OperatorGenerator(
+      static_cast<const char*>(schema_str),
+      std::forward<Func>(op),
+      alias_analysis);
+}
+
+template <typename Func>
+std::optional<Operator> OperatorGenerator(
+    torch::detail::SelectiveStr<false> schema_str,
+    Func&& op,
+    AliasAnalysisKind alias_analysis) {
+  return std::nullopt;
+}
+
+template <typename Func>
+std::optional<Operator> OperatorGenerator(
+    const std::string name,
+    const std::string overload_name,
+    const std::vector<c10::Argument> arguments,
+    const std::vector<c10::Argument> returns,
+    Func&& op,
+    AliasAnalysisKind alias_analysis) {
+  return std::optional<Operator>(Operator(
+      name,
+      overload_name,
+      arguments,
+      returns,
+      std::forward<Func>(op),
+      alias_analysis));
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/operator_options.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/operator_options.h
new file mode 100644
index 0000000000000000000000000000000000000000..ade935e052b07e79dcbd722407422538c2ea3c68
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/operator_options.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/dispatch/OperatorOptions.h>
+
+namespace torch::jit {
+
+using AliasAnalysisKind = c10::AliasAnalysisKind;
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/print_handler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/print_handler.h
new file mode 100644
index 0000000000000000000000000000000000000000..911e9bc905ce609ce146090b24c684e370d08660
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/print_handler.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+
+#include <string>
+
+namespace torch::jit {
+
+using PrintHandler = void (*)(const std::string&);
+
+TORCH_API PrintHandler getDefaultPrintHandler();
+TORCH_API PrintHandler getPrintHandler();
+TORCH_API void setPrintHandler(PrintHandler ph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/profiling_graph_executor_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..be8dd0ae4748a9e5ecfc124f72dd87563ae22cb8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
@@ -0,0 +1,87 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Flags.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/runtime/graph_executor_impl.h>
+
+TORCH_DECLARE_bool(torch_jit_static_then_dynamic);
+
+TORCH_DECLARE_bool(torch_jit_always_dynamic);
+
+C10_DECLARE_bool(torch_jit_input_independent_optimization);
+C10_DECLARE_bool(torch_jit_release_profiling_graph_after_optimization);
+C10_DECLARE_int32(torch_jit_release_profiling_graph_delay_in_seconds);
+C10_DECLARE_int64(torch_jit_num_profiled_runs);
+C10_DECLARE_int64(torch_jit_bailout_depth);
+
+namespace torch::jit {
+
+TORCH_API void runNooptPassPipeline(std::shared_ptr<Graph>& graph);
+
+struct TORCH_API ProfilingGraphExecutorImpl : public GraphExecutorImplBase {
+  ProfilingGraphExecutorImpl(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name);
+
+  const ExecutionPlan& getPlanFor(
+      Stack& stack,
+      std::optional<size_t> remaining_bailout_depth) override;
+  const ExecutionPlan& getInputIndependentPlan() override;
+  GraphExecutorState getDebugState() override;
+  ~ProfilingGraphExecutorImpl() override = default;
+
+  void debugFlushCompilationCache();
+
+  bool isOptimized() const override {
+    return optimized_plan_.has_value();
+  }
+
+ private:
+  const ExecutionPlan& getOptimizedPlanFor(
+      Stack& stack,
+      std::optional<size_t> remaining_bailout_depth);
+  // Input-independent optimization, assumes compile_mutex is held.
+  const ExecutionPlan& getInputIndependentPlanImpl();
+  void runProfilingInsensitiveOptimizations(std::shared_ptr<Graph>& graph);
+  void runProfilingOptimizations(
+      std::shared_ptr<Graph>& graph,
+      size_t remaining_depth);
+  void replaceFallbackGraphWithFallbackFunction(Block* b);
+  FusionBehavior getCurrentBehavior(size_t remaining_depth);
+  size_t getInstantiatedBailoutDepth();
+  void runNoGradOptimizations(
+      std::shared_ptr<Graph>& graph,
+      size_t remaining_bailout_depth);
+  void runFinalOptimizations(std::shared_ptr<Graph>& graph);
+
+  void clearTheGraphCompilationIntermediateGraphs();
+
+  std::unique_ptr<ProfilingRecord> pr_;
+  std::optional<ExecutionPlan>
+      profiling_plan_; // plan to run in order to profiling the code
+  std::optional<ExecutionPlan> optimized_plan_;
+  FusionStrategy fusion_strategy_;
+
+  // this plan is used if getGraphExecutorOptimize is unset
+  std::optional<ExecutionPlan> fallback_plan_;
+  // fallback functions are inserted for tensorexpr fusion groups
+  // and by specialize_autogradzero. Whenever, at runtime, input
+  // tensor don't match profiled properties, fallback functions are called
+  // They are the deoptimized version of the logic in fusion groups
+  // and/or autograd.
+  // The fallback functions are owned by a GraphExecutor instance
+  // They only exist in the optimized graph which is a private property
+  // of the GraphExecutor and only shared with InterpreterState
+  std::vector<std::unique_ptr<Function>> fallback_functions_;
+  std::optional<size_t> remaining_bailout_depth_;
+  // The time the optimized_plan_ is created.
+  int32_t time_optimized_plan_created_ = 0;
+  // Has the extra memory used by the graph for profiling is released?
+  bool is_graph_extra_memory_released_ = false;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/profiling_record.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/profiling_record.h
new file mode 100644
index 0000000000000000000000000000000000000000..a72e67bd31d3b530b6118251c78b7fdea7f69564
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/profiling_record.h
@@ -0,0 +1,211 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <list>
+#include <map>
+#include <unordered_map>
+#include <vector>
+
+// We would like to assign each position/axis of a tensor an abstract size
+// * For each `tensor` we have a profiled `Value` of a `TensorType` describing
+// the properties of the `tensor`.
+// * `TensorType` has a property called `symbolic_sizes_` to describe observed
+// `tensor.sizes()`
+// * `symbolic_sizes_` is a vector of abstract sizes (or
+// `std::vector<ShapeSymbol>`) where
+//   * `ShapeSymbol`at `symbolic_sizes_[i]`  describes the size value
+//   (`Dimension`) at `tensor.sizes()[i]`
+// * We may see the same `Dimension` at different positions `i` in
+// `tensor.sizes()` or even in different `tensor`
+//   * First, we would like associate the same `ShapeSymbol` to the same
+//   `Dimension` across **one** profiling execution or run of a TorchScript
+//   function.
+//     * The same `ShapeSymbol`s in different positions of `symbolic_shapes_` in
+//     possibly different `TensorType`s (i.e. `TensorType`s for different
+//     profiled values) form an implicit set. The elements of such a set are
+//     called *dimension locations*.
+//     * These sets allow us to track how the shapes of input arguments of some
+//     operation relate to operation's output shapes as the input and output
+//     shapes might share the same `ShapeSymbol`s
+// * For **every** profiling run, we would like to maintain the invariant that
+// *the same `ShapeSymbol` is always associated with the same `Dimension`*.
+// * To maintain this invariant we merge the profiling information from all
+// profiling runs,
+//   * For every two runs, we iterate over all `symbic_shapes_`  and compare
+//   their `ShapeSymbol`s in the same position.
+//     * if we observe that for every dimension location that has
+//     the`ShapeSymbol S1`  in run #1 there is **only one** `ShapeSymbol S2` in
+//     the same dimension location in run #2, we conclude that the invariant
+//     holds.
+//     * However, if we observe some dimension locations in run #2 have
+//     `ShapeSymbol S2` and the other ones have `ShapeSymbol S3` we would like
+//     to partition the virtual set of dimension locations associated with
+//     `ShapeSymbol S1` into two new subsets, so the invariant holds.
+//     * The partitioning works by assigning a new symbol to the dimension
+//     locations (associated with `ShapeSymbol S1`) that have `ShapeSymbol S2`
+//     and another new symbol to the dimension locations that have `ShapeSymbol
+//     S3`. In other words,
+//       * Subset #1 will consist of the dimension locations that in run #2 have
+//       `ShapeSymbol S2`  and will have `ShapeSymbol S4`  in those dimension
+//       locations
+//       * Subset #2 will consist of the dimension locations that in run #2 have
+//       `ShapeSymbol S4`  and will have `ShapeSymbol S5`  in those dimension
+//       locations
+//     * The effective result of merging the profiling information from two runs
+//     is new `TensorTypes` whose `symbolic_sizes_` /dimension locations have
+//     either `ShapeSymbol S4` or `ShapeSymbol S5`.
+//     * Partitioning can be done even before we have seen all the dimension
+//     locations associated with `ShapeSymbol S1`
+//       * We use `getSymbolInSet` of `ShapeSymbolTable` to remember all
+//       `ShapeSymbols` from run #2 we observed in the dimension locations
+//       associated with `ShapeSymbol S1` .
+//       * For every `ShapeSymbol` from run #2 in the dimension location
+//       associated with `ShapeSymbol S1`  `getSymbolInSet` returns a symbol
+//       that we assign to the dimension location in a new TensorType.
+//         * It's important to point out that the same `ShapeSymbol S2` from run
+//         #2 in two dimension locations that have different `ShapeSymbol`s in
+//         run #1 are different! These dimension locations will belong to
+//         different subsets and have different `ShapeSymbol`s after merge.
+//         * On the other hand, for the same `ShapeSymbol S2` in two dimension
+//         locations that have `ShapeSymbol S1` in run #1`getSymbolInSet` will
+//         return the same symbol.
+
+namespace torch::jit {
+
+using ::c10::TensorTypePtr;
+using Dimension = int64_t;
+
+TORCH_API void RegisterProfilingNode(
+    const std::function<bool(const Node*)>& /*func*/);
+
+struct ProfilingRecord;
+
+// `SetPartitioningHelper` is used to maintain the following invariant:
+// For **every** profiling run, *the same `ShapeSymbol` is always associated
+// with the same `Dimension`*.
+// while merging the profiling information from multiple runs.
+struct SetPartitioningHelper {
+  std::map<c10::ShapeSymbol, std::map<Dimension, c10::ShapeSymbol>>
+      sets2subsets_;
+
+  // `partitionSetByDimension` partitions a virtual set
+  // of dimension locations associated with ShapeSymbol `symbol` into subsets.
+  // Partitioning is equivalent to giving (or renaming) a particular
+  // dimension location a new `ShapeSymbol`.
+  // The same `Dimension` value in different dimension locations
+  // that used to have `symbol` will receive the same
+  // new `ShapeSymbol`, effectively forming a new set.
+  c10::ShapeSymbol partitionSetByDimension(
+      Dimension new_size,
+      c10::ShapeSymbol symbol) {
+    auto& dims2symbols = getSetForSymbol(symbol);
+
+    if (dims2symbols.count(new_size) == 0) {
+      auto new_sym = c10::ShapeSymbol::newSymbol();
+      dims2symbols[new_size] = new_sym;
+      return new_sym;
+    }
+
+    return dims2symbols[new_size];
+  }
+
+ private:
+  std::map<Dimension, c10::ShapeSymbol>& getSetForSymbol(c10::ShapeSymbol s) {
+    auto& set = sets2subsets_[s];
+    // N.B. adding a mapping { s.static_size(), s }
+    // makes sure we preserve the fact that
+    // some dimension values remain the same
+    // across all profiled runs
+    if (s.is_static()) {
+      set.insert({s.static_size(), s});
+    }
+    return set;
+  }
+};
+
+// ShapeSymbolTable is used by Interpreter
+// to assign dimension values to ShapeSymbols
+// and fail a guard if the same symbol
+// is assigned more than one dimension value.
+struct ShapeSymbolTable {
+  // N.B. we treat static symbols as always assigned
+  // to themselves
+  bool isBound(c10::ShapeSymbol s) {
+    if (s.is_static()) {
+      return true;
+    }
+    return data_.count(s) != 0;
+  }
+
+  // N.B. we treat static symbols as always assigned
+  // to themselves
+  Dimension getValue(c10::ShapeSymbol s) {
+    if (s.is_static()) {
+      return s.static_size();
+    }
+    return data_[s];
+  }
+  void assign(c10::ShapeSymbol s, Dimension v) {
+    TORCH_INTERNAL_ASSERT(!s.is_static());
+    data_[s] = v;
+  }
+  std::map<c10::ShapeSymbol, Dimension> data_;
+  // Tries to assign dimension values from `new_sizes` to
+  // `ShapeSymbol`s `sym_shapes`.
+  // Returns `true` if every dimension value from `new_sizes`
+  // can be assigned to the corresponding `ShapeSymbol` from
+  // `sym_shapes`
+  // A dimension value can be assigned to a `ShapeSymbol`
+  // * if the symbol isn't assigned yet any dimension value
+  // * if the symbol is assigned and its value is equal to
+  // the dimension value from `new_sizes`
+  bool bindSymbolicShapes(
+      at::IntArrayRef new_sizes,
+      const c10::SymbolicShape& sym_shapes);
+};
+
+struct ProfilingRecord {
+  // N.B. ProfilingRecord's copy and move c-tor are disabled, so we won't
+  // end up accidentally copying or moving ProfilingRecords whose addresses
+  // are captured in callbacks_
+  ProfilingRecord(const ProfilingRecord&) = delete;
+  ProfilingRecord(ProfilingRecord&&) noexcept = delete;
+  TORCH_API static std::unique_ptr<ProfilingRecord> instrumentGraph(
+      const std::shared_ptr<Graph>& graph);
+  TORCH_API static void removeProfilingNodes(Block* b);
+  TORCH_API static void removeProfileCounter(Block* b);
+
+  std::shared_ptr<Graph> profiled_graph_;
+  mutable std::mutex mutex_;
+  size_t profiling_count_;
+
+  bool ready() const;
+
+  std::shared_ptr<Graph> graph() const {
+    return profiled_graph_;
+  }
+
+  TORCH_API ProfileIValueOp* createProfileIValueNode(Value* in_val);
+  TORCH_API ProfileIValueOp* createProfileIValueNode(ArrayRef<Value*> inputs);
+
+ private:
+  ProfileOp* createProfileNode(
+      const std::function<void(Stack&)>& fp,
+      at::ArrayRef<Value*> inputs);
+  void instrumentBlock(Block* block);
+  void insertShapeProfile(Node* n, size_t offset, const TypePtr& input_type);
+  ProfilingRecord(std::shared_ptr<Graph> g);
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/register_ops_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/register_ops_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..4c6c441a5a9cd0ec0e90e6b13c035f35e9a3ea67
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/register_ops_utils.h
@@ -0,0 +1,888 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Context.h>
+#include <c10/core/DeviceType.h>
+#include <torch/csrc/autograd/autograd.h>
+#include <torch/csrc/autograd/edge.h>
+#include <torch/csrc/autograd/function.h>
+#include <torch/csrc/autograd/generated/variable_factories.h>
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/jit/api/compilation_unit.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/frontend/error_report.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/mobile/register_ops_common_utils.h>
+#include <torch/csrc/jit/runtime/custom_operator.h>
+#include <torch/csrc/jit/runtime/graph_executor.h>
+#include <torch/csrc/jit/runtime/jit_exception.h>
+#include <torch/csrc/jit/runtime/logging.h>
+#include <torch/csrc/jit/runtime/operator.h>
+#include <torch/csrc/jit/runtime/print_handler.h>
+#include <torch/csrc/jit/runtime/profiling_record.h>
+#include <torch/csrc/jit/runtime/vararg_functions.h>
+#include <torch/csrc/jit/serialization/pickle.h>
+
+#include <ATen/ExpandUtils.h>
+#include <ATen/Parallel.h>
+#include <ATen/WrapDimUtils.h>
+#include <ATen/core/Dict.h>
+#include <ATen/core/Generator.h>
+#include <ATen/core/ivalue.h>
+#include <c10/core/Device.h>
+#include <c10/core/thread_pool.h>
+#include <c10/util/SmallVector.h>
+#include <c10/util/irange.h>
+
+namespace torch::jit {
+constexpr inline c10::AliasAnalysisKind aliasAnalysisFromSchema() {
+  return c10::AliasAnalysisKind::FROM_SCHEMA;
+}
+
+constexpr inline c10::AliasAnalysisKind aliasAnalysisConservative() {
+  return c10::AliasAnalysisKind::CONSERVATIVE;
+}
+
+constexpr inline c10::AliasAnalysisKind aliasAnalysisSpecialCase() {
+  return c10::AliasAnalysisKind::INTERNAL_SPECIAL_CASE;
+}
+
+template <class T>
+c10::List<T> make_result_list(const TypePtr& elemType) {
+  return c10::List<T>();
+}
+
+template <>
+c10::impl::GenericList make_result_list<IValue>(const TypePtr& elemType);
+
+// As described in https://docs.python.org/3/library/functions.html#round
+// When a number is exactly halfway between two integers, python builtin round
+// function will round to even number. We use round(x/2)*2 to handle the
+// special halfway case. For positive 'x', round(x/2)*2 =
+// round((x_e + x_r)/2)*2 = x_e + round(x_r/2)*2, where x_e is an even integer,
+// x_r is either 0.5 of 1.5, round(x_r/2)*2 results a 0 or 2, so the final
+// result will always be a even number. Due to symmetricity, it also applies to
+// negative cases.
+inline double round_to_even(double a) {
+  return a - std::floor(a) == 0.5 ? (std::round(a * 0.5) * 2.0) : std::round(a);
+}
+
+// using the rules from python_arg_parser FunctionParameter::check
+// tensor cannot have grad set, tensor must be 0 dim,
+// and if the dest is an int the source must be integral type
+void checkImplicitTensorToNum(const at::Tensor& t, bool toInt);
+
+[[maybe_unused]] static int64_t floordiv(int64_t a, int64_t b) {
+  if (b == 0) {
+    throw std::runtime_error("division by 0");
+  }
+  if ((a > 0) == (b > 0)) {
+    // simple case, both have same sign
+    return a / b;
+  } else {
+    // in python division rounds down, it doesn't not truncate like in c++
+    auto r = lldiv(a, b);
+    return (r.rem) ? r.quot - 1 : r.quot;
+  }
+}
+TORCH_API void checkDoubleInRange(double a);
+[[maybe_unused]] static int64_t floor(double a) {
+  checkDoubleInRange(a);
+  return std::floor(a);
+}
+[[maybe_unused]] static int64_t ceil(double a) {
+  checkDoubleInRange(a);
+  return std::ceil(a);
+}
+
+[[maybe_unused]] static int64_t gcd(int64_t a, int64_t b) {
+  while (b != 0) {
+    int64_t r = a % b;
+    a = b;
+    b = r;
+  }
+  // in python gcd returns non-negative values
+  return std::abs(a);
+}
+
+int64_t partProduct(int n, int m);
+
+void loop(int n, int64_t& p, int64_t& r);
+
+int nminussumofbits(int v);
+
+int64_t factorial(int n);
+static const double degToRad = std::acos(-1.0) / 180.0;
+static const double radToDeg = 180.0 / std::acos(-1.0);
+double degrees(double x);
+double radians(double x);
+
+// Convert an python index (which may be negative) into an index usable for a
+// C++ container
+
+// Equivalent to list.at(idx)
+template <typename T>
+auto getItem(const c10::List<T>& list, int64_t idx) {
+  const int64_t list_size = list.size();
+  const int64_t normalized_idx = normalizeIndex(idx, list_size);
+  if (normalized_idx < 0 || normalized_idx >= list_size) {
+    throw std::out_of_range("list index out of range");
+  }
+  return list.get(normalized_idx);
+}
+
+template <typename T>
+void setItem(const c10::List<T>& list, int64_t idx, T&& value) {
+  const int64_t list_size = list.size();
+  const int64_t normalized_idx = normalizeIndex(idx, list_size);
+  if (normalized_idx < 0 || normalized_idx >= list_size) {
+    throw std::out_of_range("list index out of range");
+  }
+  list.set(normalized_idx, std::forward<T>(value));
+}
+
+void listAppend(Stack& stack);
+
+void listReverse(Stack& stack);
+
+template <typename T>
+void minList(Stack& stack) {
+  c10::List<T> a = pop(stack).to<c10::List<T>>();
+  c10::List<T> b = pop(stack).to<c10::List<T>>();
+
+  size_t min_size = std::min(a.size(), b.size());
+  for (const auto i : c10::irange(min_size)) {
+    if (a[i] == b[i]) {
+      continue;
+    }
+
+    push(stack, a[i] < b[i] ? a : b);
+    return;
+  }
+
+  push(stack, b.size() < a.size() ? b : a);
+}
+
+template <typename T>
+void maxList(Stack& stack) {
+  c10::List<T> a = pop(stack).to<c10::List<T>>();
+  c10::List<T> b = pop(stack).to<c10::List<T>>();
+
+  size_t min_size = std::min(a.size(), b.size());
+  for (const auto i : c10::irange(min_size)) {
+    if (a[i] == b[i]) {
+      continue;
+    }
+
+    push(stack, a[i] > b[i] ? a : b);
+    return;
+  }
+
+  push(stack, b.size() > a.size() ? b : a);
+}
+
+void listPopImpl(Stack& stack, const char* empty_message);
+
+void listPop(Stack& stack);
+
+void listClear(Stack& stack);
+
+void listDelete(Stack& stack);
+
+void listInsert(Stack& stack);
+
+template <typename T>
+void listRemove(Stack& stack) {
+  T elem = pop(stack).to<T>();
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+
+  auto pos = std::find(list.begin(), list.end(), elem);
+
+  if (pos != list.end()) {
+    list.erase(pos);
+  } else {
+    TORCH_CHECK(false, "list.remove(x): x not in list");
+  }
+}
+
+template <typename T>
+void listMin(Stack& stack) {
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+  size_t list_size = list.size();
+  if (list_size == 0) {
+    throw std::runtime_error("min() arg is an empty sequence");
+  }
+
+  T min_elem = list[0];
+  for (const auto i : c10::irange(1, list_size)) {
+    T elem = list[i];
+    min_elem = elem < min_elem ? elem : min_elem;
+  }
+
+  stack.push_back(min_elem);
+}
+
+template <typename T>
+void listMax(Stack& stack) {
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+  size_t list_size = list.size();
+  if (list_size == 0) {
+    throw std::runtime_error("max() arg is an empty sequence");
+  }
+
+  T max_elem = list[0];
+  for (const auto i : c10::irange(1, list_size)) {
+    T elem = list[i];
+    max_elem = elem > max_elem ? elem : max_elem;
+  }
+
+  stack.push_back(max_elem);
+}
+
+template <>
+void listRemove<at::Tensor>(Stack& stack);
+
+template <typename T>
+void listIndex(Stack& stack) {
+  T elem = pop(stack).to<T>();
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+
+  auto pos = std::find(list.begin(), list.end(), elem);
+
+  if (pos != list.end()) {
+    push(stack, static_cast<int64_t>(std::distance(list.begin(), pos)));
+  } else {
+    TORCH_CHECK(false, "'", elem, "' is not in list");
+  }
+}
+
+template <>
+void listIndex<at::Tensor>(Stack& stack);
+
+template <typename T>
+void listCount(Stack& stack) {
+  T elem = pop(stack).to<T>();
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+
+  const int64_t count = std::count(list.begin(), list.end(), elem);
+  push(stack, count);
+}
+
+template <>
+void listCount<at::Tensor>(Stack& stack);
+
+void listExtend(Stack& stack);
+
+void listCopy(Stack& stack);
+
+void listSelect(Stack& stack);
+
+void listLen(Stack& stack);
+
+template <typename T>
+void listEq(Stack& stack) {
+  c10::List<T> b = pop(stack).to<c10::List<T>>();
+  c10::List<T> a = pop(stack).to<c10::List<T>>();
+  push(stack, a == b);
+}
+
+template <typename T>
+void listNe(Stack& stack) {
+  c10::List<T> b = pop(stack).to<c10::List<T>>();
+  c10::List<T> a = pop(stack).to<c10::List<T>>();
+  push(stack, a != b);
+}
+
+inline bool tensor_list_equal(
+    const c10::List<at::Tensor>& a,
+    const c10::List<at::Tensor>& b) {
+  if (a.size() != b.size()) {
+    return false;
+  }
+
+  for (const auto i : c10::irange(a.size())) {
+    const at::Tensor& a_element = a[i];
+    const at::Tensor& b_element = b[i];
+    // This preserves Python's semantics, which uses eq() to compare two
+    // elements, then passes the result to bool().
+    // see: https://docs.python.org/3.4/reference/datamodel.html#object.__ge__
+    const auto cmp_result = a_element.eq(b_element);
+    if (!at::native::is_nonzero(cmp_result)) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+// Specialization for at::Tensor, since it doesn't define operator==
+template <>
+void listEq<at::Tensor>(Stack& stack);
+
+// Specialization for at::Tensor, since it doesn't define operator==
+template <>
+void listNe<at::Tensor>(Stack& stack);
+
+void listList(Stack& stack);
+
+template <typename T>
+void listContains(Stack& stack) {
+  auto key = pop(stack).to<T>();
+  auto list = pop(stack).to<c10::List<T>>();
+  // NOLINTNEXTLINE(performance-implicit-conversion-in-loop)
+  for (const T& item : list) {
+    if (item == key) {
+      push(stack, true);
+      return;
+    }
+  }
+  push(stack, false);
+}
+
+void listAdd(Stack& stack);
+
+void listInplaceAdd(Stack& stack);
+
+void listMulIntLeftInPlace(Stack& stack);
+
+void listMulIntLeft(Stack& stack);
+
+void listMulIntRight(Stack& stack);
+
+void listSlice(Stack& stack);
+
+template <typename T>
+void listSort(Stack& stack) {
+  bool reverse = pop(stack).toBool();
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+  std::sort(list.begin(), list.end(), [reverse](const T& a, const T& b) {
+    // FBCode errors without this check - "strict weak ordering"
+    // TODO: remove when possible, since it just slows down
+    // sorting and doesn't do anything useful
+    if (a == b) {
+      return false;
+    }
+    return (a < b) != reverse;
+  });
+}
+
+// Specialization for at::Tensor
+template <>
+void listSort<at::Tensor>(Stack& stack);
+
+template <typename T>
+void listCopyAndSort(Stack& stack) {
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+  auto list_copied = list.copy();
+  std::sort(list_copied.begin(), list_copied.end(), [](const T& a, const T& b) {
+    // "strict weak ordering" issue - see other sort
+    if (a == b) {
+      return false;
+    }
+    return a < b;
+  });
+  push(stack, list_copied);
+}
+
+// Specialization for at::Tensor
+template <>
+void listCopyAndSort<at::Tensor>(Stack& stack);
+
+void listSetItem(Stack& stack);
+
+struct OperatorGeneratorArgs {
+  const char* schema_str;
+  bool isOperationCreator;
+  union {
+    void (*operation)(Stack&);
+    OperationCreator operationCreator;
+  };
+  AliasAnalysisKind aliasAnalysis;
+
+  explicit constexpr OperatorGeneratorArgs(
+      torch::detail::SelectiveStr<true> schema_str,
+      void (*op)(Stack&),
+      AliasAnalysisKind aa)
+      : schema_str(schema_str),
+        isOperationCreator(false),
+        operation(op),
+        aliasAnalysis(aa) {}
+
+  explicit constexpr OperatorGeneratorArgs(
+      torch::detail::SelectiveStr<true> schema_str,
+      OperationCreator opCreator,
+      AliasAnalysisKind aa)
+      : schema_str(schema_str),
+        isOperationCreator(true),
+        operationCreator(opCreator),
+        aliasAnalysis(aa) {}
+
+  template <typename... Args>
+  explicit constexpr OperatorGeneratorArgs(
+      torch::detail::SelectiveStr<false> /*unused*/,
+      Args... /*unused*/)
+      : schema_str(nullptr),
+        isOperationCreator(false),
+        operation(nullptr),
+        aliasAnalysis(AliasAnalysisKind::INTERNAL_SPECIAL_CASE) {}
+};
+
+#define DEFINE_GENERIC_BINARY_OP(                                             \
+    aten_op, op, int_float_result, complex_result)                            \
+  OperatorGeneratorArgs(                                                      \
+      TORCH_SELECTIVE_SCHEMA(#aten_op                                         \
+                             ".int_int(int a, int b) -> " #int_float_result), \
+      [](Stack& stack) {                                                      \
+        int64_t a, b;                                                         \
+        pop(stack, a, b);                                                     \
+        push(stack, op);                                                      \
+      },                                                                      \
+      aliasAnalysisFromSchema()),                                             \
+      OperatorGeneratorArgs(                                                  \
+          TORCH_SELECTIVE_SCHEMA(                                             \
+              #aten_op                                                        \
+              ".float_float(float a, float b) -> " #int_float_result),        \
+          [](Stack& stack) {                                                  \
+            double a, b;                                                      \
+            pop(stack, a, b);                                                 \
+            push(stack, op);                                                  \
+          },                                                                  \
+          aliasAnalysisFromSchema()),                                         \
+      OperatorGeneratorArgs(                                                  \
+          TORCH_SELECTIVE_SCHEMA(                                             \
+              #aten_op                                                        \
+              ".complex_complex(complex a, complex b) -> " #complex_result),  \
+          [](Stack& stack) {                                                  \
+            c10::complex<double> a, b;                                        \
+            pop(stack, a, b);                                                 \
+            push(stack, op);                                                  \
+          },                                                                  \
+          aliasAnalysisFromSchema())
+
+// define implementations for primitive number ops
+#define DEFINE_GENERIC_OP(aten_op, int_op, float_op, int_result, float_result) \
+  OperatorGeneratorArgs(                                                       \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".int(int a, int b) -> " #int_result),   \
+      [](Stack& stack) {                                                       \
+        int64_t a, b;                                                          \
+        pop(stack, a, b);                                                      \
+        push(stack, int_op);                                                   \
+      },                                                                       \
+      aliasAnalysisFromSchema()),                                              \
+      OperatorGeneratorArgs(                                                   \
+          TORCH_SELECTIVE_SCHEMA(                                              \
+              #aten_op ".float(float a, float b) -> " #float_result),          \
+          [](Stack& stack) {                                                   \
+            double a, b;                                                       \
+            pop(stack, a, b);                                                  \
+            push(stack, float_op);                                             \
+          },                                                                   \
+          aliasAnalysisFromSchema())
+
+#define DEFINE_INT_FLOAT_OP(aten_op, op, result)                            \
+  OperatorGeneratorArgs(                                                    \
+      TORCH_SELECTIVE_SCHEMA(#aten_op                                       \
+                             ".int_float(int a, float b) -> " #result),     \
+      [](Stack& stack) {                                                    \
+        int64_t a;                                                          \
+        double b;                                                           \
+        pop(stack, a, b);                                                   \
+        push(stack, op);                                                    \
+      },                                                                    \
+      aliasAnalysisFromSchema()),                                           \
+      OperatorGeneratorArgs(                                                \
+          TORCH_SELECTIVE_SCHEMA(#aten_op                                   \
+                                 ".float_int(float a, int b) -> " #result), \
+          [](Stack& stack) {                                                \
+            double a;                                                       \
+            int64_t b;                                                      \
+            pop(stack, a, b);                                               \
+            push(stack, op);                                                \
+          },                                                                \
+          aliasAnalysisFromSchema())
+
+#define DEFINE_INT_OP(aten_op, op)                                  \
+  OperatorGeneratorArgs(                                            \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".int(int a, int b) -> int"), \
+      [](Stack& stack) {                                            \
+        int64_t a, b;                                               \
+        pop(stack, a, b);                                           \
+        push(stack, op); /* NOLINT(hicpp-signed-bitwise) */         \
+      },                                                            \
+      aliasAnalysisFromSchema())
+
+#define DEFINE_STR_CMP_OP(aten_op, op)                               \
+  OperatorGeneratorArgs(                                             \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".str(str a, str b) -> bool"), \
+      [](Stack& stack) {                                             \
+        auto b = pop(stack).toStringRef();                           \
+        auto a = pop(stack).toStringRef();                           \
+        push(stack, op);                                             \
+      },                                                             \
+      aliasAnalysisFromSchema())
+
+// define a primitive op over Scalar operands.
+// it's necessary to register this overload following
+// int/float variations to avoid trapping Scalar args
+// in unintended implicit conversions
+#define DEFINE_SCALAR_BINARY_OP_AVOID_COLLISION_GENERIC(          \
+    aten_op, int_op, float_op, result, string_val)                \
+  OperatorGeneratorArgs(                                          \
+      TORCH_SELECTIVE_SCHEMA(#aten_op string_val                  \
+                             "(Scalar a, Scalar b) -> " #result), \
+      [](Stack& stack) {                                          \
+        IValue x, y;                                              \
+        pop(stack, x, y);                                         \
+        if (x.isDouble()) {                                       \
+          if (y.isDouble()) {                                     \
+            double a = x.toDouble();                              \
+            double b = y.toDouble();                              \
+            push(stack, float_op);                                \
+          } else {                                                \
+            double a = x.toDouble();                              \
+            int64_t b = y.toInt();                                \
+            push(stack, float_op);                                \
+          }                                                       \
+        } else {                                                  \
+          if (y.isDouble()) {                                     \
+            int64_t a = x.toInt();                                \
+            double b = y.toDouble();                              \
+            push(stack, float_op);                                \
+          } else {                                                \
+            int64_t a = x.toInt();                                \
+            int64_t b = y.toInt();                                \
+            push(stack, int_op);                                  \
+          }                                                       \
+        }                                                         \
+      },                                                          \
+      aliasAnalysisFromSchema())
+
+#define DEFINE_SCALAR_BINARY_OP(aten_op, int_op, float_op, result) \
+  DEFINE_SCALAR_BINARY_OP_AVOID_COLLISION_GENERIC(                 \
+      aten_op, int_op, float_op, result, "")
+
+#define DEFINE_SCALAR_BINARY_OP_AVOID_COLLISION(   \
+    aten_op, int_op, float_op, result)             \
+  DEFINE_SCALAR_BINARY_OP_AVOID_COLLISION_GENERIC( \
+      aten_op, int_op, float_op, result, ".Scalar_Scalar")
+
+#define DEFINE_BINARY_OP(aten_op, op)             \
+  DEFINE_GENERIC_OP(aten_op, op, op, int, float), \
+      DEFINE_INT_FLOAT_OP(aten_op, op, float),    \
+      DEFINE_SCALAR_BINARY_OP(aten_op, op, op, Scalar)
+
+#define DEFINE_BINARY_FLOAT_OP(aten_op, op)         \
+  DEFINE_GENERIC_OP(aten_op, op, op, float, float), \
+      DEFINE_INT_FLOAT_OP(aten_op, op, float),      \
+      DEFINE_SCALAR_BINARY_OP(aten_op, op, op, float)
+
+#define DEFINE_COMPARISON_OP(aten_op, op)             \
+  DEFINE_GENERIC_OP(aten_op, op, op, bool, bool),     \
+      DEFINE_INT_FLOAT_OP(aten_op, op, bool),         \
+      DEFINE_SCALAR_BINARY_OP(aten_op, op, op, bool), \
+      DEFINE_STR_CMP_OP(aten_op, op)
+
+#define DEFINE_UNARY_INT_OP(aten_op, op, result)                  \
+  OperatorGeneratorArgs(                                          \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".int(int a) -> " #result), \
+      [](Stack& stack) {                                          \
+        int64_t a;                                                \
+        pop(stack, a);                                            \
+        push(stack, op);                                          \
+      },                                                          \
+      aliasAnalysisFromSchema())
+
+#define DEFINE_UNARY_FLOAT_OP(aten_op, op, result)                    \
+  OperatorGeneratorArgs(                                              \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".float(float a) -> " #result), \
+      [](Stack& stack) {                                              \
+        double a;                                                     \
+        pop(stack, a);                                                \
+        push(stack, op);                                              \
+      },                                                              \
+      aliasAnalysisFromSchema())
+
+#define DEFINE_UNARY_OP(aten_op, op, int_result, float_result)            \
+  DEFINE_UNARY_INT_OP(aten_op, op, int_result),                           \
+      DEFINE_UNARY_FLOAT_OP(aten_op, op, float_result),                   \
+      OperatorGeneratorArgs(                                              \
+          TORCH_SELECTIVE_SCHEMA(#aten_op ".Scalar(Scalar a) -> Scalar"), \
+          [](Stack& stack) {                                              \
+            IValue x;                                                     \
+            pop(stack, x);                                                \
+            if (x.isDouble()) {                                           \
+              double a = x.toDouble();                                    \
+              push(stack, static_cast<float_result>(op));                 \
+            } else {                                                      \
+              int64_t a = x.toInt();                                      \
+              push(stack, static_cast<int_result>(op));                   \
+            }                                                             \
+          },                                                              \
+          aliasAnalysisFromSchema())
+#define DEFINE_BOOL_OP(aten_op, op)                                     \
+  OperatorGeneratorArgs(                                                \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".bool(bool a, bool b) -> bool"), \
+      [](Stack& stack) {                                                \
+        bool a, b;                                                      \
+        pop(stack, a, b);                                               \
+        push(stack, op);                                                \
+      },                                                                \
+      aliasAnalysisFromSchema())
+#define DEFINE_STRING_OP(op_name, string_op, result)                    \
+  OperatorGeneratorArgs(                                                \
+      TORCH_SELECTIVE_SCHEMA(#op_name ".str(str a, str b) ->" #result), \
+      [](Stack& stack) {                                                \
+        auto b = pop(stack).toStringRef();                              \
+        auto a = pop(stack).toStringRef();                              \
+        push(stack, string_op);                                         \
+      },                                                                \
+      aliasAnalysisFromSchema())
+
+//-----------------------------------------------------------------------------
+//-----------------------------------------------------------------------------
+//-----------------------------------------------------------------------------
+//-----------------------------------------------------------------------------
+#define DEFINE_UNARY_COMPLEX_OP(aten_op, op, result)                      \
+  OperatorGeneratorArgs(                                                  \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".complex(complex a) -> " #result), \
+      [](Stack& stack) {                                                  \
+        c10::complex<double> a;                                           \
+        pop(stack, a);                                                    \
+        push(stack, op);                                                  \
+      },                                                                  \
+      aliasAnalysisFromSchema())
+
+// Some complex unary ops (like abs, angle) return real valued output, but most
+// other unary ops return complex valued output. So, this macro is used in the
+// former case where we can explicitly pass complex_result_cast argument, which
+// is set to c10::complex<float> in the macro `DEFINE_UNARY_OP_WITH_COMPLEX`
+// defined below.
+#define DEFINE_UNARY_OP_WITH_COMPLEX_CAST(                                \
+    aten_op,                                                              \
+    op,                                                                   \
+    int_result,                                                           \
+    float_result,                                                         \
+    complex_result,                                                       \
+    complex_result_cast)                                                  \
+  DEFINE_UNARY_INT_OP(aten_op, op, int_result),                           \
+      DEFINE_UNARY_FLOAT_OP(aten_op, op, float_result),                   \
+      DEFINE_UNARY_COMPLEX_OP(aten_op, op, complex_result),               \
+      OperatorGeneratorArgs(                                              \
+          TORCH_SELECTIVE_SCHEMA(#aten_op ".Scalar(Scalar a) -> Scalar"), \
+          [](Stack& stack) {                                              \
+            IValue x;                                                     \
+            pop(stack, x);                                                \
+            if (x.isDouble()) {                                           \
+              double a = x.toDouble();                                    \
+              push(stack, static_cast<float_result>(op));                 \
+            } else if (x.isComplexDouble()) {                             \
+              c10::complex<double> a = x.toComplexDouble();               \
+              push(stack, static_cast<complex_result_cast>(op));          \
+            } else {                                                      \
+              int64_t a = x.toInt();                                      \
+              push(stack, static_cast<int_result>(op));                   \
+            }                                                             \
+          },                                                              \
+          aliasAnalysisFromSchema())
+
+#define DEFINE_UNARY_OP_WITH_COMPLEX(aten_op, op, int_result, float_result) \
+  DEFINE_UNARY_OP_WITH_COMPLEX_CAST(                                        \
+      aten_op, op, int_result, float_result, complex, c10::complex<double>)
+
+#define DEFINE_GENERIC_OP_WITH_COMPLEX(                                       \
+    aten_op,                                                                  \
+    int_op,                                                                   \
+    float_op,                                                                 \
+    complex_op,                                                               \
+    int_result,                                                               \
+    float_result,                                                             \
+    complex_result)                                                           \
+  OperatorGeneratorArgs(                                                      \
+      TORCH_SELECTIVE_SCHEMA(#aten_op ".int(int a, int b) -> " #int_result),  \
+      [](Stack& stack) {                                                      \
+        int64_t a, b;                                                         \
+        pop(stack, a, b);                                                     \
+        push(stack, int_op);                                                  \
+      },                                                                      \
+      aliasAnalysisFromSchema()),                                             \
+      OperatorGeneratorArgs(                                                  \
+          TORCH_SELECTIVE_SCHEMA(                                             \
+              #aten_op ".complex(complex a, complex b) -> " #complex_result), \
+          [](Stack& stack) {                                                  \
+            c10::complex<double> a, b;                                        \
+            pop(stack, a, b);                                                 \
+            push(stack, complex_op);                                          \
+          },                                                                  \
+          aliasAnalysisFromSchema()),                                         \
+      OperatorGeneratorArgs(                                                  \
+          TORCH_SELECTIVE_SCHEMA(                                             \
+              #aten_op ".float(float a, float b) -> " #float_result),         \
+          [](Stack& stack) {                                                  \
+            double a, b;                                                      \
+            pop(stack, a, b);                                                 \
+            push(stack, float_op);                                            \
+          },                                                                  \
+          aliasAnalysisFromSchema())
+
+#define DEFINE_INT_COMPLEX_OP(aten_op, op, result)                          \
+  OperatorGeneratorArgs(                                                    \
+      TORCH_SELECTIVE_SCHEMA(#aten_op                                       \
+                             ".int_complex(int a, complex b) -> " #result), \
+      [](Stack& stack) {                                                    \
+        int64_t a;                                                          \
+        c10::complex<double> b;                                             \
+        pop(stack, a, b);                                                   \
+        push(stack, op);                                                    \
+      },                                                                    \
+      aliasAnalysisFromSchema()),                                           \
+      OperatorGeneratorArgs(                                                \
+          TORCH_SELECTIVE_SCHEMA(                                           \
+              #aten_op ".complex_int(complex a, int b) -> " #result),       \
+          [](Stack& stack) {                                                \
+            c10::complex<double> a;                                         \
+            int64_t b;                                                      \
+            pop(stack, a, b);                                               \
+            push(stack, op);                                                \
+          },                                                                \
+          aliasAnalysisFromSchema())
+
+#define DEFINE_FLOAT_COMPLEX_OP(aten_op, op, result)                      \
+  OperatorGeneratorArgs(                                                  \
+      TORCH_SELECTIVE_SCHEMA(                                             \
+          #aten_op ".float_complex(float a, complex b) -> " #result),     \
+      [](Stack& stack) {                                                  \
+        double a;                                                         \
+        c10::complex<double> b;                                           \
+        pop(stack, a, b);                                                 \
+        push(stack, op);                                                  \
+      },                                                                  \
+      aliasAnalysisFromSchema()),                                         \
+      OperatorGeneratorArgs(                                              \
+          TORCH_SELECTIVE_SCHEMA(                                         \
+              #aten_op ".complex_float(complex a, float b) -> " #result), \
+          [](Stack& stack) {                                              \
+            c10::complex<double> a;                                       \
+            double b;                                                     \
+            pop(stack, a, b);                                             \
+            push(stack, op);                                              \
+          },                                                              \
+          aliasAnalysisFromSchema())
+
+#define DEFINE_SCALAR_BINARY_OP_WITH_COMPLEX_AVOID_COLLISION_GENERIC( \
+    aten_op, int_op, float_op, complex_op, result, string_val)        \
+  OperatorGeneratorArgs(                                              \
+      TORCH_SELECTIVE_SCHEMA(#aten_op string_val                      \
+                             "(Scalar a, Scalar b) -> " #result),     \
+      [](Stack& stack) {                                              \
+        IValue x, y;                                                  \
+        pop(stack, x, y);                                             \
+        if (x.isComplexDouble()) {                                    \
+          c10::complex<double> a = x.toComplexDouble();               \
+          if (y.isComplexDouble()) {                                  \
+            c10::complex<double> b = y.toComplexDouble();             \
+            push(stack, complex_op);                                  \
+          } else if (y.isDouble()) {                                  \
+            double b = y.toDouble();                                  \
+            push(stack, complex_op);                                  \
+          } else {                                                    \
+            int64_t b = y.toInt();                                    \
+            push(stack, complex_op);                                  \
+          }                                                           \
+        } else if (x.isDouble()) {                                    \
+          double a = x.toDouble();                                    \
+          if (y.isComplexDouble()) {                                  \
+            c10::complex<double> b = y.toComplexDouble();             \
+            push(stack, complex_op);                                  \
+          } else if (y.isDouble()) {                                  \
+            double b = y.toDouble();                                  \
+            push(stack, float_op);                                    \
+          } else {                                                    \
+            int64_t b = y.toInt();                                    \
+            push(stack, float_op);                                    \
+          }                                                           \
+        } else {                                                      \
+          int64_t a = x.toInt();                                      \
+          if (y.isComplexDouble()) {                                  \
+            c10::complex<double> b = y.toComplexDouble();             \
+            push(stack, complex_op);                                  \
+          } else if (y.isDouble()) {                                  \
+            double b = y.toDouble();                                  \
+            push(stack, float_op);                                    \
+          } else {                                                    \
+            int64_t b = y.toInt();                                    \
+            push(stack, int_op);                                      \
+          }                                                           \
+        }                                                             \
+      },                                                              \
+      aliasAnalysisFromSchema())
+
+#define DEFINE_SCALAR_BINARY_OP_WITH_COMPLEX_WITHOUT_INT_COMPLEX_PAIR(     \
+    aten_op, int_op, float_op, complex_op, result)                         \
+  OperatorGeneratorArgs(                                                   \
+      TORCH_SELECTIVE_SCHEMA(#aten_op "(Scalar a, Scalar b) -> " #result), \
+      [](Stack& stack) {                                                   \
+        IValue x, y;                                                       \
+        pop(stack, x, y);                                                  \
+        if (x.isComplexDouble()) {                                         \
+          c10::complex<double> a = x.toComplexDouble();                    \
+          if (y.isComplexDouble()) {                                       \
+            c10::complex<double> b = y.toComplexDouble();                  \
+            push(stack, complex_op);                                       \
+          } else if (y.isDouble()) {                                       \
+            double b = y.toDouble();                                       \
+            push(stack, complex_op);                                       \
+          }                                                                \
+        } else if (x.isDouble()) {                                         \
+          double a = x.toDouble();                                         \
+          if (y.isComplexDouble()) {                                       \
+            c10::complex<double> b = y.toComplexDouble();                  \
+            push(stack, complex_op);                                       \
+          } else if (y.isDouble()) {                                       \
+            double b = y.toDouble();                                       \
+            push(stack, float_op);                                         \
+          } else {                                                         \
+            int64_t b = y.toInt();                                         \
+            push(stack, float_op);                                         \
+          }                                                                \
+        } else {                                                           \
+          int64_t a = x.toInt();                                           \
+          if (y.isDouble()) {                                              \
+            double b = y.toDouble();                                       \
+            push(stack, float_op);                                         \
+          } else if (y.isInt()) {                                          \
+            int64_t b = y.toInt();                                         \
+            push(stack, int_op);                                           \
+          }                                                                \
+        }                                                                  \
+      },                                                                   \
+      aliasAnalysisFromSchema())
+
+#define DEFINE_SCALAR_BINARY_OP_WITH_COMPLEX(                   \
+    aten_op, int_op, float_op, complex_op, result)              \
+  DEFINE_SCALAR_BINARY_OP_WITH_COMPLEX_AVOID_COLLISION_GENERIC( \
+      aten_op, int_op, float_op, complex_op, result, "")
+
+#define DEFINE_BINARY_OP_WITH_COMPLEX(aten_op, op)                          \
+  DEFINE_GENERIC_OP_WITH_COMPLEX(aten_op, op, op, op, int, float, complex), \
+      DEFINE_INT_COMPLEX_OP(aten_op, op, complex),                          \
+      DEFINE_FLOAT_COMPLEX_OP(aten_op, op, complex),                        \
+      DEFINE_INT_FLOAT_OP(aten_op, op, float),                              \
+      DEFINE_SCALAR_BINARY_OP_WITH_COMPLEX(aten_op, op, op, op, Scalar)
+
+#define DEFINE_COMPARISON_OP_WITH_COMPLEX(aten_op, op)                   \
+  DEFINE_GENERIC_OP_WITH_COMPLEX(aten_op, op, op, op, bool, bool, bool), \
+      DEFINE_INT_FLOAT_OP(aten_op, op, bool),                            \
+      DEFINE_FLOAT_COMPLEX_OP(aten_op, op, bool),                        \
+      DEFINE_SCALAR_BINARY_OP_WITH_COMPLEX_WITHOUT_INT_COMPLEX_PAIR(     \
+          aten_op, op, op, op, bool),                                    \
+      DEFINE_STR_CMP_OP(aten_op, op)
+
+TORCH_API at::Generator make_generator_for_device(
+    c10::Device device,
+    std::optional<int64_t> seed = std::nullopt);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/script_profile.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/script_profile.h
new file mode 100644
index 0000000000000000000000000000000000000000..04db7ab64ca4c0a2849ca29da6c4b7ad823d4eef
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/script_profile.h
@@ -0,0 +1,108 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <chrono>
+#include <map>
+#include <string>
+
+#include <ATen/core/ivalue.h>
+#include <c10/macros/Macros.h>
+#include <torch/csrc/jit/frontend/source_ref.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+namespace profiling {
+
+struct Datapoint {
+  using Timepoint = std::chrono::time_point<std::chrono::steady_clock>;
+  SourceRange sourceRange;
+  Timepoint start;
+  Timepoint end;
+
+  explicit Datapoint(SourceRange sr)
+      : sourceRange(std::move(sr)), start(std::chrono::steady_clock::now()) {}
+};
+
+class TORCH_API InstructionSpan {
+ public:
+  explicit InstructionSpan(Node& /*node*/);
+  ~InstructionSpan();
+  InstructionSpan(InstructionSpan&&) = delete;
+  InstructionSpan& operator=(InstructionSpan&&) = delete;
+
+ private:
+  std::unique_ptr<Datapoint> datapoint_;
+};
+
+bool TORCH_API isProfilingOngoing();
+
+} // namespace profiling
+
+struct TORCH_API InstructionStats : public CustomClassHolder {
+  int64_t count{0};
+  std::chrono::nanoseconds duration{0};
+};
+
+class TORCH_API SourceStats : public CustomClassHolder {
+ public:
+  using LineMap = c10::Dict<int64_t, c10::intrusive_ptr<InstructionStats>>;
+
+  SourceStats(SourceRef source, const LineMap& lineMap)
+      : source_(std::move(source)), lineMap_(lineMap) {}
+
+  const SourceRef& getSourceRef() const {
+    return source_;
+  }
+
+  const LineMap& getLineMap() const {
+    return lineMap_;
+  }
+
+ private:
+  SourceRef source_;
+  LineMap lineMap_;
+};
+
+/**
+ * ScriptProfile is an underlying C++ implementation for TorchScript profiling.
+ * The profiling section is specified by calling enable() and disable():
+ *
+ * ...
+ * scriptProfile.enable();
+ * ...
+ * (scripts)
+ * ...
+ * scriptProfile.disable();
+ * ...
+ *
+ * NOTE: you cannot attach the profiler while the script is running.
+ *
+ * To retrieve collected runtime data, users may call dumpStats() and do
+ * arbitrary filtering on the data they want. Note that dumpStats() should
+ * not be called inside a profiling section.
+ * In general, stats are aggregated per source function body, and then by line
+ * number.
+ */
+class TORCH_API ScriptProfile : public CustomClassHolder {
+  // Aggregates datapoints by function source id, then by line number.
+  using LineMap = std::map<int64_t, InstructionStats>;
+  using SourceMap = std::map<SourceRef, LineMap, std::less<>>;
+
+ public:
+  void enable();
+  void disable();
+  const SourceMap& dumpStats();
+  void addDatapoint(std::shared_ptr<profiling::Datapoint> /*datapoint*/);
+  ~ScriptProfile() override;
+
+ private:
+  bool enabled_{false};
+  std::vector<std::shared_ptr<profiling::Datapoint>> datapoints_;
+  SourceMap sourceMap_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/serialized_shape_function_registry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/serialized_shape_function_registry.h
new file mode 100644
index 0000000000000000000000000000000000000000..d690b7e1a4eedb8ddd14ef9f52780ea6e8940602
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/serialized_shape_function_registry.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API const std::string& GetSerializedShapeFunctions();
+
+TORCH_API const OperatorMap<std::string>& GetShapeFunctionMappings();
+
+TORCH_API const OperatorMap<std::pair<std::string, std::string>>&
+GetBoundedShapeMappings();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/shape_function_registry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/shape_function_registry.h
new file mode 100644
index 0000000000000000000000000000000000000000..c2f8ed5907bea4a2a241b9187236203ae04965f9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/shape_function_registry.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API const std::string& GetSerializedFuncs();
+
+TORCH_API const OperatorMap<std::string>& GetFuncMapping();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/simple_graph_executor_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/simple_graph_executor_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..59f5a37db6630df9811f9a95f8eb54d70a015e9c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/simple_graph_executor_impl.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Flags.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/runtime/graph_executor_impl.h>
+
+namespace torch::jit {
+
+struct TORCH_API SimpleGraphExecutorImpl : public GraphExecutorImplBase {
+  SimpleGraphExecutorImpl(
+      const std::shared_ptr<Graph>& graph,
+      std::string function_name);
+
+  const ExecutionPlan& getPlanFor(
+      Stack& stack,
+      std::optional<size_t> remaining_bailout_depth) override;
+  const ExecutionPlan& getInputIndependentPlan() override;
+  GraphExecutorState getDebugState() override;
+  ~SimpleGraphExecutorImpl() override = default;
+
+ private:
+  const ExecutionPlan& getInputIndependentPlanImpl();
+
+  std::optional<ExecutionPlan> execution_plan_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/slice_indices_adjust.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/slice_indices_adjust.h
new file mode 100644
index 0000000000000000000000000000000000000000..5c386e8728c02f2a83fd9cd97efc78663f989463
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/slice_indices_adjust.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <cstddef>
+#include <cstdint>
+
+namespace torch::jit {
+
+// Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010,
+// 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020 Python Software
+// Foundation; All Rights Reserved
+//
+// Stolen (with appropriate modifications) by @agolynski
+// (https://github.com/pytorch/pytorch/pull/33019) from cpython repo
+// Objects/sliceobject.c with comment: this is harder to get right than you
+// might think
+//
+// This adjusts indexes according to python list semantics and returns number
+// of elements in the resulting list.
+TORCH_API int64_t slice_indices_adjust(
+    int64_t length,
+    int64_t* start,
+    int64_t* stop,
+    int64_t step);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/ProcessedNodeInputs.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/ProcessedNodeInputs.h
new file mode 100644
index 0000000000000000000000000000000000000000..17742bab4f2e11bd6ddd909c3406a9c5a759bbf0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/ProcessedNodeInputs.h
@@ -0,0 +1,246 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+
+#include <memory>
+
+#include <c10/macros/Macros.h>
+#include <c10/util/Logging.h>
+
+/**
+ * Packed representation of input indices for ProcessedNode.
+ */
+class ProcessedNodeInputs {
+ private:
+  // This keeps the size usage for inputs + outputs down to 16 bytes;
+  // we use 12 bytes, and then two 2-byte integers are used to store
+  // the outputs.
+  static constexpr size_t kMaxInlineInputs = 5;
+
+ public:
+  ProcessedNodeInputs() : ProcessedNodeInputs(0) {}
+
+  explicit ProcessedNodeInputs(size_t size) {
+    TORCH_DCHECK_LT(size, (1 << 16));
+    if (size <= kMaxInlineInputs) {
+      repr_.inline_repr_.size = size;
+    } else {
+      new (&repr_.outline_repr_) HeapArrayPtr(size);
+    }
+  }
+
+  uint16_t operator[](uint16_t idx) const {
+    // NOLINTNEXTLINE(*const-cast*)
+    return (*const_cast<ProcessedNodeInputs*>(this))[idx];
+  }
+
+  uint16_t& operator[](uint16_t idx) {
+    if (C10_LIKELY(repr_.is_inline())) {
+      TORCH_DCHECK_LT(idx, repr_.inline_repr_.size);
+      return repr_.inline_repr_.inputs[idx];
+    } else {
+      return repr_.outline_repr_[idx];
+    }
+  }
+
+  [[nodiscard]] uint16_t size() const {
+    if (C10_LIKELY(repr_.is_inline())) {
+      return repr_.inline_repr_.size;
+    } else {
+      return repr_.outline_repr_.size();
+    }
+  }
+
+  [[nodiscard]] bool empty() const {
+    return size() == 0;
+  }
+
+ private:
+  class HeapArrayPtr {
+   public:
+    HeapArrayPtr() = default;
+    ~HeapArrayPtr() = default;
+
+    explicit HeapArrayPtr(uint16_t size) : array_(alloc(size)) {}
+
+    HeapArrayPtr(const HeapArrayPtr& rhs) : array_(alloc(rhs.size())) {
+      if (rhs.array_) {
+        std::memcpy(
+            array_.get(),
+            rhs.array_.get(),
+            (rhs.size() + 1) * sizeof(uint16_t));
+      }
+    }
+
+    HeapArrayPtr& operator=(const HeapArrayPtr& rhs) {
+      if (&rhs == this) {
+        return *this;
+      }
+
+      if (size() != rhs.size()) {
+        array_ = alloc(rhs.size());
+      }
+
+      if (rhs.array_) {
+        std::memcpy(
+            array_.get(),
+            rhs.array_.get(),
+            (rhs.size() + 1) * sizeof(uint16_t));
+      }
+      return *this;
+    }
+
+    HeapArrayPtr(HeapArrayPtr&&) noexcept = default;
+    HeapArrayPtr& operator=(HeapArrayPtr&&) noexcept = default;
+
+    [[nodiscard]] bool empty() const {
+      return size() != 0;
+    }
+
+    [[nodiscard]] uint16_t size() const {
+      return array_ ? array_[0] : 0;
+    }
+
+    uint16_t operator[](uint16_t idx) const {
+      TORCH_DCHECK_LT(idx, size());
+      return array_[idx + 1];
+    }
+
+    uint16_t& operator[](uint16_t idx) {
+      TORCH_DCHECK_LT(idx, size());
+      return array_[idx + 1];
+    }
+
+   private:
+    // NOLINTNEXTLINE(modernize-avoid-c-arrays)
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays)
+    std::unique_ptr<uint16_t[]> array_;
+
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays)
+    // NOLINTNEXTLINE(modernize-avoid-c-arrays)
+    static std::unique_ptr<uint16_t[]> alloc(uint16_t num_elts) {
+      if (num_elts) {
+        auto result = std::make_unique<uint16_t[]>(num_elts + 1);
+        result[0] = num_elts;
+        return result;
+      } else {
+        return nullptr;
+      }
+    }
+  };
+
+  // We want ProcessedNode to be able to pack two more `uint16_t`
+  // fields after its ProcessedNodeInputs, and we'll end up being
+  // aligned to an 8-byte boundary anyway. We could avoid this pragma
+  // at the cost of having to move ProcessedNode::outputs_offset_ and
+  // ProcessedNode::num_outputs_ into this class, which would be
+  // awkward.
+#pragma pack(push, 2)
+  union Repr {
+    [[nodiscard]] bool is_inline() const {
+      uint8_t tag = 0;
+      // Use of reinterpret_cast to pointer to char or unsigned char
+      // is defined behavior; see
+      // https://en.cppreference.com/w/cpp/language/reinterpret_cast .
+      std::memcpy(&tag, reinterpret_cast<const uint8_t*>(this), 1);
+      // HeapArrayPtr will be represented as a plain old pointer,
+      // which will have alignment to at least a 2-byte boundary
+      // (because it's uint16_t*) and more likely an 8- or 16-byte
+      // boundary because malloc will tend to just align everything to
+      // one of those. So, we just set tag to 1 when inline_repr_ is
+      // active so as to be able to differentiate the two.
+      return (tag & 1) != 0;
+    }
+
+    // NOLINTNEXTLINE(modernize-use-equals-default)
+    Repr() {}
+
+    ~Repr() {
+      destroyIfOutline();
+    }
+
+    Repr(const Repr& rhs) {
+      if (rhs.is_inline()) {
+        std::memcpy(&inline_repr_, &rhs.inline_repr_, sizeof(inline_repr_));
+      } else {
+        new (&outline_repr_) OutlineRepr(rhs.outline_repr_);
+      }
+    }
+
+    Repr& operator=(const Repr& rhs) {
+      if (&rhs == this) {
+        return *this;
+      }
+      if (rhs.is_inline()) {
+        destroyIfOutline();
+        new (&inline_repr_) InlineRepr();
+        std::memcpy(&inline_repr_, &rhs.inline_repr_, sizeof(inline_repr_));
+      } else {
+        if (is_inline()) {
+          new (&outline_repr_) OutlineRepr(rhs.outline_repr_);
+        } else {
+          outline_repr_ = rhs.outline_repr_;
+        }
+      }
+      return *this;
+    }
+
+    Repr(Repr&& rhs) noexcept {
+      if (rhs.is_inline()) {
+        std::memcpy(&inline_repr_, &rhs.inline_repr_, sizeof(inline_repr_));
+      } else {
+        new (&outline_repr_) OutlineRepr(std::move(rhs.outline_repr_));
+      }
+    }
+
+    Repr& operator=(Repr&& rhs) noexcept {
+      if (&rhs == this) {
+        return *this;
+      }
+
+      if (rhs.is_inline()) {
+        destroyIfOutline();
+        new (&inline_repr_) InlineRepr();
+        std::memcpy(&inline_repr_, &rhs.inline_repr_, sizeof(inline_repr_));
+      } else {
+        if (is_inline()) {
+          new (&outline_repr_) OutlineRepr(std::move(rhs.outline_repr_));
+        } else {
+          outline_repr_ = std::move(rhs.outline_repr_);
+        }
+      }
+
+      return *this;
+    }
+
+    struct InlineRepr {
+      uint8_t tag = 0x1;
+      uint8_t size{};
+      uint16_t inputs[kMaxInlineInputs]{};
+    };
+
+    using OutlineRepr = HeapArrayPtr;
+
+    InlineRepr inline_repr_{};
+    OutlineRepr outline_repr_;
+
+   private:
+    void destroyIfOutline() {
+      if (!is_inline()) {
+        outline_repr_.~OutlineRepr();
+      }
+    }
+  } repr_;
+#pragma pack(pop)
+};
+
+static_assert(
+    sizeof(ProcessedNodeInputs) == 12,
+    "ProcessedNodeInputs has the wrong size!");
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/fusion.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/fusion.h
new file mode 100644
index 0000000000000000000000000000000000000000..3009ca031cf304074424eb64b52e07fa0937978c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/fusion.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void fuseStaticSubgraphs(
+    std::shared_ptr<Graph> graph,
+    size_t min_size);
+
+TORCH_API void performTensorExprFusion(
+    std::shared_ptr<Graph> graph,
+    std::vector<IValue> sample_inputs);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..2cab635be12a85c2f4501985f3730448ec1d09ad
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/impl.h
@@ -0,0 +1,1152 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/ivalue.h>
+#include <ATen/core/symbol.h>
+#include <c10/core/CPUAllocator.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/FbcodeMaps.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/graph_node_list.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/constant_propagation.h>
+#include <torch/csrc/jit/passes/freeze_module.h>
+#include <torch/csrc/jit/passes/inliner.h>
+#include <torch/csrc/jit/runtime/static/ProcessedNodeInputs.h>
+#include <torch/custom_class.h>
+#include <limits>
+
+#ifdef FBCODE_CAFFE2
+#include <folly/container/F14Map.h>
+#include <folly/container/F14Set.h>
+#endif
+
+namespace torch::jit {
+
+TORCH_API bool canEnableStaticRuntime(
+    const std::shared_ptr<torch::jit::Graph>& graph);
+
+TORCH_API std::string dumpValueSet(
+    const c10::FastSet<const Value*>& value_set,
+    const char* set_name = "");
+
+inline bool doesNotHeapAllocateWhenStoredInIValue(const Type& type) {
+  switch (type.kind()) {
+    // NOTE: NumberType may allocate because it includes complex.
+    case TypeKind::NoneType:
+    case TypeKind::IntType:
+    case TypeKind::FloatType:
+    case TypeKind::BoolType:
+    case TypeKind::DeviceObjType:
+    case TypeKind::StreamObjType:
+      return true;
+    default:
+      return false;
+  }
+}
+
+inline c10::Symbol getStaticRuntimeMetadataSymbol() {
+  return Symbol::attr("static_runtime::metadata");
+}
+
+inline bool borrowsOutputs(c10::Symbol kind) {
+  static const std::array<c10::Symbol, 4> symbols_with_borrowed_outputs = {
+      c10::Symbol::fromQualString("static_runtime::select_tensor"),
+      c10::Symbol::fromQualString("static_runtime::dict_unpack"),
+      c10::Symbol::fromQualString("static_runtime::VarTupleUnpack"),
+      c10::Symbol::fromQualString("prim::IfThenElse"),
+  };
+  return std::find(
+             symbols_with_borrowed_outputs.begin(),
+             symbols_with_borrowed_outputs.end(),
+             kind) != symbols_with_borrowed_outputs.end();
+}
+
+// Group values used by `graph` into three categories:
+//
+// - output_aliases:
+//     values that are either outputs or contain aliases of outputs
+// - external_aliases:
+//     values that are inputs, constants, or their aliases.
+//     The output aliases that end up here are as a result of aliasDb failing to
+//     recognize them as outputs due to collection object (e.g., Tuple) aliasing
+//     inputs.
+// Values that don't show up in output_aliases or external_aliases are created
+// and consumed within the graph.
+class ValueGroup {
+ public:
+  explicit ValueGroup() = default;
+  void init(const Block& block, const AliasDb& db);
+
+  bool isExternalAlias(const Value* value) const {
+    return external_aliases_.find(value) != external_aliases_.end();
+  }
+
+  bool isOutputAlias(const Value* value) const {
+    return output_aliases_.find(value) != output_aliases_.end();
+  }
+
+  bool isAlwaysAlive(const Value* value) const {
+    return isExternalAlias(value) || isOutputAlias(value);
+  }
+
+  std::string toString() const {
+    return c10::str(
+        dumpValueSet(output_aliases_, "ValueGroup::output_aliases_"),
+        "\n",
+        dumpValueSet(external_aliases_, "ValueGroup::external_aliases_"));
+  }
+
+ private:
+  c10::FastSet<const Value*> output_aliases_;
+  c10::FastSet<const Value*> external_aliases_;
+};
+
+class TORCH_API ManagedTensorRanges {
+ public:
+  ManagedTensorRanges() = default;
+  ManagedTensorRanges(
+      Block& block,
+      const AliasDb& alias_db,
+      const c10::FastSet<const Value*>& managed_tensor_values);
+
+  // If true, then this node is the last use of at least one
+  // managed tensor. availableTensorValuesAfterNode(node) will return a vector
+  // of the managed tensors that are available for reuse
+  // in the nodes following this one.
+  bool nodeFreesManagedTensors(Node* node) const;
+  const std::vector<const Value*>& availableTensorValuesAfterNode(
+      Node* node) const;
+
+  // For testing. True if v1 and v2 are both mutable types and have lifetimes
+  // that overlap.
+  bool lifetimesOverlap(const Value* v1, const Value* v2) const;
+
+ private:
+  struct Lifetime {
+    Lifetime(size_t start_, size_t end_) : start(start_), end(end_) {}
+    size_t start;
+    size_t end;
+  };
+
+  // Returns nullptr if we are not tracking the lifetime of value
+  Lifetime* getLifetime(const Value* value);
+  const Lifetime* getLifetime(const Value* value) const;
+  // Collect all values in the input that have tracked lifetimes.
+  // A value's lifetime may not be tracked if it is a graph input
+  // or immutable type (containers with at least one mutable
+  // type are mutable)
+  std::vector<const Value*> collectValuesWithTrackedLifetimes(
+      at::ArrayRef<const Value*> values);
+  void extendLifetime(Value* input, size_t new_end);
+  void extendInputLifetime(Node* node, size_t new_end);
+
+  // Maps Node* to the set of managed tensors that are now available
+  // for reuse after this node.
+  c10::FastMap<Node*, std::vector<const Value*>> node_to_newly_free_tensors_;
+  // Maps each Value* to its lifetime (start node index, end node index)
+  c10::FastMap<const Value*, Lifetime> value_lifetimes_;
+};
+
+struct TORCH_API StaticModuleOptions {
+  // enabling out variant allows Static Runtime to do memory planning
+  bool enable_out_variant{true};
+  // to reuse tensor storage for tensors whose live-range do not overlap to
+  // reduce memory footprint (enable_out_variant must be true)
+  bool optimize_memory{true};
+  // to batch allocate tensor storage for output tensors of the
+  // graph, where storage is deallocated outside static runtime
+  // (enable_out_variant must be true)
+  bool manage_output_tensors{false};
+  // Gates the ReplaceWithCopy pass, which replaces ops that
+  // sometimes alias their outputs with out variants that
+  // always copy (so the output may participate in memory planning).
+  // Since replacing with copies is done after TensorExpr fusion, the
+  // resulting graph does not conform to the assumptions made in the fuser.
+  // So, even if this flag is turned on, the ReplaceWithCopy pass will not
+  // be executed if TensorExpr fusion is enabled.
+  bool use_copy_variants{true};
+  // Gates the ReplaceWithMaybeCopy pass, which replaces ops that
+  // sometimes alias their outputs with subgraphs that include an out
+  // variant.
+  // For the same reason as `use_copy_variants`, the ReplaceWithMaybeCopy pass
+  // will not be executed if TensorExpr fusion is enabled, even if this flag
+  // is turned on.
+  bool use_maybe_copy_variants{true};
+  // enable TensorExpr fusion of ops at model loading time
+  bool enable_tensorexpr_fusion{false};
+};
+
+/*
+  Responsible for plugging StaticRuntime metadata onto the
+  IR nodes. StaticRuntimeMetdata extends CustomClassHolder
+  which can be casted to IValue and attached to IR node.
+  This is needed to pass parent graph metadata to forked
+  graph in presence of prim::fork operator
+*/
+class TORCH_API StaticRuntimeMetadata : public torch::CustomClassHolder {
+ public:
+  explicit StaticRuntimeMetadata(const StaticModuleOptions& opts)
+      : opts_(opts) {}
+
+  const StaticModuleOptions& get_opts() {
+    return opts_;
+  }
+
+ private:
+  StaticModuleOptions opts_;
+};
+
+/// The static runime supports two execution modes.
+///
+/// Mode 1: single-threaded with no parallelism except for intra-op parallelism
+/// For this mode, you can do either:
+/// @code
+///   // m is a TorchScript module
+///   auto module = StaticModule(m, opts);
+///   auto output = module(args, kwargs);
+/// @endcode
+///
+/// or
+///
+/// @code
+///   // g is the TorchScript graph
+///   auto module = StaticModule(g, opts);
+///   auto output = module(args, kwargs);
+/// @endcode
+///
+/// Mode 2: similar to data parallelism, run the same model for different inputs
+/// on different threads at the same time.
+/// You should have one StaticModule per model, and one StaticRuntime instance
+/// per running thread. To avoiding creating StaticRuntimes on the fly, use a
+/// synchronized stack (i.e. boost::lockfree::stack) to cache all the
+/// StaticRuntime instances in your code.
+/// @code
+///   // initialization
+///   auto module = std::make_shared<StaticModule>(m, opts);
+///
+///   // 128 is good for most cases. Pick a number that works for you
+///   boost::lockfree::stack<std::shared_ptr<StaticRuntime>,
+///     boost::lockfree::fixed_sized<true>> pool(128);
+///
+///   // inference
+///   std::shared_ptr<StaticRuntime> runtime = nullptr;
+///   pool.pop(runtime);
+///   if (!runtime) {
+///     // holds a reference to the underlying module
+///     // but does its own memory management
+///     runtime = std::make_shared<StaticRuntime>(*module);
+///   }
+///   auto output = runtime(args, kwargs);
+///   pool.push(runtime);
+/// @endcode
+///
+class MemoryPlanner;
+class ProcessedNode;
+class StaticRuntime;
+
+using SROperator = std::function<void(ProcessedNode*)>;
+
+#ifdef FBCODE_CAFFE2
+struct TORCH_API SROperatorObserver {
+  using OperatorCallback = void (*)(const Node*);
+  OperatorCallback startCb = nullptr;
+  OperatorCallback endCb = nullptr;
+
+  static void setCurrentThreadObserver(SROperatorObserver* observer);
+  static SROperatorObserver* getCurrentThreadObserver();
+  static void onStart(const Node* name);
+  static void onEnd(const Node* name);
+};
+#endif
+
+class TORCH_API ProcessedFunction {
+ public:
+  ProcessedFunction(
+      Node* node,
+      bool enable_out_variant,
+      bool check_memory_overlap);
+
+  enum class Kind : uint8_t {
+    kOutVariant,
+    kNativeFunction,
+    kInterpreterFallback,
+  };
+
+  void run(ProcessedNode* pnode) const {
+    return f_(pnode);
+  }
+
+  Kind kind() const {
+    return kind_;
+  }
+
+  bool checkMemoryOverlap() const {
+    return check_memory_overlap_;
+  }
+
+  size_t num_outputs() const {
+    return num_outputs_;
+  }
+
+ private:
+  SROperator f_;
+  Kind kind_{ProcessedFunction::Kind::kOutVariant};
+  bool check_memory_overlap_{false};
+  size_t num_outputs_{0};
+};
+
+class TORCH_API StaticNodeInfo {
+ public:
+  StaticNodeInfo(
+      Node* n,
+      ProcessedFunction* fn,
+      ProcessedNodeInputs inputs,
+      uint16_t outputs_offset);
+
+  Node* node() const {
+    return node_;
+  }
+
+  size_t num_outputs() const {
+    DCHECK(fn_ != nullptr);
+    return fn_->num_outputs();
+  }
+
+  bool has_out_variant() const {
+    return fn_->kind() == ProcessedFunction::Kind::kOutVariant;
+  }
+
+ private:
+  friend class ProcessedNode;
+
+  Node* node_;
+  const ProcessedFunction* fn_;
+  ProcessedNodeInputs inputs_;
+  uint16_t outputs_offset_;
+};
+
+// A `BlockInfo` instance stores all of the shared state that each
+// `BlockRunner` will need to access. Most of this information is
+// read-only and shared between threads.
+// - Each `BlockInfo` corresponds to one block in the graph.
+// - Each `BlockInfo` may be used by multiple block runners (when there are many
+//   threads).
+// - All of the `BlockInfo`s are stored in a vector in the `StaticModule` and
+//   are initialized during `StaticModule` construction.
+// - Most of the information stored is used to initialize the block's memory
+//   planner.
+class BlockInfo {
+ public:
+  BlockInfo(uint32_t input_idx, Block& block);
+
+  void set_nodes(
+      std::vector<StaticNodeInfo> nodes,
+      const c10::FastMap<Node*, bool>& node_has_out_variant);
+
+  const std::vector<StaticNodeInfo>& nodes() const {
+    return nodes_;
+  }
+
+  size_t num_nodes() const;
+
+  size_t num_inputs() const {
+    return block_.inputs().size();
+  }
+
+  size_t num_outputs() const {
+    return block_.outputs().size();
+  }
+
+  graph_node_list node_ptrs() const {
+    return block_.nodes();
+  }
+
+  void set_output_indices(std::vector<uint16_t> indices) {
+    output_indices_ = std::move(indices);
+  }
+
+  const std::vector<uint16_t>& block_output_indices() const {
+    return output_indices_;
+  }
+
+  auto block_inputs_idx() const {
+    return input_idx_;
+  }
+
+  bool node_is_optimizable_container_type(const Node* node) const {
+    return node_is_optimizable_container_type_.find(node) !=
+        node_is_optimizable_container_type_.end();
+  }
+
+  bool value_is_managed_tensor(const Value* value) const {
+    return managed_tensor_values_.find(value) != managed_tensor_values_.end();
+  }
+
+  bool value_is_leaked_container(const Value* value) const {
+    return leaked_values_.find(value) != leaked_values_.end();
+  }
+
+  const ValueGroup& value_group() const {
+    return value_group_;
+  }
+
+  const ManagedTensorRanges& managed_tensor_ranges() const {
+    return managed_tensor_ranges_;
+  }
+
+  void init_value_group(const AliasDb& alias_db) {
+    value_group_.init(block_, alias_db);
+  }
+
+  void prepare_for_memory_planner(
+      const AliasDb& alias_db,
+      const StaticModuleOptions& opt);
+
+  const auto& managed_output_tensor_values() const {
+    return managed_output_tensor_values_;
+  }
+
+  const auto& managed_tensor_values() const {
+    return managed_tensor_values_;
+  }
+
+  const auto& leaked_values() const {
+    return leaked_values_;
+  }
+
+ private:
+  std::vector<StaticNodeInfo> nodes_;
+
+  ValueGroup value_group_;
+
+  c10::FastSet<const Node*> node_is_optimizable_container_type_;
+  c10::FastSet<const Value*> managed_tensor_values_;
+  c10::FastSet<const Value*> managed_output_tensor_values_;
+  c10::FastSet<const Value*> leaked_values_;
+
+  ManagedTensorRanges managed_tensor_ranges_;
+
+  // The index of this block's inputs in the shared values_ array.
+  const uint16_t input_idx_;
+  // The indices of this block's outputs in the shared values_ array.
+  std::vector<uint16_t> output_indices_;
+  Block& block_;
+};
+
+class TORCH_API StaticModule {
+ public:
+  explicit StaticModule(
+      const std::shared_ptr<torch::jit::Graph>& g,
+      const StaticModuleOptions& opts = StaticModuleOptions(),
+      std::vector<IValue> sample_inputs = {});
+
+  explicit StaticModule(
+      const torch::jit::Module& m,
+      bool is_frozen = false,
+      const StaticModuleOptions& opts = StaticModuleOptions(),
+      std::vector<IValue> sample_inputs = {});
+
+ private:
+  explicit StaticModule(
+      std::pair<std::shared_ptr<torch::jit::Graph>, std::optional<Module>>
+          graph_and_module,
+      const StaticModuleOptions& opts);
+
+ public:
+  using KeywordArgs = std::unordered_map<std::string, c10::IValue>;
+  c10::IValue operator()(
+      const std::vector<c10::IValue>& args,
+      const KeywordArgs& kwargs = KeywordArgs());
+  c10::IValue operator()(
+      std::vector<c10::IValue>&& args,
+      const KeywordArgs& kwargs = KeywordArgs());
+
+  const Graph& graph() const {
+    return *graph_;
+  }
+
+  const Module& module() const {
+    DCHECK(module_.has_value());
+    return *module_;
+  }
+
+  const StaticModuleOptions& opts() const;
+
+  size_t num_inputs() const;
+  size_t num_outputs() const;
+
+  size_t num_constants() const {
+    return constants_.size();
+  }
+
+  size_t num_intermediate_values() const {
+    return num_intermediate_values_;
+  }
+
+  size_t total_num_values() const {
+    return num_inputs() + num_constants() + num_intermediate_values();
+  }
+
+  [[nodiscard]] const std::vector<uint16_t>& output_indices() const {
+    return output_indices_;
+  }
+
+  const std::vector<IValue>& constants() const {
+    return constants_;
+  }
+
+  const BlockInfo& block_info(Block* block) const {
+    return block_infos_.at(block);
+  }
+
+  Block* root_block() const {
+    return graph_->block();
+  }
+
+ private:
+  friend class StaticRuntime;
+  friend class BlockRunner;
+
+ public:
+  auto num_nodes() const {
+    return std::accumulate(
+        block_infos_.begin(),
+        block_infos_.end(),
+        0,
+        [](size_t sum, const auto& block_and_info) {
+          auto& block_info = block_and_info.second;
+          return sum + block_info.num_nodes();
+        });
+  }
+
+  [[nodiscard]] Node* findNodeWithKindForTesting(const std::string& kind) const;
+
+  const std::optional<c10::FunctionSchema>& schema() const {
+    return schema_;
+  }
+
+  bool first_input_is_self() const {
+    return module_.has_value();
+  }
+
+  StaticRuntime& runtime();
+
+  // See [Shared values array]
+  size_t value_buffer_size() const {
+    return value_buffer_size_;
+  }
+
+ private:
+  // Recursively prepares the BlockInfo array.
+  // - Populates `value_to_index` with the indices of each intermediate value
+  // - Returns the number of Value* processed, including sub-blocks.
+  size_t prepareBlockInfo(
+      Block* block,
+      const size_t start_idx,
+      c10::FastMap<const Value*, uint32_t>& value_to_index);
+
+  void prepareFunctionsAndConstants(
+      Block* block,
+      const AliasDb& alias_db,
+      c10::FastMap<const Value*, uint32_t>& value_to_index);
+
+  // Recursively traverse the graph and attach SR metadata
+  // to the prim::fork nodes as additional attributes
+  void attachNodeMetadata(Block* block);
+
+  // Recurses on sub-blocks and populates the array of ProcessedNodes
+  // Returns (number of nodes processed, number of blocks processed)
+  size_t prepareStaticNodeInfos(
+      Block* block,
+      const c10::FastMap<const Value*, uint32_t>& value_to_index,
+      const AliasDb& alias_db,
+      size_t node_idx = 0);
+
+  // Initialize various attributes that the memory planner will need.
+  // To be called at the tail of the ctor.
+  void prepareForMemoryPlanner();
+
+  StaticModuleOptions opts_;
+  // metadata that is stored in IR nodes as attribute
+  at::intrusive_ptr<jit::StaticRuntimeMetadata> sr_metadata_;
+  std::shared_ptr<torch::jit::Graph> graph_;
+  std::optional<torch::jit::Module> module_;
+  std::optional<c10::FunctionSchema> schema_;
+  std::unique_ptr<StaticRuntime> cached_runtime_;
+
+  // Bookkeeping for creating new StaticRuntime instances
+  // IValue table (defined by prim::Constant nodes)
+  std::vector<IValue> constants_;
+  // The functions to be called by corresponding ProcessedNode.
+  std::vector<ProcessedFunction> functions_;
+  // A list of pre-processed nodes from which ProcessedNode are created per
+  // StaticRuntime instance.
+  std::vector<StaticNodeInfo> nodes_;
+  // Indices of graph outputs in the single values array.
+  std::vector<uint16_t> output_indices_;
+
+  size_t num_intermediate_values_ = 0;
+
+  // Includes self if module_ != std::nullopt.
+  // Note that we might have num_inputs_ == 0 even if the schema has a `self`
+  // argument. In this case, `self` isn't used in the graph, but the schema
+  // includes it anyways to be consistent with the JIT interpreter.
+  size_t num_inputs_;
+  // See `BlockInfo` definition. The blocks are stored in depth-first order.
+  c10::FastMap<Block*, BlockInfo> block_infos_;
+  size_t value_buffer_size_ = 0;
+};
+
+// `BlockRunner` contains the core runtime logic. Each block runner
+// corresponds to one block in the graph and has its own memory planner.
+// `StaticRuntime` will initialize all `BlockRunner`s
+// upon construction. Each block runner only directly executes nodes from its
+// block. Special ops with sub-blocks like `prim::If` may have
+// `BlockRunner`s stored in their `ProcessedNode`s; these
+// sub-blocks get executed in the op's implementation.
+// `StaticRuntime` stores a vector of IValues that all
+// `BlockRunner`s share. This vector is used to store all
+// constants, inputs, and intermediate tensors.
+class TORCH_API BlockRunner {
+ public:
+  BlockRunner(
+      const StaticModule& sm,
+      IValue* values,
+      Block* block,
+      torch::jit::TaskLauncher* launcher,
+      bool is_root_block = false);
+  BlockRunner(BlockRunner&&) noexcept;
+  BlockRunner& operator=(BlockRunner&&) = delete;
+  ~BlockRunner();
+
+  C10_DISABLE_COPY_AND_ASSIGN(BlockRunner);
+
+  using KeywordArgs = std::unordered_map<std::string, c10::IValue>;
+  c10::IValue operator()(
+      const std::vector<c10::IValue>& args,
+      const KeywordArgs& kwargs = KeywordArgs());
+  c10::IValue operator()(
+      std::vector<c10::IValue>&& args,
+      const KeywordArgs& kwargs = KeywordArgs());
+
+  c10::intrusive_ptr<c10::ivalue::Future> runAsync(
+      const std::vector<c10::IValue>& args,
+      const KeywordArgs& kwargs);
+
+  c10::intrusive_ptr<c10::ivalue::Future> runAsync(
+      std::vector<c10::IValue>&& args,
+      const KeywordArgs& kwargs);
+
+  void benchmark(
+      const std::vector<std::vector<c10::IValue>>& args_list,
+      const std::vector<KeywordArgs>& kwargs_list,
+      const uint32_t warmup_runs,
+      const uint32_t main_runs,
+      bool print_per_node_time = false,
+      bool generate_ai_pep_output = false);
+
+  struct IndividualMetrics {
+    float setup_time{0.0};
+    float memory_alloc_time{0.0};
+    float memory_dealloc_time{0.0};
+    float output_dealloc_time{0.0};
+    float first_iter_time{0.0};
+    float total_time{0.0};
+    size_t out_nodes_count{0};
+    size_t total_nodes_count{0};
+    std::vector<float> time_per_node;
+    std::unordered_map<std::string, float> time_per_node_type;
+    std::unordered_map<std::string, float> percent_per_node_type;
+    std::unordered_map<std::string, int> instances_per_node_type;
+    std::unordered_set<std::string> out_nodes;
+    std::unordered_set<std::string> native_nodes;
+  };
+
+  IndividualMetrics benchmark_individual_ops(
+      const std::vector<std::vector<c10::IValue>>& args_list,
+      const std::vector<KeywordArgs>& kwargs_list,
+      const uint32_t warmup_runs,
+      const uint32_t main_runs);
+
+  // Input is readwrite
+  IValue& Input(uint32_t i) {
+    TORCH_DCHECK_LT(i, block_info_.num_inputs());
+    return values_[i + block_info_.block_inputs_idx()];
+  }
+
+  // Output is readonly. The writing process happens inside ProcessedNodes
+  [[nodiscard]] const IValue& Output(uint32_t i) const {
+    DCHECK(i < outputs_.size());
+    return *outputs_[i];
+  }
+
+  const std::vector<IValue*> outputs() const {
+    return outputs_;
+  }
+
+  const std::vector<ProcessedNode>& nodes() const {
+    return nodes_;
+  }
+
+  std::vector<ProcessedNode>& nodes() {
+    return nodes_;
+  }
+
+  graph_node_list node_ptrs() const {
+    return block_info_.node_ptrs();
+  }
+
+  const Graph& graph() const {
+    return static_module_.graph();
+  }
+
+  const MemoryPlanner* get_memory_planner() const {
+    return planner_.get();
+  }
+
+  bool check_for_memory_leak(
+      bool output_returned = true,
+      bool recurse_on_sub_blocks = false);
+
+  // WARNING: Deallocate managed output tensors.  A client receiving Static
+  // Runtime-managed Tensors needs to be very careful to call
+  // `StaticRuntime::deallocateOutputTensors` after all references of output
+  // Tensors are gone.
+  void deallocateOutputTensors();
+
+  bool checkOutputTensorMemoryLeaks();
+
+  bool isManagedOutputTensor(const IValue& ivalue) const;
+  bool isManagedOutputTensorValue(const Value* value) const;
+
+  void disableManageOutputTensors();
+
+  // This is the fallback path taken if we can't construct the memory planner
+  // on the first iteration.
+  // IMPORTANT: Nothing here should be able to throw!!!
+  // This function can be called from the (implicitly) `noexcept` destructor
+  // of Deallocator, meaning that std::terminate will be called
+  // if any exception escapes. Even if resetMemory and ~Deallocator were
+  // `noexcept(false)`, it's possible that when ~Deallocator is called, the
+  // stack is already unwinding, so there's still danger of calling
+  // std::terminate.
+  void resetMemory() noexcept;
+
+ private:
+  // A helper object that invokes memory planner deallocation code
+  // when destructed.
+  class Deallocator {
+   public:
+    explicit Deallocator(BlockRunner& block_runner)
+        : block_runner_(block_runner) {}
+
+    Deallocator(Deallocator&&) = default;
+    Deallocator(const Deallocator&) = default;
+    Deallocator& operator=(const Deallocator&) = delete;
+    Deallocator& operator=(Deallocator&&) = delete;
+    ~Deallocator();
+
+    void setFinished() {
+      finished_ = true;
+    }
+
+   private:
+    void cleanupImpl();
+
+    bool finished_ = false;
+    BlockRunner& block_runner_;
+  };
+
+  template <typename IValueList>
+  c10::IValue run_impl(IValueList&& args, const KeywordArgs& kwargs);
+
+  template <typename IValueList>
+  c10::IValue run_impl_record_functions(
+      IValueList&& args,
+      const KeywordArgs& kwargs);
+
+  template <typename IValueList>
+  c10::intrusive_ptr<c10::ivalue::Future> run_impl_async(
+      IValueList&& args,
+      const KeywordArgs& kwargs);
+
+  template <typename IValueList>
+  c10::intrusive_ptr<c10::ivalue::Future> run_impl_record_functions_async(
+      IValueList&& args,
+      const KeywordArgs& kwargs);
+
+  // helper method for copying input args/kwargs into inputs_
+  template <typename IValueList>
+  void set_inputs(IValueList&& args, const KeywordArgs& kwargs);
+
+  // Set Input(idx) to args[idx]. Invoked by set_inputs. Copies or moves
+  // depending on overload.
+  void set_arg(const size_t idx, std::vector<IValue>&& args);
+  void set_arg(const size_t idx, const std::vector<IValue>& args);
+
+  // Set Input(idx) to arg. Always copies. Used for kwargs.
+  void set_arg(const size_t idx, const IValue& arg);
+
+  bool fast_check_and_correct_overlap_with(
+      ProcessedNode& n,
+      c10::IValue& tensor_ival);
+  void verify_and_correct_memory_overlap(ProcessedNode& n);
+
+  // clean up owning refs of input IValues
+  void clean_up_input_ivalues() noexcept {
+    for (const auto idx : c10::irange(block_info_.num_inputs())) {
+      values_[idx + inputs_begin_] = IValue();
+    }
+  }
+
+  void clean_up_intermediate_ivalues() noexcept;
+
+  IValue move_outputs_to_tuple(uint32_t num_outputs);
+
+  void create_memory_planner();
+
+  float benchmark_model(
+      const std::vector<std::vector<c10::IValue>>& args_list,
+      const std::vector<KeywordArgs>& kwargs_list,
+      const uint32_t warmup_runs,
+      const uint32_t main_runs);
+
+  void display_nodes(
+      const std::vector<c10::IValue>& args,
+      const KeywordArgs& kwargs);
+
+  const StaticModule& static_module_;
+  const BlockInfo& block_info_;
+
+  const bool is_root_block_;
+  // Cache this so we don't have to call static_module_.first_input_is_self()
+  const bool first_input_is_self_;
+  // Index of the start of this blocks inputs in the shared values_ array.
+  const uint16_t inputs_begin_;
+
+  bool manage_output_tensors_enabled_ = false;
+  std::unique_ptr<MemoryPlanner> planner_;
+  // [Shared values array]
+  // ProcessedNodes reference their inputs and outputs with
+  // offsets into this array, which saves memory.
+  // All BlockRunners share the same array. The layout is as
+  // follows:
+  // [constants][block_0][block_1]...[block_N]
+  // Note that constants from all blocks are pooled together at the start.
+  // The block ordering is depth-first.
+  // Each block is further divided into inputs and intermediates:
+  // [block_i] = [inputs_i][intermediates_i]
+  // Each BlockRunner knows where its inputs start. Each ProcessedNode
+  // knows how to find the indices of its outputs/inputs in this array.
+  IValue* values_;
+
+  std::vector<IValue*> outputs_;
+  std::vector<ProcessedNode> nodes_;
+};
+
+inline size_t BlockInfo::num_nodes() const {
+  return nodes_.size();
+}
+
+/*
+  ProcessedNodeMetadata class wraps the possible metadata
+  for ProcessedNode. Depending upon the nature of op, processedNode
+  can have one of the below possibilities of metadata:
+  - prim::If/prim::Loop ops contains block_runners_ as their metadata
+  - prim::fork op contains TaskLauncher (std::function) responsible for
+    execution of forked subgraph
+*/
+class TORCH_API ProcessedNodeMetadata {
+ public:
+  ProcessedNodeMetadata(
+      std::vector<BlockRunner> runners,
+      torch::jit::TaskLauncher* launcher)
+      : block_runners_(std::move(runners)), launcher_(launcher) {}
+
+  ProcessedNodeMetadata() : launcher_(nullptr) {}
+
+  // deleted copy ctor/assignment as standard containers (vector) always
+  // have copy constructors, but their instantiation is not well-formed
+  // if the contained type (BlockRunner) is not copyable
+  ProcessedNodeMetadata(const ProcessedNodeMetadata&) = delete;
+  ProcessedNodeMetadata& operator=(const ProcessedNodeMetadata&) = delete;
+  ProcessedNodeMetadata(ProcessedNodeMetadata&&) = delete;
+  ProcessedNodeMetadata&& operator=(ProcessedNodeMetadata&&) = delete;
+  ~ProcessedNodeMetadata() = default;
+
+  std::vector<BlockRunner>& block_runners() {
+    return block_runners_;
+  }
+
+  void set_block_runners(std::vector<BlockRunner> runners) {
+    block_runners_ = std::move(runners);
+  }
+
+  void set_launcher(torch::jit::TaskLauncher* launcher) {
+    launcher_ = launcher;
+  }
+
+  torch::jit::TaskLauncher* launcher() {
+    return launcher_;
+  }
+
+ private:
+  std::vector<BlockRunner> block_runners_;
+  torch::jit::TaskLauncher* launcher_;
+};
+
+class TORCH_API ProcessedNode {
+ public:
+  ProcessedNode() = default;
+
+  ProcessedNode(const StaticNodeInfo& other, IValue* values)
+      : node_(other.node_),
+        fn_(other.fn_),
+        inputs_(other.inputs_),
+        outputs_offset_(other.outputs_offset_),
+        values_(values),
+        metadata_(nullptr) {}
+
+  // These should be noexcept, but some Android build is failing
+  // saying the noexcept specification doesn't match the calculated
+  // one. Maybe std::variant is throwing it off?
+  ProcessedNode(ProcessedNode&&) = default;
+
+  ProcessedNode(const ProcessedNode&) = delete;
+  ProcessedNode& operator=(const ProcessedNode& other) = delete;
+  ProcessedNode& operator=(ProcessedNode&&) = default;
+  ~ProcessedNode() = default;
+
+  void run();
+
+  Node* node() const {
+    return node_;
+  }
+
+  // Input is readonly
+  [[nodiscard]] const IValue& Input(uint32_t i) const {
+    return values_[inputs_[i]];
+  }
+
+  // Output is readwrite
+  IValue& Output(uint32_t i) {
+    DCHECK(i < num_outputs());
+    return values_[outputs_offset_ + i];
+  }
+
+  [[nodiscard]] const IValue& Output(uint32_t i) const {
+    DCHECK(i < num_outputs());
+    return values_[outputs_offset_ + i];
+  }
+
+  uint32_t num_outputs() const {
+    DCHECK(fn_ != nullptr);
+    return static_cast<uint32_t>(fn_->num_outputs());
+  }
+
+  [[nodiscard]] c10::ArrayRef<const IValue> outputs() const {
+    return c10::ArrayRef<const IValue>(
+        values_ + outputs_offset_, num_outputs());
+  }
+
+  [[nodiscard]] uint16_t num_inputs() const {
+    return inputs_.size();
+  }
+
+  std::vector<IValue> inputs_ivalue_vec() const;
+
+  bool has_out_variant() const {
+    return fn_->kind() == ProcessedFunction::Kind::kOutVariant;
+  }
+
+  bool has_native() const {
+    return fn_->kind() == ProcessedFunction::Kind::kNativeFunction;
+  }
+
+#ifndef PYTORCH_DISABLE_PER_OP_PROFILING
+  const char* get_op_name() const {
+    return node_->kind().toQualString();
+  }
+#endif
+
+  bool check_outputs_for_memory_overlap() const {
+    return fn_->checkMemoryOverlap();
+  }
+
+  void set_outputs_memory_overlap_detected() {
+    overlap_detected_ = true;
+  }
+
+  bool outputs_memory_overlap_detected() {
+    return overlap_detected_;
+  }
+
+  bool check_and_correct_overlap_with(
+      const at::Tensor& input,
+      c10::IValue& output);
+  void verify_and_correct_memory_overlap();
+
+  void set_values(IValue* values) {
+    DCHECK(values_ == nullptr);
+    values_ = values;
+  }
+
+  [[nodiscard]] uint16_t output_ivalue_index(uint16_t i) const {
+    DCHECK(i < num_outputs());
+    return outputs_offset_ + i;
+  }
+  // used in debug mode
+  bool verify_no_memory_overlap(bool force_check = false) const;
+
+  // returns pointer to ProcessedNodeMetadata or nullptr if no object is owned
+  ProcessedNodeMetadata* metadata() {
+    return metadata_.get();
+  }
+
+  // attach block_runner to metadata of ProcessedNode
+  void set_metadata(std::vector<BlockRunner> block_runners) {
+    if (metadata_ == nullptr) {
+      metadata_ = std::make_unique<ProcessedNodeMetadata>();
+    }
+    metadata_->set_block_runners(std::move(block_runners));
+  }
+
+  // attach TaskLauncher to metadata of ProcessedNode
+  void set_metadata(torch::jit::TaskLauncher* launcher) {
+    if (metadata_ == nullptr) {
+      metadata_ = std::make_unique<ProcessedNodeMetadata>();
+    }
+    metadata_->set_launcher(launcher);
+  }
+
+ private:
+  [[nodiscard]] bool verify_outputs_dont_overlap_each_other() const;
+
+  [[nodiscard]] bool verify_inputs_dont_overlap_outputs(bool force_check) const;
+
+  Node* node_{nullptr};
+  const ProcessedFunction* fn_{nullptr};
+  ProcessedNodeInputs inputs_;
+  uint16_t outputs_offset_{0};
+  bool overlap_detected_{false};
+  IValue* values_ = nullptr; // unowned
+  // Metadata for ProcessedNode.
+  // 1. prim::If/Loop nodes contains sub-blocks as metadata
+  // 2. prim::fork nodes contains custom executor for async execution
+  std::unique_ptr<ProcessedNodeMetadata> metadata_;
+};
+
+// `StaticRuntime` is the owner of the array of IValues (used for constants,
+// inputs, and intermediate tensors) that all `BlockRunner`s share.
+// Upon construction, it initializes all block runners. `operator()` simply
+// forwards the inputs to the top-level block runner. Each `StaticRuntime`
+// instance corresponds to one `StaticModule`. Multiple `StaticRuntime`
+// instances can be created; this is useful for multi-threaded execution, since
+// `operator()` is not thread-safe.
+class TORCH_API StaticRuntime {
+ public:
+  explicit StaticRuntime(const StaticModule& sm);
+
+  using KeywordArgs = std::unordered_map<std::string, c10::IValue>;
+  c10::IValue operator()(
+      const std::vector<c10::IValue>& args,
+      const KeywordArgs& kwargs = KeywordArgs());
+  c10::IValue operator()(
+      std::vector<c10::IValue>&& args,
+      const KeywordArgs& kwargs = KeywordArgs());
+
+  // runAsync performs inline execution of graph on
+  // caller thread and async execution on taskLauncher
+  // If no custom taskLauncher is specified, execution is done
+  // on inter-op thread pool.
+  c10::intrusive_ptr<c10::ivalue::Future> runAsync(
+      const std::vector<c10::IValue>& args,
+      const KeywordArgs& kwargs = KeywordArgs(),
+      torch::jit::TaskLauncher taskLauncher = at::launch);
+
+  c10::intrusive_ptr<c10::ivalue::Future> runAsync(
+      std::vector<c10::IValue>&& args,
+      const KeywordArgs& kwargs = KeywordArgs(),
+      torch::jit::TaskLauncher taskLauncher = at::launch);
+
+  bool check_for_memory_leak(bool output_returned = true);
+  bool checkOutputTensorMemoryLeaks();
+
+  void deallocateOutputTensors();
+  bool isManagedOutputTensor(const IValue& ivalue) const;
+  void disableManageOutputTensors();
+
+  // Gets the top-level memory planner. Used for testing.
+  const MemoryPlanner* get_memory_planner() const;
+
+  void benchmark(
+      const std::vector<std::vector<c10::IValue>>& args_list,
+      const std::vector<KeywordArgs>& kwargs_list,
+      const uint32_t warmup_runs,
+      const uint32_t main_runs,
+      bool print_per_node_time = false,
+      bool generate_ai_pep_output = false) {
+    block_->benchmark(
+        args_list,
+        kwargs_list,
+        warmup_runs,
+        main_runs,
+        print_per_node_time,
+        generate_ai_pep_output);
+  }
+
+  using IndividualMetrics = BlockRunner::IndividualMetrics;
+
+  IndividualMetrics benchmark_individual_ops(
+      const std::vector<std::vector<c10::IValue>>& args_list,
+      const std::vector<KeywordArgs>& kwargs_list,
+      const int warmup_runs,
+      const int main_runs) {
+    return block_->benchmark_individual_ops(
+        args_list, kwargs_list, warmup_runs, main_runs);
+  }
+
+ private:
+  // An array of IValues with unchanging size/data ptr.
+  class IValueArray {
+   public:
+    IValueArray() = default;
+    explicit IValueArray(size_t size) : array_(allocate(size)), size_(size) {}
+
+    IValue* data() const {
+      return array_.get();
+    }
+
+    size_t size() const {
+      return size_;
+    }
+
+   private:
+    // NOLINTNEXTLINE(modernize-avoid-c-arrays)
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays)
+    static std::unique_ptr<IValue[]> allocate(size_t size) {
+      if (size) {
+        return std::make_unique<IValue[]>(size);
+      }
+      return nullptr;
+    }
+
+    // NOLINTNEXTLINE(modernize-avoid-c-arrays)
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays)
+    std::unique_ptr<IValue[]> array_ = nullptr;
+    size_t size_ = 0;
+  };
+
+  std::unique_ptr<BlockRunner> block_;
+  // for execution of async operations present in graph
+  torch::jit::TaskLauncher async_task_launcher_;
+  IValueArray values_;
+};
+
+} // namespace torch::jit
+C10_DECLARE_bool(static_runtime_disable_debug_memory_overlap_check);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..d39c7141cb007a050e10d9a7d7588fa8a1ca0604
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/init.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/jit/python/pybind_utils.h>
+
+namespace torch::jit {
+
+void initStaticModuleBindings(PyObject* module);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/memory_planner.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/memory_planner.h
new file mode 100644
index 0000000000000000000000000000000000000000..5c9bef359ea47839007714576ef262f1aae93861
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/memory_planner.h
@@ -0,0 +1,303 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/runtime/static/impl.h>
+
+namespace torch::jit {
+
+// A StorageGroup represents a collection of tensors that share backing storage.
+class StorageGroup {
+ public:
+  // Every storage group must contain at least one tensor.
+  explicit StorageGroup(at::Tensor* tensor) : group_{tensor} {}
+
+  void addTensor(at::Tensor* tensor) {
+    group_.push_back(tensor);
+  }
+
+  const std::vector<at::Tensor*>& group() const {
+    return group_;
+  }
+
+  size_t maxTensorSize() const {
+    return max_tensor_size_;
+  }
+
+  void setMaxTensorSize(size_t new_size) {
+    max_tensor_size_ = new_size;
+  }
+
+  size_t numManagedTensors() const {
+    return group_.size();
+  }
+
+ private:
+  // The size attribute represents the amount of memory that will be
+  // allocated for all tensors in this storage group. Initially it
+  // is zero, eventually it gets updated by the MemoryPlanner.
+  size_t max_tensor_size_ = 0;
+  std::vector<at::Tensor*> group_;
+};
+
+// A contiguous buffer of `StorageImpl`s
+class ManagedStorages {
+ public:
+  ManagedStorages();
+
+  ~ManagedStorages();
+
+  void allocate(size_t capacity);
+
+  void deallocate();
+
+  bool is_allocated() const {
+    return storages_ != nullptr;
+  }
+
+  // Append a new StorageImpl to the buffer. The new StorageImpl is given the
+  // same size and allocator as `storageImpl` argument
+  void append(at::StorageImpl& storageImpl);
+
+  at::StorageImpl& operator[](size_t idx) {
+    TORCH_INTERNAL_ASSERT(storages_ != nullptr);
+    return storages_[idx];
+  }
+
+  const at::StorageImpl& operator[](size_t idx) const {
+    TORCH_INTERNAL_ASSERT(storages_ != nullptr);
+    return storages_[idx];
+  }
+
+  size_t size() const {
+    return size_;
+  }
+
+  bool empty() const {
+    return size_ == 0;
+  }
+
+  size_t capacity() const {
+    return capacity_;
+  }
+
+ private:
+  // We will use placement-new to add new storages to this buffer
+  at::StorageImpl* storages_{nullptr};
+
+  // Current number of storages that have been placed into the storage buffer
+  size_t size_{0};
+
+  // Total allocated capacity of the storage buffer
+  size_t capacity_{0};
+};
+
+TORCH_API std::vector<StorageGroup> assignStorageToManagedTensors(
+    graph_node_list nodes,
+    const ManagedTensorRanges& ranges,
+    const c10::FastMap<const Value*, at::Tensor*>& tensor_value_to_tensor);
+
+// There are three types of ops in a processed graph in Static Runtime:
+//   1. op with _out variant
+//   2. view-producing op
+//   3. tensor-producing op (could be replaced with type 1 by adding the _out
+//      variant to Static Runtime)
+// In Static Runtime, type 2 ops are replaced with their corresponding copy
+// versions when enable_out_variant is enabled and become type 1 ops.The memory
+// planner only manages tensors that are outputs of type 1 ops. For type 3, the
+// output tensors are allocated inside the operator and can't be directly
+// managed by memory planner.
+//
+// Memory planner tries to minimize the number of memory allocations by
+// tracking the output tensors of ops with _out variants with unique DataPtr
+// (part of StorageImpl). It tries to do this in several steps:
+//   1. record the max memory usage for each Tensor with unique DataPtr at the
+//      end of each iteration
+//   2. in the next iteration, allocate the buffer for the max total usage and
+//      compute the offset of each allocation with regard to the single memory
+//      buffer, optionally reusing memory. In the first iteration, we rely on
+//      the default allocator for memory allocation.
+//   3. free the buffer at the end of each iteration
+// Steps 1 and 3 are handled by `deallocate()`, and step 2 by `allocate()`.
+// Only models with simple output types are supported, i.e. None, Tensor or
+// List/Tuple/Dict of Tensors. Complex output types such as List of Lists are
+// not supported.
+//
+// Additional Optimizations:
+//
+// [Borrowed IValue Outputs]
+// A few native ops (notably, `static_runtime::dict_unpack` and
+// `static_runtime::VarTupleUnpack`) simply unpack IValues to a bunch of
+// outputs without modification. For example, `dict_unpack` does the following:
+// for each key in inputs:
+//     output[i] = dict_input[key]
+// To avoid refcount bumps, the outputs of these ops are non-owning references.
+// This requires special logic in the memory planner - when adding an op that
+// borrows outputs, be sure that the memory planner is updated accordingly!
+//
+// [Managed Output Tensors]
+// The memory planner is able to manage output tensors if the appropriate
+// `StaticModuleOptions` are set. However, the memory planner handles output
+// tensors separately from regular intermediate tensors:
+// 1. They don't participate in memory reuse.
+// 2. The memory planner cannot reclaim their backing storage until they have
+//    been explicitly freed by the client.
+
+class MemoryPlanner {
+ public:
+  MemoryPlanner(
+      BlockRunner* block_runner,
+      const BlockInfo& block_info,
+      bool enable_out_variant,
+      bool manage_output_tensors);
+
+  // disable copying and moving
+  MemoryPlanner(const MemoryPlanner&) = delete;
+  MemoryPlanner& operator=(const MemoryPlanner&) = delete;
+  MemoryPlanner(MemoryPlanner&&) = delete;
+  MemoryPlanner& operator=(MemoryPlanner&&) = delete;
+  virtual ~MemoryPlanner() = default;
+
+  void allocate();
+  void deallocate();
+  void deallocateOutputTensors();
+
+  size_t total_num_managed_tensors() const {
+    return num_managed_tensors_;
+  }
+
+  size_t total_reused_tensors() const {
+    return reused_tensors_;
+  }
+
+  size_t total_num_managed_output_tensors() const {
+    return managed_output_tensors_.size();
+  }
+
+  [[nodiscard]] size_t total_num_unmanaged() const {
+    return num_unmanaged_non_scalars() + num_unmanaged_scalars();
+  }
+
+  [[nodiscard]] size_t num_unmanaged_non_scalars() const {
+    return unmanaged_ivalues_.size() + unmanaged_borrowed_ivalues_.size();
+  }
+
+  [[nodiscard]] size_t num_unmanaged_scalars() const {
+    return num_unmanaged_scalar_ivalues_;
+  }
+
+  size_t total_managed() const {
+    return managed_bytes_;
+  }
+
+  size_t numOutputBufferBytes() const {
+    return output_buffer_bytes_;
+  }
+
+  // Check if `ivalue` is contained as a managed tensor. Only used in DCHECK().
+  bool isManagedOutputTensor(const IValue& ivalue) const {
+    if (!output_buffer_ || // output buffer got already deallocated.
+        output_buffer_bytes_ == 0 || // memory planning is not yet initialized.
+        !ivalue.isTensor() // a non-tensor is never managed
+    ) {
+      return false;
+    }
+    const auto& tensor = ivalue.toTensor();
+    if (!tensor.has_storage() || !tensor.storage().data_ptr()) {
+      return false;
+    }
+    // TODO: Improve this once D31357486 is landed.
+    uint8_t* tensor_ptr =
+        static_cast<uint8_t*>(tensor.storage().data_ptr().get());
+    uint8_t* buffer_start = static_cast<uint8_t*>(output_buffer_.get());
+    uint8_t* buffer_end = buffer_start + output_buffer_bytes_;
+    return buffer_start <= tensor_ptr && tensor_ptr < buffer_end;
+  }
+
+  bool isManagedStorageImpl(const at::StorageImpl* impl) const {
+    if (storages_.empty()) {
+      return false;
+    }
+    // Comparing pointers that aren't within the same array is
+    // UB. We're doing fancy memory allocation stuff, so we cast to an
+    // integer type and carry on.
+    const auto impl_p = reinterpret_cast<uintptr_t>(impl);
+    const auto start = reinterpret_cast<uintptr_t>(&storages_[0]);
+    const auto end =
+        reinterpret_cast<uintptr_t>(&storages_[0] + storages_.size());
+    return impl_p >= start && impl_p < end;
+  }
+
+  bool overlapWithInternalBuffer(void* data_ptr) {
+    return buffer_start_ <= data_ptr && data_ptr < buffer_end_;
+  }
+
+ protected:
+  uint8_t* allocateBuffer(size_t num_bytes);
+
+  size_t managed_bytes_{0};
+  size_t reused_tensors_{0};
+
+  // We allocate StorageImpls ourselves so that 1) we don't have to do
+  // an extra two loads per Tensor (which will likely miss in the CPU
+  // data cache) first reading the Storage (i.e., StorageImpl pointer)
+  // from the TensorImpl object and then second dereferencing it and
+  // 2) our memory access pattern during allocate() has high locality.
+  // We don't have any guarantee that the model doesn't change the
+  // Storage for managed tensors out from under us during execution,
+  // so we have to check the StorageImpls each time we deallocate.
+  ManagedStorages storages_;
+
+  // Contains the size (in bytes) of the data to be allocated for each storage
+  std::vector<size_t> storages_nbytes_;
+
+ private:
+  // ivalues created in one run but not managed by MemoryPlanner
+  std::vector<IValue*> unmanaged_ivalues_;
+
+  // Special class of unmanaged values: some native ops create IValues
+  // in a "borrowed" state that can and must be cleaned up without a
+  // reference count decrement.
+  std::vector<IValue*> unmanaged_borrowed_ivalues_;
+
+  // Even more special class of unmanaged values: if select_tensor
+  // outputs are outputs of the graph, then they need to be restored
+  // to an ordinary "strong reference" state.
+  std::vector<IValue*> borrowed_ivalues_needing_incref_;
+
+  std::vector<std::pair<size_t, at::Tensor*>> managed_output_tensors_;
+  at::DataPtr buffer_; // allocated each time we call Run()
+  uint8_t* buffer_start_{nullptr};
+  uint8_t* buffer_end_{nullptr};
+  size_t num_managed_tensors_{0};
+  size_t num_unmanaged_scalar_ivalues_{0};
+
+  at::DataPtr output_buffer_;
+  size_t output_buffer_bytes_{0};
+
+  virtual void allocateManagedTensors() = 0;
+  virtual void deallocateManagedTensors() = 0;
+
+  void allocateOutputTensors();
+};
+
+class StandardMemoryPlanner : public MemoryPlanner {
+ public:
+  StandardMemoryPlanner(
+      BlockRunner* block_runner,
+      const BlockInfo& block_info,
+      bool enable_out_variant,
+      bool manage_output_tensors,
+      bool optimize_memory);
+
+ protected:
+  void allocateManagedTensors() override;
+  void deallocateManagedTensors() override;
+
+  std::vector<StorageGroup> managed_tensors_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..eee8b1d798cbd2d60a324d482216ec7584d08ec3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/ops.h
@@ -0,0 +1,192 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Utils.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/runtime/static/impl.h>
+
+namespace at::native {
+at::Tensor& reshape_copy_out(
+    at::Tensor& out,
+    const at::Tensor& self,
+    const at::DimVector& proposed_shape,
+    bool infer_size = true);
+at::Tensor& to_copy_out(
+    Tensor& out,
+    const Tensor& self,
+    bool non_blocking,
+    bool copy_strides,
+    std::optional<MemoryFormat> memory_format);
+} // namespace at::native
+
+namespace torch::jit {
+
+using SROpFunctor = SROperator (*)(Node* n);
+struct SROperatorFunctor {
+  virtual SROperator Generate(Node* /*unused*/) {
+    SROperator out;
+    return out;
+  }
+  virtual ~SROperatorFunctor() = default;
+};
+
+TORCH_DECLARE_REGISTRY(SROperatorRegistry, SROperatorFunctor);
+
+#define REGISTER_OPERATOR_FUNCTOR(name, id, ...)             \
+  struct SROperatorFunctor_##id : public SROperatorFunctor { \
+    SROpFunctor fn = __VA_ARGS__;                            \
+    SROperator Generate(Node* n) override {                  \
+      return fn(n);                                          \
+    }                                                        \
+  };                                                         \
+  C10_REGISTER_CLASS(SROperatorRegistry, name, SROperatorFunctor_##id)
+
+TORCH_DECLARE_REGISTRY(SRNativeOperatorRegistry, SROperatorFunctor);
+#define REGISTER_NATIVE_OPERATOR_FUNCTOR(name, id, ...)            \
+  struct SRNativeOperatorFunctor_##id : public SROperatorFunctor { \
+    SROpFunctor fn = __VA_ARGS__;                                  \
+    SROperator Generate(Node* n) override {                        \
+      return fn(n);                                                \
+    }                                                              \
+  };                                                               \
+  C10_REGISTER_CLASS(                                              \
+      SRNativeOperatorRegistry, name, SRNativeOperatorFunctor_##id)
+
+inline at::Tensor create_empty_from(const at::Tensor& t) {
+  return at::detail::empty_cpu(
+      {0},
+      c10::typeMetaToScalarType(t.dtype()),
+      t.layout(),
+      t.device(),
+      std::nullopt,
+      std::nullopt);
+}
+
+inline at::Tensor create_empty_from(
+    at::IntArrayRef sizes,
+    const at::Tensor& t) {
+  return at::detail::empty_cpu(
+      sizes,
+      c10::typeMetaToScalarType(t.dtype()),
+      t.layout(),
+      t.device(),
+      std::nullopt,
+      std::nullopt);
+}
+
+inline at::Tensor create_empty(c10::ScalarType dtype) {
+  return at::detail::empty_cpu(
+      {0}, dtype, std::nullopt, std::nullopt, std::nullopt, std::nullopt);
+}
+
+inline at::Tensor create_empty_from(
+    const at::Tensor& t,
+    c10::ScalarType dtype) {
+  return at::detail::empty_cpu(
+      {0}, dtype, t.layout(), t.device(), std::nullopt, std::nullopt);
+}
+
+inline at::Tensor create_empty_from(const at::Tensor& t, c10::Layout layout) {
+  return at::detail::empty_cpu(
+      {0},
+      c10::typeMetaToScalarType(t.dtype()),
+      layout,
+      t.device(),
+      std::nullopt,
+      std::nullopt);
+}
+
+inline at::Tensor create_empty_from(const at::Tensor& t, c10::Device device) {
+  return at::detail::empty_cpu(
+      {0},
+      c10::typeMetaToScalarType(t.dtype()),
+      t.layout(),
+      device,
+      std::nullopt,
+      std::nullopt);
+}
+
+inline at::Tensor create_empty_from(
+    const at::Tensor& t,
+    c10::MemoryFormat memory_format) {
+  return at::detail::empty_cpu(
+      {0},
+      c10::typeMetaToScalarType(t.dtype()),
+      t.layout(),
+      t.device(),
+      std::nullopt,
+      memory_format);
+}
+
+inline at::Tensor create_empty_from(
+    const at::Tensor& t,
+    c10::ScalarType dtype,
+    c10::MemoryFormat memory_format) {
+  return at::detail::empty_cpu(
+      {0}, dtype, t.layout(), t.device(), std::nullopt, memory_format);
+}
+
+inline bool checkResizedDataPtr(at::Tensor& t) {
+  auto const prev_data_ptr = t.data_ptr();
+  t.resize_({0});
+  return prev_data_ptr == t.data_ptr();
+}
+
+inline void fastResizeToZero(at::Tensor& t) {
+  t.unsafeGetTensorImpl()->set_sizes_contiguous({0});
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(checkResizedDataPtr(t));
+}
+
+// check if an op has an out variant registered in Static Runtime
+bool opIsRegistered(const c10::Symbol& op_name);
+// check if Static Runtime can run an op natively.
+// prim ops that are implemented directly in the jit interpreter are implemented
+// as native ops in Static Runtime
+bool nativeOpIsRegistered(const c10::Symbol& op_name);
+
+bool canReuseInputsOutputs(
+    Node* n,
+    const c10::FastMap<Node*, bool>& node_has_out_variant);
+bool isOptimizableContainerType(
+    Node* n,
+    const c10::FastMap<Node*, bool>& node_has_out_variant);
+
+SROperator getOutOfPlaceOperation(Node* n);
+SROperator getNativeOperation(Node* n);
+
+bool hasVarArgs(Node* n);
+
+inline std::string PrintNode(const Node* node) {
+  std::ostringstream ss;
+  node->print(ss, 0, nullptr, false);
+  return ss.str();
+}
+
+inline void LogAndDumpSchema(const Node* node) {
+  VLOG(1) << "Found schema mismatch for: " << node->schema();
+}
+
+inline bool sr_schema_check(torch::jit::Node* /*unused*/) {
+  return true;
+}
+
+template <typename Schema, typename... Schemas>
+bool sr_schema_check(
+    torch::jit::Node* node,
+    Schema&& first,
+    Schemas&&... rest) {
+  auto is_match = node->matches(first) || sr_schema_check(node, rest...);
+  if (!is_match) {
+    torch::jit::LogAndDumpSchema(node);
+  }
+  return is_match;
+}
+
+bool sr_schema_check_kind(torch::jit::Node* node, c10::Symbol node_kind);
+} // namespace torch::jit
+
+C10_DECLARE_bool(static_runtime_enable_fast_math);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/passes.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/passes.h
new file mode 100644
index 0000000000000000000000000000000000000000..2169e4775263a39000b236828336d98ebb18f271
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/passes.h
@@ -0,0 +1,96 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API void FuseInferenceOpsForSparseNN(
+    std::shared_ptr<torch::jit::Graph>& graph);
+
+TORCH_API void EliminateTrivialEquallySplit(
+    std::shared_ptr<torch::jit::Graph>& graph);
+
+TORCH_API void FuseListUnpack(std::shared_ptr<torch::jit::Graph>& graph);
+
+// If outputs_are_immutable is set to false, don't replace the view ops that
+// produce aliases of graph outputs with the copy version.
+TORCH_API void ReplaceWithCopy(
+    std::shared_ptr<torch::jit::Graph>& graph,
+    bool outputs_are_immutable = true);
+
+TORCH_API void ReplacePermuteWithCopy(
+    std::shared_ptr<torch::jit::Graph>& graph,
+    bool outputs_are_immutable = true);
+
+TORCH_API void ReplaceWithMaybeCopy(
+    std::shared_ptr<torch::jit::Graph>& graph,
+    bool outputs_are_immutable = true);
+
+TORCH_API void RemoveImmutableInputDictLookups(
+    std::shared_ptr<torch::jit::Graph>& graph);
+
+TORCH_API bool graphHasOp(std::shared_ptr<Graph>& graph, const char* op_name);
+
+TORCH_API bool forwardHasOp(const Module& module, const char* op_name);
+
+TORCH_API void FuseSignLog1P(std::shared_ptr<Graph>& graph);
+
+TORCH_API void UseVariadicTupleUnpack(const std::shared_ptr<Graph>& graph);
+
+// c10::Symbol::fromQualString is a bit long to type everywhere, and
+// we can't use a `using` statement since it's a static class function.
+inline c10::Symbol fromQualString(const std::string& qual_string) {
+  return c10::Symbol::fromQualString(qual_string);
+}
+
+// [Create owned refs for special values]
+// StaticRuntimeBlockRunner moves its outputs to the return value at the end of
+// run_impl. However, there's a corner case where this can cause problems. If
+// we return a constant, then the only reference in the constants_ array can
+// be destroyed by this move.
+// We could add special logic to handle this in run_impl. But since this is a
+// relatively rare corner case, it's simpler to just add an op that does nothing
+// but create an owned reference to its input. This owned reference can be
+// safely moved out of StaticRuntimeBlockRunner. Note that for scalars,
+// this actually does a copy.
+// Note that we have to do the same thing if we are returning a value from an
+// outer scope in a sub-block.
+TORCH_API void CreateOwnedRefsForSpecialValues(Graph& graph);
+
+// [Force non-empty outputs]
+// It is technically possible for sub-blocks to not return anything. This is
+// problematic for StaticRuntimeBlockRunner because it assumes that at least one
+// output is being returned. Rather than slowing down SR with special logic for
+// this corner case, we simply force blocks that return nothing to return None.
+TORCH_API void ForceNonEmptyOutputs(Graph& graph);
+
+TORCH_API void UseVariadicGroupedAccessor(const std::shared_ptr<Graph>& graph);
+
+TORCH_API void EliminateExtraPermuteOps(std::shared_ptr<Graph>& graph);
+
+TORCH_API void EliminateNoOpSlice(std::shared_ptr<Graph>& graph);
+
+TORCH_API void UseSplitAndSqueeze(std::shared_ptr<Graph>& graph);
+
+// [Remove unnecessary outputs]]
+// Removes outputs to reduce compute when it is not used later in the graph.
+// Currently used to remove the max_indices output of embedding_bag, which
+// isn't necessary to compute the main output.
+TORCH_API void RemoveUnnecessaryOutputs(std::shared_ptr<Graph>& graph);
+
+TORCH_API void RemoveUnnecessaryEmbeddingBagOutputs(
+    std::shared_ptr<Graph>& graph);
+
+TORCH_API void FuseClampNaNToNum(std::shared_ptr<Graph>& graph);
+
+TORCH_API void UseInPlaceGetRealInputsFromOptionalInputsV2(
+    std::shared_ptr<Graph>& graph);
+
+TORCH_API void PrepackWeights(std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+C10_DECLARE_bool(enable_clip_ranges_gather_fusions);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/processed_node_wrapper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/processed_node_wrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..634dc574e72f47fa1ea683a7eb37145f0f60f5ff
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/processed_node_wrapper.h
@@ -0,0 +1,216 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <torch/csrc/jit/runtime/static/impl.h>
+
+namespace torch::jit {
+
+// The following class facilitates code reuse between ProcessedNodeInputWrapper
+// and ProcessedNodeOutputWrapper via CRTP
+template <typename DerivedWrapper>
+class ProcessedNodeWrapperBase {
+ public:
+  class ProcessedNodeWrapperBaseIter {
+   public:
+    using iterator_category = std::forward_iterator_tag;
+    using value_type = at::Tensor;
+    using difference_type = size_t;
+    using pointer = const at::Tensor*;
+    using reference = const at::Tensor&;
+
+    ProcessedNodeWrapperBaseIter() = default;
+
+    ProcessedNodeWrapperBaseIter(
+        const DerivedWrapper* container,
+        size_t start_idx)
+        : container_(container), idx_(start_idx) {}
+
+    ProcessedNodeWrapperBaseIter& operator++() {
+      TORCH_DCHECK_NE(idx_, container_->size());
+      ++idx_;
+      return *this;
+    }
+
+    ProcessedNodeWrapperBaseIter operator++(int) {
+      ProcessedNodeWrapperBaseIter old = *this;
+      ++(*this);
+      return old;
+    }
+
+    reference operator*() const {
+      TORCH_CHECK(container_ != nullptr);
+      return (*container_)[idx_];
+    }
+
+    pointer operator->() const {
+      TORCH_CHECK(container_ != nullptr);
+      return &(*container_)[idx_];
+    }
+
+    friend bool operator==(
+        ProcessedNodeWrapperBaseIter lhs,
+        ProcessedNodeWrapperBaseIter rhs) {
+      TORCH_DCHECK_EQ(lhs.container_, rhs.container_);
+      return lhs.idx_ == rhs.idx_;
+    }
+
+    friend bool operator!=(
+        ProcessedNodeWrapperBaseIter lhs,
+        ProcessedNodeWrapperBaseIter rhs) {
+      return !(lhs == rhs);
+    }
+
+   private:
+    const DerivedWrapper* container_ = nullptr;
+    size_t idx_ = 0;
+  };
+
+  // NB: to mimic the behavior of at::ArrayRef, both iterators are
+  // the const version.
+  using iterator = ProcessedNodeWrapperBaseIter;
+  using const_iterator = ProcessedNodeWrapperBaseIter;
+  using size_type = size_t;
+  using value_type = at::Tensor;
+
+  explicit ProcessedNodeWrapperBase(ProcessedNode& pnode) : pnode_(pnode) {}
+
+  iterator begin() {
+    return ProcessedNodeWrapperBaseIter(static_cast<DerivedWrapper*>(this), 0);
+  }
+  iterator end() {
+    return ProcessedNodeWrapperBaseIter(
+        static_cast<DerivedWrapper*>(this),
+        static_cast<DerivedWrapper*>(this)->size());
+  }
+
+  const_iterator begin() const {
+    return ProcessedNodeWrapperBaseIter(
+        static_cast<const DerivedWrapper*>(this), 0);
+  }
+  const_iterator end() const {
+    return ProcessedNodeWrapperBaseIter(
+        static_cast<const DerivedWrapper*>(this),
+        static_cast<const DerivedWrapper*>(this)->size());
+  }
+
+  const_iterator cbegin() const {
+    return ProcessedNodeWrapperBaseIter(
+        static_cast<const DerivedWrapper*>(this), 0);
+  }
+  const_iterator cend() const {
+    return ProcessedNodeWrapperBaseIter(
+        static_cast<const DerivedWrapper*>(this),
+        static_cast<const DerivedWrapper*>(this)->size());
+  }
+
+  bool empty() const {
+    return static_cast<const DerivedWrapper*>(this)->size() == 0;
+  }
+
+ protected:
+  ProcessedNode& pnode_;
+};
+
+// A ProcessedNodeWrapperBase lets us use ProcessedNode directly in a context
+// where a container of IValues is expected. This trick is handy for avoiding
+// refcount bumps in perf-sensitive native ops. For example, suppose we have an
+// op that takes a list of tensors as an argument and we've turned the op into a
+// variadic variant in static runtime. To use the PyTorch library implementation
+// of the op, we would have to pack the variadic arguments into a list:
+//   std::vector<Tensor> tensor_list;
+//   tensor_list.reserve(pnode->num_outputs());
+//   for (const auto i : c10::irange(pnode->num_inputs())
+//     tensor_list.push_back(pnode->Input(i).toTensor());
+//   op_impl(tensor_list);
+// Using ProcessedNodeWrapperBase, we can avoid this round of refcount bumps.
+// All we need to do is turn `op_impl` into a template and pass it
+// ProcessedNodeInputWrapper(*pnode)!
+class ProcessedNodeInputWrapper
+    : public ProcessedNodeWrapperBase<ProcessedNodeInputWrapper> {
+ public:
+  // The last `back_elements_ignored` elements are not considered.
+  // Same for the first `front_elements_ignored` elements.
+  // This is useful for ops where
+  // only the first N elements are tensors (N < inputs.size()).
+  // For instance, the last argument to VarStack is an integer dimension.
+  explicit ProcessedNodeInputWrapper(
+      ProcessedNode& pnode,
+      size_t front_elements_ignored = 0,
+      size_t back_elements_ignored = 1)
+      : ProcessedNodeWrapperBase<ProcessedNodeInputWrapper>(pnode),
+        front_elements_ignored_(front_elements_ignored),
+        back_elements_ignored_(back_elements_ignored) {
+    TORCH_CHECK(front_elements_ignored_ <= pnode_.num_inputs());
+    TORCH_CHECK(
+        back_elements_ignored_ <=
+        pnode_.num_inputs() - front_elements_ignored_);
+  }
+
+  size_t size() const {
+    return pnode_.num_inputs() - back_elements_ignored_ -
+        front_elements_ignored_;
+  }
+
+  const at::Tensor& operator[](size_t idx) const {
+    TORCH_CHECK(idx < size());
+    return pnode_.Input(front_elements_ignored_ + idx).toTensor();
+  }
+
+  const at::Tensor& front() const {
+    TORCH_CHECK(
+        !empty(),
+        "Attempted to access front() of empty ProcessedNodeInputWrapper");
+    return pnode_.Input(front_elements_ignored_).toTensor();
+  }
+
+  const at::Tensor& back() const {
+    TORCH_CHECK(
+        !empty(),
+        "Attempted to access back() of empty ProcessedNodeInputWrapper");
+    return pnode_.Input(pnode_.num_inputs() - back_elements_ignored_ - 1)
+        .toTensor();
+  }
+
+ private:
+  size_t front_elements_ignored_;
+  size_t back_elements_ignored_;
+};
+
+// Similar to ProcessedNodeInputWrapper, but wraps outputs and allows for
+// writing.
+class ProcessedNodeOutputWrapper
+    : public ProcessedNodeWrapperBase<ProcessedNodeOutputWrapper> {
+ public:
+  using ProcessedNodeWrapperBase<
+      ProcessedNodeOutputWrapper>::ProcessedNodeWrapperBase;
+
+  size_t size() const {
+    return pnode_.num_outputs();
+  }
+
+  at::Tensor& operator[](size_t idx) const {
+    TORCH_CHECK(idx < size());
+    return pnode_.Output(idx).toTensor();
+  }
+
+  at::Tensor& front() const {
+    TORCH_CHECK(
+        !empty(),
+        "Attempted to access front() of empty ProcessedNodeOutputWrapper");
+    return pnode_.Output(0).toTensor();
+  }
+
+  at::Tensor& back() const {
+    TORCH_CHECK(
+        !empty(),
+        "Attempted to access back() of empty ProcessedNodeOutputWrapper");
+    return pnode_.Output(size() - 1).toTensor();
+  }
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/static_method.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/static_method.h
new file mode 100644
index 0000000000000000000000000000000000000000..26d9043555d6575b127dac18735da727036d8a6f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/static_method.h
@@ -0,0 +1,53 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/api/include/torch/imethod.h>
+#include <torch/csrc/jit/runtime/static/impl.h>
+
+namespace torch::jit {
+
+class StaticMethod : public torch::IMethod {
+ public:
+  StaticMethod(
+      std::shared_ptr<StaticModule> static_module,
+      std::string method_name)
+      : static_module_(std::move(static_module)),
+        method_name_(std::move(method_name)) {
+    TORCH_CHECK(static_module_);
+  }
+
+  c10::IValue operator()(
+      std::vector<IValue> args,
+      const IValueMap& kwargs = IValueMap()) const override {
+    return (*static_module_)(std::move(args), kwargs);
+  }
+
+  const std::string& name() const override {
+    return method_name_;
+  }
+
+ protected:
+  void setArgumentNames(
+      std::vector<std::string>& argument_names_out) const override {
+    const auto& schema = static_module_->schema();
+    CAFFE_ENFORCE(schema.has_value());
+    const auto& arguments = schema->arguments();
+    argument_names_out.clear();
+    argument_names_out.reserve(arguments.size());
+    std::transform(
+        arguments.begin(),
+        arguments.end(),
+        std::back_inserter(argument_names_out),
+        [](const c10::Argument& arg) -> std::string { return arg.name(); });
+  }
+
+ private:
+  std::shared_ptr<StaticModule> static_module_;
+  std::string method_name_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/te_wrapper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/te_wrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..b359ca7db890fbae81cad1a1bd34ca7cb1977d8a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/static/te_wrapper.h
@@ -0,0 +1,49 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_simplifier.h>
+#include <torch/csrc/jit/tensorexpr/llvm_codegen.h>
+#include <torch/csrc/jit/tensorexpr/loopnest.h>
+
+namespace torch::jit {
+
+class TEWrapper {
+ public:
+  TEWrapper() = default;
+  void call(const std::vector<void*>& args);
+
+  template <typename ExpectedType>
+  bool checkInput(const at::Tensor& t) {
+#ifdef TORCH_ENABLE_LLVM
+    return t.is_contiguous() && t.dtype().Match<ExpectedType>();
+#else
+    return false;
+#endif
+  }
+
+#ifdef TORCH_ENABLE_LLVM
+  void update(std::unique_ptr<tensorexpr::LLVMCodeGen>&& cg_);
+#endif
+
+ private:
+#ifdef TORCH_ENABLE_LLVM
+  std::unique_ptr<tensorexpr::LLVMCodeGen> cg;
+#endif
+};
+
+std::shared_ptr<TEWrapper> createDiv();
+std::shared_ptr<TEWrapper> createLogit();
+std::shared_ptr<TEWrapper> createRelu();
+std::shared_ptr<TEWrapper> createTanh();
+std::shared_ptr<TEWrapper> createSigmoid();
+std::shared_ptr<TEWrapper> createSignedLog1p();
+std::shared_ptr<TEWrapper> createClamp();
+std::shared_ptr<TEWrapper> createClampNanToNum();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_script.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_script.h
new file mode 100644
index 0000000000000000000000000000000000000000..48732494fef63473ce22ca94d4015a3c18dac5fb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_script.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+// This file is temporary until native_functions.yaml and derivatives.yaml are
+// merged. Ideally this should all go into native_functions.yaml
+
+#include <c10/util/StringUtil.h>
+#include <torch/csrc/jit/api/module.h>
+#include <optional>
+
+namespace torch::jit {
+struct GradientPair {
+  std::shared_ptr<Graph> forward;
+  std::shared_ptr<Graph> backward;
+};
+
+TORCH_API std::optional<GradientPair> gradientInfoForSchema(
+    const FunctionSchema& schema);
+TORCH_API bool hasGradientInfoForSchema(const FunctionSchema& schema);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_shape_registry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_shape_registry.h
new file mode 100644
index 0000000000000000000000000000000000000000..9f2ea5b6ae517a455878b178ceea1d3aaf2f86f2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_shape_registry.h
@@ -0,0 +1,74 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+// This file is temporary until native_functions.yaml and derivatives.yaml are
+// merged. Ideally this should all go into native_functions.yaml
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+/*
+ADDING A NEW SHAPE GRAPH:
+- For one node schema, there is one corresponding registered shape compute
+graph. The schema of the graph should be the same except for Tensor arguments.
+For every Tensor input in operator schema, there should be a List[int]
+corresponding to that Tensor's shape. For example: "aten::linear(Tensor input,
+Tensor weight, Tensor? bias=None) -> Tensor" ==> def linear(input: List[int],
+weight: List[int], bias: Optional[List[int]])
+
+Additionally, arguments which are unused at the end of the schema may be left
+off. This allows sharing a single graph for multiple function schemas, such as
+unary operators with different trailing arguments that do not affect the output
+shape.
+
+The shape graph should return a new, unaliased List[int] (or tuple of lists for
+multiple returns) and should not modify any input lists. This allows the shape
+graphs to be composed and executed.
+
+The shape analysis (particularly for non-complete, or symbolic shapes) works by
+partially evaluating the JIT IR. It may be possible for a Graph to be registered
+that we cannot currently partially evaluate. If this happens, please file an
+issue. There are lints registered to avoid particular known patterns (continue
+or break or early return in a loop). Those may be improved in the future, please
+file an issue if necessary.
+
+To debug (and write initially) the recommended flow is to define these functions
+in python and iterate there. Functions should be added to
+torch/jit/_shape_functions.
+
+To test operators, the preferred flow is through OpInfos, with
+`assert_jit_shape_analysis=True`. If this is not feasible, you can look at tests
+in `test_symbolic_shape_analysis.py` such as `test_adaptive_avg_pool2d`.
+
+Operators which take in a list of tensors, such as concat, are not yet
+supported. Concat has been special cased and could be generalized as needed.
+Please file an issue.
+*/
+
+struct BoundedShapeGraphs {
+  std::shared_ptr<Graph> lower_bound;
+  std::shared_ptr<Graph> upper_bound;
+};
+
+TORCH_API void RegisterShapeComputeGraphForSchema(
+    const FunctionSchema& schema,
+    const std::shared_ptr<Graph>& g);
+
+TORCH_API std::optional<std::shared_ptr<Graph>> shapeComputeGraphForSchema(
+    const FunctionSchema& schema);
+
+TORCH_API std::optional<BoundedShapeGraphs> boundedGraphsForSchema(
+    const FunctionSchema& schema);
+
+TORCH_API std::vector<const FunctionSchema*> RegisteredShapeComputeSchemas();
+
+TORCH_API void LintShapeComputeGraph(
+    const FunctionSchema* schema,
+    const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_shape_registry_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_shape_registry_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..324cfc5bf4f25f98004fe54a08dced86cfc318c9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/symbolic_shape_registry_util.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+// This file is temporary until native_functions.yaml and derivatives.yaml are
+// merged. Ideally this should all go into native_functions.yaml
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API const OperatorMap<std::string>& get_tensorexpr_elementwise_set();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/vararg_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/vararg_functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..45c5c94e9b37b9b84e5a79822258d163f2262fcf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/vararg_functions.h
@@ -0,0 +1,46 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/List.h>
+#include <ATen/core/functional.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/stack.h>
+
+namespace torch::jit {
+
+void tupleUnpack(Stack& stack);
+
+void format(Stack& stack, size_t num_inputs);
+
+void einsum(Stack& stack, size_t num_inputs);
+
+void percentFormat(Stack& stack, size_t num_inputs);
+
+void listUnpack(Stack& stack, size_t num_outputs);
+
+void tupleConstruct(Stack& stack, size_t num_inputs);
+
+void namedTupleConstruct(Stack& stack, c10::TypePtr type, size_t num_inputs);
+
+void listConstruct(Stack& stack, const c10::Type& list_type, size_t num_inputs);
+
+void dictConstruct(Stack& stack, const c10::Type& type, size_t num_inputs);
+
+// as weak_ref will create a Object with a non-owning CompilationUnit reference,
+// for use as a constant in the Graph to avoid a reference cycle
+void createObject(
+    Stack& stack,
+    const at::ClassTypePtr& type,
+    bool as_weak_ref = false);
+
+void isinstance(Stack& stack, at::ArrayRef<at::TypePtr> types);
+
+void tupleSlice(Stack& stack, size_t begin, size_t end);
+
+void dequantize(Stack& stack);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/variable_tensor_list.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/variable_tensor_list.h
new file mode 100644
index 0000000000000000000000000000000000000000..335992c91015ae68f444be5f4f98a3173fe9a317
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/runtime/variable_tensor_list.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/Tensor.h>
+
+namespace torch::jit {
+
+// a wrapper to mark places where we expect all the at::Tensors to be
+// variables
+struct variable_tensor_list : public std::vector<at::Tensor> {
+  variable_tensor_list() = default;
+  template <class InputIt>
+  variable_tensor_list(InputIt first, InputIt last)
+      : std::vector<at::Tensor>(first, last) {}
+  explicit variable_tensor_list(std::vector<at::Tensor>&& tensor)
+      : std::vector<at::Tensor>(std::move(tensor)) {}
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/callstack_debug_info_serialization.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/callstack_debug_info_serialization.h
new file mode 100644
index 0000000000000000000000000000000000000000..afccfce9cdfb397f7f57c33780a8c97d8a31890c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/callstack_debug_info_serialization.h
@@ -0,0 +1,94 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/Allocator.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/ir/scope.h>
+
+#include <ATen/core/ivalue.h>
+
+#include <vector>
+
+#include <c10/util/flat_hash_map.h>
+
+namespace c10 {
+struct IValue;
+}
+
+namespace torch::jit {
+
+class Pickler;
+class InlinedCallStackSerializer {
+ public:
+  // Serialize InlinedCallStack as
+  // SerializedInlinedCallStack =
+  // [module_info, source range tag, SerializedInlinedCallStack]
+  // module_info = [ClassType.qualifiedName, instance_name]
+  // source_range_tag = unique source range id
+  c10::IValue serialize(
+      const InlinedCallStackPtr& cs_ptr,
+      const SourceRangeTagMap& source_range_tags);
+
+ private:
+  // module_info = [ClassType.qualifiedName, instance_name]
+  c10::IValue serialize_module_instance_info(
+      const std::optional<ModuleInstanceInfo>& m);
+
+  // This caches serialized inlined callstack ptr, since many
+  // InlinedCallStackPtr can refer to the same one.
+  ska::flat_hash_map<InlinedCallStackPtr, c10::IValue>
+      serialized_inlined_callstack_;
+  // This caches serialized module instance info.
+  // There might be many nodes that are part of the same
+  // parent, grandparent etc. module.
+  ska::flat_hash_map<std::string, c10::IValue> serialized_module_instance_info_;
+};
+
+class TORCH_API CallStackDebugInfoPickler {
+ public:
+  CallStackDebugInfoPickler() = default;
+
+  std::vector<char> pickle(
+      const std::unordered_map<int64_t, DebugInfoTuple>& callstack_ptrs,
+      const SourceRangeTagMap& source_range_tags);
+
+ private:
+  InlinedCallStackSerializer css_;
+};
+
+class InlinedCallStackDeserializer {
+ public:
+  InlinedCallStackPtr deserialize(
+      const c10::IValue& iv,
+      const ska::flat_hash_map<int64_t, SourceRange>& source_range_map,
+      const std::shared_ptr<CompilationUnit>& cu);
+
+ private:
+  std::optional<ModuleInstanceInfo> deserialize_module_instance_info(
+      const c10::IValue& iv,
+      const std::shared_ptr<CompilationUnit>& cu);
+
+  ska::
+      flat_hash_map<c10::intrusive_ptr<c10::ivalue::Tuple>, InlinedCallStackPtr>
+          cached_inlined_callstacks_;
+  ska::flat_hash_map<c10::intrusive_ptr<c10::ivalue::Tuple>, ModuleInstanceInfo>
+      cached_module_instance_info_;
+};
+
+class TORCH_API CallStackDebugInfoUnpickler {
+ public:
+  ska::flat_hash_map<int64_t, DebugInfoTuple> unpickle(
+      const at::DataPtr& data,
+      size_t size,
+      const ska::flat_hash_map<int64_t, SourceRange>& source_range_map,
+      const std::shared_ptr<CompilationUnit>& cu);
+
+ private:
+  InlinedCallStackDeserializer csds_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/export.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/export.h
new file mode 100644
index 0000000000000000000000000000000000000000..d8778cb81b160e3f53680843dbe2b89fea97e0a6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/export.h
@@ -0,0 +1,283 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <caffe2/serialize/inline_container.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/serialization/export_bytecode.h>
+#include <torch/csrc/jit/serialization/flatbuffer_serializer.h>
+#include <torch/csrc/jit/serialization/python_print.h>
+#include <torch/csrc/jit/serialization/storage_context.h>
+#include <torch/csrc/jit/serialization/type_name_uniquer.h>
+#include <torch/csrc/onnx/onnx.h>
+#include <ostream>
+
+namespace ONNX_NAMESPACE {
+class ModelProto;
+}
+
+namespace torch::jit {
+
+// This map is used to keep track of parameters that should be exported
+// externally. When `defer_weight_export` is true, the returned map contains
+// kv pairs that map {external reference name} -> {at::Tensor to be exported}.
+// It is the responsibility of the caller to export these appropriately.
+//
+// For example, when exporting to a zip archive, the caller may write out files
+// for each entry in the export map, with the filename being the key and the
+// file contents being the raw tensor data.
+using RawDataExportMap = std::unordered_map<std::string, at::Tensor>;
+
+using SymbolDimMap = std::map<c10::ShapeSymbol, std::string>;
+using DimSymbolMap = std::map<std::string, c10::ShapeSymbol>;
+
+using NodeNameMap = std::unordered_map<const Node*, std::string>;
+
+// Used for modularized export settling function and node attributes.
+using NodeAttrNameMap = std::
+    unordered_map<const Node*, std::unordered_map<std::string, std::string>>;
+
+TORCH_API std::tuple<
+    std::shared_ptr<::ONNX_NAMESPACE::ModelProto>,
+    RawDataExportMap,
+    SymbolDimMap,
+    bool,
+    NodeNameMap>
+export_onnx(
+    const std::shared_ptr<Graph>& graph,
+    const std::map<std::string, at::Tensor>& initializers,
+    int64_t onnx_opset_version,
+    const std::unordered_map<
+        std::string,
+        std::unordered_map<int64_t, std::string>>& dynamic_axes,
+    bool defer_weight_export = false,
+    ::torch::onnx::OperatorExportTypes operator_export_type =
+        ::torch::onnx::OperatorExportTypes::ONNX,
+    bool strip_doc_string = true,
+    bool keep_initializers_as_inputs = true,
+    const std::map<std::string, int>& custom_opsets = {},
+    bool add_node_names = true,
+    bool use_external_data_format = false,
+    const std::string& onnx_file_path = std::string(),
+    const NodeAttrNameMap& node_attr_to_name = {});
+
+TORCH_API std::string serialize_model_proto_to_string(
+    const std::shared_ptr<::ONNX_NAMESPACE::ModelProto>& model_proto);
+
+TORCH_API void check_onnx_proto(const std::string& proto_string);
+
+// Serializer for both oldsyle and unified format TorchScript serialization
+class TORCH_API ScriptModuleSerializer {
+ public:
+  explicit ScriptModuleSerializer(
+      caffe2::serialize::PyTorchStreamWriter& export_writer)
+      : writer_(export_writer) {}
+
+  void writeFiles(const std::string& code_dir);
+  void serialize(
+      const Module& module,
+      const ExtraFilesMap& extra_files,
+      bool bytecode_format,
+      bool save_mobile_debug_info);
+  void serialize_unified_format(Module& module, uint64_t script_module_id);
+  SerializationStorageContext& storage_context();
+
+  ~ScriptModuleSerializer() = default;
+
+ private:
+  void convertNamedType(const c10::NamedTypePtr& class_type);
+  void convertTypes(const at::NamedTypePtr& root_type);
+  void writeExtraFiles(const Module& module, const ExtraFilesMap& extra_files);
+  void writeByteCode(const Module& module, bool save_mobile_debug_info);
+  void writeArchive(
+      const IValue& value,
+      const std::string& archive_name,
+      const std::string& archive_dir,
+      const std::string& tensor_dir,
+      bool use_storage_context = false,
+      bool skip_tensor_data = false);
+  void updateSourceRangeTags(const SourceRangeRecords& ranges);
+
+  caffe2::serialize::PyTorchStreamWriter& writer_;
+  std::vector<at::IValue> constant_table_;
+
+  std::unordered_set<c10::NamedTypePtr> converted_types_;
+  PrintDepsTable class_deps_;
+  TypeNameUniquer type_name_uniquer_;
+  // qualifier, e.g. '__torch__.Bar' -> PythonPrint for the file that will be
+  // created
+  OrderedDict<std::string, PythonPrint> file_streams_;
+  // Used to keep references of storages around during serialization to solve
+  // for ABA memory reuse problem hit when storages are created/destroyed
+  // during serialization process. Also used to coordinate sharing of storages
+  // between Script and eager modules in torch.package.
+  SerializationStorageContext storage_context_;
+
+  // Uniquely identifies a SourceRange in a model.
+  // SourceRanges are associated with Nodes of Graphs.
+  // However for mobile deployment we dont intend to ship
+  // full JIT with capabilities of reading code and constructing
+  // graphs.
+  // Instead we serialize the Code generated from graph of the methods.
+  // Code is serialized in bytecode format that contains instructions
+  // corresponding to the nodes of the graph. Since original graph is gone, the
+  // question is how do we identify where the ops, in serialized bytecode, come
+  // from in original model code. We do this in two parts.
+  // 1. Associate a unique tag to SourceRange.
+  // 2. Serialize this unique_tag.
+  //  2.1 Meaning save <byte_offset, source_range_tag, source range> instead of
+  //      <byte_offset, source range>
+  // 3. During serializing model for mobile, i.e. bytecode generation,
+  //    save unique tag of SourceRange corresponding to the Node.
+  // 4. During deserialization, read all the debug_pkl, to construct a map
+  //    of <unique_tag, SourceRange> and use tag saved with OPs in bytecode
+  //    to lookup the source range.
+  // Strictly speaking we will serialize InlinedCallStack directly, which
+  // contains SourceRange. This way we have access to entire callstack and not
+  // just source information about where the node is, since bytecode inlines the
+  // graph before saving it.
+  SourceRangeTagMap source_range_tags_;
+  int64_t current_source_range_tag_{0};
+};
+
+// For testing purposes
+TORCH_API std::string pretty_print_onnx(
+    const std::shared_ptr<Graph>& graph,
+    const std::map<std::string, at::Tensor>& initializers,
+    int64_t onnx_opset_version,
+    bool defer_weight_export,
+    ::torch::onnx::OperatorExportTypes operator_export_type =
+        ::torch::onnx::OperatorExportTypes::ONNX,
+    bool google_printer = false,
+    bool keep_initializers_as_inputs = true,
+    const std::map<std::string, int>& custom_opsets = {},
+    bool add_node_names = true);
+
+TORCH_API void ExportModule(
+    const Module& module,
+    std::ostream& out,
+    const ExtraFilesMap& metadata = ExtraFilesMap(),
+    bool bytecode_format = false,
+    bool save_mobile_debug_info = false,
+    bool use_flatbuffer = false);
+
+TORCH_API void ExportModule(
+    const Module& module,
+    const std::string& filename,
+    const ExtraFilesMap& metadata = ExtraFilesMap(),
+    bool bytecode_format = false,
+    bool save_mobile_debug_info = false,
+    bool use_flatbuffer = false);
+
+TORCH_API void ExportModule(
+    const Module& module,
+    const std::function<size_t(const void*, size_t)>& writer_func,
+    const ExtraFilesMap& metadata = ExtraFilesMap(),
+    bool bytecode_format = false,
+    bool save_mobile_debug_info = false,
+    bool use_flatbuffer = false);
+
+// Write the bytes of a pickle archive and the tensors referenced inside that
+// archive
+TORCH_API void writeArchiveAndTensors(
+    const std::string& archive_name,
+    const char* pickle_bytes,
+    size_t size,
+    const std::vector<at::Tensor>& tensors,
+    caffe2::serialize::PyTorchStreamWriter& out);
+
+// Surrounding system can install an additional hook to produce extra files
+// with metadata based on environment every time a module is serialized.
+using ExportModuleExtraFilesHook = std::function<ExtraFilesMap(const Module&)>;
+TORCH_API void SetExportModuleExtraFilesHook(ExportModuleExtraFilesHook hook);
+
+/**
+ * Generates new bytecode for a Script module and returns what the op list
+ * would be for a LiteScriptModule based off the current code base. If you
+ * have a LiteScriptModule and want to get the currently present
+ * list of ops call _export_operator_list instead.
+ */
+TORCH_API std::vector<std::string> export_opnames(const Module& m);
+
+struct TORCH_API BytecodeEmitMode {
+  static bool is_default_value_for_unspecified_arg_enabled();
+  static void set_default_value_for_unspecified_arg_enabled(bool enabled);
+
+  static bool is_default_args_before_out_args_enabled();
+  static void set_default_args_before_out_args_enabled(bool enabled);
+
+  static bool is_emit_promoted_ops_enabled();
+  static void set_default_emit_promoted_ops_enabled(bool enabled);
+};
+
+// RAII guard to switch the way JIT emits the bytecode for inputs.
+// default_value_for_unspecified_arg:
+// true: instruction of default argument values (like LOADC) is emitted.
+// false: instruction of default argument values are not emitted. Instead
+// they are fetched from operator schema.
+// default_args_before_out_args (to forward compatible support
+// operators allowing out arguments and default arguments):
+// true: the number of specified arguments will deserialized to (#all_args -
+// #default_args). false: the number of specified arguments will deserialized to
+// (#all_args).
+struct TORCH_API BytecodeEmitModeGuard {
+  BytecodeEmitModeGuard(
+      bool enable_default_value_for_unspecified_arg,
+      bool enable_default_args_before_out_args,
+      bool enable_emit_promoted_ops)
+      : prev_default_value_for_unspecified_arg_mode(
+            BytecodeEmitMode::is_default_value_for_unspecified_arg_enabled()),
+        prev_default_args_before_out_args(
+            BytecodeEmitMode::is_default_args_before_out_args_enabled()),
+        prev_default_emit_promoted_ops(
+            BytecodeEmitMode::is_emit_promoted_ops_enabled()) {
+    BytecodeEmitMode::set_default_value_for_unspecified_arg_enabled(
+        enable_default_value_for_unspecified_arg);
+    BytecodeEmitMode::set_default_args_before_out_args_enabled(
+        enable_default_args_before_out_args);
+    BytecodeEmitMode::set_default_emit_promoted_ops_enabled(
+        enable_emit_promoted_ops);
+  }
+  ~BytecodeEmitModeGuard() {
+    BytecodeEmitMode::set_default_value_for_unspecified_arg_enabled(
+        prev_default_value_for_unspecified_arg_mode);
+    BytecodeEmitMode::set_default_args_before_out_args_enabled(
+        prev_default_args_before_out_args);
+    BytecodeEmitMode::set_default_emit_promoted_ops_enabled(
+        prev_default_emit_promoted_ops);
+  }
+  bool prev_default_value_for_unspecified_arg_mode;
+  bool prev_default_args_before_out_args;
+  bool prev_default_emit_promoted_ops;
+};
+
+TORCH_API IValue to_tuple(std::vector<IValue> ivalues);
+TORCH_API IValue
+Table(const std::vector<std::pair<std::string, IValue>>& entries);
+
+// TODO remove these switches once interface call is rolled out.
+TORCH_API void enableMobileInterfaceCallExport();
+bool getMobileInterfaceCallExport();
+
+TORCH_API CompilationOptions getOptionsFromGlobal();
+
+TORCH_API void save_jit_module(
+    const Module& module,
+    const std::string& filename,
+    const ExtraFilesMap& extra_files = ExtraFilesMap());
+
+TORCH_API DetachedBuffer::UniqueDetachedBuffer save_jit_module_to_bytes(
+    const Module& module,
+    const ExtraFilesMap& extra_files = ExtraFilesMap());
+
+TORCH_API void save_jit_module_to_write_func(
+    const Module& module,
+    const ExtraFilesMap& extra_files,
+    bool save_mobile_debug_info,
+    const std::function<size_t(const void*, size_t)>& writer_func);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/export_bytecode.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/export_bytecode.h
new file mode 100644
index 0000000000000000000000000000000000000000..608e70c7b5f337bf0798888f1773018562cba00b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/export_bytecode.h
@@ -0,0 +1,46 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/function_schema.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/qualified_name.h>
+#include <torch/csrc/jit/backends/backend_debug_handler.h>
+#include <torch/csrc/jit/mobile/function.h>
+#include <torch/csrc/jit/mobile/module.h>
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <torch/csrc/jit/serialization/type_name_uniquer.h>
+
+namespace torch::jit {
+
+struct TORCH_API CompilationOptions {
+  bool incl_interface_call = false;
+  bool enable_default_value_for_unspecified_arg = false;
+  bool enable_default_args_before_out_args = true;
+  bool enable_emit_promoted_ops = true;
+  int model_version = caffe2::serialize::kProducedBytecodeVersion;
+};
+
+TORCH_API mobile::Module jitModuleToMobile(
+    const Module& module,
+    const CompilationOptions& options);
+
+mobile::Code compileGraphToMobileCode(
+    const std::string& name,
+    const std::shared_ptr<Graph>& graph,
+    const CompilationOptions& compilation_options,
+    BackendDebugInfoRecorder& debug_info_recorder);
+
+TORCH_API std::unique_ptr<mobile::Function> convertJitFunctionToMobileFunction(
+    const GraphFunction& function,
+    const CompilationOptions& options);
+
+TORCH_API IValue convertMobileFunctionToCodeTable(
+    const mobile::Function& func,
+    const CompilationOptions& compilation_options);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/flatbuffer_serializer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/flatbuffer_serializer.h
new file mode 100644
index 0000000000000000000000000000000000000000..cfd7d513a0c578bc7dd102ec2c98db8c9a7e582c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/flatbuffer_serializer.h
@@ -0,0 +1,97 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <ATen/core/ivalue.h>
+#include <c10/macros/Macros.h>
+#include <torch/csrc/jit/mobile/module.h>
+
+/**
+ * Defines the public API for serializing mobile modules to flatbuffer.
+ * Note that this header must not include or depend on flatbuffer-defined
+ * types, to avoid leaking those details to PyTorch clients.
+ */
+
+namespace torch::jit {
+
+/// Maps file names to file contents.
+using ExtraFilesMap = std::unordered_map<std::string, std::string>;
+
+/**
+ * Represents a span of data. Typically owned by a UniqueDetachedBuffer.
+ */
+class TORCH_API DetachedBuffer final {
+ public:
+  /// Creates a new DetachedBuffer with an optional data owner. This interface
+  /// is provided to let users create objects of this type for testing.
+  DetachedBuffer(void* data, size_t size, void* internal_data_owner = nullptr)
+      : data_(data), size_(size), data_owner_(internal_data_owner) {}
+
+  /// Returns a pointer to the data.
+  [[nodiscard]] void* data() {
+    return data_;
+  }
+  /// Returns a pointer to the data.
+  [[nodiscard]] const void* data() const {
+    return data_;
+  }
+  /// Returns the size of the data, in bytes.
+  [[nodiscard]] size_t size() const {
+    return size_;
+  }
+
+  /// Wrapper type that typically owns data_owner_.
+  using UniqueDetachedBuffer =
+      std::unique_ptr<DetachedBuffer, std::function<void(DetachedBuffer*)>>;
+
+ private:
+  /// Deletes the owner, if present, and the buf itself.
+  /// Note: we could have provided a movable type with a destructor that did
+  /// this work, but the unique wrapper was easier in practice.
+  static void destroy(DetachedBuffer* buf);
+
+  /// Provides access to destroy() for implementation and testing.
+  friend struct DetachedBufferFriend;
+  friend struct DetachedBufferTestingFriend;
+
+  /// Pointer to the data. Not owned by this class.
+  void* data_;
+  /// The size of `data_`, in bytes.
+  size_t size_;
+  /// Opaque pointer to the underlying owner of `data_`. This class
+  /// (DetachedBuffer) does not own the owner or the data. It will typically be
+  /// owned by a UniqueDetachedBuffer that knows how to delete the owner along
+  /// with this class.
+  void* data_owner_;
+};
+
+TORCH_API void save_mobile_module(
+    const mobile::Module& module,
+    const std::string& filename,
+    const ExtraFilesMap& extra_files = ExtraFilesMap(),
+    const ExtraFilesMap& jit_sources = ExtraFilesMap(),
+    const std::vector<IValue>& jit_constants = {});
+
+TORCH_API DetachedBuffer::UniqueDetachedBuffer save_mobile_module_to_bytes(
+    const mobile::Module& module,
+    const ExtraFilesMap& extra_files = ExtraFilesMap(),
+    const ExtraFilesMap& jit_sources = ExtraFilesMap(),
+    const std::vector<IValue>& jit_constants = {});
+
+TORCH_API void save_mobile_module_to_func(
+    const mobile::Module& module,
+    const std::function<size_t(const void*, size_t)>& writer_func);
+
+// TODO(qihan): delete
+TORCH_API bool register_flatbuffer_serializer();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/flatbuffer_serializer_jit.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/flatbuffer_serializer_jit.h
new file mode 100644
index 0000000000000000000000000000000000000000..fd11071344bc19fca9b1317d399d0c071928ac25
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/flatbuffer_serializer_jit.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/serialization/flatbuffer_serializer.h>
+
+namespace torch::jit {
+
+TORCH_API bool register_flatbuffer_all();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import.h
new file mode 100644
index 0000000000000000000000000000000000000000..48c18b7905f4c70fb2f765e151dd2ce6cd449fc2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import.h
@@ -0,0 +1,152 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <caffe2/serialize/inline_container.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+
+#include <istream>
+
+namespace caffe2::serialize {
+class ReadAdapterInterface;
+} // namespace caffe2::serialize
+
+namespace torch::jit {
+
+class DeserializationStorageContext;
+
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    const std::string& filename,
+    std::optional<c10::Device> device = std::nullopt,
+    bool load_debug_files = true);
+
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    std::istream& in,
+    std::optional<c10::Device> device = std::nullopt,
+    bool load_debug_files = true);
+
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    std::unique_ptr<caffe2::serialize::ReadAdapterInterface> rai,
+    std::optional<c10::Device> device = std::nullopt,
+    bool load_debug_files = true);
+
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    const std::string& filename,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    bool load_debug_files = true,
+    bool restore_shapes = false);
+
+// For reading unified serialization format from torch.Package
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    std::shared_ptr<caffe2::serialize::PyTorchStreamReader> reader,
+    std::shared_ptr<torch::jit::DeserializationStorageContext> storage_context,
+    std::optional<at::Device> device,
+    const std::string& ts_id /* torchscript identifier inside package */);
+
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    std::istream& in,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    bool load_debug_files = true,
+    bool restore_shapes = false);
+
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    std::unique_ptr<caffe2::serialize::ReadAdapterInterface> rai,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    bool load_debug_files = true);
+
+TORCH_API Module import_ir_module(
+    std::shared_ptr<CompilationUnit> cu,
+    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    bool load_debug_files = true);
+
+/// Loads a serialized `Module` from the given `istream`.
+///
+/// The istream must contain a serialized `Module`, exported via
+/// `torch::jit::ExportModule` in C++.
+TORCH_API Module load(
+    std::istream& in,
+    std::optional<c10::Device> device = std::nullopt,
+    bool load_debug_files = true);
+
+TORCH_API Module load(
+    std::istream& in,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    bool load_debug_files = true);
+
+/// Loads a serialized `Module` from the given `filename`.
+///
+/// The file stored at the location given in `filename` must contain a
+/// serialized `Module`, exported either via `ScriptModule.save()` in
+/// Python or `torch::jit::ExportModule` in C++.
+TORCH_API Module load(
+    const std::string& filename,
+    std::optional<c10::Device> device = std::nullopt,
+    bool load_debug_files = true);
+
+TORCH_API Module load(
+    const std::string& filename,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    bool load_debug_files = true);
+
+/// Loads a serialized `Module` from the given shared_ptr `rai`.
+///
+/// The reader adapter, which is for customized input stream, must contain a
+/// serialized `Module`, exported either via `ScriptModule.save()` in
+/// Python or `torch::jit::ExportModule` in C++.
+TORCH_API Module load(
+    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai,
+    std::optional<c10::Device> device = std::nullopt,
+    bool load_debug_files = true);
+
+TORCH_API Module load(
+    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai,
+    std::optional<c10::Device> device,
+    ExtraFilesMap& extra_files,
+    bool load_debug_files = true);
+
+TORCH_API Module jitModuleFromSourceAndConstants(
+    const IValue& ivalue,
+    const ExtraFilesMap& source,
+    const std::vector<IValue>& constants,
+    int32_t version);
+
+TORCH_API Module parse_and_initialize_jit_module(
+    const std::shared_ptr<char>& data,
+    size_t size,
+    ExtraFilesMap& extra_files,
+    std::optional<at::Device> device = std::nullopt);
+
+TORCH_API Module load_jit_module_from_file(
+    const std::string& filename,
+    ExtraFilesMap& extra_files,
+    std::optional<at::Device> device = std::nullopt);
+
+TORCH_API Module load_jit_module_from_stream(
+    std::istream& in,
+    ExtraFilesMap& extra_files,
+    std::optional<at::Device> device = std::nullopt);
+
+TORCH_API c10::intrusive_ptr<c10::ivalue::Object> ObjLoaderFunc(
+    const at::StrongTypePtr& type,
+    IValue input);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_constants.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_constants.h
new file mode 100644
index 0000000000000000000000000000000000000000..2722a6533036cd74343c12898c9e268c17329856
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_constants.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <cstddef>
+
+namespace torch::jit {
+constexpr size_t BYTECODE_INDEX_INSTRUCTION = 0;
+constexpr size_t BYTECODE_INDEX_OPERATOR = 1;
+constexpr size_t BYTECODE_INDEX_CONSTANT = 2;
+constexpr size_t BYTECODE_INDEX_TYPE = 3;
+constexpr size_t BYTECODE_INDEX_REGISTER_SIZE = 4;
+
+constexpr size_t BYTECODE_INDEX_SCHEMA_ARGUMENTS = 0;
+constexpr size_t BYTECODE_INDEX_SCHEMA_RETURNS = 1;
+
+constexpr size_t BYTECODE_INDEX_ARGUMENT_NAME = 0;
+constexpr size_t BYTECODE_INDEX_ARGUMENT_TYPE = 1;
+constexpr size_t BYTECODE_INDEX_ARGUMENT_DEFAULT_VALUE = 2;
+
+constexpr size_t BYTECODE_INDEX_MODULE_DEBUG_HANDLES = 0;
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..ea0af9d7782890380984ac33d376fffbb8540d78
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_functions.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/core/ivalue.h>
+
+// Functions that are used in both import and export processes
+
+namespace torch::jit {
+using c10::IValue;
+IValue expect_field(
+    c10::ivalue::TupleElements& elements,
+    const std::string& expected_name,
+    size_t entry);
+std::string operator_str(
+    const std::string& name,
+    const std::string& overloadname);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_helpers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_helpers.h
new file mode 100644
index 0000000000000000000000000000000000000000..40a4f968c1ab3994a113b5e850f3c3bdb97c0e19
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_export_helpers.h
@@ -0,0 +1,33 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <string>
+
+namespace caffe2::serialize {
+class PyTorchStreamReader;
+}
+
+namespace torch::jit {
+
+struct Source;
+
+// Convert a class type's qualifier name to the corresponding path the source
+// file it should be written to.
+//
+// Qualifier is like: foo.bar.baz
+// Returns: libs/foo/bar/baz.py
+std::string qualifierToArchivePath(
+    const std::string& qualifier,
+    const std::string& export_prefix);
+
+std::shared_ptr<Source> findSourceInArchiveFromQualifier(
+    caffe2::serialize::PyTorchStreamReader& reader,
+    const std::string& export_prefix,
+    const std::string& qualifier);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_read.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_read.h
new file mode 100644
index 0000000000000000000000000000000000000000..431430271c69db92255ab8786686b0f51c390dff
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_read.h
@@ -0,0 +1,32 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/serialization/unpickler.h>
+#include <memory>
+
+namespace caffe2::serialize {
+class PyTorchStreamReader;
+} // namespace caffe2::serialize
+
+namespace torch::jit {
+
+TORCH_API IValue readArchiveAndTensors(
+    const std::string& archive_name,
+    const std::string& pickle_prefix,
+    const std::string& tensor_prefix,
+    std::optional<TypeResolver> type_resolver,
+    std::optional<ObjLoader> obj_loader,
+    std::optional<at::Device> device,
+    caffe2::serialize::PyTorchStreamReader& stream_reader,
+    c10::TypePtr (*type_parser)(const std::string&) =
+        Unpickler::defaultTypeParser,
+    std::shared_ptr<DeserializationStorageContext> storage_context = nullptr);
+
+bool check_zip_file(
+    const std::shared_ptr<caffe2::serialize::ReadAdapterInterface>& rai);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_source.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_source.h
new file mode 100644
index 0000000000000000000000000000000000000000..f279d6cb18e0f143076ce19c328c1ebbd19082be
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/import_source.h
@@ -0,0 +1,105 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue_inl.h>
+#include <ATen/core/qualified_name.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/frontend/parser.h>
+#include <torch/csrc/jit/frontend/resolver.h>
+#include <torch/csrc/jit/frontend/script_type_parser.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/serialization/export.h>
+#include <torch/custom_class.h>
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace torch::jit {
+
+using SourceLoader = std::function<std::shared_ptr<Source>(const std::string&)>;
+
+struct SourceImporterImpl : public Resolver,
+                            std::enable_shared_from_this<SourceImporterImpl> {
+  SourceImporterImpl(
+      std::shared_ptr<CompilationUnit> cu,
+      const std::vector<at::IValue>* constant_table,
+      SourceLoader source_loader,
+      size_t version);
+  TypePtr findNamedType(const QualifiedName& name);
+  Function* findFunction(const QualifiedName& name);
+  void parseSourceIfNeeded(const std::string& qualifier);
+  void LEGACY_import_methods(
+      const Module& mod,
+      const std::shared_ptr<Source>& src);
+
+  std::shared_ptr<SugaredValue> resolveValue(
+      const std::string& name,
+      GraphFunction& m,
+      const SourceRange& loc) override;
+  TypePtr resolveType(const std::string& name, const SourceRange& loc) override;
+
+ private:
+  void importFunction(const std::string& qualifier, const Def& def);
+  void importNamedType(const std::string& qualifier, const ClassDef& class_def);
+  std::optional<Assign> attributeAssignmentSpecialHandlingHack(
+      const QualifiedName& qualified_classname,
+      const Assign& assign);
+  void importClass(
+      const QualifiedName& qualified_classname,
+      const ClassDef& class_def,
+      bool is_module);
+  void importEnum(
+      const QualifiedName& qualified_name,
+      const ClassDef& enum_def);
+  void importNamedTuple(
+      const QualifiedName& qualified_name,
+      const ClassDef& named_tuple_def);
+
+  void parsePossibleVersionNumber(Lexer& L);
+
+  void parseImports(Lexer& L);
+
+  std::shared_ptr<CompilationUnit> cu_;
+  std::unordered_map<std::string, std::shared_ptr<SugaredValue>> env_;
+  SourceLoader source_loader_;
+  std::optional<size_t> version_ = std::nullopt;
+  std::unordered_set<std::string> loaded_sources_;
+  // named types and functions loaded from a file but not yet defined because
+  // their type has not been requested yet.
+  std::unordered_map<QualifiedName, TreeRef> to_be_defined_;
+};
+
+// Given a directory of serialized TorchScript sources,
+// This class allows the loading of individual named types in source.
+// Resolves the dependencies between source files and parses
+// the source files as necessary.
+
+struct TORCH_API SourceImporter {
+  SourceImporter(
+      // The compilation unit that will own the imported source
+      std::shared_ptr<CompilationUnit> cu,
+      const std::vector<at::IValue>* constant_table,
+      SourceLoader loader,
+      size_t version);
+
+  TypePtr loadType(const QualifiedName& name) const;
+
+  // Add the methods defined in `src` to the module `mod`, using SourceImporter
+  // to resolve any classes via loadType
+  void LEGACY_import_methods(
+      const Module& mod,
+      const std::shared_ptr<Source>& src);
+  ~SourceImporter();
+
+ private:
+  std::shared_ptr<SourceImporterImpl> pImpl;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/mobile_bytecode_generated.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/mobile_bytecode_generated.h
new file mode 100644
index 0000000000000000000000000000000000000000..d2032493f671eb9e728d9a35467038b47a985a10
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/mobile_bytecode_generated.h
@@ -0,0 +1,2605 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// automatically generated by the FlatBuffers compiler, do not modify
+
+
+#ifndef FLATBUFFERS_GENERATED_MOBILEBYTECODE_TORCH_JIT_MOBILE_SERIALIZATION_H_
+#define FLATBUFFERS_GENERATED_MOBILEBYTECODE_TORCH_JIT_MOBILE_SERIALIZATION_H_
+
+#include "flatbuffers/flatbuffers.h"
+
+// Ensure the included flatbuffers.h is the same version as when this file was
+// generated, otherwise it may not be compatible.
+static_assert(FLATBUFFERS_VERSION_MAJOR == 24 &&
+              FLATBUFFERS_VERSION_MINOR == 12 &&
+              FLATBUFFERS_VERSION_REVISION == 23,
+             "Non-compatible flatbuffers version included");
+
+namespace torch {
+namespace jit {
+namespace mobile {
+namespace serialization {
+
+struct Int;
+
+struct Bool;
+
+struct Double;
+
+struct PerTensorAffineSchema;
+
+struct QuantizedSchema;
+struct QuantizedSchemaBuilder;
+
+struct TensorMetadata;
+struct TensorMetadataBuilder;
+
+struct String;
+struct StringBuilder;
+
+struct Device;
+struct DeviceBuilder;
+
+struct List;
+struct ListBuilder;
+
+struct IntList;
+struct IntListBuilder;
+
+struct DoubleList;
+struct DoubleListBuilder;
+
+struct BoolList;
+struct BoolListBuilder;
+
+struct Tuple;
+struct TupleBuilder;
+
+struct Dict;
+struct DictBuilder;
+
+struct ObjectType;
+struct ObjectTypeBuilder;
+
+struct Object;
+struct ObjectBuilder;
+
+struct ComplexDouble;
+
+struct EnumValue;
+struct EnumValueBuilder;
+
+struct Instruction;
+
+struct Operator;
+struct OperatorBuilder;
+
+struct Arg;
+struct ArgBuilder;
+
+struct Schema;
+struct SchemaBuilder;
+
+struct DebugInfo;
+struct DebugInfoBuilder;
+
+struct Function;
+struct FunctionBuilder;
+
+struct StorageData;
+struct StorageDataBuilder;
+
+struct IValue;
+struct IValueBuilder;
+
+struct ExtraFile;
+struct ExtraFileBuilder;
+
+struct Module;
+struct ModuleBuilder;
+
+enum class TypeType : uint8_t {
+  UNSET = 0,
+  CLASS_WITH_FIELD = 1,
+  CUSTOM_CLASS = 2,
+  CLASS_WITH_SETSTATE = 3,
+  NON_OBJ = 4,
+  MIN = UNSET,
+  MAX = NON_OBJ
+};
+
+inline const TypeType (&EnumValuesTypeType())[5] {
+  static const TypeType values[] = {
+    TypeType::UNSET,
+    TypeType::CLASS_WITH_FIELD,
+    TypeType::CUSTOM_CLASS,
+    TypeType::CLASS_WITH_SETSTATE,
+    TypeType::NON_OBJ
+  };
+  return values;
+}
+
+inline const char * const *EnumNamesTypeType() {
+  static const char * const names[6] = {
+    "UNSET",
+    "CLASS_WITH_FIELD",
+    "CUSTOM_CLASS",
+    "CLASS_WITH_SETSTATE",
+    "NON_OBJ",
+    nullptr
+  };
+  return names;
+}
+
+inline const char *EnumNameTypeType(TypeType e) {
+  if (::flatbuffers::IsOutRange(e, TypeType::UNSET, TypeType::NON_OBJ)) return "";
+  const size_t index = static_cast<size_t>(e);
+  return EnumNamesTypeType()[index];
+}
+
+enum class IValueUnion : uint8_t {
+  NONE = 0,
+  Int = 1,
+  Bool = 2,
+  Double = 3,
+  ComplexDouble = 4,
+  TensorMetadata = 5,
+  String = 6,
+  List = 7,
+  Tuple = 8,
+  Dict = 9,
+  Object = 10,
+  IntList = 11,
+  DoubleList = 12,
+  BoolList = 13,
+  Device = 14,
+  EnumValue = 15,
+  Function = 16,
+  MIN = NONE,
+  MAX = Function
+};
+
+inline const IValueUnion (&EnumValuesIValueUnion())[17] {
+  static const IValueUnion values[] = {
+    IValueUnion::NONE,
+    IValueUnion::Int,
+    IValueUnion::Bool,
+    IValueUnion::Double,
+    IValueUnion::ComplexDouble,
+    IValueUnion::TensorMetadata,
+    IValueUnion::String,
+    IValueUnion::List,
+    IValueUnion::Tuple,
+    IValueUnion::Dict,
+    IValueUnion::Object,
+    IValueUnion::IntList,
+    IValueUnion::DoubleList,
+    IValueUnion::BoolList,
+    IValueUnion::Device,
+    IValueUnion::EnumValue,
+    IValueUnion::Function
+  };
+  return values;
+}
+
+inline const char * const *EnumNamesIValueUnion() {
+  static const char * const names[18] = {
+    "NONE",
+    "Int",
+    "Bool",
+    "Double",
+    "ComplexDouble",
+    "TensorMetadata",
+    "String",
+    "List",
+    "Tuple",
+    "Dict",
+    "Object",
+    "IntList",
+    "DoubleList",
+    "BoolList",
+    "Device",
+    "EnumValue",
+    "Function",
+    nullptr
+  };
+  return names;
+}
+
+inline const char *EnumNameIValueUnion(IValueUnion e) {
+  if (::flatbuffers::IsOutRange(e, IValueUnion::NONE, IValueUnion::Function)) return "";
+  const size_t index = static_cast<size_t>(e);
+  return EnumNamesIValueUnion()[index];
+}
+
+template<typename T> struct IValueUnionTraits {
+  static const IValueUnion enum_value = IValueUnion::NONE;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Int> {
+  static const IValueUnion enum_value = IValueUnion::Int;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Bool> {
+  static const IValueUnion enum_value = IValueUnion::Bool;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Double> {
+  static const IValueUnion enum_value = IValueUnion::Double;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::ComplexDouble> {
+  static const IValueUnion enum_value = IValueUnion::ComplexDouble;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::TensorMetadata> {
+  static const IValueUnion enum_value = IValueUnion::TensorMetadata;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::String> {
+  static const IValueUnion enum_value = IValueUnion::String;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::List> {
+  static const IValueUnion enum_value = IValueUnion::List;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Tuple> {
+  static const IValueUnion enum_value = IValueUnion::Tuple;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Dict> {
+  static const IValueUnion enum_value = IValueUnion::Dict;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Object> {
+  static const IValueUnion enum_value = IValueUnion::Object;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::IntList> {
+  static const IValueUnion enum_value = IValueUnion::IntList;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::DoubleList> {
+  static const IValueUnion enum_value = IValueUnion::DoubleList;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::BoolList> {
+  static const IValueUnion enum_value = IValueUnion::BoolList;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Device> {
+  static const IValueUnion enum_value = IValueUnion::Device;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::EnumValue> {
+  static const IValueUnion enum_value = IValueUnion::EnumValue;
+};
+
+template<> struct IValueUnionTraits<torch::jit::mobile::serialization::Function> {
+  static const IValueUnion enum_value = IValueUnion::Function;
+};
+
+bool VerifyIValueUnion(::flatbuffers::Verifier &verifier, const void *obj, IValueUnion type);
+bool VerifyIValueUnionVector(::flatbuffers::Verifier &verifier, const ::flatbuffers::Vector<::flatbuffers::Offset<void>> *values, const ::flatbuffers::Vector<IValueUnion> *types);
+
+FLATBUFFERS_MANUALLY_ALIGNED_STRUCT(8) Int FLATBUFFERS_FINAL_CLASS {
+ private:
+  int64_t int_val_;
+
+ public:
+  Int()
+      : int_val_(0) {
+  }
+  Int(int64_t _int_val)
+      : int_val_(::flatbuffers::EndianScalar(_int_val)) {
+  }
+  int64_t int_val() const {
+    return ::flatbuffers::EndianScalar(int_val_);
+  }
+  void mutate_int_val(int64_t _int_val) {
+    ::flatbuffers::WriteScalar(&int_val_, _int_val);
+  }
+};
+FLATBUFFERS_STRUCT_END(Int, 8);
+
+FLATBUFFERS_MANUALLY_ALIGNED_STRUCT(1) Bool FLATBUFFERS_FINAL_CLASS {
+ private:
+  uint8_t bool_val_;
+
+ public:
+  Bool()
+      : bool_val_(0) {
+  }
+  Bool(bool _bool_val)
+      : bool_val_(::flatbuffers::EndianScalar(static_cast<uint8_t>(_bool_val))) {
+  }
+  bool bool_val() const {
+    return ::flatbuffers::EndianScalar(bool_val_) != 0;
+  }
+  void mutate_bool_val(bool _bool_val) {
+    ::flatbuffers::WriteScalar(&bool_val_, static_cast<uint8_t>(_bool_val));
+  }
+};
+FLATBUFFERS_STRUCT_END(Bool, 1);
+
+FLATBUFFERS_MANUALLY_ALIGNED_STRUCT(8) Double FLATBUFFERS_FINAL_CLASS {
+ private:
+  double double_val_;
+
+ public:
+  Double()
+      : double_val_(0) {
+  }
+  Double(double _double_val)
+      : double_val_(::flatbuffers::EndianScalar(_double_val)) {
+  }
+  double double_val() const {
+    return ::flatbuffers::EndianScalar(double_val_);
+  }
+  void mutate_double_val(double _double_val) {
+    ::flatbuffers::WriteScalar(&double_val_, _double_val);
+  }
+};
+FLATBUFFERS_STRUCT_END(Double, 8);
+
+FLATBUFFERS_MANUALLY_ALIGNED_STRUCT(8) PerTensorAffineSchema FLATBUFFERS_FINAL_CLASS {
+ private:
+  double q_scale_;
+  int32_t q_zero_point_;
+  int32_t padding0__;
+
+ public:
+  PerTensorAffineSchema()
+      : q_scale_(0),
+        q_zero_point_(0),
+        padding0__(0) {
+    (void)padding0__;
+  }
+  PerTensorAffineSchema(double _q_scale, int32_t _q_zero_point)
+      : q_scale_(::flatbuffers::EndianScalar(_q_scale)),
+        q_zero_point_(::flatbuffers::EndianScalar(_q_zero_point)),
+        padding0__(0) {
+    (void)padding0__;
+  }
+  double q_scale() const {
+    return ::flatbuffers::EndianScalar(q_scale_);
+  }
+  void mutate_q_scale(double _q_scale) {
+    ::flatbuffers::WriteScalar(&q_scale_, _q_scale);
+  }
+  int32_t q_zero_point() const {
+    return ::flatbuffers::EndianScalar(q_zero_point_);
+  }
+  void mutate_q_zero_point(int32_t _q_zero_point) {
+    ::flatbuffers::WriteScalar(&q_zero_point_, _q_zero_point);
+  }
+};
+FLATBUFFERS_STRUCT_END(PerTensorAffineSchema, 16);
+
+FLATBUFFERS_MANUALLY_ALIGNED_STRUCT(8) ComplexDouble FLATBUFFERS_FINAL_CLASS {
+ private:
+  double real_;
+  double imag_;
+
+ public:
+  ComplexDouble()
+      : real_(0),
+        imag_(0) {
+  }
+  ComplexDouble(double _real, double _imag)
+      : real_(::flatbuffers::EndianScalar(_real)),
+        imag_(::flatbuffers::EndianScalar(_imag)) {
+  }
+  double real() const {
+    return ::flatbuffers::EndianScalar(real_);
+  }
+  void mutate_real(double _real) {
+    ::flatbuffers::WriteScalar(&real_, _real);
+  }
+  double imag() const {
+    return ::flatbuffers::EndianScalar(imag_);
+  }
+  void mutate_imag(double _imag) {
+    ::flatbuffers::WriteScalar(&imag_, _imag);
+  }
+};
+FLATBUFFERS_STRUCT_END(ComplexDouble, 16);
+
+FLATBUFFERS_MANUALLY_ALIGNED_STRUCT(4) Instruction FLATBUFFERS_FINAL_CLASS {
+ private:
+  int8_t op_;
+  int8_t padding0__;
+  uint16_t n_;
+  int32_t x_;
+
+ public:
+  Instruction()
+      : op_(0),
+        padding0__(0),
+        n_(0),
+        x_(0) {
+    (void)padding0__;
+  }
+  Instruction(int8_t _op, uint16_t _n, int32_t _x)
+      : op_(::flatbuffers::EndianScalar(_op)),
+        padding0__(0),
+        n_(::flatbuffers::EndianScalar(_n)),
+        x_(::flatbuffers::EndianScalar(_x)) {
+    (void)padding0__;
+  }
+  int8_t op() const {
+    return ::flatbuffers::EndianScalar(op_);
+  }
+  void mutate_op(int8_t _op) {
+    ::flatbuffers::WriteScalar(&op_, _op);
+  }
+  uint16_t n() const {
+    return ::flatbuffers::EndianScalar(n_);
+  }
+  void mutate_n(uint16_t _n) {
+    ::flatbuffers::WriteScalar(&n_, _n);
+  }
+  int32_t x() const {
+    return ::flatbuffers::EndianScalar(x_);
+  }
+  void mutate_x(int32_t _x) {
+    ::flatbuffers::WriteScalar(&x_, _x);
+  }
+};
+FLATBUFFERS_STRUCT_END(Instruction, 8);
+
+struct QuantizedSchema FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef QuantizedSchemaBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_QSCHEME = 4,
+    VT_SCALE = 6,
+    VT_ZERO_POINT = 8,
+    VT_SCALES = 10,
+    VT_ZERO_POINTS = 12,
+    VT_AXIS = 14
+  };
+  int8_t qscheme() const {
+    return GetField<int8_t>(VT_QSCHEME, 0);
+  }
+  bool mutate_qscheme(int8_t _qscheme = 0) {
+    return SetField<int8_t>(VT_QSCHEME, _qscheme, 0);
+  }
+  double scale() const {
+    return GetField<double>(VT_SCALE, 0.0);
+  }
+  bool mutate_scale(double _scale = 0.0) {
+    return SetField<double>(VT_SCALE, _scale, 0.0);
+  }
+  int32_t zero_point() const {
+    return GetField<int32_t>(VT_ZERO_POINT, 0);
+  }
+  bool mutate_zero_point(int32_t _zero_point = 0) {
+    return SetField<int32_t>(VT_ZERO_POINT, _zero_point, 0);
+  }
+  const torch::jit::mobile::serialization::TensorMetadata *scales() const {
+    return GetPointer<const torch::jit::mobile::serialization::TensorMetadata *>(VT_SCALES);
+  }
+  torch::jit::mobile::serialization::TensorMetadata *mutable_scales() {
+    return GetPointer<torch::jit::mobile::serialization::TensorMetadata *>(VT_SCALES);
+  }
+  const torch::jit::mobile::serialization::TensorMetadata *zero_points() const {
+    return GetPointer<const torch::jit::mobile::serialization::TensorMetadata *>(VT_ZERO_POINTS);
+  }
+  torch::jit::mobile::serialization::TensorMetadata *mutable_zero_points() {
+    return GetPointer<torch::jit::mobile::serialization::TensorMetadata *>(VT_ZERO_POINTS);
+  }
+  int32_t axis() const {
+    return GetField<int32_t>(VT_AXIS, 0);
+  }
+  bool mutate_axis(int32_t _axis = 0) {
+    return SetField<int32_t>(VT_AXIS, _axis, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyField<int8_t>(verifier, VT_QSCHEME, 1) &&
+           VerifyField<double>(verifier, VT_SCALE, 8) &&
+           VerifyField<int32_t>(verifier, VT_ZERO_POINT, 4) &&
+           VerifyOffset(verifier, VT_SCALES) &&
+           verifier.VerifyTable(scales()) &&
+           VerifyOffset(verifier, VT_ZERO_POINTS) &&
+           verifier.VerifyTable(zero_points()) &&
+           VerifyField<int32_t>(verifier, VT_AXIS, 4) &&
+           verifier.EndTable();
+  }
+};
+
+struct QuantizedSchemaBuilder {
+  typedef QuantizedSchema Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_qscheme(int8_t qscheme) {
+    fbb_.AddElement<int8_t>(QuantizedSchema::VT_QSCHEME, qscheme, 0);
+  }
+  void add_scale(double scale) {
+    fbb_.AddElement<double>(QuantizedSchema::VT_SCALE, scale, 0.0);
+  }
+  void add_zero_point(int32_t zero_point) {
+    fbb_.AddElement<int32_t>(QuantizedSchema::VT_ZERO_POINT, zero_point, 0);
+  }
+  void add_scales(::flatbuffers::Offset<torch::jit::mobile::serialization::TensorMetadata> scales) {
+    fbb_.AddOffset(QuantizedSchema::VT_SCALES, scales);
+  }
+  void add_zero_points(::flatbuffers::Offset<torch::jit::mobile::serialization::TensorMetadata> zero_points) {
+    fbb_.AddOffset(QuantizedSchema::VT_ZERO_POINTS, zero_points);
+  }
+  void add_axis(int32_t axis) {
+    fbb_.AddElement<int32_t>(QuantizedSchema::VT_AXIS, axis, 0);
+  }
+  explicit QuantizedSchemaBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<QuantizedSchema> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<QuantizedSchema>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<QuantizedSchema> CreateQuantizedSchema(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    int8_t qscheme = 0,
+    double scale = 0.0,
+    int32_t zero_point = 0,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::TensorMetadata> scales = 0,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::TensorMetadata> zero_points = 0,
+    int32_t axis = 0) {
+  QuantizedSchemaBuilder builder_(_fbb);
+  builder_.add_scale(scale);
+  builder_.add_axis(axis);
+  builder_.add_zero_points(zero_points);
+  builder_.add_scales(scales);
+  builder_.add_zero_point(zero_point);
+  builder_.add_qscheme(qscheme);
+  return builder_.Finish();
+}
+
+struct TensorMetadata FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef TensorMetadataBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_STORAGE_LOCATION_INDEX = 4,
+    VT_SCALAR_TYPE = 6,
+    VT_STORAGE_OFFSET = 8,
+    VT_SIZES = 10,
+    VT_STRIDES = 12,
+    VT_REQUIRES_GRAD = 14,
+    VT_QUANTIZED_SCHEMA = 16
+  };
+  uint32_t storage_location_index() const {
+    return GetField<uint32_t>(VT_STORAGE_LOCATION_INDEX, 0);
+  }
+  bool mutate_storage_location_index(uint32_t _storage_location_index = 0) {
+    return SetField<uint32_t>(VT_STORAGE_LOCATION_INDEX, _storage_location_index, 0);
+  }
+  int8_t scalar_type() const {
+    return GetField<int8_t>(VT_SCALAR_TYPE, 0);
+  }
+  bool mutate_scalar_type(int8_t _scalar_type = 0) {
+    return SetField<int8_t>(VT_SCALAR_TYPE, _scalar_type, 0);
+  }
+  int32_t storage_offset() const {
+    return GetField<int32_t>(VT_STORAGE_OFFSET, 0);
+  }
+  bool mutate_storage_offset(int32_t _storage_offset = 0) {
+    return SetField<int32_t>(VT_STORAGE_OFFSET, _storage_offset, 0);
+  }
+  const ::flatbuffers::Vector<int32_t> *sizes() const {
+    return GetPointer<const ::flatbuffers::Vector<int32_t> *>(VT_SIZES);
+  }
+  ::flatbuffers::Vector<int32_t> *mutable_sizes() {
+    return GetPointer<::flatbuffers::Vector<int32_t> *>(VT_SIZES);
+  }
+  const ::flatbuffers::Vector<int32_t> *strides() const {
+    return GetPointer<const ::flatbuffers::Vector<int32_t> *>(VT_STRIDES);
+  }
+  ::flatbuffers::Vector<int32_t> *mutable_strides() {
+    return GetPointer<::flatbuffers::Vector<int32_t> *>(VT_STRIDES);
+  }
+  bool requires_grad() const {
+    return GetField<uint8_t>(VT_REQUIRES_GRAD, 0) != 0;
+  }
+  bool mutate_requires_grad(bool _requires_grad = false) {
+    return SetField<uint8_t>(VT_REQUIRES_GRAD, static_cast<uint8_t>(_requires_grad), 0);
+  }
+  const torch::jit::mobile::serialization::QuantizedSchema *quantized_schema() const {
+    return GetPointer<const torch::jit::mobile::serialization::QuantizedSchema *>(VT_QUANTIZED_SCHEMA);
+  }
+  torch::jit::mobile::serialization::QuantizedSchema *mutable_quantized_schema() {
+    return GetPointer<torch::jit::mobile::serialization::QuantizedSchema *>(VT_QUANTIZED_SCHEMA);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyField<uint32_t>(verifier, VT_STORAGE_LOCATION_INDEX, 4) &&
+           VerifyField<int8_t>(verifier, VT_SCALAR_TYPE, 1) &&
+           VerifyField<int32_t>(verifier, VT_STORAGE_OFFSET, 4) &&
+           VerifyOffset(verifier, VT_SIZES) &&
+           verifier.VerifyVector(sizes()) &&
+           VerifyOffset(verifier, VT_STRIDES) &&
+           verifier.VerifyVector(strides()) &&
+           VerifyField<uint8_t>(verifier, VT_REQUIRES_GRAD, 1) &&
+           VerifyOffset(verifier, VT_QUANTIZED_SCHEMA) &&
+           verifier.VerifyTable(quantized_schema()) &&
+           verifier.EndTable();
+  }
+};
+
+struct TensorMetadataBuilder {
+  typedef TensorMetadata Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_storage_location_index(uint32_t storage_location_index) {
+    fbb_.AddElement<uint32_t>(TensorMetadata::VT_STORAGE_LOCATION_INDEX, storage_location_index, 0);
+  }
+  void add_scalar_type(int8_t scalar_type) {
+    fbb_.AddElement<int8_t>(TensorMetadata::VT_SCALAR_TYPE, scalar_type, 0);
+  }
+  void add_storage_offset(int32_t storage_offset) {
+    fbb_.AddElement<int32_t>(TensorMetadata::VT_STORAGE_OFFSET, storage_offset, 0);
+  }
+  void add_sizes(::flatbuffers::Offset<::flatbuffers::Vector<int32_t>> sizes) {
+    fbb_.AddOffset(TensorMetadata::VT_SIZES, sizes);
+  }
+  void add_strides(::flatbuffers::Offset<::flatbuffers::Vector<int32_t>> strides) {
+    fbb_.AddOffset(TensorMetadata::VT_STRIDES, strides);
+  }
+  void add_requires_grad(bool requires_grad) {
+    fbb_.AddElement<uint8_t>(TensorMetadata::VT_REQUIRES_GRAD, static_cast<uint8_t>(requires_grad), 0);
+  }
+  void add_quantized_schema(::flatbuffers::Offset<torch::jit::mobile::serialization::QuantizedSchema> quantized_schema) {
+    fbb_.AddOffset(TensorMetadata::VT_QUANTIZED_SCHEMA, quantized_schema);
+  }
+  explicit TensorMetadataBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<TensorMetadata> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<TensorMetadata>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<TensorMetadata> CreateTensorMetadata(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    uint32_t storage_location_index = 0,
+    int8_t scalar_type = 0,
+    int32_t storage_offset = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<int32_t>> sizes = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<int32_t>> strides = 0,
+    bool requires_grad = false,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::QuantizedSchema> quantized_schema = 0) {
+  TensorMetadataBuilder builder_(_fbb);
+  builder_.add_quantized_schema(quantized_schema);
+  builder_.add_strides(strides);
+  builder_.add_sizes(sizes);
+  builder_.add_storage_offset(storage_offset);
+  builder_.add_storage_location_index(storage_location_index);
+  builder_.add_requires_grad(requires_grad);
+  builder_.add_scalar_type(scalar_type);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<TensorMetadata> CreateTensorMetadataDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    uint32_t storage_location_index = 0,
+    int8_t scalar_type = 0,
+    int32_t storage_offset = 0,
+    const std::vector<int32_t> *sizes = nullptr,
+    const std::vector<int32_t> *strides = nullptr,
+    bool requires_grad = false,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::QuantizedSchema> quantized_schema = 0) {
+  auto sizes__ = sizes ? _fbb.CreateVector<int32_t>(*sizes) : 0;
+  auto strides__ = strides ? _fbb.CreateVector<int32_t>(*strides) : 0;
+  return torch::jit::mobile::serialization::CreateTensorMetadata(
+      _fbb,
+      storage_location_index,
+      scalar_type,
+      storage_offset,
+      sizes__,
+      strides__,
+      requires_grad,
+      quantized_schema);
+}
+
+struct String FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef StringBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_DATA = 4
+  };
+  const ::flatbuffers::String *data() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_DATA);
+  }
+  ::flatbuffers::String *mutable_data() {
+    return GetPointer<::flatbuffers::String *>(VT_DATA);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_DATA) &&
+           verifier.VerifyString(data()) &&
+           verifier.EndTable();
+  }
+};
+
+struct StringBuilder {
+  typedef String Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_data(::flatbuffers::Offset<::flatbuffers::String> data) {
+    fbb_.AddOffset(String::VT_DATA, data);
+  }
+  explicit StringBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<String> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<String>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<String> CreateString(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> data = 0) {
+  StringBuilder builder_(_fbb);
+  builder_.add_data(data);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<String> CreateStringDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *data = nullptr) {
+  auto data__ = data ? _fbb.CreateString(data) : 0;
+  return torch::jit::mobile::serialization::CreateString(
+      _fbb,
+      data__);
+}
+
+struct Device FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef DeviceBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_STR = 4
+  };
+  const ::flatbuffers::String *str() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_STR);
+  }
+  ::flatbuffers::String *mutable_str() {
+    return GetPointer<::flatbuffers::String *>(VT_STR);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_STR) &&
+           verifier.VerifyString(str()) &&
+           verifier.EndTable();
+  }
+};
+
+struct DeviceBuilder {
+  typedef Device Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_str(::flatbuffers::Offset<::flatbuffers::String> str) {
+    fbb_.AddOffset(Device::VT_STR, str);
+  }
+  explicit DeviceBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Device> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Device>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Device> CreateDevice(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> str = 0) {
+  DeviceBuilder builder_(_fbb);
+  builder_.add_str(str);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Device> CreateDeviceDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *str = nullptr) {
+  auto str__ = str ? _fbb.CreateString(str) : 0;
+  return torch::jit::mobile::serialization::CreateDevice(
+      _fbb,
+      str__);
+}
+
+struct List FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef ListBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_ITEMS = 4,
+    VT_ANNOTATION_STR = 6
+  };
+  const ::flatbuffers::Vector<uint32_t> *items() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_ITEMS);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_items() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_ITEMS);
+  }
+  const ::flatbuffers::String *annotation_str() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_ANNOTATION_STR);
+  }
+  ::flatbuffers::String *mutable_annotation_str() {
+    return GetPointer<::flatbuffers::String *>(VT_ANNOTATION_STR);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_ITEMS) &&
+           verifier.VerifyVector(items()) &&
+           VerifyOffset(verifier, VT_ANNOTATION_STR) &&
+           verifier.VerifyString(annotation_str()) &&
+           verifier.EndTable();
+  }
+};
+
+struct ListBuilder {
+  typedef List Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_items(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> items) {
+    fbb_.AddOffset(List::VT_ITEMS, items);
+  }
+  void add_annotation_str(::flatbuffers::Offset<::flatbuffers::String> annotation_str) {
+    fbb_.AddOffset(List::VT_ANNOTATION_STR, annotation_str);
+  }
+  explicit ListBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<List> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<List>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<List> CreateList(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> items = 0,
+    ::flatbuffers::Offset<::flatbuffers::String> annotation_str = 0) {
+  ListBuilder builder_(_fbb);
+  builder_.add_annotation_str(annotation_str);
+  builder_.add_items(items);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<List> CreateListDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<uint32_t> *items = nullptr,
+    const char *annotation_str = nullptr) {
+  auto items__ = items ? _fbb.CreateVector<uint32_t>(*items) : 0;
+  auto annotation_str__ = annotation_str ? _fbb.CreateString(annotation_str) : 0;
+  return torch::jit::mobile::serialization::CreateList(
+      _fbb,
+      items__,
+      annotation_str__);
+}
+
+struct IntList FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef IntListBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_ITEMS = 4
+  };
+  const ::flatbuffers::Vector<int64_t> *items() const {
+    return GetPointer<const ::flatbuffers::Vector<int64_t> *>(VT_ITEMS);
+  }
+  ::flatbuffers::Vector<int64_t> *mutable_items() {
+    return GetPointer<::flatbuffers::Vector<int64_t> *>(VT_ITEMS);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_ITEMS) &&
+           verifier.VerifyVector(items()) &&
+           verifier.EndTable();
+  }
+};
+
+struct IntListBuilder {
+  typedef IntList Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_items(::flatbuffers::Offset<::flatbuffers::Vector<int64_t>> items) {
+    fbb_.AddOffset(IntList::VT_ITEMS, items);
+  }
+  explicit IntListBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<IntList> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<IntList>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<IntList> CreateIntList(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<int64_t>> items = 0) {
+  IntListBuilder builder_(_fbb);
+  builder_.add_items(items);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<IntList> CreateIntListDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<int64_t> *items = nullptr) {
+  auto items__ = items ? _fbb.CreateVector<int64_t>(*items) : 0;
+  return torch::jit::mobile::serialization::CreateIntList(
+      _fbb,
+      items__);
+}
+
+struct DoubleList FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef DoubleListBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_ITEMS = 4
+  };
+  const ::flatbuffers::Vector<double> *items() const {
+    return GetPointer<const ::flatbuffers::Vector<double> *>(VT_ITEMS);
+  }
+  ::flatbuffers::Vector<double> *mutable_items() {
+    return GetPointer<::flatbuffers::Vector<double> *>(VT_ITEMS);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_ITEMS) &&
+           verifier.VerifyVector(items()) &&
+           verifier.EndTable();
+  }
+};
+
+struct DoubleListBuilder {
+  typedef DoubleList Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_items(::flatbuffers::Offset<::flatbuffers::Vector<double>> items) {
+    fbb_.AddOffset(DoubleList::VT_ITEMS, items);
+  }
+  explicit DoubleListBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<DoubleList> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<DoubleList>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<DoubleList> CreateDoubleList(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<double>> items = 0) {
+  DoubleListBuilder builder_(_fbb);
+  builder_.add_items(items);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<DoubleList> CreateDoubleListDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<double> *items = nullptr) {
+  auto items__ = items ? _fbb.CreateVector<double>(*items) : 0;
+  return torch::jit::mobile::serialization::CreateDoubleList(
+      _fbb,
+      items__);
+}
+
+struct BoolList FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef BoolListBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_ITEMS = 4
+  };
+  const ::flatbuffers::Vector<uint8_t> *items() const {
+    return GetPointer<const ::flatbuffers::Vector<uint8_t> *>(VT_ITEMS);
+  }
+  ::flatbuffers::Vector<uint8_t> *mutable_items() {
+    return GetPointer<::flatbuffers::Vector<uint8_t> *>(VT_ITEMS);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_ITEMS) &&
+           verifier.VerifyVector(items()) &&
+           verifier.EndTable();
+  }
+};
+
+struct BoolListBuilder {
+  typedef BoolList Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_items(::flatbuffers::Offset<::flatbuffers::Vector<uint8_t>> items) {
+    fbb_.AddOffset(BoolList::VT_ITEMS, items);
+  }
+  explicit BoolListBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<BoolList> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<BoolList>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<BoolList> CreateBoolList(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint8_t>> items = 0) {
+  BoolListBuilder builder_(_fbb);
+  builder_.add_items(items);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<BoolList> CreateBoolListDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<uint8_t> *items = nullptr) {
+  auto items__ = items ? _fbb.CreateVector<uint8_t>(*items) : 0;
+  return torch::jit::mobile::serialization::CreateBoolList(
+      _fbb,
+      items__);
+}
+
+struct Tuple FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef TupleBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_ITEMS = 4
+  };
+  const ::flatbuffers::Vector<uint32_t> *items() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_ITEMS);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_items() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_ITEMS);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_ITEMS) &&
+           verifier.VerifyVector(items()) &&
+           verifier.EndTable();
+  }
+};
+
+struct TupleBuilder {
+  typedef Tuple Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_items(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> items) {
+    fbb_.AddOffset(Tuple::VT_ITEMS, items);
+  }
+  explicit TupleBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Tuple> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Tuple>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Tuple> CreateTuple(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> items = 0) {
+  TupleBuilder builder_(_fbb);
+  builder_.add_items(items);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Tuple> CreateTupleDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<uint32_t> *items = nullptr) {
+  auto items__ = items ? _fbb.CreateVector<uint32_t>(*items) : 0;
+  return torch::jit::mobile::serialization::CreateTuple(
+      _fbb,
+      items__);
+}
+
+struct Dict FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef DictBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_KEYS = 4,
+    VT_VALUES = 6,
+    VT_ANNOTATION_STR = 8
+  };
+  const ::flatbuffers::Vector<uint32_t> *keys() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_KEYS);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_keys() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_KEYS);
+  }
+  const ::flatbuffers::Vector<uint32_t> *values() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_VALUES);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_values() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_VALUES);
+  }
+  const ::flatbuffers::String *annotation_str() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_ANNOTATION_STR);
+  }
+  ::flatbuffers::String *mutable_annotation_str() {
+    return GetPointer<::flatbuffers::String *>(VT_ANNOTATION_STR);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_KEYS) &&
+           verifier.VerifyVector(keys()) &&
+           VerifyOffset(verifier, VT_VALUES) &&
+           verifier.VerifyVector(values()) &&
+           VerifyOffset(verifier, VT_ANNOTATION_STR) &&
+           verifier.VerifyString(annotation_str()) &&
+           verifier.EndTable();
+  }
+};
+
+struct DictBuilder {
+  typedef Dict Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_keys(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> keys) {
+    fbb_.AddOffset(Dict::VT_KEYS, keys);
+  }
+  void add_values(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> values) {
+    fbb_.AddOffset(Dict::VT_VALUES, values);
+  }
+  void add_annotation_str(::flatbuffers::Offset<::flatbuffers::String> annotation_str) {
+    fbb_.AddOffset(Dict::VT_ANNOTATION_STR, annotation_str);
+  }
+  explicit DictBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Dict> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Dict>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Dict> CreateDict(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> keys = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> values = 0,
+    ::flatbuffers::Offset<::flatbuffers::String> annotation_str = 0) {
+  DictBuilder builder_(_fbb);
+  builder_.add_annotation_str(annotation_str);
+  builder_.add_values(values);
+  builder_.add_keys(keys);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Dict> CreateDictDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<uint32_t> *keys = nullptr,
+    const std::vector<uint32_t> *values = nullptr,
+    const char *annotation_str = nullptr) {
+  auto keys__ = keys ? _fbb.CreateVector<uint32_t>(*keys) : 0;
+  auto values__ = values ? _fbb.CreateVector<uint32_t>(*values) : 0;
+  auto annotation_str__ = annotation_str ? _fbb.CreateString(annotation_str) : 0;
+  return torch::jit::mobile::serialization::CreateDict(
+      _fbb,
+      keys__,
+      values__,
+      annotation_str__);
+}
+
+struct ObjectType FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef ObjectTypeBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_TYPE_NAME = 4,
+    VT_TYPE = 6,
+    VT_ATTR_NAMES = 8
+  };
+  const ::flatbuffers::String *type_name() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_TYPE_NAME);
+  }
+  ::flatbuffers::String *mutable_type_name() {
+    return GetPointer<::flatbuffers::String *>(VT_TYPE_NAME);
+  }
+  torch::jit::mobile::serialization::TypeType type() const {
+    return static_cast<torch::jit::mobile::serialization::TypeType>(GetField<uint8_t>(VT_TYPE, 0));
+  }
+  bool mutate_type(torch::jit::mobile::serialization::TypeType _type = static_cast<torch::jit::mobile::serialization::TypeType>(0)) {
+    return SetField<uint8_t>(VT_TYPE, static_cast<uint8_t>(_type), 0);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *attr_names() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *>(VT_ATTR_NAMES);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *mutable_attr_names() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *>(VT_ATTR_NAMES);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_TYPE_NAME) &&
+           verifier.VerifyString(type_name()) &&
+           VerifyField<uint8_t>(verifier, VT_TYPE, 1) &&
+           VerifyOffset(verifier, VT_ATTR_NAMES) &&
+           verifier.VerifyVector(attr_names()) &&
+           verifier.VerifyVectorOfStrings(attr_names()) &&
+           verifier.EndTable();
+  }
+};
+
+struct ObjectTypeBuilder {
+  typedef ObjectType Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_type_name(::flatbuffers::Offset<::flatbuffers::String> type_name) {
+    fbb_.AddOffset(ObjectType::VT_TYPE_NAME, type_name);
+  }
+  void add_type(torch::jit::mobile::serialization::TypeType type) {
+    fbb_.AddElement<uint8_t>(ObjectType::VT_TYPE, static_cast<uint8_t>(type), 0);
+  }
+  void add_attr_names(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>>> attr_names) {
+    fbb_.AddOffset(ObjectType::VT_ATTR_NAMES, attr_names);
+  }
+  explicit ObjectTypeBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<ObjectType> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<ObjectType>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<ObjectType> CreateObjectType(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> type_name = 0,
+    torch::jit::mobile::serialization::TypeType type = torch::jit::mobile::serialization::TypeType::UNSET,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>>> attr_names = 0) {
+  ObjectTypeBuilder builder_(_fbb);
+  builder_.add_attr_names(attr_names);
+  builder_.add_type_name(type_name);
+  builder_.add_type(type);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<ObjectType> CreateObjectTypeDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *type_name = nullptr,
+    torch::jit::mobile::serialization::TypeType type = torch::jit::mobile::serialization::TypeType::UNSET,
+    const std::vector<::flatbuffers::Offset<::flatbuffers::String>> *attr_names = nullptr) {
+  auto type_name__ = type_name ? _fbb.CreateString(type_name) : 0;
+  auto attr_names__ = attr_names ? _fbb.CreateVector<::flatbuffers::Offset<::flatbuffers::String>>(*attr_names) : 0;
+  return torch::jit::mobile::serialization::CreateObjectType(
+      _fbb,
+      type_name__,
+      type,
+      attr_names__);
+}
+
+struct Object FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef ObjectBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_TYPE_INDEX = 4,
+    VT_STATE = 6,
+    VT_ATTRS = 8,
+    VT_SETSTATE_FUNC = 10
+  };
+  uint32_t type_index() const {
+    return GetField<uint32_t>(VT_TYPE_INDEX, 0);
+  }
+  bool mutate_type_index(uint32_t _type_index = 0) {
+    return SetField<uint32_t>(VT_TYPE_INDEX, _type_index, 0);
+  }
+  uint32_t state() const {
+    return GetField<uint32_t>(VT_STATE, 0);
+  }
+  bool mutate_state(uint32_t _state = 0) {
+    return SetField<uint32_t>(VT_STATE, _state, 0);
+  }
+  const ::flatbuffers::Vector<uint32_t> *attrs() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_ATTRS);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_attrs() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_ATTRS);
+  }
+  uint32_t setstate_func() const {
+    return GetField<uint32_t>(VT_SETSTATE_FUNC, 0);
+  }
+  bool mutate_setstate_func(uint32_t _setstate_func = 0) {
+    return SetField<uint32_t>(VT_SETSTATE_FUNC, _setstate_func, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyField<uint32_t>(verifier, VT_TYPE_INDEX, 4) &&
+           VerifyField<uint32_t>(verifier, VT_STATE, 4) &&
+           VerifyOffset(verifier, VT_ATTRS) &&
+           verifier.VerifyVector(attrs()) &&
+           VerifyField<uint32_t>(verifier, VT_SETSTATE_FUNC, 4) &&
+           verifier.EndTable();
+  }
+};
+
+struct ObjectBuilder {
+  typedef Object Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_type_index(uint32_t type_index) {
+    fbb_.AddElement<uint32_t>(Object::VT_TYPE_INDEX, type_index, 0);
+  }
+  void add_state(uint32_t state) {
+    fbb_.AddElement<uint32_t>(Object::VT_STATE, state, 0);
+  }
+  void add_attrs(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> attrs) {
+    fbb_.AddOffset(Object::VT_ATTRS, attrs);
+  }
+  void add_setstate_func(uint32_t setstate_func) {
+    fbb_.AddElement<uint32_t>(Object::VT_SETSTATE_FUNC, setstate_func, 0);
+  }
+  explicit ObjectBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Object> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Object>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Object> CreateObject(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    uint32_t type_index = 0,
+    uint32_t state = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> attrs = 0,
+    uint32_t setstate_func = 0) {
+  ObjectBuilder builder_(_fbb);
+  builder_.add_setstate_func(setstate_func);
+  builder_.add_attrs(attrs);
+  builder_.add_state(state);
+  builder_.add_type_index(type_index);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Object> CreateObjectDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    uint32_t type_index = 0,
+    uint32_t state = 0,
+    const std::vector<uint32_t> *attrs = nullptr,
+    uint32_t setstate_func = 0) {
+  auto attrs__ = attrs ? _fbb.CreateVector<uint32_t>(*attrs) : 0;
+  return torch::jit::mobile::serialization::CreateObject(
+      _fbb,
+      type_index,
+      state,
+      attrs__,
+      setstate_func);
+}
+
+struct EnumValue FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef EnumValueBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_TYPE_NAME = 4,
+    VT_VALUE = 6
+  };
+  const ::flatbuffers::String *type_name() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_TYPE_NAME);
+  }
+  ::flatbuffers::String *mutable_type_name() {
+    return GetPointer<::flatbuffers::String *>(VT_TYPE_NAME);
+  }
+  uint32_t value() const {
+    return GetField<uint32_t>(VT_VALUE, 0);
+  }
+  bool mutate_value(uint32_t _value = 0) {
+    return SetField<uint32_t>(VT_VALUE, _value, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_TYPE_NAME) &&
+           verifier.VerifyString(type_name()) &&
+           VerifyField<uint32_t>(verifier, VT_VALUE, 4) &&
+           verifier.EndTable();
+  }
+};
+
+struct EnumValueBuilder {
+  typedef EnumValue Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_type_name(::flatbuffers::Offset<::flatbuffers::String> type_name) {
+    fbb_.AddOffset(EnumValue::VT_TYPE_NAME, type_name);
+  }
+  void add_value(uint32_t value) {
+    fbb_.AddElement<uint32_t>(EnumValue::VT_VALUE, value, 0);
+  }
+  explicit EnumValueBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<EnumValue> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<EnumValue>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<EnumValue> CreateEnumValue(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> type_name = 0,
+    uint32_t value = 0) {
+  EnumValueBuilder builder_(_fbb);
+  builder_.add_value(value);
+  builder_.add_type_name(type_name);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<EnumValue> CreateEnumValueDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *type_name = nullptr,
+    uint32_t value = 0) {
+  auto type_name__ = type_name ? _fbb.CreateString(type_name) : 0;
+  return torch::jit::mobile::serialization::CreateEnumValue(
+      _fbb,
+      type_name__,
+      value);
+}
+
+struct Operator FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef OperatorBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_NAME = 4,
+    VT_OVERLOAD_NAME = 6,
+    VT_NUM_ARGS_SERIALIZED = 8
+  };
+  const ::flatbuffers::String *name() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_NAME);
+  }
+  ::flatbuffers::String *mutable_name() {
+    return GetPointer<::flatbuffers::String *>(VT_NAME);
+  }
+  const ::flatbuffers::String *overload_name() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_OVERLOAD_NAME);
+  }
+  ::flatbuffers::String *mutable_overload_name() {
+    return GetPointer<::flatbuffers::String *>(VT_OVERLOAD_NAME);
+  }
+  int32_t num_args_serialized() const {
+    return GetField<int32_t>(VT_NUM_ARGS_SERIALIZED, -1);
+  }
+  bool mutate_num_args_serialized(int32_t _num_args_serialized = -1) {
+    return SetField<int32_t>(VT_NUM_ARGS_SERIALIZED, _num_args_serialized, -1);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_NAME) &&
+           verifier.VerifyString(name()) &&
+           VerifyOffset(verifier, VT_OVERLOAD_NAME) &&
+           verifier.VerifyString(overload_name()) &&
+           VerifyField<int32_t>(verifier, VT_NUM_ARGS_SERIALIZED, 4) &&
+           verifier.EndTable();
+  }
+};
+
+struct OperatorBuilder {
+  typedef Operator Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_name(::flatbuffers::Offset<::flatbuffers::String> name) {
+    fbb_.AddOffset(Operator::VT_NAME, name);
+  }
+  void add_overload_name(::flatbuffers::Offset<::flatbuffers::String> overload_name) {
+    fbb_.AddOffset(Operator::VT_OVERLOAD_NAME, overload_name);
+  }
+  void add_num_args_serialized(int32_t num_args_serialized) {
+    fbb_.AddElement<int32_t>(Operator::VT_NUM_ARGS_SERIALIZED, num_args_serialized, -1);
+  }
+  explicit OperatorBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Operator> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Operator>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Operator> CreateOperator(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> name = 0,
+    ::flatbuffers::Offset<::flatbuffers::String> overload_name = 0,
+    int32_t num_args_serialized = -1) {
+  OperatorBuilder builder_(_fbb);
+  builder_.add_num_args_serialized(num_args_serialized);
+  builder_.add_overload_name(overload_name);
+  builder_.add_name(name);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Operator> CreateOperatorDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *name = nullptr,
+    const char *overload_name = nullptr,
+    int32_t num_args_serialized = -1) {
+  auto name__ = name ? _fbb.CreateString(name) : 0;
+  auto overload_name__ = overload_name ? _fbb.CreateString(overload_name) : 0;
+  return torch::jit::mobile::serialization::CreateOperator(
+      _fbb,
+      name__,
+      overload_name__,
+      num_args_serialized);
+}
+
+struct Arg FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef ArgBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_NAME = 4,
+    VT_TYPE = 6,
+    VT_DEFAULT_VALUE = 8
+  };
+  const ::flatbuffers::String *name() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_NAME);
+  }
+  ::flatbuffers::String *mutable_name() {
+    return GetPointer<::flatbuffers::String *>(VT_NAME);
+  }
+  const ::flatbuffers::String *type() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_TYPE);
+  }
+  ::flatbuffers::String *mutable_type() {
+    return GetPointer<::flatbuffers::String *>(VT_TYPE);
+  }
+  uint32_t default_value() const {
+    return GetField<uint32_t>(VT_DEFAULT_VALUE, 0);
+  }
+  bool mutate_default_value(uint32_t _default_value = 0) {
+    return SetField<uint32_t>(VT_DEFAULT_VALUE, _default_value, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_NAME) &&
+           verifier.VerifyString(name()) &&
+           VerifyOffset(verifier, VT_TYPE) &&
+           verifier.VerifyString(type()) &&
+           VerifyField<uint32_t>(verifier, VT_DEFAULT_VALUE, 4) &&
+           verifier.EndTable();
+  }
+};
+
+struct ArgBuilder {
+  typedef Arg Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_name(::flatbuffers::Offset<::flatbuffers::String> name) {
+    fbb_.AddOffset(Arg::VT_NAME, name);
+  }
+  void add_type(::flatbuffers::Offset<::flatbuffers::String> type) {
+    fbb_.AddOffset(Arg::VT_TYPE, type);
+  }
+  void add_default_value(uint32_t default_value) {
+    fbb_.AddElement<uint32_t>(Arg::VT_DEFAULT_VALUE, default_value, 0);
+  }
+  explicit ArgBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Arg> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Arg>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Arg> CreateArg(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> name = 0,
+    ::flatbuffers::Offset<::flatbuffers::String> type = 0,
+    uint32_t default_value = 0) {
+  ArgBuilder builder_(_fbb);
+  builder_.add_default_value(default_value);
+  builder_.add_type(type);
+  builder_.add_name(name);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Arg> CreateArgDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *name = nullptr,
+    const char *type = nullptr,
+    uint32_t default_value = 0) {
+  auto name__ = name ? _fbb.CreateString(name) : 0;
+  auto type__ = type ? _fbb.CreateString(type) : 0;
+  return torch::jit::mobile::serialization::CreateArg(
+      _fbb,
+      name__,
+      type__,
+      default_value);
+}
+
+struct Schema FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef SchemaBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_ARGUMENTS = 4,
+    VT_RETURNS = 6
+  };
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *arguments() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *>(VT_ARGUMENTS);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *mutable_arguments() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *>(VT_ARGUMENTS);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *returns() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *>(VT_RETURNS);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *mutable_returns() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *>(VT_RETURNS);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_ARGUMENTS) &&
+           verifier.VerifyVector(arguments()) &&
+           verifier.VerifyVectorOfTables(arguments()) &&
+           VerifyOffset(verifier, VT_RETURNS) &&
+           verifier.VerifyVector(returns()) &&
+           verifier.VerifyVectorOfTables(returns()) &&
+           verifier.EndTable();
+  }
+};
+
+struct SchemaBuilder {
+  typedef Schema Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_arguments(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>>> arguments) {
+    fbb_.AddOffset(Schema::VT_ARGUMENTS, arguments);
+  }
+  void add_returns(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>>> returns) {
+    fbb_.AddOffset(Schema::VT_RETURNS, returns);
+  }
+  explicit SchemaBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Schema> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Schema>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Schema> CreateSchema(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>>> arguments = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>>> returns = 0) {
+  SchemaBuilder builder_(_fbb);
+  builder_.add_returns(returns);
+  builder_.add_arguments(arguments);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Schema> CreateSchemaDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *arguments = nullptr,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>> *returns = nullptr) {
+  auto arguments__ = arguments ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>>(*arguments) : 0;
+  auto returns__ = returns ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::Arg>>(*returns) : 0;
+  return torch::jit::mobile::serialization::CreateSchema(
+      _fbb,
+      arguments__,
+      returns__);
+}
+
+struct DebugInfo FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef DebugInfoBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_DEBUG_HANDLE = 4
+  };
+  const ::flatbuffers::Vector<int64_t> *debug_handle() const {
+    return GetPointer<const ::flatbuffers::Vector<int64_t> *>(VT_DEBUG_HANDLE);
+  }
+  ::flatbuffers::Vector<int64_t> *mutable_debug_handle() {
+    return GetPointer<::flatbuffers::Vector<int64_t> *>(VT_DEBUG_HANDLE);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_DEBUG_HANDLE) &&
+           verifier.VerifyVector(debug_handle()) &&
+           verifier.EndTable();
+  }
+};
+
+struct DebugInfoBuilder {
+  typedef DebugInfo Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_debug_handle(::flatbuffers::Offset<::flatbuffers::Vector<int64_t>> debug_handle) {
+    fbb_.AddOffset(DebugInfo::VT_DEBUG_HANDLE, debug_handle);
+  }
+  explicit DebugInfoBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<DebugInfo> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<DebugInfo>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<DebugInfo> CreateDebugInfo(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<int64_t>> debug_handle = 0) {
+  DebugInfoBuilder builder_(_fbb);
+  builder_.add_debug_handle(debug_handle);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<DebugInfo> CreateDebugInfoDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<int64_t> *debug_handle = nullptr) {
+  auto debug_handle__ = debug_handle ? _fbb.CreateVector<int64_t>(*debug_handle) : 0;
+  return torch::jit::mobile::serialization::CreateDebugInfo(
+      _fbb,
+      debug_handle__);
+}
+
+struct Function FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef FunctionBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_QN = 4,
+    VT_INSTRUCTIONS = 6,
+    VT_OPERATORS = 8,
+    VT_CONSTANTS = 10,
+    VT_TYPE_ANNOTATIONS = 12,
+    VT_REGISTER_SIZE = 14,
+    VT_SCHEMA = 16,
+    VT_DEBUG_INFO = 18,
+    VT_CLASS_TYPE = 20
+  };
+  const ::flatbuffers::String *qn() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_QN);
+  }
+  ::flatbuffers::String *mutable_qn() {
+    return GetPointer<::flatbuffers::String *>(VT_QN);
+  }
+  const ::flatbuffers::Vector<const torch::jit::mobile::serialization::Instruction *> *instructions() const {
+    return GetPointer<const ::flatbuffers::Vector<const torch::jit::mobile::serialization::Instruction *> *>(VT_INSTRUCTIONS);
+  }
+  ::flatbuffers::Vector<const torch::jit::mobile::serialization::Instruction *> *mutable_instructions() {
+    return GetPointer<::flatbuffers::Vector<const torch::jit::mobile::serialization::Instruction *> *>(VT_INSTRUCTIONS);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>> *operators() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>> *>(VT_OPERATORS);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>> *mutable_operators() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>> *>(VT_OPERATORS);
+  }
+  const ::flatbuffers::Vector<uint32_t> *constants() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_CONSTANTS);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_constants() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_CONSTANTS);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *type_annotations() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *>(VT_TYPE_ANNOTATIONS);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *mutable_type_annotations() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>> *>(VT_TYPE_ANNOTATIONS);
+  }
+  int32_t register_size() const {
+    return GetField<int32_t>(VT_REGISTER_SIZE, 0);
+  }
+  bool mutate_register_size(int32_t _register_size = 0) {
+    return SetField<int32_t>(VT_REGISTER_SIZE, _register_size, 0);
+  }
+  const torch::jit::mobile::serialization::Schema *schema() const {
+    return GetPointer<const torch::jit::mobile::serialization::Schema *>(VT_SCHEMA);
+  }
+  torch::jit::mobile::serialization::Schema *mutable_schema() {
+    return GetPointer<torch::jit::mobile::serialization::Schema *>(VT_SCHEMA);
+  }
+  const torch::jit::mobile::serialization::DebugInfo *debug_info() const {
+    return GetPointer<const torch::jit::mobile::serialization::DebugInfo *>(VT_DEBUG_INFO);
+  }
+  torch::jit::mobile::serialization::DebugInfo *mutable_debug_info() {
+    return GetPointer<torch::jit::mobile::serialization::DebugInfo *>(VT_DEBUG_INFO);
+  }
+  uint32_t class_type() const {
+    return GetField<uint32_t>(VT_CLASS_TYPE, 0);
+  }
+  bool mutate_class_type(uint32_t _class_type = 0) {
+    return SetField<uint32_t>(VT_CLASS_TYPE, _class_type, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_QN) &&
+           verifier.VerifyString(qn()) &&
+           VerifyOffset(verifier, VT_INSTRUCTIONS) &&
+           verifier.VerifyVector(instructions()) &&
+           VerifyOffset(verifier, VT_OPERATORS) &&
+           verifier.VerifyVector(operators()) &&
+           verifier.VerifyVectorOfTables(operators()) &&
+           VerifyOffset(verifier, VT_CONSTANTS) &&
+           verifier.VerifyVector(constants()) &&
+           VerifyOffset(verifier, VT_TYPE_ANNOTATIONS) &&
+           verifier.VerifyVector(type_annotations()) &&
+           verifier.VerifyVectorOfStrings(type_annotations()) &&
+           VerifyField<int32_t>(verifier, VT_REGISTER_SIZE, 4) &&
+           VerifyOffset(verifier, VT_SCHEMA) &&
+           verifier.VerifyTable(schema()) &&
+           VerifyOffset(verifier, VT_DEBUG_INFO) &&
+           verifier.VerifyTable(debug_info()) &&
+           VerifyField<uint32_t>(verifier, VT_CLASS_TYPE, 4) &&
+           verifier.EndTable();
+  }
+};
+
+struct FunctionBuilder {
+  typedef Function Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_qn(::flatbuffers::Offset<::flatbuffers::String> qn) {
+    fbb_.AddOffset(Function::VT_QN, qn);
+  }
+  void add_instructions(::flatbuffers::Offset<::flatbuffers::Vector<const torch::jit::mobile::serialization::Instruction *>> instructions) {
+    fbb_.AddOffset(Function::VT_INSTRUCTIONS, instructions);
+  }
+  void add_operators(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>>> operators) {
+    fbb_.AddOffset(Function::VT_OPERATORS, operators);
+  }
+  void add_constants(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> constants) {
+    fbb_.AddOffset(Function::VT_CONSTANTS, constants);
+  }
+  void add_type_annotations(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>>> type_annotations) {
+    fbb_.AddOffset(Function::VT_TYPE_ANNOTATIONS, type_annotations);
+  }
+  void add_register_size(int32_t register_size) {
+    fbb_.AddElement<int32_t>(Function::VT_REGISTER_SIZE, register_size, 0);
+  }
+  void add_schema(::flatbuffers::Offset<torch::jit::mobile::serialization::Schema> schema) {
+    fbb_.AddOffset(Function::VT_SCHEMA, schema);
+  }
+  void add_debug_info(::flatbuffers::Offset<torch::jit::mobile::serialization::DebugInfo> debug_info) {
+    fbb_.AddOffset(Function::VT_DEBUG_INFO, debug_info);
+  }
+  void add_class_type(uint32_t class_type) {
+    fbb_.AddElement<uint32_t>(Function::VT_CLASS_TYPE, class_type, 0);
+  }
+  explicit FunctionBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Function> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Function>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Function> CreateFunction(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> qn = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<const torch::jit::mobile::serialization::Instruction *>> instructions = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>>> operators = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> constants = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<::flatbuffers::String>>> type_annotations = 0,
+    int32_t register_size = 0,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::Schema> schema = 0,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::DebugInfo> debug_info = 0,
+    uint32_t class_type = 0) {
+  FunctionBuilder builder_(_fbb);
+  builder_.add_class_type(class_type);
+  builder_.add_debug_info(debug_info);
+  builder_.add_schema(schema);
+  builder_.add_register_size(register_size);
+  builder_.add_type_annotations(type_annotations);
+  builder_.add_constants(constants);
+  builder_.add_operators(operators);
+  builder_.add_instructions(instructions);
+  builder_.add_qn(qn);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Function> CreateFunctionDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *qn = nullptr,
+    const std::vector<torch::jit::mobile::serialization::Instruction> *instructions = nullptr,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>> *operators = nullptr,
+    const std::vector<uint32_t> *constants = nullptr,
+    const std::vector<::flatbuffers::Offset<::flatbuffers::String>> *type_annotations = nullptr,
+    int32_t register_size = 0,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::Schema> schema = 0,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::DebugInfo> debug_info = 0,
+    uint32_t class_type = 0) {
+  auto qn__ = qn ? _fbb.CreateString(qn) : 0;
+  auto instructions__ = instructions ? _fbb.CreateVectorOfStructs<torch::jit::mobile::serialization::Instruction>(*instructions) : 0;
+  auto operators__ = operators ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::Operator>>(*operators) : 0;
+  auto constants__ = constants ? _fbb.CreateVector<uint32_t>(*constants) : 0;
+  auto type_annotations__ = type_annotations ? _fbb.CreateVector<::flatbuffers::Offset<::flatbuffers::String>>(*type_annotations) : 0;
+  return torch::jit::mobile::serialization::CreateFunction(
+      _fbb,
+      qn__,
+      instructions__,
+      operators__,
+      constants__,
+      type_annotations__,
+      register_size,
+      schema,
+      debug_info,
+      class_type);
+}
+
+struct StorageData FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef StorageDataBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_DATA = 4
+  };
+  const ::flatbuffers::Vector<uint8_t> *data() const {
+    return GetPointer<const ::flatbuffers::Vector<uint8_t> *>(VT_DATA);
+  }
+  ::flatbuffers::Vector<uint8_t> *mutable_data() {
+    return GetPointer<::flatbuffers::Vector<uint8_t> *>(VT_DATA);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_DATA) &&
+           verifier.VerifyVector(data()) &&
+           verifier.EndTable();
+  }
+};
+
+struct StorageDataBuilder {
+  typedef StorageData Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_data(::flatbuffers::Offset<::flatbuffers::Vector<uint8_t>> data) {
+    fbb_.AddOffset(StorageData::VT_DATA, data);
+  }
+  explicit StorageDataBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<StorageData> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<StorageData>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<StorageData> CreateStorageData(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint8_t>> data = 0) {
+  StorageDataBuilder builder_(_fbb);
+  builder_.add_data(data);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<StorageData> CreateStorageDataDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<uint8_t> *data = nullptr) {
+  if (data) { _fbb.ForceVectorAlignment(data->size(), sizeof(uint8_t), 16); }
+  auto data__ = data ? _fbb.CreateVector<uint8_t>(*data) : 0;
+  return torch::jit::mobile::serialization::CreateStorageData(
+      _fbb,
+      data__);
+}
+
+struct IValue FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef IValueBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_VAL_TYPE = 4,
+    VT_VAL = 6
+  };
+  torch::jit::mobile::serialization::IValueUnion val_type() const {
+    return static_cast<torch::jit::mobile::serialization::IValueUnion>(GetField<uint8_t>(VT_VAL_TYPE, 0));
+  }
+  const void *val() const {
+    return GetPointer<const void *>(VT_VAL);
+  }
+  template<typename T> const T *val_as() const;
+  const torch::jit::mobile::serialization::Int *val_as_Int() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Int ? static_cast<const torch::jit::mobile::serialization::Int *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::Bool *val_as_Bool() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Bool ? static_cast<const torch::jit::mobile::serialization::Bool *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::Double *val_as_Double() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Double ? static_cast<const torch::jit::mobile::serialization::Double *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::ComplexDouble *val_as_ComplexDouble() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::ComplexDouble ? static_cast<const torch::jit::mobile::serialization::ComplexDouble *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::TensorMetadata *val_as_TensorMetadata() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::TensorMetadata ? static_cast<const torch::jit::mobile::serialization::TensorMetadata *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::String *val_as_String() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::String ? static_cast<const torch::jit::mobile::serialization::String *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::List *val_as_List() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::List ? static_cast<const torch::jit::mobile::serialization::List *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::Tuple *val_as_Tuple() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Tuple ? static_cast<const torch::jit::mobile::serialization::Tuple *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::Dict *val_as_Dict() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Dict ? static_cast<const torch::jit::mobile::serialization::Dict *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::Object *val_as_Object() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Object ? static_cast<const torch::jit::mobile::serialization::Object *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::IntList *val_as_IntList() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::IntList ? static_cast<const torch::jit::mobile::serialization::IntList *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::DoubleList *val_as_DoubleList() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::DoubleList ? static_cast<const torch::jit::mobile::serialization::DoubleList *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::BoolList *val_as_BoolList() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::BoolList ? static_cast<const torch::jit::mobile::serialization::BoolList *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::Device *val_as_Device() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Device ? static_cast<const torch::jit::mobile::serialization::Device *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::EnumValue *val_as_EnumValue() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::EnumValue ? static_cast<const torch::jit::mobile::serialization::EnumValue *>(val()) : nullptr;
+  }
+  const torch::jit::mobile::serialization::Function *val_as_Function() const {
+    return val_type() == torch::jit::mobile::serialization::IValueUnion::Function ? static_cast<const torch::jit::mobile::serialization::Function *>(val()) : nullptr;
+  }
+  void *mutable_val() {
+    return GetPointer<void *>(VT_VAL);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyField<uint8_t>(verifier, VT_VAL_TYPE, 1) &&
+           VerifyOffset(verifier, VT_VAL) &&
+           VerifyIValueUnion(verifier, val(), val_type()) &&
+           verifier.EndTable();
+  }
+};
+
+template<> inline const torch::jit::mobile::serialization::Int *IValue::val_as<torch::jit::mobile::serialization::Int>() const {
+  return val_as_Int();
+}
+
+template<> inline const torch::jit::mobile::serialization::Bool *IValue::val_as<torch::jit::mobile::serialization::Bool>() const {
+  return val_as_Bool();
+}
+
+template<> inline const torch::jit::mobile::serialization::Double *IValue::val_as<torch::jit::mobile::serialization::Double>() const {
+  return val_as_Double();
+}
+
+template<> inline const torch::jit::mobile::serialization::ComplexDouble *IValue::val_as<torch::jit::mobile::serialization::ComplexDouble>() const {
+  return val_as_ComplexDouble();
+}
+
+template<> inline const torch::jit::mobile::serialization::TensorMetadata *IValue::val_as<torch::jit::mobile::serialization::TensorMetadata>() const {
+  return val_as_TensorMetadata();
+}
+
+template<> inline const torch::jit::mobile::serialization::String *IValue::val_as<torch::jit::mobile::serialization::String>() const {
+  return val_as_String();
+}
+
+template<> inline const torch::jit::mobile::serialization::List *IValue::val_as<torch::jit::mobile::serialization::List>() const {
+  return val_as_List();
+}
+
+template<> inline const torch::jit::mobile::serialization::Tuple *IValue::val_as<torch::jit::mobile::serialization::Tuple>() const {
+  return val_as_Tuple();
+}
+
+template<> inline const torch::jit::mobile::serialization::Dict *IValue::val_as<torch::jit::mobile::serialization::Dict>() const {
+  return val_as_Dict();
+}
+
+template<> inline const torch::jit::mobile::serialization::Object *IValue::val_as<torch::jit::mobile::serialization::Object>() const {
+  return val_as_Object();
+}
+
+template<> inline const torch::jit::mobile::serialization::IntList *IValue::val_as<torch::jit::mobile::serialization::IntList>() const {
+  return val_as_IntList();
+}
+
+template<> inline const torch::jit::mobile::serialization::DoubleList *IValue::val_as<torch::jit::mobile::serialization::DoubleList>() const {
+  return val_as_DoubleList();
+}
+
+template<> inline const torch::jit::mobile::serialization::BoolList *IValue::val_as<torch::jit::mobile::serialization::BoolList>() const {
+  return val_as_BoolList();
+}
+
+template<> inline const torch::jit::mobile::serialization::Device *IValue::val_as<torch::jit::mobile::serialization::Device>() const {
+  return val_as_Device();
+}
+
+template<> inline const torch::jit::mobile::serialization::EnumValue *IValue::val_as<torch::jit::mobile::serialization::EnumValue>() const {
+  return val_as_EnumValue();
+}
+
+template<> inline const torch::jit::mobile::serialization::Function *IValue::val_as<torch::jit::mobile::serialization::Function>() const {
+  return val_as_Function();
+}
+
+struct IValueBuilder {
+  typedef IValue Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_val_type(torch::jit::mobile::serialization::IValueUnion val_type) {
+    fbb_.AddElement<uint8_t>(IValue::VT_VAL_TYPE, static_cast<uint8_t>(val_type), 0);
+  }
+  void add_val(::flatbuffers::Offset<void> val) {
+    fbb_.AddOffset(IValue::VT_VAL, val);
+  }
+  explicit IValueBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<IValue> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<IValue>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<IValue> CreateIValue(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    torch::jit::mobile::serialization::IValueUnion val_type = torch::jit::mobile::serialization::IValueUnion::NONE,
+    ::flatbuffers::Offset<void> val = 0) {
+  IValueBuilder builder_(_fbb);
+  builder_.add_val(val);
+  builder_.add_val_type(val_type);
+  return builder_.Finish();
+}
+
+struct ExtraFile FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef ExtraFileBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_NAME = 4,
+    VT_CONTENT = 6
+  };
+  const ::flatbuffers::String *name() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_NAME);
+  }
+  ::flatbuffers::String *mutable_name() {
+    return GetPointer<::flatbuffers::String *>(VT_NAME);
+  }
+  const ::flatbuffers::String *content() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_CONTENT);
+  }
+  ::flatbuffers::String *mutable_content() {
+    return GetPointer<::flatbuffers::String *>(VT_CONTENT);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_NAME) &&
+           verifier.VerifyString(name()) &&
+           VerifyOffset(verifier, VT_CONTENT) &&
+           verifier.VerifyString(content()) &&
+           verifier.EndTable();
+  }
+};
+
+struct ExtraFileBuilder {
+  typedef ExtraFile Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_name(::flatbuffers::Offset<::flatbuffers::String> name) {
+    fbb_.AddOffset(ExtraFile::VT_NAME, name);
+  }
+  void add_content(::flatbuffers::Offset<::flatbuffers::String> content) {
+    fbb_.AddOffset(ExtraFile::VT_CONTENT, content);
+  }
+  explicit ExtraFileBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<ExtraFile> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<ExtraFile>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<ExtraFile> CreateExtraFile(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::String> name = 0,
+    ::flatbuffers::Offset<::flatbuffers::String> content = 0) {
+  ExtraFileBuilder builder_(_fbb);
+  builder_.add_content(content);
+  builder_.add_name(name);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<ExtraFile> CreateExtraFileDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const char *name = nullptr,
+    const char *content = nullptr) {
+  auto name__ = name ? _fbb.CreateString(name) : 0;
+  auto content__ = content ? _fbb.CreateString(content) : 0;
+  return torch::jit::mobile::serialization::CreateExtraFile(
+      _fbb,
+      name__,
+      content__);
+}
+
+struct Module FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef ModuleBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_BYTECODE_VERSION = 4,
+    VT_EXTRA_FILES = 6,
+    VT_METHODS = 8,
+    VT_STATE_OBJ = 10,
+    VT_IVALUES = 12,
+    VT_STORAGE_DATA_SIZE = 14,
+    VT_STORAGE_DATA = 16,
+    VT_OBJECT_TYPES = 18,
+    VT_JIT_SOURCES = 20,
+    VT_JIT_CONSTANTS = 22,
+    VT_OPERATOR_VERSION = 24,
+    VT_MOBILE_IVALUE_SIZE = 26
+  };
+  uint32_t bytecode_version() const {
+    return GetField<uint32_t>(VT_BYTECODE_VERSION, 0);
+  }
+  bool mutate_bytecode_version(uint32_t _bytecode_version = 0) {
+    return SetField<uint32_t>(VT_BYTECODE_VERSION, _bytecode_version, 0);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *extra_files() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *>(VT_EXTRA_FILES);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *mutable_extra_files() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *>(VT_EXTRA_FILES);
+  }
+  const ::flatbuffers::Vector<uint32_t> *methods() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_METHODS);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_methods() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_METHODS);
+  }
+  uint32_t state_obj() const {
+    return GetField<uint32_t>(VT_STATE_OBJ, 0);
+  }
+  bool mutate_state_obj(uint32_t _state_obj = 0) {
+    return SetField<uint32_t>(VT_STATE_OBJ, _state_obj, 0);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>> *ivalues() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>> *>(VT_IVALUES);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>> *mutable_ivalues() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>> *>(VT_IVALUES);
+  }
+  int32_t storage_data_size() const {
+    return GetField<int32_t>(VT_STORAGE_DATA_SIZE, 0);
+  }
+  bool mutate_storage_data_size(int32_t _storage_data_size = 0) {
+    return SetField<int32_t>(VT_STORAGE_DATA_SIZE, _storage_data_size, 0);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>> *storage_data() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>> *>(VT_STORAGE_DATA);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>> *mutable_storage_data() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>> *>(VT_STORAGE_DATA);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>> *object_types() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>> *>(VT_OBJECT_TYPES);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>> *mutable_object_types() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>> *>(VT_OBJECT_TYPES);
+  }
+  const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *jit_sources() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *>(VT_JIT_SOURCES);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *mutable_jit_sources() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *>(VT_JIT_SOURCES);
+  }
+  const ::flatbuffers::Vector<uint32_t> *jit_constants() const {
+    return GetPointer<const ::flatbuffers::Vector<uint32_t> *>(VT_JIT_CONSTANTS);
+  }
+  ::flatbuffers::Vector<uint32_t> *mutable_jit_constants() {
+    return GetPointer<::flatbuffers::Vector<uint32_t> *>(VT_JIT_CONSTANTS);
+  }
+  uint32_t operator_version() const {
+    return GetField<uint32_t>(VT_OPERATOR_VERSION, 0);
+  }
+  bool mutate_operator_version(uint32_t _operator_version = 0) {
+    return SetField<uint32_t>(VT_OPERATOR_VERSION, _operator_version, 0);
+  }
+  uint32_t mobile_ivalue_size() const {
+    return GetField<uint32_t>(VT_MOBILE_IVALUE_SIZE, 0);
+  }
+  bool mutate_mobile_ivalue_size(uint32_t _mobile_ivalue_size = 0) {
+    return SetField<uint32_t>(VT_MOBILE_IVALUE_SIZE, _mobile_ivalue_size, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyField<uint32_t>(verifier, VT_BYTECODE_VERSION, 4) &&
+           VerifyOffset(verifier, VT_EXTRA_FILES) &&
+           verifier.VerifyVector(extra_files()) &&
+           verifier.VerifyVectorOfTables(extra_files()) &&
+           VerifyOffset(verifier, VT_METHODS) &&
+           verifier.VerifyVector(methods()) &&
+           VerifyField<uint32_t>(verifier, VT_STATE_OBJ, 4) &&
+           VerifyOffset(verifier, VT_IVALUES) &&
+           verifier.VerifyVector(ivalues()) &&
+           verifier.VerifyVectorOfTables(ivalues()) &&
+           VerifyField<int32_t>(verifier, VT_STORAGE_DATA_SIZE, 4) &&
+           VerifyOffset(verifier, VT_STORAGE_DATA) &&
+           verifier.VerifyVector(storage_data()) &&
+           verifier.VerifyVectorOfTables(storage_data()) &&
+           VerifyOffset(verifier, VT_OBJECT_TYPES) &&
+           verifier.VerifyVector(object_types()) &&
+           verifier.VerifyVectorOfTables(object_types()) &&
+           VerifyOffset(verifier, VT_JIT_SOURCES) &&
+           verifier.VerifyVector(jit_sources()) &&
+           verifier.VerifyVectorOfTables(jit_sources()) &&
+           VerifyOffset(verifier, VT_JIT_CONSTANTS) &&
+           verifier.VerifyVector(jit_constants()) &&
+           VerifyField<uint32_t>(verifier, VT_OPERATOR_VERSION, 4) &&
+           VerifyField<uint32_t>(verifier, VT_MOBILE_IVALUE_SIZE, 4) &&
+           verifier.EndTable();
+  }
+};
+
+struct ModuleBuilder {
+  typedef Module Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_bytecode_version(uint32_t bytecode_version) {
+    fbb_.AddElement<uint32_t>(Module::VT_BYTECODE_VERSION, bytecode_version, 0);
+  }
+  void add_extra_files(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>> extra_files) {
+    fbb_.AddOffset(Module::VT_EXTRA_FILES, extra_files);
+  }
+  void add_methods(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> methods) {
+    fbb_.AddOffset(Module::VT_METHODS, methods);
+  }
+  void add_state_obj(uint32_t state_obj) {
+    fbb_.AddElement<uint32_t>(Module::VT_STATE_OBJ, state_obj, 0);
+  }
+  void add_ivalues(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>>> ivalues) {
+    fbb_.AddOffset(Module::VT_IVALUES, ivalues);
+  }
+  void add_storage_data_size(int32_t storage_data_size) {
+    fbb_.AddElement<int32_t>(Module::VT_STORAGE_DATA_SIZE, storage_data_size, 0);
+  }
+  void add_storage_data(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>>> storage_data) {
+    fbb_.AddOffset(Module::VT_STORAGE_DATA, storage_data);
+  }
+  void add_object_types(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>>> object_types) {
+    fbb_.AddOffset(Module::VT_OBJECT_TYPES, object_types);
+  }
+  void add_jit_sources(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>> jit_sources) {
+    fbb_.AddOffset(Module::VT_JIT_SOURCES, jit_sources);
+  }
+  void add_jit_constants(::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> jit_constants) {
+    fbb_.AddOffset(Module::VT_JIT_CONSTANTS, jit_constants);
+  }
+  void add_operator_version(uint32_t operator_version) {
+    fbb_.AddElement<uint32_t>(Module::VT_OPERATOR_VERSION, operator_version, 0);
+  }
+  void add_mobile_ivalue_size(uint32_t mobile_ivalue_size) {
+    fbb_.AddElement<uint32_t>(Module::VT_MOBILE_IVALUE_SIZE, mobile_ivalue_size, 0);
+  }
+  explicit ModuleBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Module> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Module>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Module> CreateModule(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    uint32_t bytecode_version = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>> extra_files = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> methods = 0,
+    uint32_t state_obj = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>>> ivalues = 0,
+    int32_t storage_data_size = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>>> storage_data = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>>> object_types = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>> jit_sources = 0,
+    ::flatbuffers::Offset<::flatbuffers::Vector<uint32_t>> jit_constants = 0,
+    uint32_t operator_version = 0,
+    uint32_t mobile_ivalue_size = 0) {
+  ModuleBuilder builder_(_fbb);
+  builder_.add_mobile_ivalue_size(mobile_ivalue_size);
+  builder_.add_operator_version(operator_version);
+  builder_.add_jit_constants(jit_constants);
+  builder_.add_jit_sources(jit_sources);
+  builder_.add_object_types(object_types);
+  builder_.add_storage_data(storage_data);
+  builder_.add_storage_data_size(storage_data_size);
+  builder_.add_ivalues(ivalues);
+  builder_.add_state_obj(state_obj);
+  builder_.add_methods(methods);
+  builder_.add_extra_files(extra_files);
+  builder_.add_bytecode_version(bytecode_version);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<Module> CreateModuleDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    uint32_t bytecode_version = 0,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *extra_files = nullptr,
+    const std::vector<uint32_t> *methods = nullptr,
+    uint32_t state_obj = 0,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>> *ivalues = nullptr,
+    int32_t storage_data_size = 0,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>> *storage_data = nullptr,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>> *object_types = nullptr,
+    const std::vector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>> *jit_sources = nullptr,
+    const std::vector<uint32_t> *jit_constants = nullptr,
+    uint32_t operator_version = 0,
+    uint32_t mobile_ivalue_size = 0) {
+  auto extra_files__ = extra_files ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>(*extra_files) : 0;
+  auto methods__ = methods ? _fbb.CreateVector<uint32_t>(*methods) : 0;
+  auto ivalues__ = ivalues ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::IValue>>(*ivalues) : 0;
+  auto storage_data__ = storage_data ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::StorageData>>(*storage_data) : 0;
+  auto object_types__ = object_types ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::ObjectType>>(*object_types) : 0;
+  auto jit_sources__ = jit_sources ? _fbb.CreateVector<::flatbuffers::Offset<torch::jit::mobile::serialization::ExtraFile>>(*jit_sources) : 0;
+  auto jit_constants__ = jit_constants ? _fbb.CreateVector<uint32_t>(*jit_constants) : 0;
+  return torch::jit::mobile::serialization::CreateModule(
+      _fbb,
+      bytecode_version,
+      extra_files__,
+      methods__,
+      state_obj,
+      ivalues__,
+      storage_data_size,
+      storage_data__,
+      object_types__,
+      jit_sources__,
+      jit_constants__,
+      operator_version,
+      mobile_ivalue_size);
+}
+
+inline bool VerifyIValueUnion(::flatbuffers::Verifier &verifier, const void *obj, IValueUnion type) {
+  switch (type) {
+    case IValueUnion::NONE: {
+      return true;
+    }
+    case IValueUnion::Int: {
+      return verifier.VerifyField<torch::jit::mobile::serialization::Int>(static_cast<const uint8_t *>(obj), 0, 8);
+    }
+    case IValueUnion::Bool: {
+      return verifier.VerifyField<torch::jit::mobile::serialization::Bool>(static_cast<const uint8_t *>(obj), 0, 1);
+    }
+    case IValueUnion::Double: {
+      return verifier.VerifyField<torch::jit::mobile::serialization::Double>(static_cast<const uint8_t *>(obj), 0, 8);
+    }
+    case IValueUnion::ComplexDouble: {
+      return verifier.VerifyField<torch::jit::mobile::serialization::ComplexDouble>(static_cast<const uint8_t *>(obj), 0, 8);
+    }
+    case IValueUnion::TensorMetadata: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::TensorMetadata *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::String: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::String *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::List: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::List *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::Tuple: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::Tuple *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::Dict: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::Dict *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::Object: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::Object *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::IntList: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::IntList *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::DoubleList: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::DoubleList *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::BoolList: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::BoolList *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::Device: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::Device *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::EnumValue: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::EnumValue *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case IValueUnion::Function: {
+      auto ptr = reinterpret_cast<const torch::jit::mobile::serialization::Function *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    default: return true;
+  }
+}
+
+inline bool VerifyIValueUnionVector(::flatbuffers::Verifier &verifier, const ::flatbuffers::Vector<::flatbuffers::Offset<void>> *values, const ::flatbuffers::Vector<IValueUnion> *types) {
+  if (!values || !types) return !values && !types;
+  if (values->size() != types->size()) return false;
+  for (::flatbuffers::uoffset_t i = 0; i < values->size(); ++i) {
+    if (!VerifyIValueUnion(
+        verifier,  values->Get(i), types->GetEnum<IValueUnion>(i))) {
+      return false;
+    }
+  }
+  return true;
+}
+
+inline const torch::jit::mobile::serialization::Module *GetModule(const void *buf) {
+  return ::flatbuffers::GetRoot<torch::jit::mobile::serialization::Module>(buf);
+}
+
+inline const torch::jit::mobile::serialization::Module *GetSizePrefixedModule(const void *buf) {
+  return ::flatbuffers::GetSizePrefixedRoot<torch::jit::mobile::serialization::Module>(buf);
+}
+
+inline Module *GetMutableModule(void *buf) {
+  return ::flatbuffers::GetMutableRoot<Module>(buf);
+}
+
+inline torch::jit::mobile::serialization::Module *GetMutableSizePrefixedModule(void *buf) {
+  return ::flatbuffers::GetMutableSizePrefixedRoot<torch::jit::mobile::serialization::Module>(buf);
+}
+
+inline const char *ModuleIdentifier() {
+  return "PTMF";
+}
+
+inline bool ModuleBufferHasIdentifier(const void *buf) {
+  return ::flatbuffers::BufferHasIdentifier(
+      buf, ModuleIdentifier());
+}
+
+inline bool SizePrefixedModuleBufferHasIdentifier(const void *buf) {
+  return ::flatbuffers::BufferHasIdentifier(
+      buf, ModuleIdentifier(), true);
+}
+
+inline bool VerifyModuleBuffer(
+    ::flatbuffers::Verifier &verifier) {
+  return verifier.VerifyBuffer<torch::jit::mobile::serialization::Module>(ModuleIdentifier());
+}
+
+inline bool VerifySizePrefixedModuleBuffer(
+    ::flatbuffers::Verifier &verifier) {
+  return verifier.VerifySizePrefixedBuffer<torch::jit::mobile::serialization::Module>(ModuleIdentifier());
+}
+
+inline void FinishModuleBuffer(
+    ::flatbuffers::FlatBufferBuilder &fbb,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::Module> root) {
+  fbb.Finish(root, ModuleIdentifier());
+}
+
+inline void FinishSizePrefixedModuleBuffer(
+    ::flatbuffers::FlatBufferBuilder &fbb,
+    ::flatbuffers::Offset<torch::jit::mobile::serialization::Module> root) {
+  fbb.FinishSizePrefixed(root, ModuleIdentifier());
+}
+
+}  // namespace serialization
+}  // namespace mobile
+}  // namespace jit
+}  // namespace torch
+
+#endif  // FLATBUFFERS_GENERATED_MOBILEBYTECODE_TORCH_JIT_MOBILE_SERIALIZATION_H_
+// @generated
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/onnx.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/onnx.h
new file mode 100644
index 0000000000000000000000000000000000000000..45d67dc480dddf573c951d42be5fc187175d8eaf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/onnx.h
@@ -0,0 +1,23 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED(
+    "-Winconsistent-missing-destructor-override")
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wsuggest-override")
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED(
+    "-Wdeprecated-dynamic-exception-spec")
+#include <onnx/onnx_pb.h>
+C10_DIAGNOSTIC_POP()
+C10_DIAGNOSTIC_POP()
+C10_DIAGNOSTIC_POP()
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit {
+
+TORCH_API std::string prettyPrint(const ::ONNX_NAMESPACE::ModelProto& model);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickle.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickle.h
new file mode 100644
index 0000000000000000000000000000000000000000..0cdf6cded25e7574958df43ea4bf338dc5bc6ba2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickle.h
@@ -0,0 +1,145 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <c10/util/ArrayRef.h>
+#include <caffe2/serialize/inline_container.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/serialization/pickler.h>
+#include <torch/csrc/jit/serialization/unpickler.h>
+
+namespace torch::jit {
+
+/// Pickle an IValue by calling a function to handle writing the data.
+///
+/// `writer` is a function that takes in a pointer to a chunk of memory and its
+/// size and consumes it.
+///
+/// See `jit::pickle` for more details.
+TORCH_API void pickle(
+    std::function<void(const char* data_start, size_t data_len)> writer,
+    const IValue& ivalue,
+    std::vector<at::Tensor>* tensor_table = nullptr);
+
+/// Save a `torch::IValue` in a format compatible with Python's `pickle` module
+///
+/// If present, `tensor_table` is a pointer to a table in which tensors that
+/// are contained within `ivalue` are stored, and the bytes returned by the
+/// pickler will only include references to these tensors in the table. This can
+/// be used to keep the binary blob size small.
+/// If not provided, tensors are stored in the same byte stream as the pickle
+/// data, similar to `torch.save()` in eager Python.
+///
+/// Pickled values can be loaded in Python and C++:
+/// \rst
+/// .. code-block:: cpp
+///
+///  torch::IValue float_value(2.3);
+///
+///  // TODO: when tensors are stored in the pickle, delete this
+///  std::vector<at::Tensor> tensor_table;
+///  auto data = torch::jit::pickle(float_value, &tensor_table);
+///
+///  std::vector<torch::IValue> ivalues =
+///      torch::jit::unpickle(data.data(), data.size());
+///
+/// .. code-block:: python
+///
+///   values = torch.load('data.pkl')
+///   print(values)
+///
+/// \endrst
+TORCH_API std::vector<char> pickle(
+    const IValue& ivalue,
+    std::vector<at::Tensor>* tensor_table = nullptr);
+
+/// Save a `torch::IValue` in a format that can be loaded by both
+/// `torch::pickle_load` in C++ and `torch.load` in Python.
+TORCH_API std::vector<char> pickle_save(const IValue& ivalue);
+
+/// Deserialize a `torch::IValue` from bytes produced by either
+/// `torch::pickle_save` in C++ or `torch.save` in Python
+TORCH_API IValue pickle_load(const std::vector<char>& data);
+
+/// Deserialize a `torch::IValue` from bytes produced by either
+/// `torch::pickle_save` in C++ or `torch.save` in Python with custom object.
+TORCH_API IValue pickle_load_obj(std::string_view data);
+
+/// `reader` is a function that takes in a size to read from some pickled
+/// binary. `reader` should remember where it last read, and return
+/// the number of bytes read.
+/// See `torch::pickle` for details.
+/// type_resolver is used to resolve any JIT type based on type str
+TORCH_API IValue unpickle(
+    std::function<size_t(char*, size_t)> reader,
+    TypeResolver type_resolver,
+    c10::ArrayRef<at::Tensor> tensor_table,
+    c10::TypePtr (*type_parser)(const std::string&) =
+        Unpickler::defaultTypeParser,
+    ObjLoader obj_loader = nullptr);
+
+/// Decode a chunk of memory containing pickled data into its `torch::IValue`s.
+///
+/// If any `torch::IValue`s in the pickled data are `Object`s, then a
+/// `class_resolver` function must be provided.
+///
+/// See `torch::pickle` for details.
+TORCH_API IValue unpickle(
+    const char* data,
+    size_t size,
+    TypeResolver type_resolver = nullptr,
+    c10::ArrayRef<at::Tensor> tensor_table = {},
+    c10::TypePtr (*type_parser)(const std::string&) =
+        Unpickler::defaultTypeParser);
+
+/// Decode a chunk of memory containing pickled data into its `torch::IValue`s.
+///
+/// If any `torch::IValue`s in the pickled data are `Object`s, then a
+/// `class_resolver` function must be provided.
+///
+/// See `torch::pickle` for details.
+TORCH_API IValue unpickle(
+    const char* data,
+    size_t size,
+    ObjLoader obj_loader,
+    TypeResolver type_resolver = nullptr,
+    c10::ArrayRef<at::Tensor> tensor_table = {},
+    c10::TypePtr (*type_parser)(const std::string&) =
+        Unpickler::defaultTypeParser);
+
+#ifndef C10_MOBILE
+class VectorReader : public caffe2::serialize::ReadAdapterInterface {
+ public:
+  VectorReader(std::vector<char> data) : data_(std::move(data)) {}
+
+  size_t size() const override {
+    return data_.size();
+  }
+
+  size_t read(uint64_t pos, void* buf, size_t n, const char* what)
+      const override;
+
+ private:
+  std::vector<char> data_;
+};
+
+class StringViewReader : public caffe2::serialize::ReadAdapterInterface {
+ public:
+  StringViewReader(std::string_view data) : data_(data) {}
+
+  size_t size() const override {
+    return data_.size();
+  }
+
+  size_t read(uint64_t pos, void* buf, size_t n, const char* what)
+      const override;
+
+ private:
+  std::string_view data_;
+};
+#endif
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickler.h
new file mode 100644
index 0000000000000000000000000000000000000000..6361fed7b397d707816fbde1dedfebf39b0f2422
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickler.h
@@ -0,0 +1,190 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include <ATen/Utils.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/jit_type.h>
+#include <ATen/core/qualified_name.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/FbcodeMaps.h>
+#include <c10/util/intrusive_ptr.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/serialization/pickler_helper.h>
+
+namespace torch::jit {
+
+using ::c10::IValue;
+
+class TORCH_API Pickler {
+  AT_DISALLOW_COPY_AND_ASSIGN(Pickler);
+
+ public:
+  Pickler(std::function<void(const char*, size_t)> writer)
+      : Pickler(std::move(writer), nullptr, nullptr, nullptr) {}
+
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+  Pickler(
+      std::function<void(const char*, size_t)> writer,
+      std::vector<at::Tensor>* tensor_table,
+      std::function<c10::QualifiedName(const c10::ClassTypePtr&)> type_renamer,
+      std::vector<c10::ClassTypePtr>* memoized_class_types,
+      std::function<std::string(const at::Tensor&)> get_tensor_id = nullptr,
+      bool tag_aggregates = true)
+      : writer_(std::move(writer)),
+        tensor_table_(tensor_table),
+        type_renamer_(std::move(type_renamer)),
+        memoized_class_types_(memoized_class_types),
+        get_tensor_id_(std::move(get_tensor_id)),
+        tag_aggregates_(tag_aggregates) {}
+  ~Pickler();
+
+  // Push protocol onto the stack
+  void protocol();
+
+  // Push STOP PickleOpCode onto the stack
+  void stop();
+
+  void pushIValue(const IValue& ivalue);
+
+  void startTuple();
+  void endTuple();
+
+  const std::vector<at::Tensor>& tensorData() {
+    return tensor_data_;
+  }
+
+  void pushEmptyDict();
+  void pushDict(const IValue& ivalue);
+  void pushInt(int64_t value);
+  void pushLong(const std::string& data);
+
+ private:
+  void pushIValueImpl(const IValue& ivalue);
+  void startTypeTag();
+  void endTypeTag(const IValue& value);
+  void pushBool(bool value);
+  void pushDouble(double value);
+  void pushComplexDouble(const IValue& value);
+  void pushGenericList(const IValue& ivalue);
+  void pushIntList(const IValue& ivalue);
+  void pushList(const IValue& ivalue);
+  void pushTensor(const IValue& ivalue);
+  void pushTensorReference(const IValue& ivalue);
+  void pushLiteralTensor(const IValue& ivalue);
+  void pushLiteralSparseTensor(const at::Tensor& tensor);
+  void pushTuple(const IValue& ivalue);
+  void pushString(const std::string& string);
+  void pushDevice(const IValue& ivalue);
+#ifdef USE_DISTRIBUTED
+  void pushRRef(const IValue& ivalue);
+#endif
+  // unmemoized version
+  void pushStringImpl(const std::string& string);
+  void pushStorageOfTensor(const at::Tensor& tensor);
+
+  void pushBinGet(uint32_t memo_id);
+  void pushSpecializedList(
+      const IValue& ivalue,
+      const char* list_name,
+      const std::function<void(const IValue&)>& item_pusher);
+  void pushGlobal(std::string_view module_name, std::string_view class_name);
+  // raw string data is appended directly to the byte stream
+  void pushBytes(const std::string& string);
+  void pushTensorData(const at::Tensor& tensor);
+
+  // Add a BINPUT op and return the memoization id used
+  size_t pushNextBinPut();
+
+  const void* getPointer(const IValue& ivalue);
+
+  // Caller checks that bufferPos_ > 0
+  void flushNonEmpty() {
+    writer_(buffer_.data(), bufferPos_);
+    bufferPos_ = 0;
+  }
+
+  void flush() {
+    if (bufferPos_ != 0) {
+      flushNonEmpty();
+    }
+  }
+
+  // These convert values to bytes and add them to the stack (NB: since T is to
+  // the left of a '::', its type cannot be deduced by the compiler so one must
+  // explicitly instantiate the template, i.e. push<int>(int) works, push(int)
+  // does not)
+  static constexpr size_t kBufferSize = 256;
+  template <typename T>
+  void push(std::common_type_t<T> value) {
+    const char* begin = reinterpret_cast<const char*>(&value);
+    if (bufferPos_ + sizeof(T) > buffer_.size()) {
+      flushNonEmpty();
+    }
+    static_assert(sizeof(T) <= kBufferSize, "Buffer size assumption");
+    memcpy(buffer_.data() + bufferPos_, begin, sizeof(T));
+    bufferPos_ += sizeof(T);
+  }
+
+  // Stream to write binary data to
+  // Code shouldn't call writer_ directly without first flushing.
+  std::function<void(const char*, size_t)> writer_;
+
+  // Buffer to avoid calling a writer_ on a per-byte basis.
+  std::array<char, kBufferSize> buffer_;
+  size_t bufferPos_{0};
+
+  // Stack of opcodes/data
+  std::vector<char> stack_;
+
+  // External table of tensors to serialize. If this is missing, then tensors
+  // are serialized directly into the pickle
+  std::vector<at::Tensor>* tensor_table_;
+
+  // TODO: only use this if necessary (add a pass to find all shared ivalues,
+  // and only memoize those)
+  uint32_t memo_id_ = 0;
+
+  // Memoization of IValues that have been written (index in table is used for
+  // BINPUT opcodes) to enable shared references
+  c10::FastMap<const void*, uint32_t> memoized_ivalue_map_;
+
+  // because we de-dup ivalues based on their raw pointer address in the above
+  // map we need to keep all the memoized values alive during the pickle.
+  // Otherwise, it is possible that a raw address gets reused for another
+  // object, and we will alias it to the old object at that address.
+  std::vector<IValue> memoized_ivalues_;
+
+  std::function<c10::QualifiedName(const c10::ClassTypePtr&)> type_renamer_;
+
+  // List of all the types that it wrote, inspect from the IValues it wrote.
+  std::vector<c10::ClassTypePtr>* memoized_class_types_;
+
+  // Function to grab next id_name for tensor storage, function is responsible
+  // for returning unique ids
+  std::function<std::string(const at::Tensor&)> get_tensor_id_;
+
+  // List of tensor storages to serialize in the same binary as the pickle data
+  // similar to ivalues, they are memoized using BINPUT
+  std::vector<at::Tensor> tensor_data_;
+  c10::FastMap<const void*, uint32_t> memoized_storage_map_;
+
+  c10::FastMap<std::string, uint32_t> memoized_globals_map_;
+  c10::FastMap<std::string, uint32_t> memoized_strings_map_;
+  c10::FastMap<std::string, uint32_t> memoized_devices_map_;
+  // when true, List and Dict objects will be wrapped in a
+  // torch.jit._pickle.restore_type_tag call to correctly set the dynamic
+  // TorchScript type for the object. When true the thing unpickling must have
+  // torch installed.
+  bool tag_aggregates_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickler_helper.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickler_helper.h
new file mode 100644
index 0000000000000000000000000000000000000000..523e912a392aab0104f6279d32b90c12c6d3bd74
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/pickler_helper.h
@@ -0,0 +1,239 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+
+#include <ATen/Utils.h>
+#include <ATen/core/ivalue.h>
+
+namespace torch::jit {
+
+// See Python's pickletools.py for a detailed description of each of these codes
+enum class PickleOpCode : char {
+  MARK = '(',
+  STOP = '.',
+  POP = '0',
+  POP_MARK = '1',
+  DUP = '2',
+  FLOAT = 'F',
+  INT = 'I',
+  BININT = 'J',
+  BININT1 = 'K',
+  LONG = 'L',
+  BININT2 = 'M',
+  NONE = 'N',
+  PERSID = 'P',
+  BINPERSID = 'Q',
+  REDUCE = 'R',
+  STRING = 'S',
+  BINSTRING = 'T',
+  SHORT_BINSTRING = 'U',
+  // NB: Avoid using UNICODE as it is a macro in the Windows API
+  UNICODE_ = 'V',
+  BINUNICODE = 'X',
+  APPEND = 'a',
+  BUILD = 'b',
+  GLOBAL = 'c',
+  DICT = 'd',
+  EMPTY_DICT = '}',
+  APPENDS = 'e',
+  GET = 'g',
+  BINGET = 'h',
+  INST = 'i',
+  LONG_BINGET = 'j',
+  LIST = 'l',
+  EMPTY_LIST = ']',
+  OBJ = 'o',
+  PUT = 'p',
+  BINPUT = 'q',
+  LONG_BINPUT = 'r',
+  SETITEM = 's',
+  TUPLE = 't',
+  EMPTY_TUPLE = ')',
+  SETITEMS = 'u',
+  BINFLOAT = 'G',
+
+  // Protocol 2
+  // NOLINTNEXTLINE(readability-redundant-inline-specifier)
+  PROTO = char('\x80'),
+  NEWOBJ = '\x81',
+  EXT1 = '\x82',
+  EXT2 = '\x83',
+  EXT4 = '\x84',
+  TUPLE1 = '\x85',
+  TUPLE2 = '\x86',
+  TUPLE3 = '\x87',
+  NEWTRUE = '\x88',
+  NEWFALSE = '\x89',
+  LONG1 = '\x8a',
+  LONG4 = '\x8b',
+
+  // Protocol 3 (Python 3.x)
+  BINBYTES = 'B',
+  SHORT_BINBYTES = 'C',
+
+  // Protocol 4
+  // NOLINTNEXTLINE(readability-redundant-inline-specifier)
+  SHORT_BINUNICODE = char('\x8c'),
+  BINUNICODE8 = '\x8d',
+  BINBYTES8 = '\x8e',
+  EMPTY_SET = '\x8f',
+  ADDITEMS = '\x90',
+  FROZENSET = '\x91',
+  NEWOBJ_EX = '\x92',
+  STACK_GLOBAL = '\x93',
+  MEMOIZE = '\x94',
+  FRAME = '\x95'
+};
+
+struct WriteableTensorData {
+  const char* data() const {
+    return static_cast<const char*>(tensor_.storage().data());
+  }
+  size_t sizeInBytes() const {
+    return size_;
+  }
+  size_t nbytes() const {
+    return tensor_.storage().nbytes();
+  }
+  bool storageHasDeleter() const {
+    return tensor_.storage().data_ptr().get_context() != nullptr;
+  }
+
+ private:
+  friend TORCH_API WriteableTensorData
+  getWriteableTensorData(const at::Tensor& tensor, bool to_cpu);
+  at::Tensor tensor_;
+  uint64_t size_;
+};
+
+// returns a (tensor, record_size) for a tensor, converting it to a CPU tensor
+// if it was CUDA and to_cpu is True.
+TORCH_API WriteableTensorData
+getWriteableTensorData(const at::Tensor& tensor, bool to_cpu = true);
+
+// if the cls has __getstate__/__setstate__
+// assert they have the right schema and return true,
+// otherwise return false
+bool checkHasValidSetGetState(const c10::ClassType& cls);
+
+// Declare BackendMeta serialization and deserialization function pointer types.
+using BackendMetaPtr = std::function<
+    void(const at::Tensor&, std::unordered_map<std::string, bool>&)>;
+
+// A allowlist of device type, currently available is PrivateUse1
+TORCH_API std::unordered_set<c10::DeviceType>& GetBackendMetaAllowlist();
+
+// Dynamically obtain serialization function pairs
+// that require the corresponding backend.
+TORCH_API std::array<
+    std::optional<std::pair<BackendMetaPtr, BackendMetaPtr>>,
+    at::COMPILE_TIME_MAX_DEVICE_TYPES>&
+GetBackendMetaSerialization();
+
+// Return a map of Tensor Metadata which including BackendMetaData for
+// serialization. For now, it only takes care of `conj` and `neg` bit.
+inline std::unordered_map<std::string, bool> getTensorMetadata(
+    const at::Tensor& t) {
+  // We don't support serializing `ZeroTensor` as it is not public
+  // facing yet.
+  TORCH_CHECK(
+      !t._is_zerotensor(),
+      "ZeroTensor is not serializable,",
+      " please file an issue if required.");
+  std::unordered_map<std::string, bool> metadata{};
+
+  // Only add meta-data if the value is not default.
+  if (t.is_conj()) {
+    metadata["conj"] = true;
+  }
+  if (t.is_neg()) {
+    metadata["neg"] = true;
+  }
+  // Only add BackendMetaData for custom backend if the function pointer is
+  // registered.
+  int device_type = static_cast<int>(t.device().type());
+  const auto& BackendMetaSerialization = GetBackendMetaSerialization();
+  if (BackendMetaSerialization[device_type].has_value()) {
+    // Pass the tensor and metadata map references as parameters to the custom
+    // serialization function.
+    BackendMetaPtr fptr = BackendMetaSerialization[device_type].value().first;
+    fptr(t, metadata);
+  }
+  return metadata;
+}
+
+// set Tensor Metadata based on the map.
+// Refer: getTensorMetadata
+inline void setTensorMetadata(
+    const at::Tensor& t,
+    std::unordered_map<std::string, bool> metadata) {
+  auto iter_end = metadata.end();
+  auto iter_temp = metadata.find("conj");
+  if (iter_temp != iter_end) {
+    t._set_conj(true);
+    metadata.erase(iter_temp);
+  }
+  iter_temp = metadata.find("neg");
+  if (iter_temp != iter_end) {
+    t._set_neg(true);
+    metadata.erase(iter_temp);
+  }
+  // Only set BackendMetaData for custom backend if the function pointer is
+  // registered.
+  int device_type = static_cast<int>(t.device().type());
+  const auto& BackendMetaSerialization = GetBackendMetaSerialization();
+  if (BackendMetaSerialization[device_type].has_value()) {
+    // Pass the tensor and metadata map references as parameters to the custom
+    // deserialization function.
+    BackendMetaPtr fptr = BackendMetaSerialization[device_type].value().second;
+    fptr(t, metadata);
+  }
+}
+
+// set Tensor metadata based on the map.
+// NOTE: This overload is required by unpickler.cpp
+inline void setTensorMetadata(
+    const at::Tensor& t,
+    const c10::Dict<c10::IValue, c10::IValue>& metadata_idict) {
+  std::unordered_map<std::string, bool> metadata;
+  for (auto& pair : metadata_idict) {
+    auto key = *pair.key().toString();
+    metadata[key] = pair.value().toBool();
+  }
+  setTensorMetadata(t, std::move(metadata));
+}
+
+// Register function pointer of Tensor BackendMetadata for serialization.
+inline void TensorBackendMetaRegistry(
+    c10::DeviceType t,
+    const BackendMetaPtr& get_fptr,
+    const BackendMetaPtr& set_fptr) {
+  // allowlist verification
+  // Only if the devicetype is in the allowlist,
+  // we allow the serialization extension to be registered for backendmeta data.
+  const auto& DeviceTypeAllowlist = GetBackendMetaAllowlist();
+  TORCH_CHECK(
+      DeviceTypeAllowlist.find(t) != DeviceTypeAllowlist.end(),
+      "It is not allowed to register the serialization method ",
+      "of backendMeta data for PrivateUse1. ",
+      "If you have related serialization requirements, ",
+      "please expand the allowlist");
+  // Register function pointer
+  int device_type = static_cast<int>(t);
+  auto& BackendMetaSerialization = GetBackendMetaSerialization();
+  TORCH_CHECK(
+      !BackendMetaSerialization[device_type].has_value(),
+      "The tensor BackendMeta serialization function pointer for ",
+      t,
+      " has been registered.");
+  BackendMetaSerialization[device_type] =
+      std::optional<std::pair<BackendMetaPtr, BackendMetaPtr>>(
+          std::make_pair(get_fptr, set_fptr));
+}
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/python_print.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/python_print.h
new file mode 100644
index 0000000000000000000000000000000000000000..7cefbe091fa581766880dcde6a7ae49a4a56705d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/python_print.h
@@ -0,0 +1,61 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <vector>
+
+namespace torch::jit {
+
+struct Method;
+struct Module;
+struct PythonPrintImpl;
+
+struct PrintDepsTable {
+  void add(const c10::NamedTypePtr& type);
+
+  size_t size() const {
+    return table_.size();
+  }
+
+  const c10::NamedTypePtr& operator[](size_t index) const {
+    return table_[index];
+  }
+
+ private:
+  std::vector<c10::NamedTypePtr> table_;
+  std::unordered_set<c10::NamedTypePtr> non_unique_;
+};
+
+struct TORCH_API PythonPrint {
+  PythonPrint(
+      std::vector<IValue>& constant_table,
+      PrintDepsTable& deps_table,
+      c10::TypePrinter type_printer = nullptr,
+      bool enforce_importable = false);
+
+  void printNamedType(const c10::NamedTypePtr& classType);
+  void printFunction(const Function& callee);
+  void printMethod(const Function& callee);
+
+  std::string str() const;
+  const SourceRangeRecords& ranges() const;
+  uint64_t minVersion() const;
+
+ private:
+  std::shared_ptr<PythonPrintImpl> pImpl;
+};
+
+TORCH_API bool printerHasSpecialCaseFor(c10::Symbol sym);
+
+TORCH_API void jitModuleToPythonCodeAndConstants(
+    const Module& module,
+    ExtraFilesMap* jit_sources, // output
+    std::vector<IValue>* constants // output
+);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/source_range_serialization.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/source_range_serialization.h
new file mode 100644
index 0000000000000000000000000000000000000000..3ff838cc62409bd0a5b9f14c9a60ea1e744c223f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/source_range_serialization.h
@@ -0,0 +1,71 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/Allocator.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+
+#include <ATen/core/ivalue.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace c10 {
+struct IValue;
+}
+
+namespace torch::jit {
+
+class Pickler;
+class SourceRangeSerializer;
+static constexpr size_t kByteOffsetIndex = 0;
+static constexpr size_t kSourceRangeIndex = 1;
+static constexpr size_t kSourceRangeTagIndex = 2;
+constexpr std::string_view kFormatWithStringTable = "FORMAT_WITH_STRING_TABLE";
+
+class SourceRangePickler {
+ public:
+  SourceRangePickler();
+
+  std::vector<char> pickle(
+      const SourceRangeRecords& ranges,
+      const SourceRangeTagMap& source_range_tags);
+
+ private:
+  std::shared_ptr<SourceRangeSerializer> srs;
+};
+
+class SourceRangeDeserializer {
+ public:
+  SourceRangeDeserializer() = default;
+  explicit SourceRangeDeserializer(const c10::IValue& text_table) {
+    for (const auto& x : text_table.toTuple()->elements()) {
+      text_table_.emplace_back(std::make_shared<std::string>(x.toStringRef()));
+    }
+  }
+  SourceRange deserialize(const c10::IValue& iv);
+
+ private:
+  std::shared_ptr<Source> deserialize_source(const c10::IValue& iv);
+  std::unordered_map<
+      c10::intrusive_ptr<c10::ivalue::Tuple>,
+      std::shared_ptr<Source>>
+      cached_sources;
+  std::vector<std::shared_ptr<std::string>> text_table_;
+};
+
+class SourceRangeUnpickler {
+ public:
+  virtual std::optional<SourceRange> findSourceRangeThatGenerated(
+      const SourceRange& range) = 0;
+
+  virtual ~SourceRangeUnpickler() = default;
+};
+
+TORCH_API void setShouldUseFormatWithStringTable(
+    bool should_use_format_with_string_table);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/source_range_serialization_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/source_range_serialization_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..35161796290f4cd756de6f789d8bd645828da80f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/source_range_serialization_impl.h
@@ -0,0 +1,33 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/serialization/source_range_serialization.h>
+
+namespace torch::jit {
+
+// Do this clownyness with virtual functions because of the split
+// between ATen core and torch
+
+class ConcreteSourceRangeUnpickler : public SourceRangeUnpickler {
+ public:
+  ConcreteSourceRangeUnpickler(at::DataPtr&& data, size_t size);
+
+  std::optional<SourceRange> findSourceRangeThatGenerated(
+      const SourceRange& range) override;
+
+ private:
+  at::DataPtr data;
+  size_t size;
+
+  void unpickle();
+
+  std::mutex mutex;
+  std::shared_ptr<SourceRangeDeserializer> deserializer;
+  std::shared_ptr<SourceRangeRecords> unpickled_records;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/storage_context.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/storage_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..fac2b53c200ebeba2daf2affbafc3a6b15fe6fd6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/storage_context.h
@@ -0,0 +1,88 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+
+namespace torch::jit {
+
+// Used in torch.package and TorchScript serialization to coordinate
+// sharing of storages between models. Also used to create deterministic
+// naming for storages.
+class TORCH_API SerializationStorageContext {
+ public:
+  explicit SerializationStorageContext() = default;
+  SerializationStorageContext operator=(const SerializationStorageContext&) =
+      delete;
+  SerializationStorageContext(const SerializationStorageContext&) = delete;
+
+  uint64_t getOrAddStorage(const c10::Storage& storage) {
+    if (!hasStorage(storage)) {
+      uint64_t size = storage_id_map_.size();
+      storage_id_map_[storage] = size;
+    }
+    return storage_id_map_[storage];
+  }
+
+  bool hasStorage(const c10::Storage& storage) {
+    return storage_id_map_.find(storage) != storage_id_map_.end();
+  }
+
+  ~SerializationStorageContext() = default;
+
+ private:
+  class StorageSerializationHash {
+   public:
+    size_t operator()(const c10::Storage& storage) const {
+      return std::hash<void*>()(
+          reinterpret_cast<void*>(storage.unsafeGetStorageImpl()));
+    }
+  };
+
+  class StorageSerializationEqual {
+   public:
+    bool operator()(const c10::Storage& lhs, const c10::Storage& rhs) const {
+      return lhs.unsafeGetStorageImpl() == rhs.unsafeGetStorageImpl();
+    }
+  };
+
+  std::unordered_map<
+      c10::Storage,
+      uint64_t,
+      StorageSerializationHash,
+      StorageSerializationEqual>
+      storage_id_map_;
+};
+
+// Used in torch.package and TorchScript deserialization to coordinate
+// sharing of storages between models.
+class TORCH_API DeserializationStorageContext {
+ public:
+  explicit DeserializationStorageContext() = default;
+  DeserializationStorageContext operator=(
+      const DeserializationStorageContext&) = delete;
+  DeserializationStorageContext(const DeserializationStorageContext&) = delete;
+
+  void addStorage(std::string name, c10::Storage storage) {
+    TORCH_INTERNAL_ASSERT(!hasStorage(name));
+    name_storage_map_.emplace(std::move(name), std::move(storage));
+  }
+
+  bool hasStorage(const std::string& name) {
+    return name_storage_map_.find(name) != name_storage_map_.end();
+  }
+
+  c10::Storage getStorage(const std::string& name) {
+    TORCH_INTERNAL_ASSERT(hasStorage(name));
+    return name_storage_map_.find(name)->second;
+  }
+  ~DeserializationStorageContext() = default;
+
+ private:
+  std::unordered_map<std::string, c10::Storage> name_storage_map_;
+};
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/type_name_uniquer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/type_name_uniquer.h
new file mode 100644
index 0000000000000000000000000000000000000000..62747c3514fa4fba667368b2830b78c3f8355dc3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/type_name_uniquer.h
@@ -0,0 +1,36 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/frontend/name_mangler.h>
+#include <torch/csrc/jit/ir/type_hashing.h>
+
+namespace torch::jit {
+
+/**
+ * class TypeNameUniquer
+ *
+ * Generates a unique name for every type `t` passed in. Types that compare
+ * equal with EqualType will receive the same unique name.
+ *
+ * This is used during Module::save(), to resolve type name collisions during
+ * serialization.
+ */
+class TORCH_API TypeNameUniquer {
+ public:
+  c10::QualifiedName getUniqueName(c10::ConstNamedTypePtr t);
+
+ private:
+  NameMangler mangler_;
+  std::unordered_set<c10::QualifiedName> used_names_;
+  std::unordered_map<
+      c10::ConstNamedTypePtr,
+      c10::QualifiedName,
+      HashType,
+      EqualType>
+      name_map_;
+};
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/unpickler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/unpickler.h
new file mode 100644
index 0000000000000000000000000000000000000000..02fc9ba8d543c72132359bb1ccc846402c498574
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/serialization/unpickler.h
@@ -0,0 +1,211 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <c10/util/ArrayRef.h>
+#include <caffe2/serialize/inline_container.h>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/script_type_parser.h>
+#include <torch/csrc/jit/serialization/pickler_helper.h>
+
+namespace torch::jit {
+
+using TypeResolver =
+    std::function<c10::StrongTypePtr(const c10::QualifiedName&)>;
+
+using ObjLoader = std::function<
+    c10::intrusive_ptr<c10::ivalue::Object>(const at::StrongTypePtr&, IValue)>;
+
+class DeserializationStorageContext;
+
+// [unpickler refactor] there is some cruft around PickleOpCode::BUILD,
+// PickleOpCode::NEWOBJ, and the last_opcode_ member below that should be
+// deleted at some point, the Pickler doesn't produce it and it's only around to
+// support models saved before 1.1
+class TORCH_API Unpickler {
+  AT_DISALLOW_COPY_AND_ASSIGN(Unpickler);
+
+  using TypeParserT = c10::TypePtr (*)(const std::string&);
+
+ public:
+  // tensors inside the pickle are references to the tensor_table.
+  // class_resolver is to resolve strong class type, type_resolver_ is
+  // to resolve any JIT type. class_resolver and type_resolver are not merged
+  // here because some use cases need to get strong class type that
+  // type_resolver_ can not return.
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+  Unpickler(
+      std::function<size_t(char*, size_t)> reader,
+      TypeResolver type_resolver,
+      c10::ArrayRef<at::Tensor> tensor_table,
+      TypeParserT type_parser = defaultTypeParser)
+      : reader_(std::move(reader)),
+        tensor_table_(tensor_table),
+        type_resolver_(std::move(type_resolver)),
+        use_storage_device_(false),
+        type_parser_(type_parser),
+        version_(caffe2::serialize::kProducedFileFormatVersion) {}
+
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+  Unpickler(
+      std::function<size_t(char*, size_t)> reader,
+      TypeResolver type_resolver,
+      c10::ArrayRef<at::Tensor> tensor_table,
+      ObjLoader obj_loader,
+      TypeParserT type_parser = defaultTypeParser)
+      : reader_(std::move(reader)),
+        tensor_table_(tensor_table),
+        type_resolver_(std::move(type_resolver)),
+        obj_loader_(std::move(obj_loader)),
+        use_storage_device_(false),
+        type_parser_(type_parser),
+        version_(caffe2::serialize::kProducedFileFormatVersion) {}
+
+  // tensors inside the pickle contain meta-data, the raw tensor
+  // dead is retrieved by calling `read_record`.
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+  Unpickler(
+      std::function<size_t(char*, size_t)> reader,
+      TypeResolver type_resolver,
+      ObjLoader obj_loader,
+      std::function<at::DataPtr(const std::string&)> read_record,
+      std::optional<at::Device> device,
+      bool use_storage_device = false,
+      TypeParserT type_parser = defaultTypeParser,
+      std::shared_ptr<DeserializationStorageContext> storage_context = nullptr)
+      : reader_(std::move(reader)),
+        type_resolver_(std::move(type_resolver)),
+        obj_loader_(std::move(obj_loader)),
+        read_record_(std::move(read_record)),
+        device_(device),
+        use_storage_device_(use_storage_device),
+        type_parser_(type_parser),
+        storage_context_(std::move(storage_context)),
+        version_(caffe2::serialize::kProducedFileFormatVersion) {}
+
+  Unpickler(Unpickler&&) = delete;
+  Unpickler& operator=(Unpickler&&) = delete;
+  ~Unpickler() = default;
+
+  // consume the pickle stream, producing an IValue from the contents.
+  // Type Tags: the pickler will restore the type tags on
+  // List and Dict objects when possible IValue is an Object.
+  // Otherwise, Dict and List objects will end up with Any as their tag.
+  // If you know the type of the ivalue, tags can be restored with
+  // restoreAccurateTypeTags
+  IValue parse_ivalue();
+
+  // [type tag serialization]
+  // This is used to determine whether to restore type tags be recursively
+  // descending into the returned stack object (if version_number <= 2), or
+  // if version_number >= 3, to use the type strings included in the pickle
+  // archive for container types. By default this is set to
+  // `kProducedFileFormatVersion` so unless you're loading a pickle file
+  // from alongside a corresponding `version` file, you don't need to set
+  // the version manually.
+  void set_version(uint64_t version_number) {
+    version_ = version_number;
+  }
+
+  static c10::TypePtr defaultTypeParser(const std::string& str) {
+    ScriptTypeParser parser;
+    return parser.parseType(str);
+  }
+
+ private:
+  // No arguments ensures that a template argument must be specified
+  // so that the number of bytes read / type read is explicit
+  template <typename T>
+  T read() {
+    T item;
+    if (sizeof(T) <= buffer_remaining_) {
+      // Fast path: entirely from buffer.
+      memcpy(&item, buffer_.data() + buffer_pos_, sizeof(T));
+      buffer_remaining_ -= sizeof(T);
+      buffer_pos_ += sizeof(T);
+    } else {
+      // Don't over-template the slow path, to avoid code size bloat.
+      readSlowWithBuffer(reinterpret_cast<char*>(&item), sizeof(T));
+    }
+    return item;
+  }
+  void readSlowWithBuffer(char* dest, size_t sz);
+  std::string readBytes(size_t num_bytes);
+
+  double readFloat();
+  void readGlobal(
+      const std::string& module_name,
+      const std::string& class_name);
+  void rebuildTensor(bool quantized);
+  void rebuildParameter();
+  void rebuildTensorFromTypeV2();
+  void rebuildSparseTensor();
+#ifdef USE_DISTRIBUTED
+  void rebuildRRef();
+#endif
+  PickleOpCode readInstruction();
+  PickleOpCode readOpCode() {
+    return static_cast<PickleOpCode>(read<uint8_t>());
+  }
+  std::string readString();
+  void readList(IValue list_ivalue);
+  void readListElements(IValue list_ivalue, size_t start);
+  void setInput(size_t memo_id);
+  void run();
+
+  // Returns the number of bytes read. This should statefully
+  // remember the position. Don't call reader_ directly.
+  std::function<size_t(char*, size_t)> reader_;
+  // Small buffer to avoid calling reader_ on a per-byte basis.
+  std::array<char, 256> buffer_;
+  size_t buffer_pos_{0};
+  size_t buffer_remaining_{0};
+
+  std::vector<IValue> stack_;
+
+  // globals are represented on the stack as IValue integer indices
+  // into this list
+  std::vector<std::function<void(void)>> globals_;
+  std::vector<IValue> memo_table_;
+  std::vector<size_t> marks_;
+  c10::ArrayRef<at::Tensor> tensor_table_;
+
+  // When deserializing types on lists and dicts, cache the type here
+  // so we don't have to parse the same type multiple times. Strings
+  // are already de-duplicated and replaced with BINGETs in the
+  // pickler, so we can just use the actual data pointer of each string.
+  std::unordered_map<std::string, c10::TypePtr> type_cache_;
+
+  // optionally nullptr, needs to be present for creating classes
+  TypeResolver type_resolver_;
+  ObjLoader obj_loader_;
+  IValue empty_tuple_;
+
+  std::function<at::DataPtr(const std::string&)> read_record_;
+  std::optional<at::Device> device_;
+  // When set to true, Unpickler will ignore the pickled device and use the
+  // device of the DataPtr returned by the read_record_ function. The default
+  // value of this flag is false.
+  const bool use_storage_device_;
+
+  TypeParserT type_parser_{defaultTypeParser};
+
+  // Used for torch.package to enable sharing of storages across
+  // ScriptModules and eager modules
+  std::shared_ptr<DeserializationStorageContext> storage_context_;
+
+  // See [type tag serialization]
+  uint64_t version_;
+
+  // See [NOTE] skip_next_read_global
+  uint8_t skip_next_read_global = 0;
+};
+
+void restoreAccurateTypeTags(const IValue& root, const c10::TypePtr& type_tag);
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/analysis.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/analysis.h
new file mode 100644
index 0000000000000000000000000000000000000000..d0b79418577b6a7bbabd2acbb418175ac5542924
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/analysis.h
@@ -0,0 +1,403 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/stmt.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+#include <utility>
+
+namespace torch::jit::tensorexpr {
+class HasRand : public IRVisitor {
+ public:
+  HasRand(StmtPtr stmt) : stmt_(std::move(stmt)) {
+    stmt_->accept(this);
+  }
+
+  bool has_rand() const {
+    return has_rand_;
+  }
+
+ private:
+  void visit(const IntrinsicsPtr& v) override {
+    if (v->op_type() == IntrinsicsOp::kRand) {
+      has_rand_ = true;
+    } else {
+      IRVisitor::visit(v);
+    }
+  }
+  StmtPtr stmt_;
+  bool has_rand_ = false;
+};
+
+template <typename Op>
+class NodeFinder : public IRVisitor {
+ public:
+  void visit(const NodePtr<Op>& v) override {
+    nodes.push_back((NodePtr<Op>)v);
+    IRVisitor::visit(v);
+  }
+
+  static std::vector<NodePtr<Op>> find(const StmtPtr& s) {
+    NodeFinder<Op> nf;
+    s->accept(&nf);
+    return nf.nodes;
+  }
+
+  static std::vector<NodePtr<Op>> find(const ExprPtr& e) {
+    NodeFinder<Op> nf;
+    e->accept(&nf);
+    return nf.nodes;
+  }
+
+  std::vector<NodePtr<Op>> nodes;
+};
+
+class VarFinder : public IRVisitor {
+ public:
+  void visit(const VarPtr& v) override {
+    vars_.insert(v);
+    IRVisitor::visit(v);
+  }
+
+  static std::unordered_set<VarPtr> find(const StmtPtr& s) {
+    VarFinder nf;
+    s->accept(&nf);
+    return nf.vars();
+  }
+
+  static std::unordered_set<VarPtr> find(const ExprPtr& e) {
+    VarFinder nf;
+    e->accept(&nf);
+    return nf.vars();
+  }
+
+  const std::unordered_set<VarPtr>& vars() {
+    return vars_;
+  }
+
+ private:
+  std::unordered_set<VarPtr> vars_;
+};
+
+class BufFinder : public IRVisitor {
+ public:
+  void visit(const BufPtr& v) override {
+    bufs_.insert(v);
+    IRVisitor::visit(v);
+  }
+
+  static std::unordered_set<BufPtr> find(const StmtPtr& s) {
+    BufFinder nf;
+    s->accept(&nf);
+    return nf.bufs();
+  }
+
+  static std::unordered_set<BufPtr> find(const ExprPtr& e) {
+    BufFinder nf;
+    e->accept(&nf);
+    return nf.bufs();
+  }
+
+  const std::unordered_set<BufPtr>& bufs() {
+    return bufs_;
+  }
+
+ private:
+  std::unordered_set<BufPtr> bufs_;
+};
+
+// Finds all kinds of write operations to the provided Buf.
+class WritesToBuf : public IRVisitor {
+ public:
+  WritesToBuf(BufPtr target) : target_(std::move(target)) {}
+
+  std::vector<StmtPtr> writes() {
+    return writes_;
+  }
+
+  static std::vector<StmtPtr> find(const StmtPtr& s, BufPtr b) {
+    WritesToBuf finder(std::move(b));
+    s->accept(&finder);
+    return finder.writes();
+  }
+
+ private:
+  void visit(const StorePtr& v) override {
+    if (v->buf() == target_) {
+      writes_.push_back(v);
+    }
+  }
+
+  void visit(const AtomicAddPtr& v) override {
+    if (v->buf() == target_) {
+      writes_.push_back(v);
+    }
+  }
+
+  BufPtr target_;
+  std::vector<StmtPtr> writes_;
+};
+
+class StmtsReadingBuf : public IRVisitor {
+ public:
+  StmtsReadingBuf(BufPtr target) : target_(std::move(target)) {}
+
+  std::vector<StmtPtr> reads() {
+    return reads_;
+  }
+
+  static std::vector<StmtPtr> find(const StmtPtr& s, BufPtr b) {
+    StmtsReadingBuf finder(std::move(b));
+    s->accept(&finder);
+    return finder.reads();
+  }
+
+ private:
+  bool readsBuffer(const StmtPtr& s) {
+    auto loads = NodeFinder<Load>::find(s);
+    for (const auto& l : loads) {
+      if (l->buf() == target_) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  void visit(const StorePtr& v) override {
+    if (readsBuffer(v)) {
+      reads_.push_back(v);
+    }
+  }
+
+  void visit(const LetPtr& v) override {
+    if (readsBuffer(v)) {
+      reads_.push_back(v);
+    }
+  }
+
+  void visit(const CondPtr& v) override {
+    if (readsBuffer(v)) {
+      reads_.push_back(v);
+    }
+  }
+
+  void visit(const AtomicAddPtr& v) override {
+    if (readsBuffer(v)) {
+      reads_.push_back(v);
+    }
+  }
+
+  BufPtr target_;
+  std::vector<StmtPtr> reads_;
+};
+
+class ExternalAllocBufFinder : public IRVisitor {
+ public:
+  void visit(const ExternalCallWithAllocPtr& v) override {
+    const auto& bufs_out = v->buf_out_args();
+    bufs_.insert(bufs_out.begin(), bufs_out.end());
+    IRVisitor::visit(v);
+  }
+
+  static std::unordered_set<BufPtr> find(const StmtPtr& s) {
+    ExternalAllocBufFinder f;
+    s->accept(&f);
+    return f.bufs();
+  }
+
+  static std::unordered_set<BufPtr> find(const ExprPtr& e) {
+    ExternalAllocBufFinder f;
+    e->accept(&f);
+    return f.bufs();
+  }
+
+  const std::unordered_set<BufPtr>& bufs() {
+    return bufs_;
+  }
+
+ private:
+  std::unordered_set<BufPtr> bufs_;
+};
+
+// Traverses the IR to determine if a particular Var is modified within it.
+class ModifiesVarChecker : public IRVisitor {
+ public:
+  ModifiesVarChecker(VarPtr v) : var_(std::move(v)) {}
+
+  static bool check(const StmtPtr& s, VarPtr v) {
+    ModifiesVarChecker checker(std::move(v));
+    s->accept(&checker);
+    return checker.found();
+  }
+
+  bool found() {
+    return found_;
+  }
+
+ private:
+  void visit(const StorePtr& v) override {
+    if (v->buf()->base_handle() == var_) {
+      found_ = true;
+      return;
+    }
+    IRVisitor::visit(v);
+  }
+
+  void visit(const AtomicAddPtr& v) override {
+    if (v->buf()->base_handle() == var_) {
+      found_ = true;
+      return;
+    }
+    IRVisitor::visit(v);
+  }
+
+  void visit(const LetPtr& v) override {
+    if (v->var() == var_) {
+      found_ = true;
+      return;
+    }
+    IRVisitor::visit(v);
+  }
+
+  void visit(const ForPtr& v) override {
+    if (v->var() == var_) {
+      found_ = true;
+      return;
+    }
+    IRVisitor::visit(v);
+  }
+
+  VarPtr var_;
+  bool found_{false};
+};
+
+// Traverse the Block stmt to identify the live range of the specified buf. The
+// live range, indicated by a pair of integers, specifies the first and last
+// stmt in block stmts that access to the buf.
+class BufLiveRange : public IRVisitor {
+ public:
+  BufLiveRange(BufPtr b) : buf_(std::move(b)) {}
+
+  static std::tuple<int32_t, int32_t> liveRange(const StmtPtr& s, BufPtr b) {
+    BlockPtr block = to<Block>(s);
+    // We Only analyze buffer live ranges for block stmts.
+    if (!block) {
+      return std::make_tuple(0, 0);
+    }
+
+    BufLiveRange analyzer(std::move(b));
+    block->accept(&analyzer);
+    return analyzer.getLiveRange();
+  }
+
+ private:
+  std::tuple<int32_t, int32_t> getLiveRange() {
+    return std::make_tuple(begin_, end_);
+  }
+
+  bool hasBufReads(const StmtPtr& s) {
+    auto loads1 = NodeFinder<Load>::find(s);
+    for (const auto& l : loads1) {
+      if (l->buf() == buf_) {
+        return true;
+      }
+    }
+    auto loads2 = NodeFinder<ExternalCall>::find(s);
+    for (const auto& l : loads2) {
+      for (const auto& lb : l->buf_args()) {
+        if (lb == buf_) {
+          return true;
+        }
+      }
+    }
+    auto loads3 = NodeFinder<ExternalCallWithAlloc>::find(s);
+    for (const auto& l : loads3) {
+      for (const auto& lb : l->buf_args()) {
+        if (lb == buf_) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
+  bool hasBufWrites(const StmtPtr& s) {
+    auto writes1 = NodeFinder<Store>::find(s);
+    for (const auto& w : writes1) {
+      if (w->buf() == buf_) {
+        return true;
+      }
+    }
+    auto writes2 = NodeFinder<ExternalCall>::find(s);
+    for (const auto& w : writes2) {
+      if (w->buf() == buf_) {
+        return true;
+      }
+    }
+    auto writes3 = NodeFinder<ExternalCallWithAlloc>::find(s);
+    for (const auto& w : writes3) {
+      for (const auto& wb : w->buf_out_args()) {
+        if (wb == buf_) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
+  void findAccAndUpdateLiveRange(const StmtPtr& s) {
+    bool has_reads = hasBufReads(s), has_writes = hasBufWrites(s);
+    if (has_reads || has_writes) {
+      if (begin_ == -1) {
+        begin_ = curr_index_;
+      };
+      end_ = curr_index_;
+    }
+  }
+
+  void visit(const BlockPtr& v) override {
+    for (const StmtPtr& s : *v) {
+      curr_index_ += 1;
+      findAccAndUpdateLiveRange(s);
+    }
+  }
+
+  BufPtr buf_;
+  int32_t begin_ = -1;
+  int32_t end_ = -1;
+  int32_t curr_index_ = -1;
+};
+
+// A class that analyzes the given program relevant for Block backend
+// It creates a map of multi dim buffers and their flat versions
+class CreateBufferMap : public IRVisitor {
+ public:
+  const std::unordered_map<std::string, BufPtr>& getBufferMap() const {
+    return map_input_to_tensor_bufs_;
+  }
+
+ private:
+  void visit(const StorePtr& v) override {
+    auto load_node = to<Load>(v->value());
+    if (load_node) {
+      auto t_buf = load_node->buf();
+      map_input_to_tensor_bufs_.emplace(t_buf->name_hint(), v->buf());
+    } else {
+      auto add_node = to<Add>(v->value());
+      auto mul_node = to<Mul>(v->value());
+      // This means for now, v->value() can be Add or Mul
+      TORCH_INTERNAL_ASSERT(add_node || mul_node, buildErrorMessage());
+      map_input_to_tensor_bufs_.emplace(v->buf()->name_hint(), v->buf());
+    }
+    v->value()->accept(this);
+  }
+  std::unordered_map<std::string, BufPtr> map_input_to_tensor_bufs_;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/block_codegen.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/block_codegen.h
new file mode 100644
index 0000000000000000000000000000000000000000..3dabc4cef5692cd72a0bdc2d79c454d8afa320fb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/block_codegen.h
@@ -0,0 +1,151 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
+#include <ATen/ATen.h>
+#include <torch/csrc/jit/resource_guard.h>
+#include <torch/csrc/jit/tensorexpr/analysis.h>
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_printer.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/unique_name_manager.h>
+
+namespace torch::jit::tensorexpr {
+
+// A class that analyzes the given program relevant for Block backend.
+class BlockAnalysis : public IRVisitor {
+ public:
+  bool is_buf_store_target(const BufPtr& buf) const {
+    return store_targets_.count(buf) > 0;
+  }
+
+  const std::unordered_set<BufPtr>& loads() const {
+    return loads_;
+  }
+
+  const std::unordered_set<BufPtr>& stores() const {
+    return store_targets_;
+  }
+
+  int64_t block_size() const {
+    return block_size_;
+  }
+
+  bool areBufsInMap(const std::unordered_set<BufPtr>& bufs) const;
+
+  BufPtr getMultiDimBuf(const BufPtr& buf) const;
+
+  std::string getInputName(const BufPtr& buf) const;
+
+  std::string getFlatInputName(const BufPtr& buf) const {
+    return getInputName(buf) + "_flat";
+  }
+
+  std::unordered_map<std::string, BufPtr> getBufferMap() const {
+    return map_input_to_tensor_bufs_;
+  }
+
+ private:
+  void visit(const StorePtr& v) override;
+  void visit(const LoadPtr& v) override;
+  void visit(const ForPtr& v) override;
+
+  std::unordered_map<std::string, BufPtr> map_input_to_tensor_bufs_;
+  std::unordered_set<BufPtr> store_targets_;
+  std::unordered_set<BufPtr> loads_;
+  int64_t block_size_ = 32;
+};
+
+// A class that overrides the underlying IRPrinter to produce Block.
+class BlockPrinter : public IRPrinter {
+ public:
+  BlockPrinter(std::ostream* os, BlockAnalysis* block_analysis)
+      : IRPrinter(*os), block_analysis_(block_analysis) {}
+
+  using IRPrinter::name_manager;
+  using IRPrinter::visit;
+
+ private:
+  BlockAnalysis* block_analysis_;
+  std::unordered_map<std::string, int> dim_values_map;
+  std::vector<std::string> dim_names = {"N", "H", "W", "C"};
+  std::vector<std::string> flat_dim_names = {"N", "NH", "NHW", "NHWC"};
+  void PrintTensorInfo(const std::unordered_set<BufPtr>& bufs);
+  void PrintArguments(const std::unordered_set<BufPtr>& bufs);
+  void PrintBufferInfo(const std::unordered_set<BufPtr>& bufs);
+  void PrintDistribution(const std::unordered_set<BufPtr>& bufs);
+  void PrintLoop(const std::unordered_set<BufPtr>& bufs, bool block_idx = true);
+  void PrintReshapeInfo(
+      const std::unordered_set<BufPtr>& bufs,
+      bool reverse = false);
+  void PrintDMAs(const std::unordered_set<BufPtr>& bufs);
+  void PrintAdjustBuffers(const std::unordered_set<BufPtr>& bufs);
+
+  void visit(const ForPtr& v) override;
+  void visit(const LoadPtr& v) override;
+  void visit(const StorePtr& v) override;
+  void visit(const BlockPtr& v) override;
+  void visit(const AddPtr& v) override;
+  void visit(const MulPtr& v) override;
+};
+
+class TORCH_API BlockCodeGen : public CodeGen {
+ public:
+  template <typename... Ts>
+  /* implicit */
+  BlockCodeGen(StmtPtr stmt, Ts... ts)
+      : CodeGen(
+            stmt,
+            std::vector<BufferArg>({BufferArg(ts)...}),
+            at::Device(at::kCPU)) {
+    Initialize();
+  }
+
+  BlockCodeGen(
+      StmtPtr stmt,
+      const std::vector<BufferArg>& buffer_args,
+      at::Device device = at::Device(at::kCPU),
+      const std::string& kernel_func_name = "func")
+      : CodeGen(std::move(stmt), buffer_args, device, kernel_func_name) {
+    Initialize();
+  }
+
+  ~BlockCodeGen() override;
+
+  void call(const std::vector<CallArg>& args) override;
+  void call_raw(const std::vector<void*>& args) override;
+
+  void Initialize();
+
+  std::string getCodeText(const std::string& attr = "") override {
+    return oss_.str();
+  }
+
+ private:
+  UniqueNameManager* name_manager() {
+    if (!printer_) {
+      throw std::runtime_error("Null IRPrinter is not expected");
+    }
+    return printer_->name_manager();
+  }
+
+  std::ostream& os() {
+    return printer_->os();
+  }
+
+  std::ostringstream oss_;
+  std::unique_ptr<BlockPrinter> printer_;
+  std::unique_ptr<BlockAnalysis> block_analysis_;
+
+  std::string GetUniqueFuncName(const std::string& func_prefix);
+};
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/bounds_inference.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/bounds_inference.h
new file mode 100644
index 0000000000000000000000000000000000000000..7d8ca6785aeeae1e886d2e46319bfca72cd3d564
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/bounds_inference.h
@@ -0,0 +1,80 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <unordered_map>
+#include <vector>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/tensorexpr/mem_dependency_checker.h>
+
+namespace torch::jit::tensorexpr {
+
+class Expr;
+class Buf;
+class Stmt;
+
+enum C10_API_ENUM TensorAccessKind { kLoad, kStore, kMutate };
+
+struct TORCH_API TensorAccessBoundsInfo {
+  TensorAccessKind kind;
+  std::vector<ExprPtr> start;
+  std::vector<ExprPtr> stop;
+};
+
+using BoundsInfo =
+    std::unordered_map<BufPtr, std::vector<TensorAccessBoundsInfo>>;
+
+TORCH_API BoundsInfo
+inferBounds(const StmtPtr& s, bool distinctAccessKinds = true);
+
+// Bounds inference caching the analysis. The MemDependencyChecker must already
+// have been run.
+TORCH_API BoundsInfo getInferredBounds(
+    analysis::MemDependencyChecker& analyzer,
+    const StmtPtr& s,
+    bool distinctAccessKinds = true);
+TORCH_API BoundsInfo getInferredBounds(
+    analysis::MemDependencyChecker& analyzer,
+    const ExprPtr& e,
+    bool distinctAccessKinds = true);
+
+TORCH_API void printBoundsInfo(const BoundsInfo& v);
+
+TORCH_API std::vector<ExprPtr> getBoundExtents(
+    const std::vector<TensorAccessBoundsInfo>& infos);
+
+// The kind of dependency found, in increasing order of exclusivity.
+enum class HazardKind {
+  ReadAfterWrite,
+  WriteAfterRead,
+  WriteAfterWrite,
+  NoDependency,
+};
+TORCH_API HazardKind getPotentialHazards(
+    analysis::MemDependencyChecker& analyzer,
+    const StmtPtr& A,
+    const StmtPtr& B);
+
+// Returns true if there is a conflicting overlap between accesses in
+// statements A and B. A conflicting overlap is an overlap in buffer accesses
+// where at least one of the accesses is a Store.
+TORCH_API bool hasConflictingOverlap(
+    analysis::MemDependencyChecker& analyzer,
+    const StmtPtr& A,
+    const StmtPtr& B);
+// Same as above, between accesses in stores S1 and S2.
+TORCH_API bool isOverlapping(
+    analysis::MemDependencyChecker& analyzer,
+    const StorePtr& S1,
+    const StorePtr& S2);
+// Same as above, between accesses in store S and load L.
+TORCH_API bool isOverlapping(
+    analysis::MemDependencyChecker& analyzer,
+    const StorePtr& S,
+    const LoadPtr& L);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/bounds_overlap.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/bounds_overlap.h
new file mode 100644
index 0000000000000000000000000000000000000000..e19bc7b4752d5f13bac53fb3e56d229daf3e42d8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/bounds_overlap.h
@@ -0,0 +1,126 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/expr.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+
+#include <utility>
+#include <vector>
+
+namespace torch::jit::tensorexpr::analysis {
+
+// A simple class containing the start and end of a range in a single dimension.
+struct TORCH_API Bound {
+  ExprPtr start{nullptr};
+  ExprPtr end{nullptr};
+
+  // This stores whether or not the start and end of this Bound have previously
+  // been swapped. This occurs when the bound is in a loop with a negative
+  // stride.
+  bool swapped{false};
+
+  Bound() = default;
+  Bound(ExprPtr s, ExprPtr e) : start(std::move(s)), end(std::move(e)) {}
+
+  void print() const;
+  bool equals(const Bound& other) const;
+
+  // The comparison operators are conservative. If the compare operator returns
+  // true, it means that all the elements satisfy the logical expression. But
+  // the false does not mean the opposite comparison is satisfied. It could be
+  // but not always.
+  bool operator==(const Bound& other) const;
+  bool operator!=(const Bound& other) const;
+  bool operator<(const Bound& other) const;
+  bool operator<=(const Bound& other) const;
+  bool operator>(const Bound& other) const;
+  bool operator>=(const Bound& other) const;
+
+  void swap() noexcept {
+    std::swap(start, end);
+    swapped = !swapped;
+  }
+};
+
+struct BoundHash {
+  size_t operator()(const Bound& b) const {
+    return std::hash<ExprPtr>()(b.start) ^ std::hash<ExprPtr>()(b.end);
+  }
+};
+
+// The type of overlap found. Each condition is true only if none of the
+// previous conditions hold.
+//     ContainedOrEqual: All elements in the Bound A are in the Bound B (this
+//                       includes the case where the bounds are equal).
+//     Contains: All elements in the Bound B are in the Bound B.
+//     PartialOverlap: Any elements in the Bound B are in the Bound A.
+//     NoOverlap: No elements in the Bound A are in the bound B.
+enum class OverlapKind {
+  ContainedOrEqual,
+  Contains,
+  PartialOverlap,
+  NoOverlap
+};
+
+// The Bound comparison result.
+//     True: Every Bound element always satisfies the given comparison operator
+//     False: Every Bound element always does NOT satisfy the given comparison
+//     operator
+//     NotDetermined: Some elements satisfy the given comparison operator and
+//     some elements not
+enum class CmpEvalResult { True, False, NotDetermined };
+
+// Returns the kind of overlap between Bound A and Bound A in a single
+// dimension.
+OverlapKind TORCH_API boundOverlap(const Bound& A, const Bound& B);
+
+// The comparison is conservative and the compare result is deterministic.
+// It means that every element of the Bound to be compared needs to satisfy
+// the given comparison operator.
+CmpEvalResult TORCH_API compareBound(
+    const Bound& a,
+    const Bound& b,
+    const CompareSelectOperation& cmp_op);
+
+// A multi dimensional bound representing the bound of a set of indices.
+using IndexBounds = std::vector<Bound>;
+
+// Returns true if two IndexBounds are equivalent.
+bool TORCH_API indexBoundsEquals(const IndexBounds& A, const IndexBounds& B);
+
+// Flattens a multi dimensional bound to a single dimension. The IndexBounds "a"
+// *must* encapsulate the entire range of the buffer.
+Bound TORCH_API flattenBounds(const IndexBounds& a);
+
+// Determines the kind of overlap in X dimensions.
+OverlapKind TORCH_API overlaps(const IndexBounds& a, const IndexBounds& b);
+
+// Returns the Bound slices created by subtracing bound B from bound A.
+// Multiple Bounds can be returned in the case where B slices A into two
+// distinct regions with no overlap.
+//
+// For example:
+//    subtractBound((0, 10), (2, 4)) => [(0, 1), (5, 10)]
+//       bound A: (0, 10)
+//       bound B: (2, 4)
+//       If we remove slice (2, 4) from the slice (0, 10), we will be left
+//       with 2 slices, one at the start (0, 1), and one at the end (5, 10).
+//       So, the result of this subtraction is [(0, 1), (5, 10)].
+//
+// Note: this doesn't use IndexBounds because the Bounds returned do not
+// represent multiple different dimensions.
+std::vector<Bound> TORCH_API subtractBound(const Bound& a, const Bound& b);
+
+// Returns the bound slices created by subtracting the IndexBounds B from A.
+std::vector<IndexBounds> TORCH_API subtractIndicesBounds(
+    const IndexBounds& A,
+    const IndexBounds& B,
+    OverlapKind overlap);
+std::vector<IndexBounds> TORCH_API
+subtractIndicesBounds(const IndexBounds& A, const IndexBounds& B);
+
+} // namespace torch::jit::tensorexpr::analysis
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/codegen.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/codegen.h
new file mode 100644
index 0000000000000000000000000000000000000000..9ad1cf9615936065658aea77059188995b88e57b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/codegen.h
@@ -0,0 +1,282 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+#include <utility>
+
+namespace torch::jit::tensorexpr {
+
+template <typename T>
+class PaddedBuffer;
+
+class TORCH_API CodeGen {
+ public:
+  class BufferArg;
+  class CallArg;
+
+  template <typename... Ts>
+  CodeGen(StmtPtr stmt, Ts... ts)
+      : stmt_(std::move(stmt)), buffer_args_({BufferArg(ts)...}) {}
+
+  CodeGen(
+      StmtPtr stmt,
+      std::vector<BufferArg> buffer_args,
+      at::Device device = at::kCPU,
+      std::string kernel_func_name = "func");
+
+  CodeGen(const CodeGen& rhs);
+
+  CodeGen(CodeGen&& rhs);
+
+  virtual ~CodeGen();
+
+  CodeGen& operator=(const CodeGen& rhs);
+
+  CodeGen& operator=(CodeGen&& rhs);
+
+  StmtPtr stmt() const {
+    return stmt_;
+  }
+
+  void set_stmt(StmtPtr s) {
+    stmt_ = std::move(s);
+  }
+
+  void apply_mutator(IRMutator* mutator) {
+    stmt_ = stmt_->accept_mutator(mutator);
+  }
+
+  void apply_visitor(IRVisitor* visitor) {
+    stmt_->accept(visitor);
+  }
+
+  std::vector<BufferArg>& buffer_args() {
+    return buffer_args_;
+  }
+
+  const std::vector<BufferArg>& buffer_args() const {
+    return buffer_args_;
+  }
+
+  at::Device device() {
+    return device_;
+  }
+
+  // This function returns the generated code as
+  // a string.
+  virtual std::string getCodeText(
+      const std::string& attr [[maybe_unused]] = "") {
+    return "";
+  }
+
+  // TODO: Figure out how to unify these call interfaces.
+
+  /// Call a function with a vector of CallArgs, which are tagged
+  /// unions that properly type the arguments.
+  virtual void call(const std::vector<CallArg>& args) = 0;
+
+  /// Call a function faster than a regular `call` by assuming that
+  /// the generated kernel already knows the type of the arguments, so
+  /// they can be type-punned with `void*`s.
+  virtual void call_raw(const std::vector<void*>& args) = 0;
+
+  /// Call a function even faster than a regular call, by assuming
+  /// that the number of thread blocks can be derived from `numel` via
+  /// a simple division, rather than evaluating an expression.
+  virtual void call_with_numel(void** args, int64_t numel);
+
+  virtual at::Tensor empty_strided(
+      c10::IntArrayRef size,
+      c10::IntArrayRef stride,
+      std::optional<c10::ScalarType> dtype_opt,
+      std::optional<c10::Layout> layout_opt,
+      std::optional<c10::Device> device_opt,
+      std::optional<bool> pin_memory_opt) {
+    return at::empty_strided(
+        size, stride, dtype_opt, layout_opt, device_opt, pin_memory_opt);
+  }
+
+  const std::string& kernel_func_name() const {
+    return kernel_func_name_;
+  }
+
+  void allocIntermediateBufs();
+
+ protected:
+  static void* argToPtr(const BufferArg& bufferArg, const CallArg& callArg);
+
+ private:
+  StmtPtr stmt_;
+  std::vector<BufferArg> buffer_args_;
+  at::Device device_ = at::kCPU;
+  std::string kernel_func_name_ = "func";
+};
+
+class TORCH_API ExtCallMemoryReuse : public IRMutator {
+  static std::unordered_map<std::string, std::string> makeExtCallFuncNameMap();
+  static const std::unordered_map<std::string, std::string> extCallFuncNameMap_;
+
+ public:
+  explicit ExtCallMemoryReuse(
+      const std::vector<CodeGen::BufferArg>& bufferArgs);
+  ~ExtCallMemoryReuse() override = default;
+  StmtPtr mutate(const ExternalCallPtr& v) override;
+
+ private:
+  std::unordered_set<BufPtr> bufferArgs_;
+};
+
+class CodeGen::BufferArg {
+ public:
+  BufferArg(const Tensor& tensor) : buf_(tensor.buf()) {}
+  BufferArg(const VarHandle& var) : var_(var.node()), isVar_(true) {}
+  BufferArg(const BufHandle& buf) : buf_(buf.node()) {}
+  BufferArg(BufPtr buf) : buf_(std::move(buf)) {}
+
+  VarPtr var() const {
+    return isVar_ ? var_ : buf_->base_handle();
+  }
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  bool isVar() const {
+    return isVar_;
+  }
+
+  Dtype dtype() const {
+    return isVar_ ? var_->dtype() : buf_->dtype();
+  }
+
+ private:
+  VarPtr var_ = nullptr;
+  BufPtr buf_ = nullptr;
+  bool isVar_ = false;
+};
+
+class CodeGen::CallArg {
+ public:
+  template <typename T>
+  CallArg(const PaddedBuffer<T>& buffer);
+
+  template <typename T>
+  CallArg(const std::vector<T>& buffer)
+      : data_(const_cast<T*>(buffer.data())) {}
+
+  CallArg(void* ptr) : data_(ptr) {}
+
+#define ARG_TYPE_CTOR(Type, Name)      \
+  CallArg(Type v) {                    \
+    memcpy(buffer_, &v, sizeof(Type)); \
+    data_ = (void*)buffer_;            \
+  }
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, ARG_TYPE_CTOR)
+#undef ARG_TYPE_CTOR
+
+  void* data() const {
+    return data_;
+  }
+
+  CallArg(const CallArg& rhs) {
+    if (rhs.data_ == rhs.buffer_) {
+      memcpy(this->buffer_, rhs.buffer_, sizeof(rhs.buffer_));
+      this->data_ = (void*)(this->buffer_);
+    } else {
+      this->data_ = rhs.data_;
+    }
+  }
+
+  CallArg& operator=(const CallArg& rhs) {
+    if (this == &rhs) {
+      return *this;
+    }
+    if (rhs.data_ == rhs.buffer_) {
+      memcpy(this->buffer_, rhs.buffer_, sizeof(rhs.buffer_));
+      this->data_ = (void*)(this->buffer_);
+    } else {
+      this->data_ = rhs.data_;
+    }
+    return *this;
+  }
+
+#define ARG_PTR_DEFINE(Type, Name)                  \
+  Type* Name##Ptr() const {                         \
+    TORCH_INTERNAL_ASSERT(data_ == (void*)buffer_); \
+    return (Type*)data_;                            \
+  }
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, ARG_PTR_DEFINE)
+#undef ARG_PTR_DEFINE
+
+ private:
+  void* data_;
+  // Regarding a scalar value, CallArg uses void**=&data_ to store it. But the
+  // bit width of a pointer is 32bit on a 32bit platform. It cannot store the
+  // scalar if the bit width of the scalar is larger than 32bit, such as double
+  // and long. Hence, we add 8 bytes buffer dedicated to storing the scalar
+  // value regardless its bit width is less or greater than 32bits.
+  char buffer_[8] = {0}; // 64bits
+};
+
+class RegisterCodeGenList {
+ public:
+  TORCH_API static RegisterCodeGenList& GetInstance();
+
+  using StmtFactoryMethod = std::function<std::unique_ptr<CodeGen>(
+      StmtPtr stmt,
+      const std::vector<CodeGen::BufferArg>&,
+      at::Device device,
+      const std::string& kernel_func_name)>;
+
+  TORCH_API StmtFactoryMethod FindStmtFactoryMethod(const std::string& name);
+  RegisterCodeGenList(const RegisterCodeGenList&) = delete;
+  RegisterCodeGenList& operator=(const RegisterCodeGenList&) = delete;
+
+ private:
+  template <class CodeGenType>
+  friend class RegisterCodeGen;
+  RegisterCodeGenList() = default;
+  TORCH_API void AddStmtFactoryMethod(
+      const std::string& name,
+      const StmtFactoryMethod& stmt_factory_method);
+
+  std::unordered_map<std::string, StmtFactoryMethod> stmt_factory_methods_;
+};
+
+template <class CodeGenType>
+class RegisterCodeGen {
+ public:
+  explicit RegisterCodeGen(const std::string& name) {
+    RegisterCodeGenList& codegen_list = RegisterCodeGenList::GetInstance();
+    codegen_list.AddStmtFactoryMethod(
+        name,
+        [](const StmtPtr& stmt,
+           const std::vector<CodeGen::BufferArg>& params,
+           at::Device device,
+           const std::string& kernel_func_name) {
+          return std::make_unique<CodeGenType>(
+              stmt, params, device, kernel_func_name);
+        });
+  }
+};
+
+TORCH_API std::unique_ptr<CodeGen> CreateCodeGen(
+    const std::string& name,
+    StmtPtr stmt,
+    const std::vector<CodeGen::BufferArg>& params,
+    at::Device device = at::kCPU,
+    const std::string& kernel_func_name = "func");
+
+class TORCH_API GenericIntrinsicsExpander : public IRMutator {
+ protected:
+  ExprPtr mutate(const IntrinsicsPtr& v) override;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cpp_codegen.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cpp_codegen.h
new file mode 100644
index 0000000000000000000000000000000000000000..825fec57e7204f2d6ac41a9352f36b56c093c3e6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cpp_codegen.h
@@ -0,0 +1,103 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/ir_printer.h>
+
+namespace torch::jit::tensorexpr {
+
+class CppVarNameRewriter;
+
+// Generates C++ code from the IR.
+//
+// Vector operations are unrolled.
+// For example:
+// C[Ramp(0, 1, 3)] = A[Ramp(0, 2, 3)] + B[Ramp(0, 3, 3)];
+// is unrolled into:
+// C[0] = A[0] + B[0];
+// C[1] = A[2] + B[3];
+// C[2] = A[4] + B[6];
+class TORCH_API CppPrinter : public IRPrinter {
+ public:
+  explicit CppPrinter(std::ostream* os);
+  ~CppPrinter() override;
+
+  void printPrologue();
+
+  using IRPrinter::visit;
+
+  // Binary expressions.
+  void visit(const ModPtr& /*v*/) override;
+  void visit(const MaxPtr& /*v*/) override;
+  void visit(const MinPtr& /*v*/) override;
+
+  // Conditional expressions.
+  void visit(const CompareSelectPtr& /*v*/) override;
+  void visit(const IfThenElsePtr& /*v*/) override;
+
+  // Tensor operations.
+  void visit(const AllocatePtr& /*v*/) override;
+  void visit(const FreePtr& /*v*/) override;
+  void visit(const LoadPtr& /*v*/) override;
+  void visit(const StorePtr& /*v*/) override;
+
+  // Casts.
+  void visit(const CastPtr& /*v*/) override;
+  void visit(const BitCastPtr& /*v*/) override;
+
+  // Calls.
+  void visit(const IntrinsicsPtr& /*v*/) override;
+  void visit(const ExternalCallPtr& /*v*/) override;
+
+  // Vars.
+  void visit(const LetPtr& /*v*/) override;
+  void visit(const VarPtr& /*v*/) override;
+
+  // Vector data types.
+  void visit(const RampPtr& /*v*/) override;
+  void visit(const BroadcastPtr& /*v*/) override;
+
+ private:
+  int lane_{0};
+  std::unordered_map<VarPtr, ExprPtr> vector_vars_;
+};
+
+class TORCH_API CppCodeGen : public CodeGen {
+ public:
+  CppCodeGen(
+      StmtPtr stmt,
+      const std::vector<BufferArg>& buffer_args,
+      at::Device device = at::kCPU,
+      const std::string& kernel_func_name = "func");
+
+  ~CppCodeGen() override;
+
+  void call(const std::vector<CallArg>& args) override;
+  void call_raw(const std::vector<void*>& args) override;
+
+  template <typename... Ts>
+  void operator()(const Ts&... ts) {
+    call(std::vector<CallArg>({CallArg(ts)...}));
+  }
+
+  std::string getCodeText(const std::string& attr = "") override {
+    return oss_.str();
+  }
+
+ private:
+  void init();
+
+  std::ostream& os() {
+    return printer_->os();
+  }
+
+  std::ostringstream oss_;
+  std::unique_ptr<CppPrinter> printer_;
+  std::unique_ptr<CppVarNameRewriter> var_name_rewriter_;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cpp_intrinsics.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cpp_intrinsics.h
new file mode 100644
index 0000000000000000000000000000000000000000..0d977fc7f1920907d1779793a6a2ef6884fe3c98
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cpp_intrinsics.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::jit::tensorexpr {
+
+constexpr auto cpp_intrinsics_definition = R"(
+namespace std {
+
+template <typename T,
+          std::enable_if_t<std::is_floating_point_v<T>, int> = 0>
+T rsqrt(T v) {
+  return 1.0f / std::sqrt(v);
+}
+
+template <typename T,
+          std::enable_if_t<std::is_floating_point_v<T>, int> = 0>
+T frac(T v) {
+  T intpart;
+  return std::modf(v, &intpart);
+}
+
+template <typename From, typename To>
+To bitcast(const From& v) {
+  assert(sizeof(To) == sizeof(From));
+  To res;
+  std::memcpy(&res, &v, sizeof(From));
+  return res;
+}
+
+} // namespace std
+)";
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cuda_codegen.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cuda_codegen.h
new file mode 100644
index 0000000000000000000000000000000000000000..2c94659c6dab2386c8aaa140ca14bb70d4ca8f86
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cuda_codegen.h
@@ -0,0 +1,291 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <unordered_set>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/nvrtc_stub/ATenNVRTC.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/csrc/jit/resource_guard.h>
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/eval.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_printer.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/llvm_codegen.h>
+#include <torch/csrc/jit/tensorexpr/unique_name_manager.h>
+
+namespace torch::jit::tensorexpr {
+
+// A class that analyzes the given program relevant for Cuda backends.
+class CudaAnalysis : public IRVisitor {
+ public:
+  CudaAnalysis() {
+    gpu_block_extents_ = {alloc<IntImm>(1), alloc<IntImm>(1), alloc<IntImm>(1)};
+    gpu_thread_extents_ = {
+        alloc<IntImm>(1), alloc<IntImm>(1), alloc<IntImm>(1)};
+  }
+  bool is_buf_store_target(const BufPtr& buf) const {
+    return store_targets_.count(buf) > 0;
+  }
+
+  const std::unordered_set<VarPtr>& thread_local_bufs() const {
+    return thread_local_bufs_;
+  }
+
+  const std::unordered_set<VarPtr>& cross_block_bufs() const {
+    return cross_block_bufs_;
+  }
+
+  const std::vector<ExprPtr>& gpu_block_extents() const {
+    return gpu_block_extents_;
+  }
+
+  const std::vector<ExprPtr>& gpu_thread_extents() const {
+    return gpu_thread_extents_;
+  }
+
+ private:
+  void visit(const StorePtr& v) override {
+    store_targets_.insert(v->buf());
+  }
+
+  void visit(const AllocatePtr& v) override;
+  void visit(const FreePtr& v) override;
+  void visit(const PlacementAllocatePtr& v) override;
+  void visit(const ForPtr& v) override;
+
+  std::unordered_set<BufPtr> store_targets_;
+  std::unordered_set<VarPtr> thread_local_bufs_;
+  std::unordered_set<VarPtr> cross_block_bufs_;
+
+  std::vector<ExprPtr> gpu_block_extents_;
+  std::vector<ExprPtr> gpu_thread_extents_;
+};
+
+// An IRMutator that replaces binding loop options with Cuda metavars, and masks
+// statements blocks which should execute with less reach than the launch
+// parameter extent.
+//
+// We do this by segmenting each block into chunks which should have the same
+// execution parameters, then if those params differ from the max mask each dim.
+class GPUMetaVarRewriter : public IRMutator {
+ public:
+  explicit GPUMetaVarRewriter(const CudaAnalysis* cuda_analysis)
+      : cuda_analysis_(cuda_analysis) {
+    gpu_block_vars_ = {
+        alloc<Var>("blockIdx.x", kInt),
+        alloc<Var>("blockIdx.y", kInt),
+        alloc<Var>("blockIdx.z", kInt)};
+    gpu_thread_vars_ = {
+        alloc<Var>("threadIdx.x", kInt),
+        alloc<Var>("threadIdx.y", kInt),
+        alloc<Var>("threadIdx.z", kInt)};
+
+    current_block_reach_ = {
+        alloc<IntImm>(1), alloc<IntImm>(1), alloc<IntImm>(1)};
+    current_thread_reach_ = {
+        alloc<IntImm>(1), alloc<IntImm>(1), alloc<IntImm>(1)};
+  }
+
+  StmtPtr mutate(const ForPtr& v) override;
+  StmtPtr mutate(const BlockPtr& v) override;
+
+  const std::vector<VarPtr>& gpu_block_vars() const {
+    return gpu_block_vars_;
+  }
+
+  const std::vector<VarPtr>& gpu_thread_vars() const {
+    return gpu_thread_vars_;
+  }
+
+  const std::vector<ExprPtr>& gpu_block_extents() const {
+    return cuda_analysis_->gpu_block_extents();
+  }
+
+  const std::vector<ExprPtr>& gpu_thread_extents() const {
+    return cuda_analysis_->gpu_thread_extents();
+  }
+
+ private:
+  // When processing a block, stores the contents of each sub-segment.
+  class Segment {
+   public:
+    void reset(bool mask) {
+      stmts_.clear();
+      mask_ = mask;
+    }
+
+    bool empty() const {
+      return stmts_.empty();
+    }
+
+    std::vector<StmtPtr>& stmts() {
+      return stmts_;
+    }
+    bool mask() {
+      return mask_;
+    }
+
+   private:
+    std::vector<StmtPtr> stmts_;
+    bool mask_{true};
+  };
+
+  // Returns true if the current execution scope is equivalent to the launch
+  // parameters.
+  bool isFullExtent();
+
+  std::vector<VarPtr> gpu_block_vars_;
+  std::vector<VarPtr> gpu_thread_vars_;
+
+  std::vector<ExprPtr> current_block_reach_;
+  std::vector<ExprPtr> current_thread_reach_;
+
+  const CudaAnalysis* cuda_analysis_;
+};
+
+// A class that overrides the underlying IRPrinter to produce Cuda C.
+class CudaPrinter : public IRPrinter {
+ public:
+  explicit CudaPrinter(
+      std::ostream* os,
+      const CudaAnalysis* cuda_analysis,
+      bool has_random)
+      : IRPrinter(*os), cuda_analysis_(cuda_analysis) {
+    if (has_random) {
+      rand_func_ = alloc<Var>("rand", kHandle);
+    }
+  }
+
+  void visit(const CastPtr& v) override;
+  void visit(const IntrinsicsPtr& v) override;
+  void visit(const ForPtr& v) override;
+
+  void visit(const LoadPtr& v) override;
+  void visit(const StorePtr& v) override;
+  void visit(const AtomicAddPtr& v) override;
+  void visit(const MaxPtr& v) override;
+  void visit(const MinPtr& v) override;
+  void visit(const IfThenElsePtr& v) override;
+  void visit(const BlockPtr& v) override;
+  void visit(const AllocatePtr& v) override;
+  void visit(const FreePtr& v) override;
+  void visit(const LetPtr& v) override;
+
+  void visit(const ExternalCallPtr& v) override;
+
+  VarPtr rand_func() const {
+    return rand_func_;
+  }
+
+  std::string dtypeToCppString(const Dtype& dtype) override;
+
+  using IRPrinter::name_manager;
+  using IRPrinter::visit;
+
+ private:
+  VarPtr rand_func_;
+  const CudaAnalysis* cuda_analysis_;
+
+  void print_flat_alloc(const AllocatePtr& alloc);
+};
+
+// Construct Cuda C from the buffer and tensor input, and invoke the
+// kernel when real arguments are provided.
+class TORCH_CUDA_CU_API CudaCodeGen : public CodeGen {
+ public:
+  template <typename... Ts>
+  CudaCodeGen(StmtPtr stmt, Ts... ts)
+      : CodeGen(
+            stmt,
+            std::vector<BufferArg>({BufferArg(ts)...}),
+            at::Device(at::kCUDA, at::cuda::current_device())) {
+    Initialize();
+  }
+
+  CudaCodeGen(
+      StmtPtr stmt,
+      const std::vector<BufferArg>& buffer_args,
+      at::Device device = at::Device(at::kCUDA, at::cuda::current_device()),
+      const std::string& kernel_func_name = "func")
+      : CodeGen(std::move(stmt), buffer_args, device, kernel_func_name) {
+    Initialize();
+  }
+
+  ~CudaCodeGen() override;
+
+  void call(const std::vector<CallArg>& args) override;
+  void call_raw(const std::vector<void*>& args) override;
+  void call_with_numel(void** args, int64_t numel) override;
+
+  template <typename... Ts>
+  void operator()(const Ts&... ts) {
+    call(std::vector<CallArg>({CallArg(ts)...}));
+  }
+
+  at::Tensor empty_strided(
+      c10::IntArrayRef size,
+      c10::IntArrayRef stride,
+      std::optional<c10::ScalarType> dtype_opt,
+      std::optional<c10::Layout> layout_opt,
+      std::optional<c10::Device> device_opt,
+      std::optional<bool> pin_memory_opt) override;
+
+  const std::vector<ExprPtr>& gpu_block_extents() const {
+    return cuda_analysis_->gpu_block_extents();
+  }
+
+  const std::vector<ExprPtr>& gpu_thread_extents() const {
+    return cuda_analysis_->gpu_thread_extents();
+  }
+
+  std::string getCodeText(const std::string& attr = "") override {
+    return oss_.str();
+  }
+
+ private:
+  void Initialize();
+
+  void CompileToNVRTC(const std::string& code, const std::string& func_name);
+
+  UniqueNameManager* name_manager() {
+    if (!printer_) {
+      throw std::runtime_error("Null IRPrinter is not expected");
+    }
+    return printer_->name_manager();
+  }
+
+  std::ostream& os() {
+    return printer_->os();
+  }
+
+  std::ostringstream oss_;
+  std::unique_ptr<CudaPrinter> printer_;
+  std::unique_ptr<CudaAnalysis> cuda_analysis_;
+  std::unique_ptr<GPUMetaVarRewriter> metavar_rewriter_;
+  std::unordered_set<std::string> taken_func_names;
+  std::mutex eval_lock_;
+  CUfunction function_{nullptr};
+  bool has_random_ = false;
+  int thread_block_size_ = -1;
+
+  std::vector<bool> arg_pos_in_extents_;
+#ifdef TORCH_ENABLE_LLVM
+  std::vector<ExprEval<LLVMCodeGen>> block_extents_eval_;
+  std::vector<ExprEval<LLVMCodeGen>> thread_extents_eval_;
+#else
+  std::vector<ExprEval<SimpleIREvaluator>> block_extents_eval_;
+  std::vector<ExprEval<SimpleIREvaluator>> thread_extents_eval_;
+#endif
+
+  std::string GetUniqueFuncName(const std::string& func_prefix);
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cuda_random.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cuda_random.h
new file mode 100644
index 0000000000000000000000000000000000000000..d0ba0493e2eaf79e75179370279d05d37c7ab7e9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/cuda_random.h
@@ -0,0 +1,105 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::jit::tensorexpr {
+
+constexpr auto philox_random_string = R"(
+
+class Philox {
+public:
+  __device__ inline Philox(unsigned long long seed,
+                           unsigned long long subsequence,
+                           unsigned long long offset) {
+    key.x = (unsigned int)seed;
+    key.y = (unsigned int)(seed >> 32);
+    counter = make_uint4(0, 0, 0, 0);
+    counter.z = (unsigned int)(subsequence);
+    counter.w = (unsigned int)(subsequence >> 32);
+    STATE = 0;
+    incr_n(offset / 4);
+  }
+
+  __device__ inline unsigned long operator()() {
+    if(STATE == 0) {
+      uint4 counter_ = counter;
+      uint2 key_ = key;
+      for(int i = 0; i < 9; i++) {
+        counter_ = single_round(counter_, key_);
+        key_.x += (kPhilox10A); key_.y += (kPhilox10B);
+      }
+      output = single_round(counter_, key_);
+      incr();
+    }
+    unsigned long ret;
+    switch(STATE) {
+      case 0: ret = output.x; break;
+      case 1: ret = output.y; break;
+      case 2: ret = output.z; break;
+      case 3: ret = output.w; break;
+    }
+    STATE = (STATE + 1) % 4;
+    return ret;
+  }
+
+private:
+  uint4 counter;
+  uint4 output;
+  uint2 key;
+  unsigned int STATE;
+  __device__ inline void incr_n(unsigned long long n) {
+    unsigned int nlo = (unsigned int)(n);
+    unsigned int nhi = (unsigned int)(n >> 32);
+    counter.x += nlo;
+    if (counter.x < nlo)
+      nhi++;
+    counter.y += nhi;
+    if (nhi <= counter.y)
+      return;
+    if (++counter.z)
+      return;
+    ++counter.w;
+  }
+  __device__ inline void incr() {
+    if (++counter.x)
+      return;
+    if (++counter.y)
+      return;
+    if (++counter.z)
+      return;
+    ++counter.w;
+  }
+  __device__ unsigned int mulhilo32(unsigned int a, unsigned int b,
+                                    unsigned int *result_high) {
+    *result_high = __umulhi(a, b);
+    return a*b;
+  }
+
+  __device__ inline uint4 single_round(uint4 ctr, uint2 key) {
+    unsigned int hi0;
+    unsigned int hi1;
+    unsigned int lo0 = mulhilo32(kPhiloxSA, ctr.x, &hi0);
+    unsigned int lo1 = mulhilo32(kPhiloxSB, ctr.z, &hi1);
+
+    uint4 ret = {hi1 ^ ctr.y ^ key.x, lo1, hi0 ^ ctr.w ^ key.y, lo0};
+    return ret;
+  }
+
+  static const unsigned long kPhilox10A = 0x9E3779B9;
+  static const unsigned long kPhilox10B = 0xBB67AE85;
+  static const unsigned long kPhiloxSA = 0xD2511F53;
+  static const unsigned long kPhiloxSB = 0xCD9E8D57;
+};
+
+// Inverse of 2^32.
+#define M_RAN_INVM32 2.3283064e-10f
+__device__  __inline__ float Uint32ToFloat(unsigned int x) {
+  return x * M_RAN_INVM32;
+}
+
+)";
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/eval.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/eval.h
new file mode 100644
index 0000000000000000000000000000000000000000..bacdb882de08365383e5548ecdfd80f8253fa359
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/eval.h
@@ -0,0 +1,330 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cmath>
+#include <cstring>
+#include <utility>
+#include <vector>
+
+#include <c10/macros/Macros.h>
+#include <c10/util/Logging.h>
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/exceptions.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_printer.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+#include <torch/csrc/jit/tensorexpr/types.h>
+#include <torch/csrc/jit/tensorexpr/var_substitutor.h>
+
+namespace torch::jit::tensorexpr {
+
+class InterpValue {
+ public:
+  InterpValue() : dtype_(kInt) {
+    Intvalues.push_back(0);
+  }
+
+  template <typename T>
+  InterpValue(Dtype dtype, T v) : dtype_(dtype) {
+#define TYPE_CASE(Type, Name)  \
+  if (dtype == k##Name) {      \
+    Name##values.push_back(v); \
+    return;                    \
+  }
+    AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TYPE_CASE)
+#undef TYPE_CASE
+    throw unsupported_dtype();
+  }
+
+#define VALUE_CTOR(Type, Name)            \
+  InterpValue(Type v) : dtype_(k##Name) { \
+    Name##values.push_back(v);            \
+  }
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, VALUE_CTOR)
+#undef VALUE_CTOR
+
+  explicit InterpValue(c10::quint8 v) : dtype_(kQUInt8) {
+    QUInt8values.emplace_back(v.val_);
+  }
+
+  explicit InterpValue(c10::qint8 v) : dtype_(kQInt8) {
+    QInt8values.emplace_back(v.val_);
+  }
+
+#define VALUE_VEC_CTOR(Type, Name)        \
+  InterpValue(const std::vector<Type>& v) \
+      : dtype_(Dtype(k##Name, v.size())), Name##values(v) {}
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, VALUE_VEC_CTOR)
+  VALUE_VEC_CTOR(c10::quint8, QUInt8)
+  VALUE_VEC_CTOR(c10::qint8, QInt8)
+#undef VALUE_VEC_CTOR
+
+  template <typename T>
+  T as() const;
+
+  template <typename T>
+  const std::vector<T>& as_vec() const;
+
+  int64_t intValue() const;
+
+  Dtype dtype() const {
+    return dtype_;
+  }
+
+ private:
+  Dtype dtype_;
+
+#define VALUE_STORAGE(Type, Name) std::vector<Type> Name##values;
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, VALUE_STORAGE)
+  VALUE_STORAGE(c10::qint8, QInt8)
+  VALUE_STORAGE(c10::quint8, QUInt8)
+#undef VALUE_STORAGE
+  void* ptr{nullptr};
+};
+
+#define VALUE_AS_DISPATCH(Type, Name)         \
+  template <>                                 \
+  inline Type InterpValue::as<Type>() const { \
+    if (dtype_ != k##Name) {                  \
+      throw unsupported_dtype();              \
+    }                                         \
+    return Name##values[0];                   \
+  }
+AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, VALUE_AS_DISPATCH)
+VALUE_AS_DISPATCH(c10::quint8, QUInt8)
+VALUE_AS_DISPATCH(c10::qint8, QInt8)
+#undef VALUE_AS_DISPATCH
+
+#define VALUE_AS_VEC_DISPATCH(Type, Name)                             \
+  template <>                                                         \
+  inline const std::vector<Type>& InterpValue::as_vec<Type>() const { \
+    if (dtype_.scalar_type() != ScalarType::Name) {                   \
+      throw unsupported_dtype();                                      \
+    }                                                                 \
+    return Name##values;                                              \
+  }
+AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, VALUE_AS_VEC_DISPATCH)
+VALUE_AS_VEC_DISPATCH(c10::quint8, QUInt8)
+VALUE_AS_VEC_DISPATCH(c10::qint8, QInt8)
+#undef VALUE_AS_VEC_DISPATCH
+
+template <typename Type>
+auto underlyingValue(Type x) {
+  return x;
+}
+
+template <>
+inline auto underlyingValue<c10::quint8>(c10::quint8 x) {
+  return x.val_;
+}
+
+template <>
+inline auto underlyingValue<c10::qint8>(c10::qint8 x) {
+  return x.val_;
+}
+
+template <typename To, typename From>
+To raw_bitcast(const From& src) {
+  TORCH_CHECK(sizeof(To) == sizeof(From), "Invalid bitcast invocation");
+  To storage;
+  std::memcpy(&storage, &src, sizeof(To));
+  return storage;
+}
+
+class SimpleIREvaluatorImpl;
+class TORCH_API SimpleIREvaluator : public CodeGen {
+ public:
+  SimpleIREvaluator(
+      StmtPtr stmt,
+      const std::vector<BufferArg>& buffer_args,
+      at::Device device = at::kCPU,
+      const std::string& kernel_func_name = "func");
+
+  ~SimpleIREvaluator() override;
+
+  void call(const std::vector<CallArg>& args) override;
+  void call_raw(const std::vector<void*>& args) override;
+
+  template <typename... Ts>
+  void operator()(const Ts&... ts) {
+    std::vector<CallArg> args({CallArg(ts)...});
+    call(args);
+  }
+
+  void bindVar(const VarPtr& v, const ExprPtr& e);
+  InterpValue value() const;
+
+ private:
+  void bindArg(const BufferArg& buf, void* data);
+  void expand_intrinsics() {
+    GenericIntrinsicsExpander intrinsics_expander;
+    apply_mutator(&intrinsics_expander);
+  }
+
+  std::unique_ptr<SimpleIREvaluatorImpl> impl_;
+};
+
+template <class CodeGenType>
+class ExprEval {
+ public:
+  using BufferArg = CodeGen::BufferArg;
+  using CallArg = CodeGen::CallArg;
+
+  template <typename... Ts>
+  ExprEval(const ExprHandle& expr, Ts... ts)
+      : ExprEval(expr, {BufferArg(ts)...}) {}
+
+  ExprEval(const ExprHandle& expr, const std::vector<BufferArg>& buffer_args)
+      : dtype_(expr.dtype()) {
+    std::vector<BufferArg> buffer_args_extended = buffer_args;
+    BufHandle ret_buf("ret_val", {1}, dtype_);
+    std::vector<ExprHandle> indices;
+    ExprHandle zero = IntImm::make(0);
+    indices.reserve(ret_buf.ndim());
+    for (size_t i = 0; i < ret_buf.ndim(); i++) {
+      indices.push_back(zero);
+    }
+    StmtPtr store_stmt = Store::make(ret_buf, indices, expr);
+    buffer_args_extended.emplace_back(ret_buf);
+    codegen_.reset(new CodeGenType(store_stmt, buffer_args_extended));
+  }
+
+  template <typename... Ts>
+  void operator()(Ts... ts) {
+    call(ts...);
+  }
+
+  void operator()(const std::vector<CallArg>& call_args) {
+    call(call_args);
+  }
+
+  void bindVar(VarPtr v, ExprPtr e) {
+    codegen_->bindVar(v, e);
+  }
+
+  void bindVar(const VarHandle& v, const ExprHandle& e) {
+    codegen_->bindVar(v.node(), e.node());
+  }
+
+  template <typename... Ts>
+  void call(Ts... ts) {
+    call({CallArg(ts)...});
+  }
+
+  void call(const std::vector<CallArg>& call_args) {
+    std::vector<CallArg> call_args_extended = call_args;
+    switch (dtype_.scalar_type()) {
+#define TYPE_CASE(Type, Name)                     \
+  case ScalarType::Name: {                        \
+    std::vector<Type> ret_val_arg(1);             \
+    call_args_extended.emplace_back(ret_val_arg); \
+    codegen_->call(call_args_extended);           \
+    ret_value_ = InterpValue(ret_val_arg[0]);     \
+  } break;
+      AT_FORALL_SCALAR_TYPES_AND2(Half, BFloat16, TYPE_CASE);
+      TYPE_CASE(c10::quint8, QUInt8);
+      TYPE_CASE(c10::qint8, QInt8);
+#undef TYPE_CASE
+      case ScalarType::Bool: {
+        std::vector<unsigned char> ret_val_arg(1);
+        call_args_extended.emplace_back(ret_val_arg.data());
+        codegen_->call(call_args_extended);
+        ret_value_ = InterpValue((bool)ret_val_arg[0]);
+      } break;
+      default:
+        throw unsupported_dtype();
+    }
+  }
+
+  void call_raw(const std::vector<void*>& args) {
+    std::vector<void*> args_extended = args;
+    switch (dtype_.scalar_type()) {
+#define TYPE_CASE(Type, Name)                    \
+  case ScalarType::Name: {                       \
+    std::vector<Type> ret_val_arg(1);            \
+    args_extended.push_back(ret_val_arg.data()); \
+    codegen_->call_raw(args_extended);           \
+    ret_value_ = InterpValue(ret_val_arg[0]);    \
+  } break;
+      AT_FORALL_SCALAR_TYPES_AND2(Half, BFloat16, TYPE_CASE);
+      TYPE_CASE(c10::quint8, QUInt8);
+      TYPE_CASE(c10::qint8, QInt8);
+#undef TYPE_CASE
+      case ScalarType::Bool: {
+        std::vector<unsigned char> ret_val_arg(1);
+        args_extended.push_back(ret_val_arg.data());
+        codegen_->call_raw(args_extended);
+        ret_value_ = InterpValue((bool)ret_val_arg[0]);
+      } break;
+      default:
+        throw unsupported_dtype();
+    }
+  }
+
+  template <typename T>
+  T value(const std::vector<void*>& args) {
+    call_raw(args);
+    return ret_value_.as<T>();
+  }
+
+  template <typename T, typename... Ts>
+  T value(Ts... ts) {
+    call(std::forward<Ts>(ts)...);
+    return ret_value_.as<T>();
+  }
+
+  Dtype dtype() {
+    return dtype_;
+  }
+
+ private:
+  Dtype dtype_;
+  std::unique_ptr<CodeGenType> codegen_;
+  InterpValue ret_value_;
+};
+
+// Evaluates the given expression and returns an int64_t value if the result of
+// the given expression is int64_t.
+std::optional<int64_t> evalInt(ExprPtr e);
+
+// Substitutes the given vars with their corresponding expressions in the input
+// expression.
+inline ExprPtr Substitute(const ExprPtr& expr, const VarMapping& var_mapping) {
+  VarSubMutator var_sub(var_mapping);
+  return expr->accept_mutator(&var_sub);
+}
+
+// Substitutes the given vars with their corresponding expressions in the input
+// statement.
+inline StmtPtr Substitute(const StmtPtr& stmt, const VarMapping& var_mapping) {
+  VarSubMutator var_sub(var_mapping);
+  return stmt->accept_mutator(&var_sub);
+}
+
+// Creates a clone of the input expression and substitutes the given vars with
+// their corresponding expressions in the clone.
+// NOTE: This works because cloning reuses variables and does not create new
+// ones, and `VarMapping` input has variables as the key.
+inline ExprPtr SubstituteInClone(
+    const ExprPtr& expr,
+    const VarMapping& var_mapping) {
+  VarSubMutator var_sub(var_mapping);
+  return Expr::clone(expr)->accept_mutator(&var_sub);
+}
+
+// Creates a clone of the input statement and substitutes the given vars with
+// their corresponding expressions in the clone.
+// NOTE: This works because cloning reuses variables and does not create new
+// ones, and `VarMapping` input has variables as the key.
+inline StmtPtr SubstituteInClone(
+    const StmtPtr& stmt,
+    const VarMapping& var_mapping) {
+  VarSubMutator var_sub(var_mapping);
+  return Stmt::clone(stmt)->accept_mutator(&var_sub);
+}
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/exceptions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/exceptions.h
new file mode 100644
index 0000000000000000000000000000000000000000..a131ec071c32c3a6731691fd1f2328559d10ab99
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/exceptions.h
@@ -0,0 +1,90 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+
+#include <stdexcept>
+
+// Forward declarations of types
+
+namespace torch::jit::tensorexpr {
+class Expr;
+class Stmt;
+} // namespace torch::jit::tensorexpr
+
+// Forward declarations of functions
+namespace std {
+TORCH_API std::string to_string(
+    const torch::jit::tensorexpr::ExprPtr& /*expr*/);
+TORCH_API std::string to_string(
+    const torch::jit::tensorexpr::StmtPtr& /*stmt*/);
+} // namespace std
+
+namespace torch::jit::tensorexpr {
+
+class unsupported_dtype : public std::runtime_error {
+ public:
+  explicit unsupported_dtype() : std::runtime_error("UNSUPPORTED DTYPE") {}
+  explicit unsupported_dtype(const std::string& err)
+      : std::runtime_error("UNSUPPORTED DTYPE: " + err) {}
+};
+
+class out_of_range_index : public std::runtime_error {
+ public:
+  explicit out_of_range_index() : std::runtime_error("OUT OF RANGE INDEX") {}
+  explicit out_of_range_index(const std::string& err)
+      : std::runtime_error("OUT OF RANGE INDEX: " + err) {}
+};
+
+class unimplemented_lowering : public std::runtime_error {
+ public:
+  explicit unimplemented_lowering()
+      : std::runtime_error("UNIMPLEMENTED LOWERING") {}
+  explicit unimplemented_lowering(const ExprPtr& expr)
+      : std::runtime_error("UNIMPLEMENTED LOWERING: " + std::to_string(expr)) {}
+  explicit unimplemented_lowering(const StmtPtr& stmt)
+      : std::runtime_error("UNIMPLEMENTED LOWERING: " + std::to_string(stmt)) {}
+};
+
+class malformed_input : public std::runtime_error {
+ public:
+  explicit malformed_input() : std::runtime_error("MALFORMED INPUT") {}
+  explicit malformed_input(const std::string& err)
+      : std::runtime_error("MALFORMED INPUT: " + err) {}
+  explicit malformed_input(const ExprPtr& expr)
+      : std::runtime_error("MALFORMED INPUT: " + std::to_string(expr)) {}
+  explicit malformed_input(const std::string& err, const ExprPtr& expr)
+      : std::runtime_error(
+            "MALFORMED INPUT: " + err + " - " + std::to_string(expr)) {}
+  explicit malformed_input(const StmtPtr& stmt)
+      : std::runtime_error("MALFORMED INPUT: " + std::to_string(stmt)) {}
+  explicit malformed_input(const std::string& err, const StmtPtr& stmt)
+      : std::runtime_error(
+            "MALFORMED INPUT: " + err + " - " + std::to_string(stmt)) {}
+};
+
+class malformed_ir : public std::runtime_error {
+ public:
+  explicit malformed_ir() : std::runtime_error("MALFORMED IR") {}
+  explicit malformed_ir(const std::string& err)
+      : std::runtime_error("MALFORMED IR: " + err) {}
+  explicit malformed_ir(const ExprPtr& expr)
+      : std::runtime_error("MALFORMED IR: " + std::to_string(expr)) {}
+  explicit malformed_ir(const std::string& err, const ExprPtr& expr)
+      : std::runtime_error(
+            "MALFORMED IR: " + err + " - " + std::to_string(expr)) {}
+  explicit malformed_ir(const StmtPtr& stmt)
+      : std::runtime_error("MALFORMED IR: " + std::to_string(stmt)) {}
+  explicit malformed_ir(const std::string& err, const StmtPtr& stmt)
+      : std::runtime_error(
+            "MALFORMED IR: " + err + " - " + std::to_string(stmt)) {}
+};
+
+TORCH_API std::string buildErrorMessage(const std::string& s = "");
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/expr.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/expr.h
new file mode 100644
index 0000000000000000000000000000000000000000..40ba9800b9ce9eb06ddcf9bbdc8384e6480dd2ec
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/expr.h
@@ -0,0 +1,498 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * This file implements the core classes for Tensor Expressions.
+ *
+ * The structure of the expressions is inspired by Halide/TVM IR.
+ */
+#pragma once
+
+#include <c10/core/MemoryFormat.h>
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/types.h>
+#include <optional>
+
+#include <utility>
+
+namespace torch::jit::tensorexpr {
+
+enum IRNodeType {
+  kPrimitive,
+  kAdd,
+  kSub,
+  kMul,
+  kDiv,
+  kMod,
+  kMax,
+  kMin,
+  kAnd,
+  kOr,
+  kLshift,
+  kRshift,
+  kXor,
+  kCompareSelect,
+  kCast,
+  kBitCast,
+  kOther,
+};
+
+// The common base between all expression node.
+class TORCH_API Expr : public std::enable_shared_from_this<Expr> {
+ public:
+  explicit Expr(Dtype dtype, IRNodeType expr_type = kOther)
+      : dtype_(dtype), expr_type_(expr_type) {}
+  virtual ~Expr() = default;
+  Dtype dtype() const {
+    return dtype_;
+  }
+  virtual void accept(IRVisitor* visitor) = 0;
+  virtual ExprPtr accept_mutator(IRMutator* mutator) = 0;
+
+  IRNodeType expr_type() const {
+    return expr_type_;
+  }
+  // Is this a fixed (constant) immediate value.
+  virtual bool isConstant() const {
+    return false;
+  }
+
+  void set_dtype(Dtype dtype) {
+    dtype_ = dtype;
+  }
+
+  /*
+   * Make a deep copy of the given expression.
+   *
+   * All sub-expressions inside the given expressions are also cloned. Note
+   * that the variables are not deep-copied since they are immutable.
+   */
+  static ExprPtr clone(const ExprPtr& s);
+
+ protected:
+  std::shared_ptr<Expr> getptr() {
+    return shared_from_this();
+  }
+
+ private:
+  Dtype dtype_;
+  IRNodeType expr_type_;
+};
+
+// A CRTP pattern to accept visitors for children class,
+// and dispatch back to the children.
+template <class Op, class Base = Expr>
+class ExprNode : public Base {
+ public:
+  using ExprNodeBase = ExprNode<Op>;
+  void accept(IRVisitor* visitor) override {
+    visitor->visit(static_to<Op>(Base::getptr()));
+  }
+  ExprPtr accept_mutator(IRMutator* mutator) override;
+  // pass the constructor to the base class
+  using Base::Base;
+};
+
+// A wrapper object to the underlying ExprNode.
+// Also serves the primary way to build and operate on other expressions.
+class TORCH_API ExprHandle {
+ public:
+  ExprHandle() = default;
+  explicit ExprHandle(ExprPtr node) : base_expr_node_(std::move(node)) {}
+
+  ExprPtr node() {
+    return base_expr_node_;
+  }
+
+  ExprPtr node() const {
+    return base_expr_node_;
+  }
+
+  bool empty() const {
+    return base_expr_node_ == nullptr;
+  }
+
+#define IMM_EXPR_DECLARE(Type, Name) ExprHandle(Type v);
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_EXPR_DECLARE)
+#undef IMM_EXPR_DECLARE
+
+  template <class Op>
+  NodePtr<Op> AsNode() {
+    return to<Op>(this->node());
+  }
+
+  template <class Op>
+  NodePtr<Op> AsNode() const {
+    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
+    return const_cast<ExprHandle*>(this)->AsNode<Op>();
+  }
+
+  Dtype dtype() const {
+    return node()->dtype();
+  }
+
+  // Handling the math operators.
+  ExprHandle operator+(const ExprHandle& other) const;
+  ExprHandle operator-(const ExprHandle& other) const;
+  ExprHandle operator*(const ExprHandle& other) const;
+  ExprHandle operator/(const ExprHandle& other) const;
+  ExprHandle operator%(const ExprHandle& other) const;
+  ExprHandle operator==(const ExprHandle& other) const;
+  ExprHandle operator!=(const ExprHandle& other) const;
+  ExprHandle operator>(const ExprHandle& other) const;
+  ExprHandle operator>=(const ExprHandle& other) const;
+  ExprHandle operator<(const ExprHandle& other) const;
+  ExprHandle operator<=(const ExprHandle& other) const;
+  ExprHandle operator&(const ExprHandle& other) const;
+  ExprHandle operator|(const ExprHandle& other) const;
+  ExprHandle operator&&(const ExprHandle& other) const;
+  ExprHandle operator||(const ExprHandle& other) const;
+  ExprHandle operator^(const ExprHandle& other) const;
+  ExprHandle operator<<(const ExprHandle& other) const;
+  ExprHandle operator>>(const ExprHandle& other) const;
+
+ private:
+  ExprPtr base_expr_node_ = nullptr;
+};
+
+// The underlying representation node to a Var.
+// Currently, each Var object represents a unique variable, even though the
+// names might be the same. We should consider add a unique_name as well.
+class TORCH_API Var : public ExprNode<Var> {
+ public:
+  static ExprHandle make(const std::string& name_hint, Dtype dtype) {
+    return ExprHandle(alloc<Var>(name_hint, dtype));
+  }
+  static ExprHandle make(Dtype dtype) {
+    return ExprHandle(alloc<Var>("", dtype));
+  }
+
+  // TODO: unique_name
+  const std::string& name_hint() const {
+    return name_hint_;
+  }
+
+  void set_name_hint(const std::string& name) {
+    name_hint_ = name;
+  }
+
+  void set_name_hint(std::string&& name) {
+    name_hint_ = std::move(name);
+  }
+
+  Var(std::string name_hint, Dtype dtype)
+      : ExprNodeBase(dtype, kPrimitive), name_hint_(std::move(name_hint)) {}
+
+ private:
+  std::string name_hint_;
+};
+
+TORCH_API std::vector<ExprPtr> make_contiguous_strides(
+    const std::vector<ExprHandle>& dims);
+TORCH_API std::vector<ExprPtr> make_channels_last_strides(
+    const std::vector<ExprHandle>& dims);
+
+class TORCH_API Buf : public ExprNode<Buf> {
+ public:
+  static BufHandle make(const std::vector<ExprHandle>& dims, Dtype dtype);
+
+  static BufHandle make(
+      const std::string& name_hint,
+      const std::vector<ExprHandle>& dims,
+      const std::vector<ExprHandle>& strides,
+      Dtype dtype);
+
+  static BufHandle make(
+      const std::string& name_hint,
+      const std::vector<ExprHandle>& dims,
+      Dtype dtype,
+      std::optional<ExprHandle> initializer = std::nullopt,
+      const std::optional<std::vector<ExprHandle>>& strides = std::nullopt,
+      std::optional<ExprHandle> qscale = std::nullopt,
+      std::optional<ExprHandle> qzero = std::nullopt);
+
+  // TODO: unique_name
+  VarPtr base_handle() const {
+    return base_handle_;
+  }
+  void set_base_handle(VarPtr base_handle) {
+    base_handle_ = std::move(base_handle);
+  }
+
+  const std::string& name_hint() const {
+    return base_handle_->name_hint();
+  }
+  void set_name_hint(const std::string& name_hint) {
+    base_handle_->set_name_hint(name_hint);
+  }
+
+  Buf(const std::string& name_hint,
+      const std::vector<ExprPtr>& dims,
+      Dtype dtype,
+      ExprPtr initializer = nullptr,
+      std::optional<std::vector<ExprPtr>> strides = std::nullopt,
+      ExprPtr qscale = nullptr,
+      ExprPtr qzero = nullptr)
+      : Buf(alloc<Var>(name_hint, kHandle),
+            dims,
+            dtype,
+            std::move(initializer),
+            std::move(strides),
+            std::move(qscale),
+            std::move(qzero)) {}
+
+  Buf(const VarPtr& var,
+      std::vector<ExprPtr> dims,
+      Dtype dtype,
+      ExprPtr initializer = nullptr,
+      std::optional<std::vector<ExprPtr>> strides = std::nullopt,
+      ExprPtr qscale = nullptr,
+      ExprPtr qzero = nullptr);
+
+  size_t ndim() const {
+    return dims_.size();
+  }
+  ExprPtr dim(size_t index) const {
+    if (index >= ndim()) {
+      throw out_of_range_index();
+    }
+    return dims_[index];
+  }
+  std::vector<ExprPtr> dims() const {
+    return dims_;
+  }
+  void set_dims(std::vector<ExprPtr> dims) {
+    dims_ = std::move(dims);
+  }
+
+  std::vector<ExprPtr> strides() const {
+    return strides_;
+  }
+
+  void set_strides(std::vector<ExprPtr> strides) {
+    strides_ = std::move(strides);
+  }
+
+  ExprPtr initializer() const {
+    return initializer_;
+  }
+
+  ExprPtr qzero() const {
+    return qzero_;
+  }
+
+  ExprPtr qscale() const {
+    return qscale_;
+  }
+
+  void set_qzero(ExprPtr qzero) {
+    qzero_ = std::move(qzero);
+  }
+
+  void set_qscale(ExprPtr qscale) {
+    qscale_ = std::move(qscale);
+  }
+
+  bool hasConstantDims() const {
+    for (const auto& d : dims_) {
+      if (!d->isConstant()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  bool is_contiguous(
+      at::MemoryFormat memory_format = at::MemoryFormat::Contiguous) const;
+
+  // The channels-last 1d can benefit the performance of some operators like
+  // conv1d. But the MemoryFormat enum has not covered this layout yet. Hence,
+  // we abstract a dedicated function to check channels-last 1d contiguous.
+  //
+  // Channels-last 1d:
+  //   dims:              n   c    l
+  //   strides(nlc):    c*l   1    c
+  bool is_channels_last_1d_contiguous() const {
+    if (dims_.size() != 3) {
+      return false;
+    }
+    return is_stride_one(1) && is_cont_with(2, 1) && is_cont_with(0, 2);
+  }
+
+ private:
+  bool is_cont_with(int cur_dim, int adjacent_dim) const;
+  bool is_stride_one(int cur_dim) const;
+
+  VarPtr base_handle_;
+  std::vector<ExprPtr> dims_;
+  std::vector<ExprPtr> strides_;
+  ExprPtr initializer_;
+  // qscale_ and qzero_ are used only for quantized dtypes Bufs: kQUInt8, kQInt8
+  ExprPtr qscale_;
+  ExprPtr qzero_;
+};
+
+class TORCH_API BufHandle : public ExprHandle {
+ public:
+  BufHandle(
+      const std::string& name_hint,
+      const std::vector<ExprHandle>& dims,
+      Dtype dtype)
+      : ExprHandle(Buf::make(name_hint, dims, dtype)) {}
+
+  BufHandle(
+      const std::string& name_hint,
+      const std::vector<ExprHandle>& dims,
+      const std::vector<ExprHandle>& strides,
+      Dtype dtype)
+      : ExprHandle(Buf::make(name_hint, dims, strides, dtype)) {}
+
+  BufHandle(const std::vector<ExprHandle>& dims, Dtype dtype)
+      : ExprHandle(Buf::make("_", dims, dtype)) {}
+
+  explicit BufHandle(Dtype dtype) : ExprHandle(Buf::make("_", {}, dtype)) {}
+
+  explicit BufHandle(BufPtr node) : ExprHandle(std::move(node)) {}
+  BufPtr node() const {
+    return static_to<Buf>(ExprHandle::node());
+  }
+  BufPtr node() {
+    return static_to<Buf>(ExprHandle::node());
+  }
+
+  template <typename... Ts>
+  inline ExprHandle load(const Ts&... ts) const;
+
+  template <typename T>
+  inline ExprHandle load(const std::vector<T>& args) const;
+
+  inline ExprHandle load(const std::vector<ExprHandle>& args) const;
+
+  StorePtr store(const std::vector<ExprHandle>& args, const ExprHandle& val)
+      const;
+
+  bool operator==(const BufHandle& other) const {
+    return this->node() == other.node();
+  }
+  bool operator!=(const BufHandle& other) const {
+    return !(*this == other);
+  }
+
+  const std::string& name_hint() const {
+    return this->node()->name_hint();
+  }
+
+  bool empty() const {
+    return (this->node() == nullptr);
+  }
+
+  size_t ndim() const {
+    return node()->ndim();
+  }
+
+  std::vector<ExprHandle> dims() const;
+
+  ExprHandle dim(size_t index) const {
+    return ExprHandle(node()->dim(index));
+  }
+
+  bool is_contiguous(
+      at::MemoryFormat memory_format = at::MemoryFormat::Contiguous) const {
+    return node()->is_contiguous(memory_format);
+  }
+
+  bool is_channels_last_1d_contiguous() const {
+    return node()->is_channels_last_1d_contiguous();
+  }
+};
+
+// An expression to construct the underlying variable node.
+// Note: do not store any info here, since it is often possible to slice this
+// object. For example: VarHandle x('x'); ExprHandle x2 = x;
+class TORCH_API VarHandle : public ExprHandle {
+ public:
+  // Creates an empty VarHandle whose base Var is set to nullptr.
+  VarHandle() = default;
+
+  explicit VarHandle(Dtype dtype) : ExprHandle(Var::make(dtype)) {}
+
+  VarHandle(const std::string& name_hint, Dtype dtype)
+      : ExprHandle(Var::make(name_hint, dtype)) {}
+
+  explicit VarHandle(VarPtr node) : ExprHandle(std::move(node)) {}
+
+  VarPtr node() const {
+    return static_to<Var>(ExprHandle::node());
+  }
+  bool operator==(const VarHandle& other) const {
+    return this->node() == other.node();
+  }
+  bool operator!=(const VarHandle& other) const {
+    return !(*this == other);
+  }
+
+  const std::string& name_hint() const {
+    return this->node()->name_hint();
+  }
+  bool empty() const {
+    return (this->node() == nullptr);
+  }
+};
+
+template <class Op, class Base>
+ExprPtr ExprNode<Op, Base>::accept_mutator(IRMutator* mutator) {
+  return mutator->mutate(static_to<Op>(Base::getptr()));
+}
+
+inline bool same_node(const ExprHandle& expr1, const ExprHandle& expr2) {
+  return expr1.AsNode<Expr>() == expr2.AsNode<Expr>();
+}
+
+TORCH_API ExprHandle sin(const ExprHandle& v);
+TORCH_API ExprHandle cos(const ExprHandle& v);
+TORCH_API ExprHandle tan(const ExprHandle& v);
+TORCH_API ExprHandle asin(const ExprHandle& v);
+TORCH_API ExprHandle acos(const ExprHandle& v);
+TORCH_API ExprHandle atan(const ExprHandle& v);
+TORCH_API ExprHandle sinh(const ExprHandle& v);
+TORCH_API ExprHandle cosh(const ExprHandle& v);
+TORCH_API ExprHandle tanh(const ExprHandle& v);
+TORCH_API ExprHandle sigmoid(const ExprHandle& v);
+TORCH_API ExprHandle exp(const ExprHandle& v);
+TORCH_API ExprHandle expm1(const ExprHandle& v);
+TORCH_API ExprHandle abs(const ExprHandle& v);
+TORCH_API ExprHandle log(const ExprHandle& v);
+TORCH_API ExprHandle fast_tanh(const ExprHandle& v);
+TORCH_API ExprHandle fast_sigmoid(const ExprHandle& v);
+TORCH_API ExprHandle fast_log(const ExprHandle& v);
+TORCH_API ExprHandle log_vml(const ExprHandle& v);
+TORCH_API ExprHandle log2(const ExprHandle& v);
+TORCH_API ExprHandle log10(const ExprHandle& v);
+TORCH_API ExprHandle log1p(const ExprHandle& v);
+TORCH_API ExprHandle erf(const ExprHandle& v);
+TORCH_API ExprHandle erfc(const ExprHandle& v);
+TORCH_API ExprHandle sqrt(const ExprHandle& v);
+TORCH_API ExprHandle rsqrt(const ExprHandle& v);
+TORCH_API ExprHandle ceil(const ExprHandle& v);
+TORCH_API ExprHandle floor(const ExprHandle& v);
+TORCH_API ExprHandle round(const ExprHandle& v);
+TORCH_API ExprHandle trunc(const ExprHandle& v);
+TORCH_API ExprHandle frac(const ExprHandle& v);
+TORCH_API ExprHandle lgamma(const ExprHandle& v);
+TORCH_API ExprHandle atan2(const ExprHandle& v1, const ExprHandle& v2);
+TORCH_API ExprHandle pow(const ExprHandle& v1, const ExprHandle& v2);
+TORCH_API ExprHandle fmod(const ExprHandle& v1, const ExprHandle& v2);
+TORCH_API ExprHandle remainder(const ExprHandle& v1, const ExprHandle& v2);
+TORCH_API ExprHandle isnan(const ExprHandle& v1);
+TORCH_API ExprHandle Relu(const ExprHandle& v1);
+
+TORCH_API ExprHandle
+ifThenElse(const ExprHandle& c, const ExprHandle& t, const ExprHandle& f);
+
+TORCH_API ExprHandle expr_to_vec(const ExprHandle& v, int lanes);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..9d2ff4b04b8c68e1d7c2071860c2c5938bed58ea
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions.h
@@ -0,0 +1,116 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Config.h>
+#include <ATen/Functions.h>
+#include <c10/macros/Macros.h>
+#include <torch/csrc/Export.h>
+#include <cstdint>
+#include <vector>
+
+#define FOR_ALL_EXTERNAL_FUNCTIONS(_)   \
+  _(nnc_aten_adaptive_avg_pool2d)       \
+  _(nnc_aten_addmm)                     \
+  _(nnc_aten_conv2d)                    \
+  _(nnc_aten_conv1d)                    \
+  _(nnc_aten_conv1d_out)                \
+  _(nnc_aten_dequantize)                \
+  _(nnc_aten_dequantize_out)            \
+  _(nnc_aten_embedding)                 \
+  _(nnc_aten_matmul)                    \
+  _(nnc_aten_mv)                        \
+  _(nnc_aten_mm)                        \
+  _(nnc_aten_mean)                      \
+  _(nnc_aten_max_red)                   \
+  _(nnc_aten_max_red_out)               \
+  _(nnc_aten_quantized_conv1d)          \
+  _(nnc_aten_quantized_conv1d_out)      \
+  _(nnc_aten_quantized_conv2d)          \
+  _(nnc_aten_quantized_conv2d_out)      \
+  _(nnc_aten_quantized_conv2d_relu)     \
+  _(nnc_aten_quantized_conv2d_relu_out) \
+  _(nnc_aten_quantized_linear)          \
+  _(nnc_aten_quantized_linear_out)      \
+  _(nnc_aten_quantized_linear_relu)     \
+  _(nnc_aten_quantized_add)             \
+  _(nnc_aten_quantized_cat)             \
+  _(nnc_aten_quantized_mul)             \
+  _(nnc_aten_quantized_mul_out)         \
+  _(nnc_aten_quantized_mul_scalar)      \
+  _(nnc_aten_quantized_mul_scalar_out)  \
+  _(nnc_aten_quantized_relu)            \
+  _(nnc_aten_quantized_sigmoid)         \
+  _(nnc_aten_quantized_sigmoid_out)     \
+  _(nnc_aten_quantize_per_tensor)       \
+  _(nnc_aten_quantize_per_tensor_out)   \
+  _(nnc_aten_triangular_solve)          \
+  _(nnc_aten_upsample_nearest2d)        \
+  _(nnc_aten_upsample_nearest2d_out)    \
+  _(nnc_prepacked_conv2d_clamp_run)     \
+  _(nnc_prepacked_linear_clamp_run)
+
+#define DECLARE_EXTERNAL_FUNCTION(NAME) \
+  TORCH_API void NAME(                  \
+      int64_t bufs_num,                 \
+      void** buf_data,                  \
+      int64_t* buf_ranks,               \
+      int64_t* buf_dims,                \
+      int64_t* buf_strides,             \
+      int8_t* buf_dtypes,               \
+      int64_t args_num,                 \
+      int64_t* extra_args);
+
+namespace torch::jit::tensorexpr {
+struct QIData final {
+  double scale;
+  int64_t zero;
+  c10::ScalarType scalarType;
+};
+std::vector<at::Tensor> constructTensors(
+    int64_t bufs_num,
+    void** buf_data,
+    int64_t* buf_ranks,
+    int64_t* buf_dims,
+    int64_t* buf_strides,
+    int8_t* buf_dtypes,
+    std::optional<std::vector<std::pair<size_t, QIData>>> qdataArg =
+        std::nullopt);
+
+std::vector<at::Tensor> constructTensors2(
+    int64_t bufs_in_num,
+    void** buf_data,
+    int64_t* buf_ranks,
+    int64_t* buf_dims,
+    int64_t* buf_strides,
+    int8_t* buf_dtypes,
+    std::optional<std::vector<std::pair<size_t, QIData>>> qdataArg =
+        std::nullopt,
+    size_t bufs_out_num = 0);
+
+#ifdef C10_MOBILE
+extern "C" {
+#endif
+void DispatchParallel(
+    int8_t* func,
+    int64_t start,
+    int64_t stop,
+    int8_t* packed_data) noexcept;
+
+FOR_ALL_EXTERNAL_FUNCTIONS(DECLARE_EXTERNAL_FUNCTION)
+#if AT_MKLDNN_ENABLED()
+DECLARE_EXTERNAL_FUNCTION(nnc_mkldnn_prepacked_conv_run)
+#endif
+
+TORCH_API void nnc_aten_free(size_t bufs_num, void** ptrs) noexcept;
+
+#ifdef C10_MOBILE
+} // extern "C"
+#endif
+
+} // namespace torch::jit::tensorexpr
+
+#undef DECLARE_EXTERNAL_FUNCTION
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions_core.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions_core.h
new file mode 100644
index 0000000000000000000000000000000000000000..0977fd06b09fa72ac98d15a01b71f5abd841af96
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions_core.h
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Parallel.h>
+#include <torch/csrc/Export.h>
+#include <cstdint>
+
+namespace torch::jit::tensorexpr {
+
+#ifdef C10_MOBILE
+extern "C" {
+#endif
+void DispatchParallel(
+    int8_t* func,
+    int64_t start,
+    int64_t stop,
+    int8_t* packed_data) noexcept;
+
+TORCH_API void nnc_aten_free(size_t bufs_num, void** ptrs) noexcept;
+
+#ifdef C10_MOBILE
+} // extern "C"
+#endif
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions_registry.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions_registry.h
new file mode 100644
index 0000000000000000000000000000000000000000..d638dfa5911bd01af113d94cba196825f61fc532
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/external_functions_registry.h
@@ -0,0 +1,62 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+
+namespace torch::jit::tensorexpr {
+
+// The external functions that could be called from NNC must have the same
+// signature defined by `NNCExternalFunction`.
+//
+// Why this signature?
+// It was picked for two reasons: 1) it should be generic enough to represent
+// most of the ops we might want to call, 2) it should be possible to generate a
+// code for this call in LLVM codegen.
+// The first 5 parameters allow to pass any number of contiguous CPU tensors in
+// case we need to run aten ops (TODO: support different devices). The first
+// buffer in the array is assumed to be the output buffer. We couldn't use
+// `at::Tensor` (or `c10::IValue`) type there directly as it would mean that
+// we'd need to declare it in LLVM codegen in LLVM IR form, which would be very
+// cumbersome and hard to maintain. Note that the dimensions of all tensors are
+// concatenated into a single array buf_dims. We do not need to pass its length,
+// since it can be deduced from total number of buffers and their ranks.
+//
+// The last 2 arguments allow to pass any non-tensor arguments encoded as an
+// array of int64_t values. The way they are encoded is not specified and could
+// be arbitrary - whatever the most convenient for the specific bridge function
+// is.
+//
+// The bridge functions must not throw exceptions - properly propagating them
+// from the generated code is too cumbersome, and thus all calls to functions
+// that could throw must be wrapped with try-catch blocks.
+using NNCExternalFunction = void (*)(
+    int64_t bufs_num,
+    void** buf_data,
+    int64_t* buf_ranks,
+    int64_t* buf_dims,
+    int64_t* buf_strides,
+    int8_t* buf_dtypes,
+    int64_t args_num,
+    int64_t* extra_args);
+
+// Return a global map "function-name" -> "function-pointer" for all registered
+// in NNC external functions
+TORCH_API std::unordered_map<std::string, NNCExternalFunction>&
+getNNCFunctionRegistry();
+
+// To register a new external function in NNC one needs to create an instance of
+// this struct
+struct RegisterNNCExternalFunction {
+  RegisterNNCExternalFunction(const std::string& name, NNCExternalFunction fn) {
+    getNNCFunctionRegistry()[name] = fn;
+  }
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/fwd_decls.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/fwd_decls.h
new file mode 100644
index 0000000000000000000000000000000000000000..c58dff538fab2fc7aafd5c3cc4a361703c921fda
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/fwd_decls.h
@@ -0,0 +1,130 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/ScalarType.h>
+#include <memory>
+
+namespace torch::jit::tensorexpr {
+
+template <typename Node>
+using NodePtr = std::shared_ptr<Node>;
+
+template <typename To, typename From>
+NodePtr<To> to(const NodePtr<From>& x) {
+  return std::dynamic_pointer_cast<To>(x);
+}
+
+template <typename To, typename From>
+NodePtr<To> static_to(NodePtr<From> x) {
+  return std::static_pointer_cast<To>(x);
+}
+
+template <typename Node, typename... Args>
+NodePtr<Node> alloc(Args&&... args) {
+  return std::make_shared<Node>(std::forward<Args>(args)...);
+}
+
+class Buf;
+class Expr;
+class Stmt;
+class Var;
+
+using BufPtr = NodePtr<Buf>;
+using ExprPtr = NodePtr<Expr>;
+using StmtPtr = NodePtr<Stmt>;
+using VarPtr = NodePtr<Var>;
+
+class ExprHandle;
+class VarHandle;
+class BufHandle;
+
+class Add;
+class And;
+class BitCast;
+class Broadcast;
+class Cast;
+class CompareSelect;
+class Div;
+class IfThenElse;
+class Intrinsics;
+class Let;
+class Load;
+class Lshift;
+class Max;
+class MaxTerm;
+class Min;
+class MinTerm;
+class Mod;
+class Mul;
+class Or;
+class Polynomial;
+class Ramp;
+class ReduceOp;
+class RoundOff;
+class Rshift;
+class Store;
+class Sub;
+class Term;
+class Xor;
+using AddPtr = NodePtr<Add>;
+using AndPtr = NodePtr<And>;
+using BitCastPtr = NodePtr<BitCast>;
+using BroadcastPtr = NodePtr<Broadcast>;
+using CastPtr = NodePtr<Cast>;
+using CompareSelectPtr = NodePtr<CompareSelect>;
+using DivPtr = NodePtr<Div>;
+using IfThenElsePtr = NodePtr<IfThenElse>;
+using IntrinsicsPtr = NodePtr<Intrinsics>;
+using LetPtr = NodePtr<Let>;
+using LoadPtr = NodePtr<Load>;
+using LshiftPtr = NodePtr<Lshift>;
+using MaxPtr = NodePtr<Max>;
+using MaxTermPtr = NodePtr<MaxTerm>;
+using MinPtr = NodePtr<Min>;
+using MinTermPtr = NodePtr<MinTerm>;
+using ModPtr = NodePtr<Mod>;
+using MulPtr = NodePtr<Mul>;
+using OrPtr = NodePtr<Or>;
+using PolynomialPtr = NodePtr<Polynomial>;
+using RampPtr = NodePtr<Ramp>;
+using ReduceOpPtr = NodePtr<ReduceOp>;
+using RoundOffPtr = NodePtr<RoundOff>;
+using RshiftPtr = NodePtr<Rshift>;
+using StorePtr = NodePtr<Store>;
+using SubPtr = NodePtr<Sub>;
+using TermPtr = NodePtr<Term>;
+using XorPtr = NodePtr<Xor>;
+
+class Allocate;
+class AtomicAdd;
+class Block;
+class Cond;
+class ExternalCall;
+class ExternalCallWithAlloc;
+class For;
+class Free;
+class FreeExt;
+class PlacementAllocate;
+class SyncThreads;
+using AllocatePtr = NodePtr<Allocate>;
+using AtomicAddPtr = NodePtr<AtomicAdd>;
+using BlockPtr = NodePtr<Block>;
+using CondPtr = NodePtr<Cond>;
+using ExternalCallPtr = NodePtr<ExternalCall>;
+using ExternalCallWithAllocPtr = NodePtr<ExternalCallWithAlloc>;
+using ForPtr = NodePtr<For>;
+using FreePtr = NodePtr<Free>;
+using FreeExtPtr = NodePtr<FreeExt>;
+using PlacementAllocatePtr = NodePtr<PlacementAllocate>;
+using SyncThreadsPtr = NodePtr<SyncThreads>;
+
+#define IMM_DECLARE(Type, Name) \
+  class Name##Imm;              \
+  using Name##ImmPtr = NodePtr<Name##Imm>;
+AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_DECLARE)
+#undef IMM_DECLARE
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/graph_opt.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/graph_opt.h
new file mode 100644
index 0000000000000000000000000000000000000000..27427893f3f71dad8da98a25080b7e17995a636c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/graph_opt.h
@@ -0,0 +1,116 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch::jit::tensorexpr {
+
+// Optimize aten::cat ops in the given subgraph.
+//
+// Moving users of cat to its inputs.
+//    Cat ops get lowered into multiple loops, one per input. When the result
+//    of cat is used by some other op, it results in a situation where inlining
+//    of cat does not happen. This in turn results in intermediate buffers
+//    being created for the result of cat, since it is not inlined.
+//
+//    For example, consider the following graph:
+//       graph(%x : Float(10, strides=[1], device=cpu),
+//             %y : Float(20, strides=[1], device=cpu)):
+//         %dim : int = prim::Constant[value=0]()
+//         %xy_list : Tensor[] = prim::ListConstruct(%x, %y)
+//         %cat : Float(60, strides=[1], device=cpu) = aten::cat(%xy_list, %dim)
+//         %5 : Float(60, strides=[1], device=cpu) = aten::log(%cat)
+//         return (%5))IR";
+//
+//     This will get lowered into:
+//         Allocate(aten_cat);
+//         for (...)
+//           aten_cat[...] = x[...]
+//         for (...)
+//           aten_cat[...] = y[...]
+//         for (...)
+//           aten_log[...] = log(aten_cat[...])
+//         Free(aten_cat);
+//     Note that aten_cat is not inlined into aten_log and it results in
+//     an intermediate buffer allocation as well.
+//
+//     Optimization:
+//        We move the ops that use the result of `cat` into its inputs whenever
+//     possible.
+//
+//     The graph above will be transformed to:
+//        graph(%x : Float(10, strides=[1], device=cpu),
+//              %y : Float(20, strides=[1], device=cpu)):
+//          %3 : int = prim::Constant[value=0]()
+//          %7 : Float(10, strides=[1], device=cpu) = aten::log(%x)
+//          %8 : Float(20, strides=[1], device=cpu) = aten::log(%y)
+//          %9 : Tensor[] = prim::ListConstruct(%7, %8)
+//          %10 : Float(60, strides=[1], device=cpu) = aten::cat(%9, %3)
+//          return (%10)
+//
+//     This will get lowered into:
+//         for (...)
+//           aten_cat[...] = log(x[...])
+//         for (...)
+//           aten_cat[...] = log(y[...])
+//     aten_cat is the output buffer here.
+
+bool OptimizeCat(const std::shared_ptr<Graph>& graph);
+
+TORCH_API void annotateInputShapes(
+    const std::shared_ptr<Graph>& graph,
+    const std::vector<std::optional<at::Tensor>>& example_inputs);
+TORCH_API std::shared_ptr<Graph> removeUnusedSelfArgument(
+    const std::shared_ptr<Graph>& graph);
+TORCH_API std::shared_ptr<Graph> removeGraphOutput(
+    const std::shared_ptr<Graph>& graph,
+    size_t idx);
+TORCH_API std::shared_ptr<Graph> replaceListOutputWithTuple(
+    const std::shared_ptr<Graph>& graph);
+
+// Perform \p ITERS rounds of "trimming" for the given \p GRAPH.
+//
+// Trimming means that we try to remove a small portion of the graph while
+// keeping it valid. This is useful for debugging when we try to find a minimal
+// example reproducing the issue at hand. When ITERS is 0, the graph remains
+// unchanged, when ITERS is a big number, the graph usually becomes empty.
+TORCH_API std::shared_ptr<Graph> trimGraph(
+    const std::shared_ptr<Graph>& graph,
+    int64_t iters);
+
+// Scan all values in the given graph and replace each dimension with a size Xi
+// present in \p SIZES with a symbolic shape Yi. Return a vector of symbol
+// values [Y0, Y1, .., Yn].
+//
+// For example:
+// Input:
+// graph(%x : Float(10, 20, 30, 40)):
+//   %y : Float(10, 20, 30, 40) = aten::relu(%x)
+//   return %y
+//
+// If we run makeShapesSymbolic(graph, {20, 40}), then we'll get:
+//
+// graph(%x : Float(10, SS(-3), 30, SS(-5))):
+//   %y : Float(10, SS(-3), 30, SS(-5)) = aten::relu(%x)
+//   return %y
+//
+// and get {-3, -5} as the return value.
+TORCH_API std::vector<int64_t> makeShapesSymbolic(
+    std::shared_ptr<Graph>& graph,
+    const std::vector<int64_t>& sizes);
+
+// Inspect the graph and report whether it can be converted to TE IR.
+// TODO: add error reporting for graphs that can't be converted.
+TORCH_API bool isGraphCompilable(const std::shared_ptr<Graph>& graph);
+
+// Examine the graph and (hackily) fill in missing tensor type info, such as
+// scalar type, device, and strides. Ideally, this should be done by a proper
+// dtype/device/shape propagation passes, but until they are ready we can use
+// this, not always correct, workaround pass.
+TORCH_API void fixupMissingShapeInfo(const std::shared_ptr<Graph>& graph);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/half_support.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/half_support.h
new file mode 100644
index 0000000000000000000000000000000000000000..2d14d0b6aba003711f461c265d7be59a11b661fd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/half_support.h
@@ -0,0 +1,217 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+namespace torch::jit::tensorexpr {
+
+// Walk the Statement looking for Half size loads/stores.
+class HalfChecker : public IRVisitor {
+ public:
+  HalfChecker(const std::vector<CodeGen::BufferArg>& args) {
+    for (const auto& BA : args) {
+      hasHalf_ |= BA.dtype().scalar_type() == ScalarType::Half;
+    }
+  }
+
+  bool hasHalf() const {
+    return hasHalf_;
+  }
+
+  bool hasBFloat16() const {
+    return hasBFloat16_;
+  }
+
+  void visit(const LoadPtr& v) override {
+    hasHalf_ |= v->dtype().scalar_type() == ScalarType::Half;
+    hasBFloat16_ |= v->dtype().scalar_type() == ScalarType::BFloat16;
+    IRVisitor::visit(v);
+  }
+
+  void visit(const StorePtr& v) override {
+    hasHalf_ |= v->buf()->dtype().scalar_type() == ScalarType::Half;
+    hasBFloat16_ |= v->buf()->dtype().scalar_type() == ScalarType::BFloat16;
+    IRVisitor::visit(v);
+  }
+
+  void visit(const HalfImmPtr& v) override {
+    hasHalf_ = true;
+  }
+
+  void visit(const BFloat16ImmPtr& v) override {
+    hasBFloat16_ = true;
+  }
+
+  void visit(const CastPtr& v) override {
+    hasHalf_ |= v->dtype().scalar_type() == ScalarType::Half;
+    hasBFloat16_ |= v->dtype().scalar_type() == ScalarType::BFloat16;
+    IRVisitor::visit(v);
+  }
+
+ private:
+  bool hasHalf_{false};
+  bool hasBFloat16_{false};
+};
+
+class HalfRewriter : public IRMutator {
+  ExprPtr mutate(const LoadPtr& v) override {
+    ExprPtr child = IRMutator::mutate(v);
+    if (!isHalf(child)) {
+      return child;
+    }
+
+    ExprPtr ret = alloc<Cast>(
+        child->dtype().cloneWithScalarType(ScalarType::Float), child);
+
+    inserted_half_casts_.insert(ret);
+    return ret;
+  }
+
+  StmtPtr mutate(const StorePtr& v) override {
+    // Since mutation changes the `value()` expression in-place, we need to
+    // get the dtype of the `value()` before that is mutated.
+    auto newType = v->value()->dtype();
+    ExprPtr new_val = v->value()->accept_mutator(this);
+    auto bufType = v->buf()->dtype();
+
+    if (isHalf(newType.scalar_type())) {
+      new_val = alloc<Cast>(newType, new_val);
+      inserted_half_casts_.insert(new_val);
+    }
+
+    // The scalar_type of value is not Half while the buf is Half
+    if (!isHalf(newType.scalar_type()) && isHalf(bufType.scalar_type())) {
+      new_val = alloc<Cast>(
+          newType.cloneWithScalarType(bufType.scalar_type()), new_val);
+      inserted_half_casts_.insert(new_val);
+    }
+
+    v->set_value(new_val);
+    return v;
+  }
+
+  ExprPtr mutate(const HalfImmPtr& v) override {
+    return alloc<Cast>(kFloat, v);
+  }
+
+  ExprPtr mutate(const BFloat16ImmPtr& v) override {
+    return alloc<Cast>(kFloat, v);
+  }
+
+  ExprPtr mutate(const CastPtr& v) override {
+    ExprPtr child = v->src_value()->accept_mutator(this);
+
+    // just don't allow half casts we didn't insert.
+    if (isHalf(v)) {
+      if (inserted_half_casts_.count(v) < 1) {
+        v->set_src_value(child);
+        v->set_dtype(v->dtype().cloneWithScalarType(c10::kFloat));
+        return v;
+      }
+    }
+
+    // Remove Half(Float()) and friends.
+    CastPtr cast_child = to<Cast>(child);
+    if (cast_child) {
+      auto cast_to_double = v->dtype().scalar_type() == ScalarType::Double;
+      auto from_half = isHalf(cast_child->src_value());
+      // Cannot simplify the double(float(half)) to double(half) as NNC does
+      // not support cast BF16 to double directly.
+      auto not_cast_half_to_doulbe = !(cast_to_double && from_half);
+      if (v->dtype().is_floating_point() &&
+          cast_child->dtype().is_floating_point() && not_cast_half_to_doulbe) {
+        return alloc<Cast>(v->dtype(), cast_child->src_value());
+      }
+    }
+
+    if (child == v->src_value()) {
+      return v;
+    }
+
+    return alloc<Cast>(v->dtype(), child);
+  }
+
+  StmtPtr mutate(const LetPtr& v) override {
+    if (isHalf(v->var()->dtype().scalar_type())) {
+      VarPtr load_new_var = alloc<Var>(v->var()->name_hint(), kFloat);
+      ExprPtr new_value = alloc<Cast>(
+          v->var()->dtype().cloneWithScalarType(ScalarType::Float),
+          v->value()->accept_mutator(this));
+      var_map[v->var()] = load_new_var;
+
+      return alloc<Let>(load_new_var, new_value);
+    }
+
+    return IRMutator::mutate(v);
+  }
+
+  ExprPtr mutate(const VarPtr& v) override {
+    auto it = var_map.find(v);
+    if (it != var_map.end()) {
+      return it->second;
+    }
+
+    return v;
+  }
+
+  template <typename T>
+  ExprPtr mutateArithmetic(T v) {
+    IRMutator::mutate(v);
+    if (isHalf(v)) {
+      v->set_dtype(v->dtype().cloneWithScalarType(c10::kFloat));
+    }
+    return v;
+  }
+
+  ExprPtr mutate(const AddPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const SubPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const MulPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const DivPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const MaxPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const MinPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const CompareSelectPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const BroadcastPtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const IfThenElsePtr& v) override {
+    return mutateArithmetic(v);
+  }
+  ExprPtr mutate(const IntrinsicsPtr& v) override {
+    return mutateArithmetic(v);
+  }
+
+ private:
+  static bool isHalf(ScalarType st) {
+    return st == ScalarType::Half || st == ScalarType::BFloat16;
+  }
+
+  static bool isHalf(const ExprPtr& v) {
+    return isHalf(v->dtype().scalar_type());
+  }
+
+  std::unordered_set<ExprPtr> inserted_half_casts_;
+  std::unordered_map<VarPtr, VarPtr> var_map;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/hash_provider.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/hash_provider.h
new file mode 100644
index 0000000000000000000000000000000000000000..e574ad7f71d2aa6845dc3ce127432d6d5ea775f7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/hash_provider.h
@@ -0,0 +1,287 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_printer.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+#include <utility>
+
+namespace torch::jit::tensorexpr {
+
+struct TORCH_API SimplifierHashType {
+  SimplifierHashType() = default;
+  explicit SimplifierHashType(size_t s) : _h(s) {}
+
+  bool operator==(const SimplifierHashType& other) const;
+  bool operator!=(const SimplifierHashType& other) const;
+  bool operator<(const SimplifierHashType& other) const;
+  bool operator==(const size_t other) const;
+  bool operator!=(const size_t other) const;
+
+  size_t _h{0};
+};
+
+} // namespace torch::jit::tensorexpr
+
+namespace std {
+template <>
+struct hash<torch::jit::tensorexpr::SimplifierHashType> {
+  size_t operator()(
+      const torch::jit::tensorexpr::SimplifierHashType& k) const noexcept {
+    return k._h;
+  }
+};
+
+} // namespace std
+
+namespace torch::jit::tensorexpr {
+
+#define CACHE_GUARD()  \
+  if (cachedHash(v)) { \
+    return;            \
+  }
+
+class Term;
+class Polynomial;
+
+/* Expression hasher providing comparable values representing sub-exprs.
+ * Uses memoization to avoid excessive recursion. */
+class TORCH_API HashProvider : public IRVisitor {
+ public:
+  template <class T>
+  SimplifierHashType hash(T e) {
+    e->accept(this);
+    return hashOf(e);
+  }
+
+  bool cachedHash(const ExprPtr& e) {
+    return exprToHash_.find(e) != exprToHash_.end();
+  }
+  bool cachedHash(const StmtPtr& s) {
+    return stmtToHash_.find(s) != stmtToHash_.end();
+  }
+
+  void clearCache() {
+    exprToHash_.clear();
+    stmtToHash_.clear();
+  }
+
+  void visit(const AddPtr& v) override;
+  void visit(const SubPtr& v) override;
+  void visit(const MulPtr& v) override;
+  void visit(const DivPtr& v) override;
+  void visit(const ModPtr& v) override;
+  void visit(const RoundOffPtr& v) override;
+  void visit(const MaxPtr& v) override;
+  void visit(const MinPtr& v) override;
+  void visit(const AndPtr& v) override;
+  void visit(const OrPtr& v) override;
+  void visit(const XorPtr& v) override;
+  void visit(const LshiftPtr& v) override;
+  void visit(const RshiftPtr& v) override;
+  void visit(const CompareSelectPtr& v) override;
+
+#define IMM_VISIT(Type, Name)                    \
+  void visit(const Name##ImmPtr& v) override {   \
+    CACHE_GUARD();                               \
+    putHash(v, hash_combine(#Name, v->value())); \
+  }
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_VISIT)
+#undef IMM_VISIT
+
+  void visit(const CastPtr& v) override;
+  void visit(const VarPtr& v) override;
+  void visit(const RampPtr& v) override;
+  void visit(const LoadPtr& v) override;
+  void visit(const StorePtr& v) override;
+  void visit(const BlockPtr& v) override;
+  void visit(const ForPtr& v) override;
+  void visit(const BroadcastPtr& v) override;
+  void visit(const IfThenElsePtr& v) override;
+  void visit(const IntrinsicsPtr& v) override;
+  void visit(const AllocatePtr& v) override;
+  void visit(const FreePtr& v) override;
+  void visit(const CondPtr& v) override;
+  void visit(const TermPtr& v) override;
+  void visit(const PolynomialPtr& v) override;
+  void visit(const MaxTermPtr& v) override;
+  void visit(const MinTermPtr& v) override;
+
+  template <typename... Types>
+  SimplifierHashType hash_combine(const Types&... args) {
+    SimplifierHashType seed;
+    _hash_combine(seed, args...);
+    return seed;
+  }
+
+ private:
+  SimplifierHashType hashOf(const ExprPtr& e) {
+    auto it = exprToHash_.find(e);
+    if (it != exprToHash_.end()) {
+      return it->second;
+    }
+
+    // As a failsafe fall back to IRPrinter.
+    std::stringstream ss;
+    IRPrinter printer(ss);
+    e->accept(&printer);
+    SimplifierHashType hash = SimplifierHashType(te_hash(ss.str()));
+    putHash(e, hash);
+
+    return hash;
+  }
+
+  SimplifierHashType hashOf(const StmtPtr& s) {
+    auto it = stmtToHash_.find(s);
+    if (it != stmtToHash_.end()) {
+      return it->second;
+    }
+
+    // As a failsafe fall back to IRPrinter.
+    std::stringstream ss;
+    IRPrinter printer(ss);
+    s->accept(&printer);
+    SimplifierHashType hash = SimplifierHashType(te_hash(ss.str()));
+    putHash(s, hash);
+
+    return hash;
+  }
+
+  // Hash funcs for various types, numbers are random.
+  template <typename T>
+  void _hash_combine(SimplifierHashType& seed, const T& val) {
+    seed._h ^= te_hash(val) + 0x1f752c19 + (seed._h << 7) + (seed._h >> 4);
+  }
+
+  void _hash_combine(SimplifierHashType& seed, const char* val) {
+    seed._h ^= te_hash(val) + 0x1f752c19 + (seed._h << 7) + (seed._h >> 4);
+  }
+
+  // at:::Half doesn't have a prime_number_hash, so cast to short.
+  void _hash_combine(SimplifierHashType& seed, const at::Half& val) {
+    seed._h ^=
+        te_hash((uint16_t)val) + 0x1f752c19 + (seed._h << 7) + (seed._h >> 4);
+  }
+
+  void _hash_combine(SimplifierHashType& seed, const Dtype& val) {
+    seed._h ^= te_hash(val.ToCppString()) + 0x1f752c19 + (seed._h << 7) +
+        (seed._h >> 4);
+  }
+
+  void _hash_combine(SimplifierHashType& seed, ExprPtr e) {
+    _hash_combine(seed, hash(std::move(e)));
+  }
+
+  template <typename T, typename... Types>
+  void _hash_combine(
+      SimplifierHashType& seed,
+      const T& val,
+      const Types&... args) {
+    _hash_combine(seed, val);
+    _hash_combine(seed, args...);
+  }
+
+  void putHash(const ExprPtr& e, SimplifierHashType h) {
+    auto res = exprToHash_.emplace(e, h);
+    if (res.second == false) {
+      // This is always a logic bug since we should check the cache first.
+      throw std::runtime_error("hash collision");
+    }
+  }
+  void putHash(const StmtPtr& s, SimplifierHashType h) {
+    auto res = stmtToHash_.emplace(s, h);
+    if (res.second == false) {
+      // This is always a logic bug since we should check the cache first.
+      throw std::runtime_error("hash collision");
+    }
+  }
+
+  std::unordered_map<ExprPtr, SimplifierHashType> exprToHash_;
+  std::unordered_map<StmtPtr, SimplifierHashType> stmtToHash_;
+  UniqueNameManager name_manager_;
+
+  size_t te_hash(SimplifierHashType val) {
+    return val._h;
+  }
+
+  size_t te_hash(int64_t val) {
+    // put the thing down.
+    size_t h = val ^ 0x647AA4D20C0B;
+    // bit flip it.
+    size_t h2 = ~h;
+    // and reverse byte order.
+    size_t h3 = 0;
+    for (unsigned int i = 0; i < 64; i += 8) {
+      h3 |= ((h2 >> i) & 0xFF) << (64 - i - 8);
+    }
+    return h3;
+  }
+
+  size_t te_hash(int32_t val) {
+    int64_t v2 = val;
+    return te_hash(v2);
+  }
+
+  size_t te_hash(uint32_t val) {
+    int64_t v2 = val;
+    return te_hash(v2);
+  }
+
+  size_t te_hash(uint64_t val) {
+    int64_t v2 = val;
+    return te_hash(v2);
+  }
+
+  size_t te_hash(int16_t val) {
+    int64_t v2 = val;
+    return te_hash(v2);
+  }
+
+  size_t te_hash(std::string val) {
+    size_t hash{0};
+    int64_t intval{0};
+    int64_t s = val.size() - 1;
+    while (s >= 0) {
+      for (unsigned int i = 0; i < 8; ++i) {
+        if (s < 0)
+          break;
+        int64_t c = val[s];
+        intval |= (c << (i * 8));
+
+        s--;
+      }
+      hash ^= te_hash(intval);
+      intval = 0;
+    }
+
+    return hash;
+  }
+
+  size_t te_hash(double d) {
+    int64_t* n = reinterpret_cast<int64_t*>(&d);
+    return te_hash(*n);
+  }
+
+  size_t te_hash(float d) {
+    int32_t* n = reinterpret_cast<int32_t*>(&d);
+    return te_hash(*n);
+  }
+
+  size_t te_hash(at::Half d) {
+    int16_t* n = reinterpret_cast<int16_t*>(&d);
+    return te_hash(*n);
+  }
+
+  size_t te_hash(at::BFloat16 d) {
+    int16_t* n = reinterpret_cast<int16_t*>(&d);
+    return te_hash(*n);
+  }
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/intrinsic_symbols.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/intrinsic_symbols.h
new file mode 100644
index 0000000000000000000000000000000000000000..ba5b5ff6728bec6e50e9151c29a9e1d2885693e2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/intrinsic_symbols.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef TORCH_ENABLE_LLVM
+#include <c10/util/ArrayRef.h>
+
+namespace torch {
+namespace jit {
+namespace tensorexpr {
+
+struct SymbolAddress {
+  const char* symbol;
+  void* address;
+
+  SymbolAddress(const char* sym, void* addr) : symbol(sym), address(addr) {}
+};
+
+c10::ArrayRef<SymbolAddress> getIntrinsicSymbols();
+
+} // namespace tensorexpr
+} // namespace jit
+} // namespace torch
+#endif // TORCH_ENABLE_LLVM
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir.h
new file mode 100644
index 0000000000000000000000000000000000000000..4ec8f9bbf8d3a9e71f90e1aa712e2d75f336fcab
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir.h
@@ -0,0 +1,921 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <torch/csrc/jit/tensorexpr/exceptions.h>
+#include <torch/csrc/jit/tensorexpr/expr.h>
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+#include <torch/csrc/jit/tensorexpr/stmt.h>
+
+#include <ATen/core/ivalue.h>
+
+namespace torch::jit::tensorexpr {
+
+enum CompareSelectOperation {
+  kEQ = 0,
+  kGT,
+  kGE,
+  kLT,
+  kLE,
+  kNE,
+};
+
+enum CompareSelectBias {
+  kUnbiased,
+  kLikely,
+  kUnlikely,
+};
+
+inline int getPrecedence(IRNodeType ty) {
+  // Match C++ operator precedence rules, since some pretty-print expressions to
+  // C++. SEE: https://en.cppreference.com/w/cpp/language/operator_precedence
+  switch (ty) {
+    case kPrimitive:
+      return 0;
+    case kCast:
+    case kBitCast:
+      return 2;
+    case kAdd:
+    case kSub:
+      return 6;
+    case kMul:
+    case kDiv:
+    case kMod:
+      return 5;
+    case kMax:
+    case kMin:
+      return 99;
+    case kAnd:
+      return 11;
+    case kOr:
+      return 13;
+    case kLshift:
+    case kRshift:
+      return 7;
+    case kXor:
+      return 12;
+    case kCompareSelect:
+      return 16;
+    default:
+      return 99;
+  }
+}
+
+class TORCH_API Cast : public ExprNode<Cast> {
+ public:
+  ExprPtr src_value() const {
+    return src_value_;
+  }
+
+  void set_src_value(ExprPtr src_value) {
+    src_value_ = std::move(src_value);
+  }
+
+  static ExprHandle make(Dtype dtype, const ExprHandle& src_value) {
+    return ExprHandle(alloc<Cast>(dtype, src_value.node()));
+  }
+  Cast(Dtype dtype, ExprPtr src_value)
+      : ExprNodeBase(dtype, kCast), src_value_(std::move(src_value)) {}
+
+  bool isConstant() const override {
+    return src_value_->isConstant();
+  }
+
+ private:
+  ExprPtr src_value_;
+};
+
+template <typename T>
+ExprHandle cast(const ExprHandle& src_value) {
+  return Cast::make(Dtype(ToDtype<T>(), src_value.dtype().lanes()), src_value);
+}
+
+// This is a bitwise cast, akin to bitcast in LLVM
+class TORCH_API BitCast : public ExprNode<BitCast> {
+ public:
+  ExprPtr src_value() const {
+    return src_value_;
+  }
+
+  void set_src_value(ExprPtr src_value) {
+    src_value_ = std::move(src_value);
+  }
+
+  static ExprHandle make(Dtype dtype, const ExprHandle& src_value) {
+    return ExprHandle(alloc<BitCast>(dtype, src_value.node()));
+  }
+  BitCast(Dtype dtype, ExprPtr src_value)
+      : ExprNodeBase(dtype, kBitCast), src_value_(std::move(src_value)) {
+    TORCH_CHECK(src_value_->dtype().byte_size() == dtype.byte_size());
+  }
+
+  bool isConstant() const override {
+    return src_value_->isConstant();
+  }
+
+ private:
+  ExprPtr src_value_;
+};
+
+template <typename T>
+ExprHandle bitcast(const ExprHandle& src_value) {
+  return BitCast::make(
+      Dtype(ToDtype<T>(), src_value.dtype().lanes()), src_value);
+}
+
+// Represent the expression node for binary operators.
+// A CRTP pattern to share common code among the operators.
+template <typename Op>
+class BinaryOpNode : public ExprNode<Op> {
+ public:
+  ExprPtr lhs() const {
+    return this->lhs_;
+  }
+  ExprPtr rhs() const {
+    return this->rhs_;
+  }
+
+  void set_lhs(ExprPtr lhs) {
+    lhs_ = std::move(lhs);
+  }
+
+  void set_rhs(ExprPtr rhs) {
+    rhs_ = std::move(rhs);
+  }
+
+  static ExprHandle make(const ExprHandle& lhs, const ExprHandle& rhs) {
+    return ExprHandle(alloc<Op>(lhs.node(), rhs.node()));
+  }
+
+  BinaryOpNode(
+      ExprPtr lhs_v,
+      ExprPtr rhs_v,
+      IRNodeType expr_type,
+      ScalarType ret_type = ScalarType::Undefined)
+      : ExprNode<Op>(
+            BinaryOpDtype(lhs_v->dtype(), rhs_v->dtype(), ret_type),
+            expr_type),
+        lhs_(CastIfNeeded(std::move(lhs_v), ExprNode<Op>::dtype())),
+        rhs_(CastIfNeeded(std::move(rhs_v), ExprNode<Op>::dtype())) {}
+
+ private:
+  static ExprPtr CastIfNeeded(ExprPtr expr, Dtype dst_dtype) {
+    if (expr->dtype() == dst_dtype) {
+      return expr;
+    }
+    return Cast::make(dst_dtype, ExprHandle(std::move(expr))).node();
+  }
+
+  ExprPtr lhs_;
+  ExprPtr rhs_;
+};
+
+namespace detail {
+template <typename T>
+void bin_op_deducer(BinaryOpNode<T>);
+bool bin_op_deducer(...);
+} // namespace detail
+
+class TORCH_API Add : public BinaryOpNode<Add> {
+ public:
+  Add(ExprPtr lhs, ExprPtr rhs)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kAdd) {}
+};
+
+class TORCH_API Sub : public BinaryOpNode<Sub> {
+ public:
+  Sub(ExprPtr lhs, ExprPtr rhs)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kSub) {}
+};
+
+class TORCH_API Mul : public BinaryOpNode<Mul> {
+ public:
+  Mul(ExprPtr lhs, ExprPtr rhs)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kMul) {}
+};
+
+class TORCH_API Div : public BinaryOpNode<Div> {
+ public:
+  Div(ExprPtr lhs, ExprPtr rhs)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kDiv) {}
+};
+
+class TORCH_API Mod : public BinaryOpNode<Mod> {
+ public:
+  Mod(ExprPtr lhs, ExprPtr rhs)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kMod) {}
+};
+
+template <typename Op>
+class BitwiseOpNode : public BinaryOpNode<Op> {
+ public:
+  BitwiseOpNode(ExprPtr lhs, ExprPtr rhs, IRNodeType type)
+      : BinaryOpNode<Op>(std::move(lhs), std::move(rhs), type) {}
+
+  static ExprHandle make(const ExprHandle& lhs, const ExprHandle& rhs) {
+    if (!lhs.dtype().is_integral()) {
+      throw unsupported_dtype();
+    }
+    if (lhs.dtype() != rhs.dtype()) {
+      throw malformed_input("lhs/rhs dtype mismatch");
+    }
+    return BinaryOpNode<Op>::make(lhs, rhs);
+  }
+};
+
+class TORCH_API And : public BitwiseOpNode<And> {
+ public:
+  And(ExprPtr lhs, ExprPtr rhs)
+      : BitwiseOpNode(std::move(lhs), std::move(rhs), IRNodeType::kAnd) {}
+};
+
+class TORCH_API Or : public BitwiseOpNode<Or> {
+ public:
+  Or(ExprPtr lhs, ExprPtr rhs)
+      : BitwiseOpNode(std::move(lhs), std::move(rhs), IRNodeType::kOr) {}
+};
+
+class TORCH_API Xor : public BitwiseOpNode<Xor> {
+ public:
+  Xor(ExprPtr lhs, ExprPtr rhs)
+      : BitwiseOpNode(std::move(lhs), std::move(rhs), IRNodeType::kXor) {}
+};
+
+class TORCH_API Lshift : public BitwiseOpNode<Lshift> {
+ public:
+  Lshift(ExprPtr lhs, ExprPtr rhs)
+      : BitwiseOpNode(std::move(lhs), std::move(rhs), IRNodeType::kLshift) {}
+};
+
+class TORCH_API Rshift : public BitwiseOpNode<Rshift> {
+ public:
+  Rshift(ExprPtr lhs, ExprPtr rhs)
+      : BitwiseOpNode(std::move(lhs), std::move(rhs), IRNodeType::kRshift) {}
+};
+
+// TODO: add TORCH_API
+// Currently adding it results in a compilation error on Windows
+class Max : public BinaryOpNode<Max> {
+ private:
+  bool propagate_nans_;
+
+ public:
+  Max(ExprPtr lhs, ExprPtr rhs, bool propagate_nans)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kMax),
+        propagate_nans_(propagate_nans) {}
+
+  bool propagate_nans() const {
+    return propagate_nans_;
+  }
+
+  static ExprHandle make(const ExprHandle& lhs, const ExprHandle& rhs) = delete;
+  static ExprHandle make(
+      const ExprHandle& lhs,
+      const ExprHandle& rhs,
+      bool propagate_nans) {
+    return ExprHandle(alloc<Max>(lhs.node(), rhs.node(), propagate_nans));
+  }
+};
+
+// TODO: add TORCH_API
+// Currently adding it results in a compilation error on Windows
+class Min : public BinaryOpNode<Min> {
+ private:
+  bool propagate_nans_;
+
+ public:
+  Min(ExprPtr lhs, ExprPtr rhs, bool propagate_nans)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kMin),
+        propagate_nans_(propagate_nans) {}
+
+  bool propagate_nans() const {
+    return propagate_nans_;
+  }
+
+  static ExprHandle make(const ExprHandle& lhs, const ExprHandle& rhs) = delete;
+  static ExprHandle make(
+      const ExprHandle& lhs,
+      const ExprHandle& rhs,
+      bool propagate_nans) {
+    return ExprHandle(alloc<Min>(lhs.node(), rhs.node(), propagate_nans));
+  }
+};
+
+// Encode typed immediate values e.g. IntImm, FloatImm.
+#define IMM_DECLARE(Type, Name)                               \
+  class TORCH_API Name##Imm : public ExprNode<Name##Imm> {    \
+   public:                                                    \
+    Name##Imm(Type value)                                     \
+        : ExprNodeBase(k##Name, kPrimitive), value_(value) {} \
+    bool isConstant() const override {                        \
+      return true;                                            \
+    }                                                         \
+    Type value() const {                                      \
+      return value_;                                          \
+    }                                                         \
+    static ExprHandle make(Type value) {                      \
+      return ExprHandle(alloc<Name##Imm>(value));             \
+    }                                                         \
+                                                              \
+   private:                                                   \
+    Type value_;                                              \
+  };
+AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_DECLARE)
+#undef IMM_DECLARE
+
+// Get immediate by ScalarType.
+template <typename T>
+ExprPtr getImmediateByType(ScalarType immType, T initialVal) {
+  switch (immType) {
+#define TYPE_CASE(Type, Name) \
+  case ScalarType::Name:      \
+    return alloc<Name##Imm>(Type(initialVal));
+    AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TYPE_CASE)
+#undef TYPE_CASE
+    default:
+      throw unsupported_dtype();
+  }
+  return nullptr;
+}
+
+template <typename T>
+ExprPtr getImmediateByType(Dtype dtype, T initialVal) {
+  return getImmediateByType<T>(dtype.scalar_type(), initialVal);
+}
+
+template <typename T>
+ExprPtr immLike(const ExprPtr& e, T v) {
+  return getImmediateByType<T>(e->dtype(), v);
+}
+
+template <typename T>
+ExprPtr immLike(const ExprHandle& e, T v) {
+  return immLike(e.node(), v);
+}
+
+inline std::optional<int64_t> intValue(const ExprPtr& e) {
+#define TYPE_CASE(Type, Name)      \
+  if (auto v = to<Name##Imm>(e)) { \
+    return v->value();             \
+  }
+  AT_FORALL_INT_TYPES(TYPE_CASE);
+#undef TYPE_CASE
+  return std::nullopt;
+}
+
+inline std::optional<int64_t> intValue(const ExprHandle& e) {
+  return intValue(e.node());
+}
+
+template <typename T>
+T immediateAs(const ExprPtr& e) {
+#define TYPE_CASE(Type, Name)                \
+  if (Name##ImmPtr imm = to<Name##Imm>(e)) { \
+    return imm->value();                     \
+  }
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TYPE_CASE)
+#undef TYPE_CASE
+  throw unsupported_dtype();
+  return 0;
+}
+
+template <typename T>
+T immediateAs(const ExprHandle& e) {
+  return immediateAs<T>(e.node());
+}
+
+template <typename T>
+bool immediateEquals(const ExprPtr& e, T val) {
+#define TYPE_CASE(Type, Name)                \
+  if (Name##ImmPtr imm = to<Name##Imm>(e)) { \
+    return imm->value() == val;              \
+  }
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TYPE_CASE)
+#undef TYPE_CASE
+  throw unsupported_dtype();
+  return false;
+}
+
+TORCH_API bool immediateIsNegative(const ExprPtr& e);
+
+TORCH_API bool immediateIsPositive(const ExprPtr& e);
+
+TORCH_API bool immediateIsZero(const ExprPtr& e);
+
+// Represents a ramp vector node:
+//     [base, base + 1 * stride, ... , base + (lanes - 1) * stride]
+class TORCH_API Ramp : public ExprNode<Ramp> {
+ public:
+  ExprPtr base() const {
+    return base_;
+  }
+  ExprPtr stride() const {
+    return stride_;
+  }
+
+  void set_base(ExprPtr base) {
+    base_ = std::move(base);
+  }
+
+  void set_stride(ExprPtr stride) {
+    stride_ = std::move(stride);
+  }
+
+  static ExprHandle make(
+      const ExprHandle& base,
+      const ExprHandle& stride,
+      int64_t lanes) {
+    if (stride.dtype() != base.dtype()) {
+      throw malformed_input("Bad stride in Ramp");
+    }
+    return ExprHandle(alloc<Ramp>(base.node(), stride.node(), lanes));
+  }
+  int64_t lanes() const {
+    return lanes_;
+  }
+
+  Ramp(ExprPtr base, ExprPtr stride, int64_t lanes)
+      : ExprNodeBase(Dtype(base->dtype(), lanes)),
+        base_(std::move(base)),
+        stride_(std::move(stride)),
+        lanes_(lanes) {}
+
+ private:
+  ExprPtr base_;
+  ExprPtr stride_;
+  int64_t lanes_;
+};
+
+class TORCH_API Load : public ExprNode<Load> {
+ public:
+  VarPtr base_handle() const {
+    return buf_->base_handle();
+  }
+  std::vector<ExprPtr> indices() const {
+    return indices_;
+  }
+  ExprPtr flat_index() const {
+    TORCH_CHECK(indices_.size() == 1, "Indices haven't been flattened.");
+    return indices_[0];
+  }
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  void set_buf(BufPtr buf) {
+    buf_ = std::move(buf);
+  }
+
+  void set_indices(std::vector<ExprPtr> indices) {
+    indices_ = std::move(indices);
+  }
+
+  static ExprHandle make(
+      Dtype dtype,
+      const BufHandle& buf,
+      const std::vector<ExprHandle>& indices);
+  static ExprHandle make(
+      const BufHandle& buf,
+      const std::vector<ExprHandle>& indices);
+
+  Load(Dtype dtype, BufPtr base_handle, std::vector<ExprPtr> indices);
+  Load(const BufPtr& base_handle, const std::vector<ExprPtr>& indices);
+
+ private:
+  BufPtr buf_;
+  std::vector<ExprPtr> indices_;
+};
+
+class TORCH_API Broadcast : public ExprNode<Broadcast> {
+ public:
+  ExprPtr value() const {
+    return value_;
+  }
+
+  void set_value(ExprPtr value) {
+    value_ = std::move(value);
+  }
+
+  int64_t lanes() const {
+    return lanes_;
+  }
+  static ExprHandle make(const ExprHandle& value, int64_t lanes) {
+    return ExprHandle(alloc<Broadcast>(value.node(), lanes));
+  }
+  Broadcast(ExprPtr value, int64_t lanes)
+      : ExprNodeBase(Dtype(value->dtype(), lanes)),
+        value_(std::move(value)),
+        lanes_(lanes) {}
+
+ private:
+  ExprPtr value_;
+  int64_t lanes_;
+};
+
+class TORCH_API IfThenElse : public ExprNode<IfThenElse> {
+ public:
+  ExprPtr condition() const {
+    return condition_;
+  }
+
+  // Lazily evaluated only if condition is true
+  ExprPtr true_value() const {
+    return true_;
+  }
+
+  // Lazily evaluated only if condition is false
+  ExprPtr false_value() const {
+    return false_;
+  }
+
+  void set_condition(ExprPtr condition) {
+    condition_ = std::move(condition);
+  }
+
+  void set_true_value(ExprPtr true_value) {
+    true_ = std::move(true_value);
+  }
+
+  void set_false_value(ExprPtr false_value) {
+    false_ = std::move(false_value);
+  }
+
+  static ExprHandle make(
+      const ExprHandle& c,
+      const ExprHandle& t,
+      const ExprHandle& f) {
+    if (!c.dtype().is_integral()) {
+      throw unsupported_dtype();
+    }
+    if (c.dtype().lanes() != 1) {
+      throw unsupported_dtype();
+    }
+    if (t.dtype() != f.dtype()) {
+      throw malformed_input("Bad dtype in IfThenElse");
+    }
+    return ExprHandle(alloc<IfThenElse>(c.node(), t.node(), f.node()));
+  }
+
+  IfThenElse(ExprPtr c, ExprPtr t, ExprPtr f)
+      : ExprNodeBase(t->dtype()),
+        condition_(std::move(c)),
+        true_(std::move(t)),
+        false_(std::move(f)) {}
+
+ private:
+  ExprPtr condition_;
+  ExprPtr true_;
+  ExprPtr false_;
+};
+
+class TORCH_API CompareSelect : public ExprNode<CompareSelect> {
+ public:
+  CompareSelectOperation compare_select_op() const {
+    return compare_op_;
+  }
+  ExprPtr lhs() const {
+    return this->lhs_;
+  }
+  ExprPtr rhs() const {
+    return this->rhs_;
+  }
+  ExprPtr ret_val1() const {
+    return this->ret_val1_;
+  }
+  ExprPtr ret_val2() const {
+    return this->ret_val2_;
+  }
+
+  void set_lhs(ExprPtr lhs) {
+    lhs_ = std::move(lhs);
+  }
+
+  void set_rhs(ExprPtr rhs) {
+    rhs_ = std::move(rhs);
+  }
+
+  void set_ret_val1(ExprPtr ret_val1) {
+    ret_val1_ = std::move(ret_val1);
+  }
+
+  void set_ret_val2(ExprPtr ret_val2) {
+    ret_val2_ = std::move(ret_val2);
+  }
+
+  CompareSelectBias bias() const {
+    return bias_;
+  }
+
+  static ExprHandle make(
+      const ExprHandle& lhs,
+      const ExprHandle& rhs,
+      CompareSelectOperation cmp_op,
+      CompareSelectBias bias = kUnbiased) {
+    if (lhs.dtype() != rhs.dtype()) {
+      throw malformed_input("bad dtype in CompareSelect");
+    }
+    return ExprHandle(alloc<CompareSelect>(
+        lhs.node(),
+        rhs.node(),
+        IntImm::make(1).node(),
+        IntImm::make(0).node(),
+        cmp_op,
+        bias));
+  }
+
+  static ExprHandle make(
+      const ExprHandle& lhs,
+      const ExprHandle& rhs,
+      const ExprHandle& ret_val1,
+      const ExprHandle& ret_val2,
+      CompareSelectOperation cmp_op,
+      CompareSelectBias bias = kUnbiased) {
+    if (lhs.dtype() != rhs.dtype() || ret_val1.dtype() != ret_val2.dtype()) {
+      throw malformed_input("bad dtype in CompareSelect");
+    }
+    return ExprHandle(alloc<CompareSelect>(
+        lhs.node(),
+        rhs.node(),
+        ret_val1.node(),
+        ret_val2.node(),
+        cmp_op,
+        bias));
+  }
+
+  CompareSelect(
+      ExprPtr lhs,
+      ExprPtr rhs,
+      ExprPtr ret_val1,
+      ExprPtr ret_val2,
+      CompareSelectOperation cmp_op,
+      CompareSelectBias bias = kUnbiased)
+      : ExprNodeBase(ret_val1->dtype()),
+        lhs_(std::move(lhs)),
+        rhs_(std::move(rhs)),
+        ret_val1_(std::move(ret_val1)),
+        ret_val2_(std::move(ret_val2)),
+        compare_op_(cmp_op),
+        bias_(bias) {}
+
+  CompareSelect(
+      ExprPtr lhs,
+      ExprPtr rhs,
+      CompareSelectOperation cmp_op,
+      CompareSelectBias bias = kUnbiased)
+      : ExprNodeBase(kInt),
+        lhs_(std::move(lhs)),
+        rhs_(std::move(rhs)),
+        ret_val1_(alloc<IntImm>(1)),
+        ret_val2_(alloc<IntImm>(0)),
+        compare_op_(cmp_op),
+        bias_(bias) {}
+
+ private:
+  ExprPtr lhs_;
+  ExprPtr rhs_;
+  ExprPtr ret_val1_;
+  ExprPtr ret_val2_;
+  CompareSelectOperation compare_op_;
+  CompareSelectBias bias_;
+};
+
+enum IntrinsicsOp {
+  kSin,
+  kCos,
+  kTan,
+  kAsin,
+  kAcos,
+  kAtan,
+  kAtan2,
+  kSinh,
+  kCosh,
+  kTanh,
+  kSigmoid,
+  kExp,
+  kExpm1,
+  kAbs,
+  kLog,
+  kLog2,
+  kLog10,
+  kLog1p,
+  kErf,
+  kErfc,
+  kSqrt,
+  kRsqrt,
+  kPow,
+  kCeil,
+  kFloor,
+  kRound,
+  kTrunc,
+  kFmod,
+  kRemainder,
+  kLgamma,
+  kFrac,
+  kIsNan,
+  kRand, // We need more discussions on this. Should we consider stateful?
+  kMaxIntrinsicsOp,
+};
+
+class TORCH_API Intrinsics : public ExprNode<Intrinsics> {
+ public:
+  static ExprHandle make(IntrinsicsOp op_type, const ExprHandle& v1) {
+    return ExprHandle(alloc<Intrinsics>(op_type, v1.node()));
+  }
+
+  static ExprHandle make(
+      IntrinsicsOp op_type,
+      const ExprHandle& v1,
+      const ExprHandle& v2) {
+    return ExprHandle(alloc<Intrinsics>(op_type, v1.node(), v2.node()));
+  }
+
+  static ExprHandle make(
+      IntrinsicsOp op_type,
+      const std::vector<ExprHandle>& params) {
+    std::vector<ExprPtr> params_nodes(params.size());
+    for (size_t i = 0; i < params.size(); i++) {
+      params_nodes[i] = params[i].node();
+    }
+    return ExprHandle(alloc<Intrinsics>(op_type, params_nodes));
+  }
+
+  static ExprHandle make(IntrinsicsOp op_type, Dtype dtype) {
+    return ExprHandle(alloc<Intrinsics>(op_type, dtype));
+  }
+
+  IntrinsicsOp op_type() const {
+    return op_type_;
+  }
+
+  std::string func_name() const {
+    switch (op_type()) {
+      case kSin:
+        return "sin";
+      case kCos:
+        return "cos";
+      case kTan:
+        return "tan";
+      case kAsin:
+        return "asin";
+      case kAcos:
+        return "acos";
+      case kAtan:
+        return "atan";
+      case kAtan2:
+        return "atan2";
+      case kSinh:
+        return "sinh";
+      case kCosh:
+        return "cosh";
+      case kTanh:
+        return "tanh";
+      case kSigmoid:
+        return "sigmoid";
+      case kExp:
+        return "exp";
+      case kAbs:
+        return "abs";
+      case kLog:
+        return "log";
+      case kLog2:
+        return "log2";
+      case kLog10:
+        return "log10";
+      case kLog1p:
+        return "log1p";
+      case kErf:
+        return "erf";
+      case kSqrt:
+        return "sqrt";
+      case kRsqrt:
+        return "rsqrt";
+      case kPow:
+        return "pow";
+      case kCeil:
+        return "ceil";
+      case kFloor:
+        return "floor";
+      case kRound:
+        return "round";
+      case kTrunc:
+        return "trunc";
+      case kRand:
+        return "rand";
+      case kFmod:
+        return "fmod";
+      case kRemainder:
+        return "remainder";
+      case kLgamma:
+        return "lgamma";
+      case kExpm1:
+        return "expm1";
+      case kErfc:
+        return "erfc";
+      case kFrac:
+        return "frac";
+      case kIsNan:
+        return "isnan";
+      default:
+        throw std::runtime_error(
+            "invalid op_type: " + std::to_string(op_type()));
+    }
+  }
+
+  Intrinsics(IntrinsicsOp op_type, Dtype dtype)
+      : ExprNodeBase(IntrinsicsDtype(op_type, dtype)),
+        params_({}),
+        op_type_(op_type) {
+    if (OpArgCount(op_type) != 0) {
+      throw malformed_input("bad arg count in Intrinsics");
+    }
+  }
+
+  Intrinsics(IntrinsicsOp op_type, ExprPtr v1)
+      : ExprNodeBase(IntrinsicsDtype(op_type, v1->dtype())),
+        params_({std::move(v1)}),
+        op_type_(op_type) {
+    if (OpArgCount(op_type) != 1) {
+      throw malformed_input("bad arg count in Intrinsics");
+    }
+  }
+
+  Intrinsics(IntrinsicsOp op_type, ExprPtr v1, ExprPtr v2)
+      : ExprNodeBase(IntrinsicsDtype(op_type, v1->dtype(), v2->dtype())),
+        params_({std::move(v1), std::move(v2)}),
+        op_type_(op_type) {
+    if (OpArgCount(op_type) != 2) {
+      throw malformed_input("bad arg count in Intrinsics");
+    }
+  }
+
+  Intrinsics(IntrinsicsOp op_type, const std::vector<ExprPtr>& params)
+      : ExprNodeBase(IntrinsicsDtype(op_type, params)),
+        params_(params),
+        op_type_(op_type) {
+    if (OpArgCount(op_type) != nparams()) {
+      throw malformed_input("bad arg count in Intrinsics");
+    }
+  }
+
+  Intrinsics(IntrinsicsOp op_type, Dtype dtype, std::vector<ExprPtr> params)
+      : ExprNodeBase(IntrinsicsDtype(op_type, dtype)),
+        params_(std::move(params)),
+        op_type_(op_type) {
+    if (OpArgCount(op_type) != nparams()) {
+      throw malformed_input("bad arg count in Intrinsics");
+    }
+  }
+
+  bool isPure() const {
+    return op_type_ != kRand;
+  }
+
+  size_t nparams() const {
+    return params_.size();
+  }
+
+  ExprPtr param(size_t index) const {
+    return params_[index];
+  }
+  const std::vector<ExprPtr>& params() const {
+    return params_;
+  }
+
+  void set_params(std::vector<ExprPtr> params) {
+    params_ = std::move(params);
+  }
+
+  static size_t OpArgCount(IntrinsicsOp op_type);
+
+ private:
+  static Dtype IntrinsicsDtype(IntrinsicsOp op_type, Dtype dt1);
+  static Dtype IntrinsicsDtype(IntrinsicsOp op_type, Dtype dt1, Dtype dt2);
+  static Dtype IntrinsicsDtype(
+      IntrinsicsOp op_type,
+      const std::vector<ExprPtr>& params);
+
+  std::vector<ExprPtr> params_;
+  IntrinsicsOp op_type_;
+};
+
+TORCH_API std::vector<ExprPtr> ExprHandleVectorToExprVector(
+    const std::vector<ExprHandle>& /*v*/);
+TORCH_API std::vector<ExprHandle> ExprVectorToExprHandleVector(
+    const std::vector<ExprPtr>& /*v*/);
+TORCH_API std::vector<VarPtr> VarHandleVectorToVarVector(
+    const std::vector<VarHandle>& /*v*/);
+TORCH_API std::vector<VarHandle> VarVectorToVarHandleVector(
+    const std::vector<VarPtr>& /*v*/);
+TORCH_API ExprPtr flatten_index(
+    const std::vector<ExprPtr>& dims,
+    const std::vector<ExprPtr>& indices,
+    const std::vector<ExprPtr>& strides);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_cloner.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_cloner.h
new file mode 100644
index 0000000000000000000000000000000000000000..9c337bb1b6fa633b28be70d7497cdd149a4c45aa
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_cloner.h
@@ -0,0 +1,67 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/Export.h>
+#include <vector>
+
+#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
+
+namespace torch::jit::tensorexpr {
+
+class TORCH_API IRCloner : public IRMutator {
+ public:
+  ~IRCloner() override = default;
+  ExprPtr mutate(const AddPtr& v) override;
+  ExprPtr mutate(const SubPtr& v) override;
+  ExprPtr mutate(const MulPtr& v) override;
+  ExprPtr mutate(const DivPtr& v) override;
+  ExprPtr mutate(const ModPtr& v) override;
+  ExprPtr mutate(const MaxPtr& v) override;
+  ExprPtr mutate(const MinPtr& v) override;
+  ExprPtr mutate(const AndPtr& v) override;
+  ExprPtr mutate(const OrPtr& v) override;
+  ExprPtr mutate(const XorPtr& v) override;
+  ExprPtr mutate(const LshiftPtr& v) override;
+  ExprPtr mutate(const RshiftPtr& v) override;
+  ExprPtr mutate(const CompareSelectPtr& v) override;
+#define IMM_MUTATE_DECLARE(Type, Name) \
+  ExprPtr mutate(const Name##ImmPtr& v) override;
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_MUTATE_DECLARE)
+#undef IMM_MUTATE_DECLARE
+  ExprPtr mutate(const CastPtr& v) override;
+  ExprPtr mutate(const BitCastPtr& v) override;
+  ExprPtr mutate(const VarPtr& v) override;
+  ExprPtr mutate(const BufPtr& v) override;
+  ExprPtr mutate(const RampPtr& v) override;
+  ExprPtr mutate(const LoadPtr& v) override;
+  ExprPtr mutate(const BroadcastPtr& v) override;
+  ExprPtr mutate(const IfThenElsePtr& v) override;
+  ExprPtr mutate(const IntrinsicsPtr& v) override;
+
+  ExprPtr mutate(const TermPtr& v) override;
+  ExprPtr mutate(const PolynomialPtr& v) override;
+  ExprPtr mutate(const RoundOffPtr& v) override;
+  ExprPtr mutate(const MaxTermPtr& v) override;
+  ExprPtr mutate(const MinTermPtr& v) override;
+
+  ExprPtr mutate(const ReduceOpPtr& v) override;
+
+  StmtPtr mutate(const ForPtr& v) override;
+  StmtPtr mutate(const BlockPtr& v) override;
+  StmtPtr mutate(const StorePtr& v) override;
+  StmtPtr mutate(const AtomicAddPtr& v) override;
+  StmtPtr mutate(const SyncThreadsPtr& v) override;
+  StmtPtr mutate(const ExternalCallPtr& v) override;
+  StmtPtr mutate(const ExternalCallWithAllocPtr& v) override;
+
+  StmtPtr mutate(const AllocatePtr& v) override;
+  StmtPtr mutate(const FreePtr& v) override;
+  StmtPtr mutate(const LetPtr& v) override;
+  StmtPtr mutate(const CondPtr& v) override;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_mutator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_mutator.h
new file mode 100644
index 0000000000000000000000000000000000000000..826122019d43d9542f27cbe8337c513dc5c6884b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_mutator.h
@@ -0,0 +1,67 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+
+namespace torch::jit::tensorexpr {
+
+class TORCH_API IRMutator {
+ public:
+  virtual ~IRMutator() = default;
+  virtual ExprPtr mutate(const AddPtr& v);
+  virtual ExprPtr mutate(const SubPtr& v);
+  virtual ExprPtr mutate(const MulPtr& v);
+  virtual ExprPtr mutate(const DivPtr& v);
+  virtual ExprPtr mutate(const ModPtr& v);
+  virtual ExprPtr mutate(const MaxPtr& v);
+  virtual ExprPtr mutate(const MinPtr& v);
+  virtual ExprPtr mutate(const AndPtr& v);
+  virtual ExprPtr mutate(const OrPtr& v);
+  virtual ExprPtr mutate(const XorPtr& v);
+  virtual ExprPtr mutate(const LshiftPtr& v);
+  virtual ExprPtr mutate(const RshiftPtr& v);
+  virtual ExprPtr mutate(const CompareSelectPtr& v);
+#define IMM_MUTATE_DECLARE(Type, Name) \
+  virtual ExprPtr mutate(const Name##ImmPtr& v);
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_MUTATE_DECLARE)
+#undef IMM_MUTATE_DECLARE
+  virtual ExprPtr mutate(const CastPtr& v);
+  virtual ExprPtr mutate(const BitCastPtr& v);
+  virtual ExprPtr mutate(const VarPtr& v);
+  virtual ExprPtr mutate(const BufPtr& v);
+  virtual ExprPtr mutate(const RampPtr& v);
+  virtual ExprPtr mutate(const LoadPtr& v);
+  virtual ExprPtr mutate(const BroadcastPtr& v);
+  virtual ExprPtr mutate(const IfThenElsePtr& v);
+  virtual ExprPtr mutate(const IntrinsicsPtr& v);
+
+  virtual ExprPtr mutate(const TermPtr& v);
+  virtual ExprPtr mutate(const PolynomialPtr& v);
+  virtual ExprPtr mutate(const RoundOffPtr& v);
+  virtual ExprPtr mutate(const MaxTermPtr& v);
+  virtual ExprPtr mutate(const MinTermPtr& v);
+
+  virtual ExprPtr mutate(const ReduceOpPtr& v);
+
+  virtual StmtPtr mutate(const ForPtr& v);
+  virtual StmtPtr mutate(const BlockPtr& v);
+  virtual StmtPtr mutate(const StorePtr& v);
+  virtual StmtPtr mutate(const AtomicAddPtr& v);
+  virtual StmtPtr mutate(const SyncThreadsPtr& v);
+  virtual StmtPtr mutate(const ExternalCallPtr& v);
+  virtual StmtPtr mutate(const ExternalCallWithAllocPtr& v);
+
+  virtual StmtPtr mutate(const AllocatePtr& v);
+  virtual StmtPtr mutate(const FreePtr& v);
+  virtual StmtPtr mutate(const FreeExtPtr& v);
+  virtual StmtPtr mutate(const PlacementAllocatePtr& v);
+  virtual StmtPtr mutate(const LetPtr& v);
+  virtual StmtPtr mutate(const CondPtr& v);
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_printer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_printer.h
new file mode 100644
index 0000000000000000000000000000000000000000..033244123919bc33842771e479c989221daf897b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_printer.h
@@ -0,0 +1,137 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ostream>
+
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/unique_name_manager.h>
+
+namespace torch::jit::tensorexpr {
+
+class Tensor;
+
+class TORCH_API IRPrinter : public IRVisitor {
+ public:
+  explicit IRPrinter(std::ostream& os) : printer_os_(this, os) {}
+
+  void print(ExprHandle /*expr*/);
+  void print(Expr& /*expr*/);
+  void print(Stmt& /*stmt*/);
+  void visit(const AddPtr& v) override;
+  void visit(const SubPtr& v) override;
+  void visit(const MulPtr& v) override;
+  void visit(const DivPtr& v) override;
+  void visit(const ModPtr& v) override;
+  void visit(const MaxPtr& v) override;
+  void visit(const MinPtr& v) override;
+  void visit(const AndPtr& v) override;
+  void visit(const OrPtr& v) override;
+  void visit(const XorPtr& v) override;
+  void visit(const LshiftPtr& v) override;
+  void visit(const RshiftPtr& v) override;
+  void visit(const CompareSelectPtr& v) override;
+#define IMM_PRINT_VISIT(Type, Name) void visit(const Name##ImmPtr& v) override;
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_PRINT_VISIT)
+#undef IMM_PRINT_VISIT
+  void visit(const CastPtr& v) override;
+  void visit(const BitCastPtr& v) override;
+  void visit(const VarPtr& v) override;
+  void visit(const BufPtr& v) override;
+  void visit(const RampPtr& v) override;
+  void visit(const LoadPtr& v) override;
+  void visit(const BroadcastPtr& v) override;
+  void visit(const IfThenElsePtr& v) override;
+  void visit(const IntrinsicsPtr& v) override;
+  void visit(const TermPtr& v) override;
+  void visit(const PolynomialPtr& v) override;
+  void visit(const RoundOffPtr& v) override;
+  void visit(const MaxTermPtr& v) override;
+  void visit(const MinTermPtr& v) override;
+  void visit(const ReduceOpPtr& v) override;
+
+  void visit(const AtomicAddPtr& v) override;
+  void visit(const SyncThreadsPtr& v) override;
+  void visit(const ExternalCallPtr& v) override;
+  void visit(const ExternalCallWithAllocPtr& v) override;
+  void visit(const StorePtr& v) override;
+  void visit(const ForPtr& v) override;
+  void visit(const CondPtr& v) override;
+  void visit(const BlockPtr& v) override;
+  void visit(const AllocatePtr& v) override;
+  void visit(const FreePtr& v) override;
+  void visit(const FreeExtPtr& v) override;
+  void visit(const PlacementAllocatePtr& v) override;
+  void visit(const LetPtr& v) override;
+
+  // A child class may have a difference rule for generating dtype
+  // string, e.g. CUDA needs int64_t to be generated as long long.
+  virtual std::string dtypeToCppString(const Dtype& dtype);
+
+  std::ostream& os() {
+    return printer_os_;
+  }
+
+  class PrinterStream : public std::ostream {
+   public:
+    PrinterStream(IRPrinter* printer, std::ostream& os)
+        : std::ostream(os.rdbuf()), printer_(printer) {
+      initialize_imbue();
+    }
+
+    void initialize_imbue();
+
+    IRPrinter* printer() {
+      return printer_;
+    }
+
+   private:
+    IRPrinter* printer_ = nullptr;
+  };
+
+ protected:
+  std::string to_string(CompareSelectOperation op);
+
+  UniqueNameManager* name_manager() {
+    return &name_manager_;
+  }
+  void emitIndent();
+
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  int indent_ = 0;
+
+ private:
+  PrinterStream printer_os_;
+  UniqueNameManager name_manager_;
+};
+
+TORCH_API std::ostream& operator<<(std::ostream& stream, const Expr& /*expr*/);
+TORCH_API std::ostream& operator<<(
+    std::ostream& stream,
+    const ExprHandle& /*expr*/);
+TORCH_API std::ostream& operator<<(std::ostream& stream, const Stmt& /*stmt*/);
+TORCH_API std::ostream& operator<<(std::ostream& stream, const Tensor& /*t*/);
+
+TORCH_API void print(const ExprPtr& expr);
+TORCH_API void print(const StmtPtr& stmt);
+TORCH_API void print(const Tensor& t);
+
+} // namespace torch::jit::tensorexpr
+
+namespace std {
+
+using torch::jit::tensorexpr::Expr;
+using torch::jit::tensorexpr::ExprPtr;
+using torch::jit::tensorexpr::Stmt;
+using torch::jit::tensorexpr::StmtPtr;
+using torch::jit::tensorexpr::Tensor;
+
+TORCH_API std::string to_string(const ExprPtr& expr);
+TORCH_API std::string to_string(const StmtPtr& stmt);
+TORCH_API std::string to_string(const Tensor& t);
+} // namespace std
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_simplifier.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_simplifier.h
new file mode 100644
index 0000000000000000000000000000000000000000..62310b4e42739135e835a886dd69daec18fe7f75
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_simplifier.h
@@ -0,0 +1,551 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/bounds_overlap.h>
+#include <torch/csrc/jit/tensorexpr/eval.h>
+#include <torch/csrc/jit/tensorexpr/hash_provider.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/types.h>
+
+#include <utility>
+
+/* IR Simplification
+ *
+ * Simplifies expressions in two stages:
+ *  1. Recursively traverse the map combining similar operations into Terms
+ * (interacted via Multiplication) and Polynomials (interacted via Addition). We
+ * reorder the components of each Term or Polynomial into a consistent order to
+ * allow combination or cancelling of like terms.
+ *  2. Once the format of the tree is minimal, expand each Term into a sequence
+ * of Muls, and each Polynomial into a sequence of Ads.
+ */
+
+namespace torch::jit::tensorexpr {
+
+// A bunch of helpers for determine the Dtype of the output of a multi argument
+// Term or Polynomial.
+template <class ExprType>
+Dtype promoteTypesVec(const ExprPtr& s, const std::vector<ExprType>& v) {
+  Dtype t = s->dtype();
+  bool first = true;
+
+  for (const auto& e : v) {
+    if (first) {
+      t = Dtype(t.scalar_type(), e->dtype().lanes());
+      first = false;
+    }
+    t = promoteTypes(t, e->dtype());
+  }
+  return t;
+}
+
+template <class ExprType>
+Dtype promoteTypesVec(const std::vector<ExprType>& v) {
+  if (v.empty()) {
+    throw malformed_input("empty list of types");
+  }
+
+  Dtype t = v[0]->dtype();
+  for (const auto& e : v) {
+    t = promoteTypes(t, e->dtype());
+  }
+  return t;
+}
+
+template <class ExprType>
+Dtype promoteTypesMap(
+    const ExprPtr& s,
+    std::unordered_map<SimplifierHashType, ExprType>& m) {
+  Dtype t = s->dtype();
+  bool first = true;
+  for (auto& e : m) {
+    if (first) {
+      t = Dtype(t.scalar_type(), e.second->dtype().lanes());
+      first = false;
+    }
+    t = promoteTypes(t, e.second->dtype());
+  }
+  return t;
+}
+
+template <class ExprType>
+Dtype promoteTypesVar(ExprType e) {
+  return e->dtype();
+}
+
+template <class ExprType, class... Args>
+Dtype promoteTypesVar(ExprType e, Args... es) {
+  Dtype lhs = e->dtype();
+  Dtype rhs = promoteTypesVar(es...);
+  if (e->isConstant()) {
+    lhs = Dtype(lhs.scalar_type(), rhs.lanes());
+  }
+
+  return promoteTypes(lhs, rhs);
+}
+
+// Uses the evaluator to fold an Expression with constant terms.
+// E.g. evaluateOp(Add(3, 4)) => 7.
+// Expr v must not have any unbound Vars.
+inline ExprPtr evaluateOp(const ExprPtr& v) {
+  ExprHandle handle(v);
+  ExprEval<SimpleIREvaluator> eval(handle);
+
+  switch (v->dtype().scalar_type()) {
+#define TYPE_CASE(Type, Name)                                 \
+  case ScalarType::Name: {                                    \
+    Type val = eval.value<Type>();                            \
+    return getImmediateByType(v->dtype().scalar_type(), val); \
+  }
+    AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TYPE_CASE)
+#undef TYPE_CASE
+    default:
+      LOG(FATAL) << "Unsupported datatype: " << v->dtype();
+      return nullptr;
+  }
+  return nullptr;
+}
+
+// A Term represents a grouping of Exprs through multiplication.
+// E.g. product(scalar, *variables).
+class Term : public ExprNode<Term> {
+ public:
+  template <class... Args>
+  Term(HashProvider& hasher, ExprPtr s, Args... ts)
+      : ExprNodeBase(promoteTypesVar(s, ts...)), scalar_(s), hasher_(hasher) {
+    CHECK(s->isConstant());
+    addComponent(ts...);
+    sort();
+  }
+
+  Term(HashProvider& hasher, ExprPtr s, std::vector<ExprPtr> v)
+      : ExprNodeBase(promoteTypesVec(s, v)),
+        variables_(std::move(v)),
+        scalar_(std::move(s)),
+        hasher_(hasher) {
+    sort();
+  }
+
+  // Convenience constructor from a map of hash -> var, used when merging Terms.
+  Term(
+      HashProvider& hasher,
+      const ExprPtr& s,
+      std::unordered_map<SimplifierHashType, ExprPtr> varmap)
+      : ExprNodeBase(promoteTypesMap(s, varmap)), scalar_(s), hasher_(hasher) {
+    for (auto& p : varmap) {
+      addComponent(p.second);
+    }
+    sort();
+  }
+
+  ExprPtr scalar() const {
+    return scalar_;
+  }
+  const std::vector<ExprPtr>& variables() const {
+    return variables_;
+  }
+  HashProvider& hasher() const {
+    return hasher_;
+  }
+
+  // Produce a hash of just the variable components of this term, to determine
+  // if it can be combined with another term.
+  SimplifierHashType hashVars() const;
+
+ private:
+  std::vector<ExprPtr> variables_;
+  ExprPtr scalar_;
+  HashProvider& hasher_;
+
+  void addComponent() {}
+  void addComponent(ExprPtr e) {
+    variables_.push_back(std::move(e));
+  }
+  template <class... Es>
+  void addComponent(ExprPtr e, Es&&... es) {
+    addComponent(std::move(e));
+    addComponent(std::forward<Es>(es)...);
+  }
+
+  // Sort by hash to normalize order of components.
+  void sort();
+};
+
+// Polynomial represents a grouping of Exprs by addition.
+// E.g. sum(*variables, scalar).
+// This would better be called Expression, but, naming conflict...
+class Polynomial : public ExprNode<Polynomial> {
+ public:
+  template <class... Args>
+  Polynomial(HashProvider& hasher, ExprPtr s, Args... ts)
+      : ExprNodeBase(promoteTypesVar(s, ts...)), scalar_(s), hasher_(hasher) {
+    CHECK(s->isConstant());
+    addTerm(ts...);
+    sort();
+  }
+
+  Polynomial(HashProvider& hasher, const ExprPtr& s, std::vector<TermPtr> v)
+      : ExprNodeBase(promoteTypesVec(s, v)),
+        variables_(std::move(v)),
+        scalar_(s),
+        hasher_(hasher) {
+    sort();
+  }
+
+  // Helper constructor for list of terms with no scalar component.
+  Polynomial(HashProvider& hasher, std::vector<TermPtr> terms)
+      : ExprNodeBase(promoteTypesVec(terms)),
+        variables_(std::move(terms)),
+        scalar_(getImmediateByType(dtype(), 0)),
+        hasher_(hasher) {
+    sort();
+  }
+
+  // Convenience constructor for map of hash -> var, used when merging
+  // Polynomials.
+  Polynomial(
+      HashProvider& hasher,
+      const ExprPtr& s,
+      std::unordered_map<SimplifierHashType, TermPtr> varmap)
+      : ExprNodeBase(promoteTypesMap(s, varmap)), scalar_(s), hasher_(hasher) {
+    for (auto& p : varmap) {
+      addTerm(p.second);
+    }
+    sort();
+  }
+
+  ExprPtr scalar() const {
+    return scalar_;
+  }
+  const std::vector<TermPtr>& variables() const {
+    return variables_;
+  }
+  HashProvider& hasher() const {
+    return hasher_;
+  }
+
+  SimplifierHashType hashVars() const;
+
+ private:
+  std::vector<TermPtr> variables_;
+  ExprPtr scalar_;
+  HashProvider& hasher_;
+
+  void addTerm(TermPtr t) {
+    variables_.push_back(std::move(t));
+  }
+  template <class... Ts>
+  void addTerm(TermPtr t, Ts&&... ts) {
+    addTerm(std::move(t));
+    addTerm(std::forward<Ts>(ts)...);
+  }
+
+  // Sort by hash to normalize order of terms.
+  void sort();
+};
+
+class RoundOff : public BinaryOpNode<RoundOff> {
+ public:
+  RoundOff(ExprPtr lhs, ExprPtr rhs)
+      : BinaryOpNode(std::move(lhs), std::move(rhs), IRNodeType::kOther) {}
+};
+
+class MaxTerm : public ExprNode<MaxTerm> {
+ public:
+  template <class... Args>
+  MaxTerm(HashProvider& hasher, ExprPtr s, bool p, Args... ts)
+      : ExprNodeBase(s ? promoteTypesVar(s, ts...) : promoteTypesVar(ts...)),
+        scalar_(s),
+        hasher_(hasher),
+        propagate_nans_(p) {
+    addComponent(ts...);
+    uniquefy();
+  }
+
+  MaxTerm(
+      HashProvider& hasher,
+      const ExprPtr& s,
+      bool p,
+      std::vector<ExprPtr> v)
+      : ExprNodeBase(s ? promoteTypesVec(s, v) : promoteTypesVec(v)),
+        variables_(std::move(v)),
+        scalar_(s),
+        hasher_(hasher),
+        propagate_nans_(p) {
+    uniquefy();
+  }
+
+  bool propagate_nans() const {
+    return propagate_nans_;
+  }
+
+  ExprPtr scalar() const {
+    return scalar_;
+  }
+  const std::vector<ExprPtr>& variables() const {
+    return variables_;
+  }
+  HashProvider& hasher() const {
+    return hasher_;
+  }
+
+ private:
+  std::vector<ExprPtr> variables_;
+  ExprPtr scalar_;
+  HashProvider& hasher_;
+  bool propagate_nans_;
+
+  void addComponent() {}
+  void addComponent(ExprPtr e) {
+    variables_.push_back(std::move(e));
+  }
+  template <class... Es>
+  void addComponent(ExprPtr e, Es&&... es) {
+    addComponent(std::move(e));
+    addComponent(std::forward<Es>(es)...);
+  }
+
+  // Uniquefy the terms using their hash.
+  void uniquefy();
+};
+
+class MinTerm : public ExprNode<MinTerm> {
+ public:
+  template <class... Args>
+  MinTerm(HashProvider& hasher, ExprPtr s, bool p, Args... ts)
+      : ExprNodeBase(s ? promoteTypesVar(s, ts...) : promoteTypesVar(ts...)),
+        scalar_(s),
+        hasher_(hasher),
+        propagate_nans_(p) {
+    addComponent(ts...);
+    uniquefy();
+  }
+
+  MinTerm(
+      HashProvider& hasher,
+      const ExprPtr& s,
+      bool p,
+      std::vector<ExprPtr> v)
+      : ExprNodeBase(s ? promoteTypesVec(s, v) : promoteTypesVec(v)),
+        variables_(std::move(v)),
+        scalar_(s),
+        hasher_(hasher),
+        propagate_nans_(p) {
+    uniquefy();
+  }
+
+  bool propagate_nans() const {
+    return propagate_nans_;
+  }
+
+  ExprPtr scalar() const {
+    return scalar_;
+  }
+  const std::vector<ExprPtr>& variables() const {
+    return variables_;
+  }
+  HashProvider& hasher() const {
+    return hasher_;
+  }
+
+ private:
+  std::vector<ExprPtr> variables_;
+  ExprPtr scalar_;
+  HashProvider& hasher_;
+  bool propagate_nans_;
+
+  void addComponent() {}
+  void addComponent(ExprPtr e) {
+    variables_.push_back(std::move(e));
+  }
+  template <class... Es>
+  void addComponent(ExprPtr e, Es&&... es) {
+    addComponent(std::move(e));
+    addComponent(std::forward<Es>(es)...);
+  }
+
+  // Uniquefy the terms using their hash.
+  void uniquefy();
+};
+
+// Context-sensitive IR simplification
+using VarBoundInfo = std::unordered_map<VarPtr, analysis::Bound>;
+
+class TORCH_API SimplifierUnderContext : public IRMutator {
+ public:
+  ~SimplifierUnderContext() override = default;
+  // Add boundary info for index variables in for-loops
+  StmtPtr mutate(const ForPtr& v) override;
+
+  ExprPtr mutate(const DivPtr& v) override;
+  ExprPtr mutate(const ModPtr& v) override;
+  ExprPtr mutate(const CompareSelectPtr& v) override;
+  ExprPtr mutate(const IfThenElsePtr& v) override;
+
+ protected:
+  bool getLoopBoundInfo(const ExprPtr& expr, analysis::Bound* loop_bound_info);
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  HashProvider hasher_;
+  VarBoundInfo var_bound_info_;
+};
+
+// Stmt simplification should occur in both modes.
+class TORCH_API PolynomialBase : public IRMutator {
+ public:
+  ~PolynomialBase() override = default;
+
+  StmtPtr mutate(const BlockPtr& v) override;
+
+  StmtPtr mutate(const CondPtr& v) override;
+
+  StmtPtr mutate(const ForPtr& v) override;
+
+  // Trivially factorize terms by GCD of scalar components.
+  TermPtr factorizePolynomial(const PolynomialPtr& poly);
+
+  HashProvider& hasher() {
+    return hasher_;
+  }
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  HashProvider hasher_;
+};
+
+// Simplify the IR by combining arithmetic expressions over common terms.
+class TORCH_API PolynomialTransformer : public PolynomialBase {
+ public:
+  using PolynomialBase::mutate;
+  // Inserts term into the provided map, in the case of a hash collision
+  // combines the term with the existing and updates the map.
+  void addOrUpdateTerm(
+      std::unordered_map<SimplifierHashType, TermPtr>& varmap,
+      const TermPtr& term);
+
+  // Add Polynomial expressions, combining Terms representing the same
+  // variables.
+  ExprPtr addPolynomials(const PolynomialPtr& lhs, const PolynomialPtr& rhs);
+
+  // Insert a new Term into the provided polynomial. If the new term has
+  // common variables to an existing term it is combined.
+  ExprPtr insertTerm(const PolynomialPtr& poly, const TermPtr& term);
+
+  // Merge and simplify addition.
+  ExprPtr mutate(const AddPtr& v) override;
+
+  // Subtract one term from another, cancelling if necessary.
+  ExprPtr subTerms(const TermPtr& lhs, TermPtr rhs, bool negated);
+
+  // Subtract the RHS Polynomial from the LHS Polynomial, cancelling out where
+  // possible.
+  ExprPtr subPolynomials(const PolynomialPtr& lhs, const PolynomialPtr& rhs);
+
+  // Merge and simplify subtraction.
+  ExprPtr mutate(const SubPtr& v) override;
+
+  // Multiply two terms together, usually creating a new term with the variable
+  // lists concatenated.
+  TermPtr mulTerms(const TermPtr& lhs, const TermPtr& rhs);
+
+  // Multiply a Polynomial by a Term.
+  ExprPtr polyByTerm(const PolynomialPtr& poly, const TermPtr& term);
+
+  // Match a rounding pattern and create a RoundOff if found.
+  ExprPtr isRoundOff(const ExprPtr& lhs, const ExprPtr& rhs);
+
+  // Inserts a new component into a term, simplifying if possible.
+  ExprPtr insertIntoTerm(const TermPtr& term, const ExprPtr& expr);
+
+  // Merge and simplify multiplication.
+  ExprPtr mutate(const MulPtr& v) override;
+
+  ExprPtr mutate(const DivPtr& v) override;
+
+  ExprPtr mutate(const ModPtr& v) override;
+
+  ExprPtr mutate(const AndPtr& v) override;
+
+  ExprPtr mutate(const XorPtr& v) override;
+
+  ExprPtr mutate(const LshiftPtr& v) override;
+
+  ExprPtr mutate(const RshiftPtr& v) override;
+
+  ExprPtr mutate(const MaxPtr& v) override;
+
+  ExprPtr mutate(const MinPtr& v) override;
+
+  ExprPtr mutate(const CompareSelectPtr& v) override;
+
+  ExprPtr mutate(const IntrinsicsPtr& v) override;
+
+  ExprPtr mutate(const CastPtr& v) override;
+
+  ExprPtr mutate(const IfThenElsePtr& v) override;
+
+  static ExprPtr simplify(ExprPtr e);
+  static ExprHandle simplify(const ExprHandle& e);
+  static StmtPtr simplify(StmtPtr e);
+};
+
+// Expands Terms and Polynomial expressions into primitive operations.
+// Does some simple factorization and reordering.
+class TORCH_API TermExpander : public PolynomialBase {
+  PolynomialTransformer* simplifier_;
+  std::set<VarPtr> eliminated_allocations_;
+
+ public:
+  using PolynomialBase::mutate;
+  TermExpander(PolynomialTransformer* simplifier) : simplifier_(simplifier) {}
+  bool check_safe() {
+    return eliminated_allocations_.empty();
+  }
+
+  // Expand Terms out to a series of Muls.
+  ExprPtr mutate(const TermPtr& v) override;
+
+  // Expand Polynomials out to a series of Adds.
+  ExprPtr mutate(const PolynomialPtr& v) override;
+
+  // Expand MaxTerms to a series of Max ops.
+  ExprPtr mutate(const MaxTermPtr& v) override;
+
+  // Expand MinTerms to a series of Min ops.
+  ExprPtr mutate(const MinTermPtr& v) override;
+
+  // Expand RoundOff to it's component: Mul(Div(lhs, rhs), rhs).
+  ExprPtr mutate(const RoundOffPtr& v) override;
+
+  // Eliminate zero length allocations.
+  StmtPtr mutate(const AllocatePtr& v) override;
+  StmtPtr mutate(const FreePtr& v) override;
+
+  // Override to enable condition fusing.
+  BlockPtr fuseConditions(BlockPtr v);
+  StmtPtr fuseSyncThreads(BlockPtr block);
+  StmtPtr mutate(const BlockPtr& v) override;
+};
+
+class TORCH_API IRSimplifier {
+ public:
+  static StmtPtr simplify(StmtPtr s);
+  static ExprPtr simplify(ExprPtr e);
+  static ExprHandle simplify(const ExprHandle& e) {
+    return ExprHandle(simplify(e.node()));
+  }
+};
+
+// Flattens the buf and performs the simplifier on the flattened dims.
+ExprPtr buf_flat_size(const BufPtr& v);
+// Returns true if expressions A and B can be simplified to an equal expression.
+TORCH_API bool exprEquals(const ExprPtr& A, const ExprPtr& B);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_verifier.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_verifier.h
new file mode 100644
index 0000000000000000000000000000000000000000..bd88ed8314ae8b5c971815348aaffc8218a8f62f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_verifier.h
@@ -0,0 +1,59 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+
+namespace torch::jit::tensorexpr {
+
+class Expr;
+class ExprHandle;
+class Mod;
+class And;
+class Or;
+class Xor;
+class Lshift;
+class Rshift;
+class CompareSelect;
+class Ramp;
+class Load;
+class IfThenElse;
+class Intrinsics;
+
+class Stmt;
+class ExternalCall;
+class Store;
+class For;
+class Block;
+
+class TORCH_API IRVerifier : public IRVisitor {
+ public:
+  IRVerifier() = default;
+
+  void visit(const ModPtr& v) override;
+  void visit(const AndPtr& v) override;
+  void visit(const OrPtr& v) override;
+  void visit(const XorPtr& v) override;
+  void visit(const LshiftPtr& v) override;
+  void visit(const RshiftPtr& v) override;
+  void visit(const CompareSelectPtr& v) override;
+  void visit(const RampPtr& v) override;
+  void visit(const LoadPtr& v) override;
+  void visit(const IfThenElsePtr& v) override;
+  void visit(const IntrinsicsPtr& v) override;
+
+  void visit(const ExternalCallPtr& v) override;
+  void visit(const StorePtr& v) override;
+  void visit(const ForPtr& v) override;
+  void visit(const BlockPtr& v) override;
+};
+
+TORCH_API void verify(const StmtPtr& /*s*/);
+TORCH_API void verify(const ExprPtr& /*e*/);
+TORCH_API void verify(const ExprHandle& /*e*/);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_visitor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_visitor.h
new file mode 100644
index 0000000000000000000000000000000000000000..d93ef37828c62407dbd27ff43a4fab3ffcf465e1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/ir_visitor.h
@@ -0,0 +1,65 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+
+namespace torch::jit::tensorexpr {
+
+class TORCH_API IRVisitor {
+ public:
+  virtual ~IRVisitor() = default;
+  virtual void visit(const AddPtr& v);
+  virtual void visit(const SubPtr& v);
+  virtual void visit(const MulPtr& v);
+  virtual void visit(const DivPtr& v);
+  virtual void visit(const ModPtr& v);
+  virtual void visit(const MaxPtr& v);
+  virtual void visit(const MinPtr& v);
+  virtual void visit(const AndPtr& v);
+  virtual void visit(const OrPtr& v);
+  virtual void visit(const XorPtr& v);
+  virtual void visit(const LshiftPtr& v);
+  virtual void visit(const RshiftPtr& v);
+  virtual void visit(const CompareSelectPtr& v);
+
+#define IMM_PRINT_VISIT(Type, Name) virtual void visit(const Name##ImmPtr& v);
+
+  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_PRINT_VISIT)
+#undef IMM_PRINT_VISIT
+
+  virtual void visit(const CastPtr& v);
+  virtual void visit(const BitCastPtr& v);
+  virtual void visit(const VarPtr& v);
+  virtual void visit(const BufPtr& v);
+  virtual void visit(const RampPtr& v);
+  virtual void visit(const LoadPtr& v);
+  virtual void visit(const ForPtr& v);
+  virtual void visit(const BlockPtr& v);
+  virtual void visit(const StorePtr& v);
+  virtual void visit(const BroadcastPtr& v);
+  virtual void visit(const IfThenElsePtr& v);
+  virtual void visit(const IntrinsicsPtr& v);
+  virtual void visit(const AllocatePtr& v);
+  virtual void visit(const FreePtr& v);
+  virtual void visit(const FreeExtPtr& v);
+  virtual void visit(const PlacementAllocatePtr& v);
+  virtual void visit(const LetPtr& v);
+  virtual void visit(const CondPtr& v);
+  virtual void visit(const TermPtr& v);
+  virtual void visit(const PolynomialPtr& v);
+  virtual void visit(const RoundOffPtr& v);
+  virtual void visit(const MaxTermPtr& v);
+  virtual void visit(const MinTermPtr& v);
+  virtual void visit(const ReduceOpPtr& v);
+  virtual void visit(const AtomicAddPtr& v);
+  virtual void visit(const SyncThreadsPtr& v);
+  virtual void visit(const ExternalCallPtr& v);
+  virtual void visit(const ExternalCallWithAllocPtr& v);
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/kernel.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/kernel.h
new file mode 100644
index 0000000000000000000000000000000000000000..b40b542412a76cdc7161b6a3f4bc97f83bf04d79
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/kernel.h
@@ -0,0 +1,383 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/symbolic_shape_runtime_fusion.h>
+#include <torch/csrc/jit/passes/utils/subgraph_utils.h>
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <torch/csrc/jit/tensorexpr/analysis.h>
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/lowerings.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+namespace torch::jit::tensorexpr {
+
+struct SmallSizeTPairHash {
+ public:
+  std::size_t operator()(const std::pair<size_t, size_t>& x) const {
+    // hashing input index and then dim index
+    return x.first * 128 + x.second;
+  }
+};
+
+// Returns true if the TE fuser supports this conv2d.
+bool conv2dIsSupportedJit(const Node* node);
+// Returns true if the TE fuser supports this conv2d with mkldnn prepacked conv.
+bool mkldnnPrepackedConvIsSupportedJit(const Node* node);
+// Returns true if the TE _convolution node is Conv2d.
+bool isConv2d(const Node* node);
+// Returns true if the TE fuser supports this matmul.
+bool matmulIsSupported(const Node* node);
+template <typename T>
+inline std::vector<int64_t> bufferSizes(const T& t) {
+  std::vector<int64_t> sizes;
+  for (size_t i = 0; i < t->ndim(); i++) {
+    sizes.push_back(*intValue(t->dim(i)));
+  }
+  return sizes;
+}
+
+// Get the dimensions of a value.
+std::vector<ExprHandle> valueShape(const ArgValue& v);
+
+// If v is a tensor, broadcast it to match the shape of axes, or return
+// directly if v is a constant.
+ExprHandle tensorOrConstant(
+    const ArgValue& v,
+    const std::vector<ExprHandle>& axes);
+
+int64_t normalizeAndCheckIndex(int64_t idx, int64_t list_size);
+
+ExprHandle broadcast(const BufHandle& b, const std::vector<ExprHandle>& axes);
+
+ExprHandle constant(const ArgValue& v);
+
+std::vector<ExprHandle> computeIndicesToBroadcast(
+    const std::vector<ExprHandle>& outputAxes,
+    const std::vector<ExprHandle>& inputSizes);
+
+inline std::string getArgValueName(const ArgValue& a) {
+  if (std::holds_alternative<tensorexpr::BufHandle>(a)) {
+    return "BufHandle";
+  } else if (std::holds_alternative<tensorexpr::VarHandle>(a)) {
+    return "VarHandle";
+  } else if (std::holds_alternative<double>(a)) {
+    return "double";
+  } else if (std::holds_alternative<int64_t>(a)) {
+    return "int64_t";
+  } else if (std::holds_alternative<bool>(a)) {
+    return "bool";
+  } else if (std::holds_alternative<BufList>(a)) {
+    return "BufList";
+  } else if (std::holds_alternative<DoubleList>(a)) {
+    return "DoubleList";
+  } else if (std::holds_alternative<IntList>(a)) {
+    return "IntList";
+  } else if (std::holds_alternative<ArgNone>(a)) {
+    return "None";
+  } else {
+    throw std::runtime_error("ArgValue type not handled in string conversion");
+  }
+}
+
+template <class T>
+std::vector<T> convertVecArgValue(const std::vector<ArgValue>& v) {
+  std::vector<T> res;
+  for (auto& x : v) {
+    auto val = std::get_if<T>(&x);
+    if (val) {
+      res.push_back(*val);
+    } else {
+      throw std::runtime_error(
+          "vector type not homogeneous - found " + getArgValueName(x) +
+          ", expected " + getArgValueName(v[0]));
+    }
+  }
+  return res;
+}
+
+class TORCH_API TensorExprKernel {
+  struct ConstantDescr {
+    BufPtr buf;
+    // Only one of ptr and node is used at a time
+    // 1) ptr for the constant tensors
+    // 2) node for the constant custom class objects
+    void* ptr = nullptr;
+    Node* node = nullptr;
+  };
+
+ public:
+  // Constructor Params:
+  //  * subgraph
+  //      - the graph that needs to be compiled.
+  //  * kernel_func_name
+  //      - the name that should be used for the generated kernel.
+  //  * custom_lowerings
+  //      - map that represents custom lowering definitions for a set of ops.
+  //  * symbolic_shape_inputs
+  //      - a list of symbolic graph inputs that represent the symbolic dims of
+  //        the input tensors.
+  //  * pre_alloc
+  //      - a flag to control pre-allocation of buffers.
+  explicit TensorExprKernel(
+      const std::shared_ptr<Graph>& subgraph,
+      std::string kernel_func_name,
+      std::unordered_map<c10::Symbol, NNCLoweringFunction> custom_lowerings =
+          {},
+      std::vector<int64_t> symbolic_shape_inputs = {},
+      bool pre_alloc = false,
+      std::unordered_map<
+          const torch::jit::Value*,
+          std::vector<torch::jit::StrideInput>> symbolic_strides = {});
+
+  explicit TensorExprKernel(
+      const std::shared_ptr<Graph>& subgraph,
+      std::unordered_map<c10::Symbol, NNCLoweringFunction> custom_lowerings =
+          {},
+      std::vector<int64_t> symbolic_shape_inputs = {},
+      bool pre_alloc = false,
+      std::unordered_map<
+          const torch::jit::Value*,
+          std::vector<torch::jit::StrideInput>> symbolic_strides = {})
+      : TensorExprKernel(
+            subgraph,
+            SubgraphUtils::generateNameForGraph(subgraph),
+            std::move(custom_lowerings),
+            std::move(symbolic_shape_inputs),
+            pre_alloc,
+            std::move(symbolic_strides)) {}
+
+  void run(Stack& stack) const;
+  void runFast(
+      const std::vector<void*>& inputs,
+      const std::vector<void*>& outputs) const;
+  // Expected format of stack:
+  //  ... <outputs> <inputs>
+  // i.e., output IValues must be below the input IValues in the stack.
+  void runWithAllocatedOutputs(Stack& stack) const;
+
+  void fallback(Stack& stack) const {
+    InterpreterState(code_).run(stack);
+  }
+  void recompile();
+
+  StmtPtr getCodeGenStmt();
+
+  std::string getCodeText(const std::string& attr = "") {
+    return codegen_->getCodeText(attr);
+  }
+
+  const std::shared_ptr<Graph> graph() {
+    return graph_;
+  }
+
+  const std::vector<ConstantDescr>& getConstantDescriptors() const {
+    return constants_;
+  }
+
+  const std::vector<CodeGen::BufferArg>& getBufferArgs() const {
+    return bufferArgs_;
+  }
+
+  const std::string& getKernelName() const {
+    return (codegen_ ? codegen_->kernel_func_name() : kernel_func_name_);
+  }
+
+  const std::vector<int64_t>& getSymbolicShapeInputs() const {
+    return symbolic_shape_inputs_;
+  }
+
+ private:
+  enum BackendType {
+    kUninitialized,
+    kSimpleIREval,
+    kLLVMCodeGen,
+    kCudaCodeGen,
+    kBlockCodeGen,
+  };
+
+  enum MemoryLayoutPolicy {
+    kContiguous,
+    kChannelsLastNdContiguous,
+  };
+
+  void compile();
+  void genInputDebugNames();
+  void runKernel(Stack& stack) const;
+
+  std::vector<ExprHandle> sizesForValue(const torch::jit::Value* v);
+
+  // These functions broadcast shape and also store a `hasBroadcast_` variable.
+  std::vector<ExprHandle> broadcastShapesMut(
+      const std::vector<ExprHandle>& a,
+      const std::vector<ExprHandle>& b);
+  std::vector<ExprHandle> broadcastShapesMut(
+      std::vector<std::vector<ExprHandle>> shapes);
+
+  ArgValue toArg(const torch::jit::Value* v) const;
+  ExprHandle constant(const torch::jit::Value* v);
+
+  Tensor computeValue(const torch::jit::Value* v);
+
+  void bindConstant(const torch::jit::Value* v);
+
+  StmtPtr transformLoops(BackendType backendType, StmtPtr st);
+
+  std::string getCodeGenName(BackendType backendType);
+
+  void getStaticOutputSizesAndStrides(
+      const at::ArrayRef<IValue>& inputs,
+      std::vector<std::vector<int64_t>>* static_sizes,
+      std::vector<std::vector<int64_t>>* static_strides) const;
+
+  std::vector<CodeGen::CallArg> prepareRunArgs(
+      const at::ArrayRef<IValue>& inputs,
+      std::vector<at::Tensor>& outputs) const;
+  BackendType inferBackendTypeFromDevice(at::Device device);
+
+  Tensor bindInput(const torch::jit::Value* input);
+  BlockPtr bindAllInputs();
+
+  // Deduce the memory layout policy to be propagated within
+  // NNC fusion group. The memory layout policy could be `kContiguous`
+  // or `kChannelsLastNdContiguous`.
+  //    `kContiguous`: Always convert the non-contiguous input tensors and
+  //        internal buffers to contiguous.
+  //    `kChannelsLastNdContiguous`: Always convert the input tensors and
+  //        internal buffers to channels-last contiguous.
+  // Currently, the rule is simple.
+  //    If all the input and out tensors of NNC fusion group are channels-last
+  //    contiguous, the policy is `kChannelsLastNdContiguous`. Otherwise, it
+  //    is always `kContiguous`.
+  void deduceMemoryLayoutPolicy();
+
+  Tensor convertSymbolicOutputToCorrectStrides(torch::jit::Value* v);
+  Tensor convertStaticShapeOutputToCorrectStrides(torch::jit::Value* v);
+  Tensor convertSymbolicOutputToCorrectStrides(
+      const std::vector<ExprHandle>& sizes,
+      const std::vector<size_t>& sorted_stride_indices_descending,
+      const std::vector<ExprPtr>& strides,
+      BufPtr& buf);
+
+  NNCLoweringFunction getCustomLoweringFor(c10::Symbol op) const;
+  std::unordered_map<c10::Symbol, NNCLoweringFunction> getCustomLowerings()
+      const {
+    return custom_lowerings_;
+  }
+
+  // Allocate memory for intermediate buffers at compile time.
+  // Specifically, we pre-allocate memory for intermediate buffers with static
+  // size and manage these buffers in the way we manage JIT constant tensors:
+  // push the buf args into the stack so NNC IR can access them at runtime.
+  std::vector<BufPtr> preAllocIntermediateBufs(
+      const std::vector<BufPtr>& interm_bufs);
+
+  struct UnpackedTensorOptions {
+    std::optional<c10::ScalarType> dtype;
+    std::optional<c10::Layout> layout;
+    std::optional<c10::Device> device;
+    std::optional<bool> pinned_memory;
+
+    UnpackedTensorOptions(const c10::TensorOptions& opts)
+        : dtype(c10::optTypeMetaToScalarType(opts.dtype_opt())),
+          layout(opts.layout_opt()),
+          device(opts.device_opt()),
+          pinned_memory(opts.pinned_memory_opt()) {}
+  };
+
+  ExprHandle getVarForShape(const c10::ShapeSymbol& ss);
+  std::vector<ExprHandle> computeInputTensorDims(
+      const torch::jit::Value* input);
+  ExprHandle getStrideArg(size_t tensor_input, size_t stride_index);
+  std::vector<ExprHandle> sizesFromSymbolicShape(
+      const c10::SymbolicShape& shape);
+  std::vector<ExprHandle> getInputStrides(
+      const torch::jit::Value* input,
+      const std::vector<ExprHandle>& inputTensorDims);
+  std::vector<torch::jit::StrideInput>& getSymbolicStrideDesc(
+      const torch::jit::Value* value);
+
+  // Apply the optimizations to the graph owned by the current fusion group,
+  // like concatenation optimization, post-op fusion, and some other graph-level
+  // optimizations.
+  void optimizeOwningGraph();
+
+  int64_t nInputs_ = 0;
+  int64_t nOutputs_ = 0;
+  std::vector<CodeGen::BufferArg> bufferArgs_;
+  std::vector<std::vector<int64_t>> tensorOutputSizes_;
+  std::vector<std::vector<int64_t>> tensorOutputStrides_;
+  std::vector<torch::jit::StrideInput> tensorOutputStrideDesc_;
+  std::vector<bool> isOutputScalar_;
+  std::vector<UnpackedTensorOptions> tensorOutputTensorOptions_;
+  std::unordered_set<BufPtr> bufOutputs_;
+  std::unordered_set<BufPtr> bufsToBeParallelized_;
+  std::unordered_map<const torch::jit::Value*, BufPtr> bufs_;
+  std::unordered_map<const torch::jit::Value*, VarHandle> scalars_;
+  std::unordered_map<const torch::jit::Value*, std::string> input_name_map_;
+  std::unique_ptr<CodeGen> codegen_;
+  at::Device device_ = at::kCPU;
+  std::shared_ptr<Graph> graph_;
+  Code code_;
+  bool allow_fallback_{false};
+  bool use_fallback_{false};
+  bool hasRandom_{false};
+  bool hasBroadcast_{false};
+  std::unordered_map<const torch::jit::Value*, std::vector<ExprHandle>>
+      known_sizes_;
+
+  std::vector<std::vector<ExprHandle>> tensorOutputSymbolicSizes_;
+  // A map from ShapeSymbol.value() to the corresponding Var.
+  std::unordered_map<int64_t, VarHandle> shapeSymbolToVar_;
+  std::unordered_map<ExprPtr, size_t> shapeSymbolInputPos_;
+  // List of values corresponding to the ShapeSymbols that are inputs to
+  // kernel being compiled. The order of these values correspond to the order
+  // of the symbolic inputs at the end of the list of inputs to the kernel.
+  std::vector<int64_t> symbolic_shape_inputs_;
+  bool has_symbolic_shapes_{false};
+
+  std::vector<at::Tensor> unpacked_constant_tensors_;
+  std::vector<ConstantDescr> constants_;
+
+  std::unordered_map<c10::Symbol, NNCLoweringFunction> custom_lowerings_;
+  StmtPtr stmt_ = nullptr;
+  bool pre_alloc_{false};
+  std::string kernel_func_name_;
+
+  // index of stack, stride index of tensor that will be appended as a codegen
+  // arg
+  std::vector<std::pair<size_t, size_t>> input_stride_args_;
+  // map from <input index, tensor dimension> to stride as arg VarHandle
+  std::unordered_map<std::pair<size_t, size_t>, VarHandle, SmallSizeTPairHash>
+      strideArgToVar_;
+  std::unordered_map<
+      const torch::jit::Value*,
+      std::vector<torch::jit::StrideInput>>
+      symbolic_strides_;
+
+  // Memory layout to be propagated with fusion group
+  MemoryLayoutPolicy memory_layout_policy_ = MemoryLayoutPolicy::kContiguous;
+};
+
+TORCH_API int& getTECudaPointwiseLoopLevels();
+TORCH_API int& getTECudaPointwiseBlockCount();
+TORCH_API int& getTECudaPointwiseBlockSize();
+TORCH_API bool& getTEGenerateBlockCode();
+TORCH_API bool& getTEMustUseLLVMOnCPU();
+TORCH_API bool fallbackAllowed();
+TORCH_API bool setFallbackAllowed(bool value);
+TORCH_API bool& getCatWoConditionals();
+TORCH_API bool& getOptConditionals();
+
+TORCH_API std::optional<at::Device> pickDeviceType(
+    const at::ArrayRef<torch::jit::Value*>& inputs);
+
+bool isContiguous(
+    const torch::jit::Value* v,
+    at::MemoryFormat memory_format = at::MemoryFormat::Contiguous);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/llvm_codegen.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/llvm_codegen.h
new file mode 100644
index 0000000000000000000000000000000000000000..f253224710956595b35d147f0687e0763ee74fd1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/llvm_codegen.h
@@ -0,0 +1,148 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef TORCH_ENABLE_LLVM
+#include <torch/csrc/Export.h>
+
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+
+#include <optional>
+
+#include <unordered_map>
+#include <vector>
+
+namespace torch {
+namespace jit {
+namespace tensorexpr {
+
+class LLVMCodeGenImpl;
+class LLVMCodeGenCallee;
+
+class TORCH_API LLVMCodeGen : public CodeGen {
+ public:
+  explicit LLVMCodeGen(
+      StmtPtr stmt,
+      const std::vector<BufferArg>& args,
+      at::Device device = at::kCPU,
+      const std::string& kernel_func_name = "func",
+      Dtype dtype = kInt,
+      std::optional<std::string> triple = std::nullopt,
+      std::optional<std::string> cpu = std::nullopt,
+      std::optional<std::string> attrs = std::nullopt);
+  explicit LLVMCodeGen(StmtPtr stmt);
+
+  LLVMCodeGen() = delete;
+  ~LLVMCodeGen() override;
+
+  // Cleans up all the memory used during LLVM code generation pass except
+  // the generated kernel. After calling this method, users should not call
+  // methods like `getCodeText` that require the LLVMCodeGenImpl data. However,
+  // users can continue to call this kernel using `call` and `call_raw`.
+  void cleanup_memory();
+
+  TORCH_API void call(const std::vector<CallArg>& args) override;
+  TORCH_API void call_raw(const std::vector<void*>& args) override;
+  TORCH_API void call_with_numel(void** args, int64_t numel) override;
+
+  at::Tensor empty_strided(
+      c10::IntArrayRef size,
+      c10::IntArrayRef stride,
+      std::optional<c10::ScalarType> dtype_opt,
+      std::optional<c10::Layout> layout_opt,
+      std::optional<c10::Device> device_opt,
+      std::optional<bool> pin_memory_opt) override;
+
+  template <typename T>
+  T value() {
+    return value<T>(nullptr);
+  }
+
+  template <typename T>
+  T value(std::vector<void*>& args) {
+    return value<T>(args.data());
+  }
+
+  template <typename T>
+  T value(void** args) {
+    T (*fp)(void**) = (T(*)(void**))getKernelAddress(callee_.get());
+    T rv = fp(args);
+    return rv;
+  }
+
+  std::string getCodeText(const std::string& attr = "") override;
+
+ private:
+  void* getKernelAddress(LLVMCodeGenCallee* callee);
+
+  std::unique_ptr<LLVMCodeGenCallee> callee_;
+  std::unique_ptr<LLVMCodeGenImpl> impl_;
+};
+
+struct TORCH_API LLVMCodeGenBuilder {
+  using BufferArg = CodeGen::BufferArg;
+
+  LLVMCodeGenBuilder(StmtPtr stmt, std::vector<BufferArg> args)
+      : stmt_(stmt), args_(std::move(args)) {}
+
+  LLVMCodeGenBuilder& device(at::Device device) {
+    device_ = device;
+    return *this;
+  }
+
+  LLVMCodeGenBuilder& kernelFuncName(std::string name) {
+    kernelFuncName_ = std::move(name);
+    return *this;
+  }
+
+  LLVMCodeGenBuilder& dtype(Dtype d) {
+    dtype_ = d;
+    return *this;
+  }
+
+  LLVMCodeGenBuilder& triple(std::string triple) {
+    triple_ = std::move(triple);
+    return *this;
+  }
+
+  LLVMCodeGenBuilder& cpu(std::string cpu) {
+    cpu_ = std::move(cpu);
+    return *this;
+  }
+
+  LLVMCodeGenBuilder& attrs(std::string attrs) {
+    attrs_ = std::move(attrs);
+    return *this;
+  }
+
+  std::unique_ptr<LLVMCodeGen> build() {
+    return std::make_unique<LLVMCodeGen>(
+        stmt_, args_, device_, kernelFuncName_, dtype_, triple_, cpu_, attrs_);
+  }
+
+ private:
+  StmtPtr stmt_;
+  std::vector<BufferArg> args_;
+  at::Device device_ = at::kCPU;
+  std::string kernelFuncName_ = "func";
+  Dtype dtype_ = kInt;
+  std::optional<std::string> triple_ = std::nullopt;
+  std::optional<std::string> cpu_ = std::nullopt;
+  std::optional<std::string> attrs_ = std::nullopt;
+};
+
+TORCH_API std::optional<std::string>& LLVMTargetTriple();
+TORCH_API std::optional<std::string>& LLVMTargetCPU();
+TORCH_API std::optional<std::string>& LLVMTargetAttrs();
+TORCH_API bool& LLVMAOTWorkflow();
+
+} // namespace tensorexpr
+} // namespace jit
+} // namespace torch
+
+#endif // TORCH_ENABLE_LLVM
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/llvm_jit.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/llvm_jit.h
new file mode 100644
index 0000000000000000000000000000000000000000..deece5fe694540855646f8a33f723ed5babf5bbe
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/llvm_jit.h
@@ -0,0 +1,84 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef TORCH_ENABLE_LLVM
+#include <c10/macros/Macros.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/Export.h>
+#include <optional>
+
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wsuggest-override")
+#include <llvm/ExecutionEngine/JITSymbol.h>
+C10_DIAGNOSTIC_POP()
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wextra-semi")
+#include <llvm/ExecutionEngine/Orc/Core.h>
+#include <llvm/ExecutionEngine/Orc/ThreadSafeModule.h>
+#include <llvm/Target/TargetMachine.h>
+C10_DIAGNOSTIC_POP()
+
+#include <memory>
+#include <string>
+
+namespace torch {
+namespace jit {
+namespace tensorexpr {
+
+inline std::string formatError(llvm::Error&& err, const char* msg) {
+  static constexpr const char* defaultErrorMsg =
+      "Unexpected failure in LLVM JIT";
+  std::string errorMsg(msg ? msg : defaultErrorMsg);
+  llvm::raw_string_ostream ss(errorMsg);
+  ss << ": " << err;
+  return ss.str();
+}
+
+template <typename T>
+T assertSuccess(llvm::Expected<T> valOrErr, const char* msg = nullptr) {
+  TORCH_INTERNAL_ASSERT(valOrErr, formatError(valOrErr.takeError(), msg));
+  return std::move(*valOrErr);
+}
+
+inline void assertSuccess(llvm::Error err, const char* msg = nullptr) {
+  TORCH_INTERNAL_ASSERT(!err, formatError(std::move(err), msg));
+}
+
+} // namespace tensorexpr
+} // namespace jit
+} // namespace torch
+
+namespace llvm {
+namespace orc {
+
+class PytorchLLVMJITImpl;
+
+class TORCH_API PytorchLLVMJIT {
+ public:
+  PytorchLLVMJIT(
+      std::optional<std::string> triple,
+      std::optional<std::string> cpu,
+      std::optional<std::string> attrs);
+  ~PytorchLLVMJIT();
+
+  void addModule(std::unique_ptr<Module> M, std::unique_ptr<LLVMContext> C);
+
+  JITSymbol findSymbol(const std::string Name);
+
+  bool hasSymbol(const std::string& Name);
+
+  TargetMachine& getTargetMachine();
+
+  const DataLayout& getDataLayout();
+
+ private:
+  // Use the PImpl idiom here to hide the no-rtti parts of the JIT structure.
+  std::unique_ptr<PytorchLLVMJITImpl> impl_;
+};
+
+} // end namespace orc
+} // end namespace llvm
+
+#endif // ENABLE LLVM
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/loopnest.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/loopnest.h
new file mode 100644
index 0000000000000000000000000000000000000000..fe3f27c0861dc6b643f5ffbd3cceb53517642cec
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/loopnest.h
@@ -0,0 +1,622 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+
+namespace torch::jit::tensorexpr {
+
+class Expr;
+class Var;
+class Buf;
+class Tensor;
+class Function;
+class Stmt;
+class For;
+class Block;
+class Store;
+class Dtype;
+
+class TORCH_API LoopNest {
+ public:
+  // A constructor for building a LoopNest from a list of Tensors
+  LoopNest(
+      const std::vector<Tensor>& output_tensors,
+      const std::vector<Tensor>& tensors_to_compute);
+
+  // A convenience constructor for the case when all tensors are output tensors
+  LoopNest(const std::vector<Tensor>& output_tensors);
+
+  // A constructor for building a LoopNest from an Stmt and a list of output
+  // buffers.
+  LoopNest(StmtPtr stmt, std::unordered_set<BufPtr> output_bufs);
+
+  // A constructor for building a LoopNest from another loopnest. It clones the
+  // other loopnest's stmt.
+  LoopNest(const LoopNest& other);
+
+  StmtPtr root_stmt() const {
+    return root_stmt_;
+  }
+
+  std::vector<ForPtr> getLoopStmtsFor(const Tensor& /*t*/) const;
+  std::vector<ForPtr> getLoopStmtsFor(const BufPtr& /*buf*/) const;
+  std::vector<ForPtr> getLoopStmtsFor(StmtPtr /*s*/) const;
+  StmtPtr getLoopBodyFor(const Tensor& /*t*/) const;
+  StmtPtr getLoopBodyFor(BufPtr /*buf*/) const;
+
+  // Returns the For stmt indexed by 'indices' in the 'root' For stmt.
+  //'indices' indicates the path to the returned loop from 'root' in AST, e.g.,
+  //
+  // root: for(int i...){
+  // j_loop: for (int j...){
+  // k1_loop:  for (int k1...){
+  //            A[i, j, k1] = ....
+  //          }
+  //          B[i, j] = ...
+  // k2_loop:  for (int k2...){
+  //            A[i, j, k2] = ...
+  //          }
+  //        }
+  //      }
+  //
+  // the path from 'root' to 'j_loop' is [0]
+  // the path from 'root' to 'k1_loop' is [0, 0]
+  // the path from 'root' to 'k2_loop' is [0, 2]
+  ForPtr getLoopAt(ForPtr root, const std::vector<int>& indices) const;
+
+  // Returns the For stmt that is immediately enclosing the given stmt.
+  static ForPtr getParentLoop(const StmtPtr& st);
+
+  // Returns the list of For stmts corresponding to the loopnest that is
+  // enclosing the given stmt.
+  static std::vector<ForPtr> getEnclosingLoopNest(const StmtPtr& st);
+
+  // Returns a list of all Stmts that write to the given buf.
+  std::vector<StmtPtr> getAllWritesToBuf(BufPtr /*buf*/) const;
+
+  // The following methods return the For loops that contain writes to
+  // the given buf.
+  //
+  // For example, consider the following code:
+  //   for i1
+  //     for j1
+  //       a[i1,j1] =
+  //   for i2
+  //     for j2
+  //       for k2
+  //         a[i2,j2] =
+  //     for j3
+  //       a[i2,j3] =
+
+  // Returns a list of For loops which directly contain a Stmt that writes
+  // to buf.
+  // For the above example:
+  //   getAllInnermostLoopsWritingToBuf(a) => {j1, k2, j3}
+  std::vector<ForPtr> getAllInnermostLoopsWritingToBuf(BufPtr /*buf*/) const;
+
+  // Returns a list of For loopnests which contain a Stmt that writes to
+  // the given buf. Each loopnest here is a vector For loops.
+  // For the above example:
+  //   getAllLoopNestsWritingToBuf(a) => {{i1,j1}, {i2,j2,k2}, {i2,j3}}
+  std::vector<std::vector<ForPtr>> getAllLoopNestsWritingToBuf(
+      BufPtr /*buf*/) const;
+
+  StmtPtr simplify();
+
+  // Sanitize variables and buffer names.
+  // The pass assigns predefined names for loop index variables
+  // (i,j,k,l,m,n,o,p,i1,j1,k1,...) and ensures these names are not conflicting
+  // anywhere. It also removes duplicates from other Buf nad Var names as well
+  // as replaces illegal characters in them with underscores.
+  //
+  // Note: since it's currently technically possible to use the same variable
+  // as index in two different loops, this transformation finds such cases and
+  // introduces new variables to avoid duplication.
+  static StmtPtr sanitizeNames(StmtPtr s);
+
+  bool computeInline(const StmtPtr& s);
+  bool computeInline(const BufPtr& b);
+  void inlineIntermediateBufs(bool allow_duplicated_work);
+
+  // Optimizes conditionals.
+  //
+  // Currently, only the following pattern of conditionals is optimized.
+  // This corresponds to the conditional format that is generated to handle
+  // `aten::cat` op.
+  //
+  //   for (int i = 0; i < 20; i++) {
+  //     A[i] = IfThenElse(i<5 ? 1 : 0, B[i], C[i-5])
+  //   }
+  //
+  // Constraints that must be satisfied for this optimization:
+  //   * All conditions should be of the form "var < expr".
+  //   * All conditions should have the same variable, say v.
+  //   * The condition variable found should be the same as the inner-most
+  //     loop variable. TODO: Remove this constraint.
+  //   * If there are multiple stores that contain conditionals using the same
+  //     loop variable, only the first conditional will be optimized.
+  //     TODO: Remove this constraint.
+  bool optimizeConditionals();
+
+  // Splits the given loop into 2 nested loops with the given factor as the
+  // inner loop bound. If the factor does not evenly divide the loop bound,
+  // then the remaining iterations are extracted into a tail loop that is
+  // added after the given loop.
+  //
+  // For example, consider the following code:
+  //   for (int i = 0; i < 100; ++i) {
+  //     A[i] =
+  //   }
+  //
+  // splitWithTail(i, 8, ...) will result in:
+  //   for (int i_outer = 0; i_outer < 12; ++i_outer) {
+  //     for (int i_inner = 0; i_inner < 8; ++i_inner) {
+  //       A[i_outer * 8 + i_inner] =
+  //     }
+  //   }
+  //   for (int i_tail = 0; i_tail < 4; ++i_tail) {
+  //     A[i_tail + 96] =
+  //   }
+  //
+  // The given loop will be transformed to the outer loop after splitting.
+  // So, the pointer to the input loop should be valid after splitting and
+  // will point to the outer loop. The `inner` and `tail` parameters will be
+  // set to point to the inner and tail loops that are generated.
+  static void splitWithTail(
+      const ForPtr& f,
+      int factor,
+      ForPtr* inner,
+      ForPtr* tail);
+  // A convenience wrapper when the caller does not need to access the
+  // split loops.
+  static void splitWithTail(const ForPtr& f, int factor);
+
+  // Splits the given loop into 2 nested loops with the given factor as the
+  // inner loop bound. If the factor does not evenly divide the loop bound,
+  // then a conditional is inserted into the body to handle the remaining
+  // iterations appropriately.
+  //
+  // For example, consider the following code:
+  //   for (int i = 0; i < 100; ++i) {
+  //     A[i] =
+  //   }
+  //
+  // splitWithMask(i, 8, ...) will result in:
+  //   for (int i_outer = 0; i_outer < 13; ++i_outer) {
+  //     for (int i_inner = 0; i_inner < 8; ++i_inner) {
+  //       if (i_outer * 8 + i_inner < 100) {
+  //         A[i_outer * 8 + i_inner] =
+  //       }
+  //     }
+  //   }
+  //
+  // The given loop will be transformed to the outer loop after splitting.
+  // So, the pointer to the input loop should be valid after splitting and
+  // will point to the outer loop. The `inner` parameter will be set to point
+  // to the inner loop that is generated.
+  static void splitWithMask(const ForPtr& f, int factor, ForPtr* inner);
+  // A convenience wrapper when the caller does not need to access the
+  // split loops.
+  static void splitWithMask(const ForPtr& f, int factor);
+
+  // The following methods support loop distribution.
+  // For example, consider the following code. This will be used to
+  // demonstrate the methods below.
+  //
+  // S0:  for m
+  // S1:    for i
+  // S2:      A[i] = 0
+  // S3:      for j
+  // S4:        A[i] = A[i] +
+  // S5:      B[i] = A[i]
+  // S6:      for k
+  // S7:        B[i] = B[i] +
+
+  // This method distributes the given loop over its body by splitting
+  // after every given pivot stmt.
+  //
+  // NOTE: Pivot stmts that are not in the given loop's body will be ignored.
+  //
+  // For the above example:
+  //   distributeLoop(S1, {S3, S5})
+  // will result in:
+  // S0:  for m
+  // S1:    for i
+  // S2:      A[i] = 0
+  // S3:      for j
+  // S4:        A[i] = A[i] +
+  //   :    for i
+  // S5:      B[i] = A[i]
+  //   :    for i
+  // S6:      for k
+  // S7:        B[i] = B[i] +
+  static std::vector<ForPtr> distributeLoop(
+      const ForPtr& loop,
+      const std::unordered_set<StmtPtr>& pivots);
+
+  // This method distributes the given loop over every stmt in its body.
+  //
+  // For the above example:
+  //   distributeLoop(S1)
+  // will result in:
+  // S0:  for m
+  // S1:    for i
+  // S2:      A[i] = 0
+  //   :    for i
+  // S3:      for j
+  // S4:        A[i] = A[i] +
+  //   :    for i
+  // S5:      B[i] = A[i]
+  //   :    for i
+  // S6:      for k
+  // S7:        B[i] = B[i] +
+  static std::vector<ForPtr> distributeLoop(const ForPtr& loop);
+  // Same as above, but also distribute parent loops.
+  // Returns the result of distributing the outermost loop.
+  //
+  // For the above example:
+  //   distributeLoopAndParents(S1) will result in:
+  // S0:  for m
+  // S1:    for i
+  // S2:      A[i] = 0
+  //   :  for m
+  //   :    for i
+  // S3:      for j
+  // S4:        A[i] = A[i] +
+  //   :  for m
+  //   :    for i
+  // S5:      B[i] = A[i]
+  //   :  for m
+  //   :    for i
+  // S6:      for k
+  // S7:        B[i] = B[i] +
+  static std::vector<ForPtr> distributeLoopAndParents(const ForPtr& loop);
+
+  // This method distributes the given loop over its body by splitting
+  // after every For stmt in its body.
+  //
+  // For the above example:
+  //   distributeLoopOverInnerLoops(S1)
+  // will result in:
+  // S0:  for m
+  // S1:    for i
+  // S2:      A[i] = 0
+  // S3:      for j
+  // S4:        A[i] = A[i] +
+  //   :    for i
+  // S5:      B[i] = A[i]
+  // S6:      for k
+  // S7:        B[i] = B[i] +
+  static std::vector<ForPtr> distributeLoopOverInnerLoops(const ForPtr& loop);
+  // Same as above, but also distribute parent loops.
+  // Returns the result of distributing the outermost loop.
+  //
+  // For the above example:
+  //   distributeLoopAndParentsOverInnerLoops(S1)
+  // will result in:
+  // S0:  for m
+  // S1:    for i
+  // S2:      A[i] = 0
+  // S3:      for j
+  // S4:        A[i] = A[i] +
+  //   :  for m
+  //   :    for i
+  // S5:      B[i] = A[i]
+  // S6:      for k
+  // S7:        B[i] = B[i] +
+  static std::vector<ForPtr> distributeLoopAndParentsOverInnerLoops(
+      const ForPtr& loop);
+
+  // This method performs loop fusion.
+  // For example, consider the following code.
+  //
+  // S1:  for m
+  // S2:    A[m] = 0
+  // S3:    for j
+  // S4:      A[m] = A[m] +
+  // S5:  for n
+  // S5:    B[n] = A[n]
+  // S6:    for k
+  // S7:      B[n] = B[n] +
+  //
+  // fuseLoops({S1, S5}), will return the following loop:
+  // S1:  for m
+  // S2:    A[m] = 0
+  // S3:    for j
+  // S4:      A[m] = A[m] +
+  // S5:    B[m] = A[m]
+  // S6:    for k
+  // S7:      B[m] = B[m] +
+  //
+  // This transformation is unsafe as it simply add all loops into the body of
+  // the first loop for fusion without correctness checks.
+  //
+  // Below are the two requirements to apply unsafeFuseLoops:
+  //  * All the loops have the same parent.
+  //  * There are no statements between these loops in their parent body.
+  static bool unsafeFuseLoops(const std::vector<ForPtr>& loops, ForPtr* fused);
+
+  // Loop fusion is done only when all the conditions below are satisfied.
+  //  * All the loops have the same parent.
+  //  * There are no statements between these loops in their parent body.
+  //  * The start bounds are the same for all loops.
+  //  * The stop bounds are the same for all loops.
+  //  * Fusing the loops does not violate or add any dependencies.
+  static bool fuseLoops(const std::vector<ForPtr>& loops, ForPtr* fused);
+
+  static void reorderAxis(const ForPtr& a, const ForPtr& b);
+
+  // Reorder the given list of loops according to the permutation specified.
+  // Here `permutation[i]` represents the position of the loop in the input
+  // which will end up at position `i` after the reorder.
+  //
+  // For example, consider the following code:
+  //   for p
+  //     for q
+  //       for r
+  //         for s
+  //           A[p,q,r,s] =
+  //
+  // reorder({p, q, r, s}, {2, 3, 0, 1}) will return the list of loops in the
+  // following form:
+  //    for r
+  //      for s
+  //        for p
+  //          for q
+  //            A[p,q,r,s] =
+  static std::vector<ForPtr> reorder(
+      const std::vector<ForPtr>& loops,
+      const std::vector<size_t>& permutation);
+
+  // Tile takes a 2d domain (x, y) and splits it into small rectangular blocks
+  // each with shape (x_factor, y_factor). The traversal over the domain turns
+  // into an outer iteration over the blocks and an inner traversal over all
+  // points in the block.
+  // Note that if x dim % x_factor or y dim % y_factor does not equal to 0, the
+  // loop body will generate corresponding tailing loops.
+  // The transformation is in-place and returns 'xtail'.
+  //
+  // For example, consider the following code:
+  //   for i: [0, 64)
+  //     for j: [0, 64)
+  //       for k: [0, 32)
+  //         A[i, j] = B[i, k] + C[j, k]
+  //
+  // tile(i, j, 4, 8) will transform "i" for-stmt into the following nested
+  // loop:
+  //   for i_outer: [0, 16)
+  //     for j_outer: [0, 8)
+  //       for i_inner: [0, 4)
+  //         for j_inner: [0, 8)
+  //           for k: [0, 32)
+  //             A[i_outer * 4 + i_inner, j_outer * 8 + j_inner] =
+  //             B[i_outer * 4 + i_inner, k] + C[j_outer * 8 + j_inner, k]
+  //
+  // tile(i, j, 4, 9) will transform "i" for-stmt into the following nested
+  // loop:
+  //   for i_outer: [0, 16)
+  //     for j_outer: [0, 7)
+  //       for i_inner: [0, 4)
+  //         for j_inner: [0, 9)
+  //           for k: (0, 32)
+  //             A[i_outer * 4 + i_inner, j_outer * 9 + j_inner] =
+  //             B[i_outer * 4 + i_inner, k] + C[j_outer * 9 + j_inner, k]
+  //     for j_tail: [0, 1)
+  //       for i_inner: [0, 4)
+  //         for k: (0, 32)
+  //           A[i_outer * 4 + i_inner, 7 * 9 + j_tail] =
+  //           B[i_outer * 4 + i_inner, k] + C[7 * 9 + j_tail, k]
+  ForPtr tile(const ForPtr& x, const ForPtr& y, int x_factor, int y_factor);
+
+  // Returns true if the given loops are perfectly nested, i.e., every loop
+  // (except the innermost) should have exactly one statement in its body
+  // and that statement must be the next inner loop.
+  static bool areLoopsPerfectlyNested(const std::vector<ForPtr>& loops);
+
+  // Returns true if the given loop has a loop-carried dependence.
+  static bool hasLoopCarriedDependence(const ForPtr& loop);
+
+  // Unrolls all the iterations of the given loop.
+  // Requires that the loop bounds are constant.
+  static void fullUnroll(const ForPtr& f, StmtPtr* unrolled);
+  static void fullUnroll(const ForPtr& f);
+
+  // Unrolls the given loop for the specified factor.
+  // This does not require constant bounds for the loop being unrolled.
+  static void unroll(const ForPtr& f, int factor, ForPtr* tail);
+  static void unroll(const ForPtr& f, int factor);
+
+  static bool normalize(const ForPtr& f);
+  static bool isNormalized(const ForPtr& f);
+
+  static bool flatten(const std::vector<ForPtr>& f, ForPtr* flattened);
+  static bool flatten(const std::vector<ForPtr>& f);
+
+  // Compresses the given buffer based on its use in the given Stmts.
+  //
+  // NOTE: This API assumes that there are no accesses to the given buffer
+  // outside the given statement. So, this should be called with the entire
+  // kernel statement to avoid incorrect buffer compressions.
+  //
+  // For example, given the input:
+  //
+  // for (int i = 0; i < 100; ++i) {
+  //   for (int j = 0; j < 200; ++j) {
+  //     A[i,j] = sin(i*j)
+  //   }
+  //   for (int j = 0; j < 199; ++j) {
+  //     B[i,j] = A[i,j] + A[i, j+1]
+  //   }
+  // }
+  //
+  // compressBuffer(A, ...) will compress buffer A from
+  // [100, 200] to [1, 200] and modify the code as follows:
+  //
+  // for (int i = 0; i < 100; ++i) {
+  //   for (int j = 0; j < 200; ++j) {
+  //     A[0,j] = sin(i*j)
+  //   }
+  //   for (int j = 0; j < 199; ++j) {
+  //     B[i,j] = A[0,j] + A[0, j+1]
+  //   }
+  // }
+  static void compressBuffer(const BufPtr& buf, const StmtPtr& stmt);
+
+  // Compresses all buffers in the given statement.
+  //
+  // NOTE: This API assumes that there are no accesses to buffers outside
+  // the given statement. So, this should be called with the entire
+  // kernel statement to avoid incorrect buffer compressions.
+  //
+  // TODO: Add an IR verifier check to detect invalidly compressed buffers.
+  static void compressAllBuffers(const StmtPtr& stmt);
+
+  // Get 'num' loops from the loopnest starting at 'f'.
+  static std::vector<ForPtr> getLoopStmtsInLoopNest(
+      const ForPtr& f,
+      size_t num);
+
+  // LoopOptions are propagated to tail.
+  static void sliceHead(
+      const ForPtr& f,
+      int factor,
+      ForPtr* head,
+      ForPtr* tail);
+  static void sliceHead(const ForPtr& f, int factor);
+  // LoopOptions are propagated to head.
+  static void sliceTail(
+      const ForPtr& f,
+      int factor,
+      ForPtr* head,
+      ForPtr* tail);
+  static void sliceTail(const ForPtr& f, int factor);
+
+  using AccessResult = std::pair<BufPtr, StmtPtr>;
+  // Insert a cache for the consumer's usages of the buffer produced in
+  // consumer, and redirect reads and writes in the consumer to that cache.
+  // Returns a pair of the new cache buffer, and the new rewritten consumer.
+  static AccessResult cacheAccesses(
+      const BufPtr& producer,
+      const std::string& name,
+      const StmtPtr& consumer);
+
+  // Insert a temporary computation of statement S in the scope of loop AT.
+  // S is assumed to be a Store or a Block containing a Store. Along with the
+  // computation itself, this transformation inserts Alloc/Free statements for
+  // the temporary buffer used in the computation.
+  static void computeAt(const StmtPtr& s, const ForPtr& at);
+
+  // Rfactor a reduction axis into a normal axis.
+  //
+  // Requirements:
+  //  * S is the reduction store
+  //  * S is the only statement in the innermost loop
+  //  * There is at least two reduction arguments in S
+  //  * OUTER_REDUCTION_FOR loop corresponds to the outermost reduction variable
+  //  used in the store and all other reduction variables are index variables of
+  //  children loops of OUTER_REDUCTION_FOR
+  //  * OUTER_REDUCTION_FOR is a perfect loop nest, i.e. it has only loops
+  //  corresponding to the other reduction variables and the store, nested into
+  //  each other
+  //
+  // What it does:
+  //   * Introduce a new buffer with an extra dimension of a size equal to the
+  //   span of the loop OUTER_REDUCTION_FOR (the new buffer is returned via
+  //   RFAC_BUF_PTR)
+  //   * Insert an initialization store for the new buffer in
+  //   OUTER_REDUCTION_FOR before its nested loop
+  //   * Replace the reduction store to the original buffer with the reduction
+  //   store to the temp buffer, removing the index var of OUTER_REDUCTION_FOR
+  //   from reduction arguments
+  //   * Insert a final reduction store over the extra dimension of the new
+  //   buffer to the original buffer
+  //   * Returns TRUE if the transformation succeeded and FALSE otherwise
+  //
+  // Example:
+  // Original IR:
+  // S1: for i      # normal axis
+  // S2:   X[i] = 0
+  // S3:   for j    # reduction axis
+  // S4:     for k  # reduction axis
+  // S5:       X[i] = ReduceOp(X[i] + Y[i,j,k], reduce_axis={j,k})
+  //
+  // After RFACTOR(S5, S3)
+  // S1: for i               # normal axis
+  // S2:   X[i] = 0
+  // S3:   for j             # reduction axis for X, normal axis for X_rfac
+  //         X_rfac[i,j] = 0
+  // S4:     for k           # reduction axis
+  //           X_rfac[i,j] = ReduceOp(X_rfac[i,j] + Y[i,j,k], reduce_axis={k})
+  //         X[i] = ReduceOp(X[i] + X_rfac[i,j], reduce_axis={j})
+  static bool rfactor(const StmtPtr& s, const ForPtr& outer_reduction_for);
+  static bool rfactor(
+      const StmtPtr& s,
+      const ForPtr& outer_reduction_for,
+      BufPtr* rfac_buf_ptr);
+
+  // Vectorize the given loop. This method requires that the given loop
+  // does not perform a reduction.
+  // It returns true if vectorization is successful and false otherwise.
+  static bool vectorize(const ForPtr& /*f*/);
+
+  // Find the inner-most loops and vectorize them. Currently, this only works
+  // for the LLVM backend, when no reductions are involved.
+  void vectorizeInnerLoops();
+
+  void eliminateDeadStores();
+
+  void prepareForCodegen();
+
+  const std::unordered_set<BufPtr> getInputBufs() const;
+  const std::unordered_set<BufPtr> getOutputBufs() const {
+    return output_bufs_;
+  }
+  std::vector<BufPtr> getIntermediateBufs() const;
+
+  // Finds which is the outer For between a and b for loops. If neither of the 2
+  // Fors is an ancestor of the other, it returns nullptr.
+  static ForPtr findOuterFor(ForPtr a, ForPtr b);
+
+ private:
+  void initialize(
+      const std::vector<Tensor>& output_tensors,
+      const std::vector<Tensor>& tensors_to_compute);
+
+  StmtPtr root_stmt_;
+
+  std::unordered_set<BufPtr> output_bufs_;
+};
+
+TORCH_API StmtPtr FlattenIndexes(const StmtPtr& s);
+
+// TODO: Revisit this once we decide on how dependencies analysis should look
+// like. Maybe we would choose to use a different API and BufUse would be
+// removed, or if we decide to keep it we need to properly document its API.
+struct BufLoadOrStoreUse {
+  StmtPtr s;
+  bool isStore;
+};
+
+/*
+ * Returns a map ( Buf -> uses of this Buf), uses are represented as vectors of
+ * BufUse elements, which are StmtPtr and a bool isStore flag. The order of uses
+ * in the vectors reflects the order in which the uses appear in the given
+ * statement.
+ */
+std::unordered_map<BufPtr, std::vector<BufLoadOrStoreUse>> findLoadOrStoreUses(
+    const StmtPtr& s);
+
+// replaces all invalid characters with underscore
+TORCH_API std::string sanitizeName(const std::string& input_name);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/loopnest_randomization.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/loopnest_randomization.h
new file mode 100644
index 0000000000000000000000000000000000000000..674fd337b0ea8ceb9cb6d09f122489eaa56cdc30
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/loopnest_randomization.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::jit::tensorexpr {
+
+// Applies a series of loop optimizations chosen randomly. This is only for
+// testing purposes. This allows automatic stress testing of NNC loop
+// transformations.
+void loopnestRandomization(int64_t seed, LoopNest& l);
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/lowerings.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/lowerings.h
new file mode 100644
index 0000000000000000000000000000000000000000..79a75576d509dbe9b143daf2611cde26d7149d7d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/lowerings.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// This file defines classes for registering standard lowerings from JIT to TE
+// IR.
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <torch/csrc/jit/tensorexpr/analysis.h>
+#include <torch/csrc/jit/tensorexpr/codegen.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+namespace torch::jit::tensorexpr {
+
+using ArgNone = std::monostate;
+using BufList = std::vector<tensorexpr::BufHandle>;
+using DoubleList = std::vector<double>;
+using IntList = std::vector<int64_t>;
+using ArgValue = std::variant<
+    tensorexpr::BufHandle,
+    tensorexpr::VarHandle,
+    double,
+    int64_t,
+    bool,
+    BufList,
+    DoubleList,
+    IntList,
+    std::string,
+    ArgNone>;
+
+using NNCLoweringFunction = std::function<Tensor(
+    const std::vector<ArgValue>&,
+    const std::vector<ExprHandle>&,
+    const std::vector<ExprHandle>&,
+    const std::optional<ScalarType>&,
+    at::Device)>;
+
+TORCH_API FunctionSchemaMap<NNCLoweringFunction>& getNNCLoweringRegistry();
+TORCH_API NNCLoweringFunction getStandardLoweringFor(const std::string& op);
+
+struct RegisterNNCLoweringsFunction {
+  RegisterNNCLoweringsFunction(
+      const std::vector<std::string>& schemas,
+      const NNCLoweringFunction& fn);
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/mem_dependency_checker.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/mem_dependency_checker.h
new file mode 100644
index 0000000000000000000000000000000000000000..8e5b9e292d862bf890e19729b0a4ba206d4a85bb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/mem_dependency_checker.h
@@ -0,0 +1,414 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/Export.h>
+#include <utility>
+#include <vector>
+
+#include <torch/csrc/jit/tensorexpr/bounds_overlap.h>
+#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
+#include <torch/csrc/jit/tensorexpr/ir_simplifier.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/stmt.h>
+
+namespace torch::jit::tensorexpr::analysis {
+
+enum class AccessType {
+  Input,
+  Output,
+  Load,
+  Store,
+  Call,
+  AtomicAdd,
+  Alloc,
+  Free
+};
+const char* AccessToString(AccessType a);
+
+class AccessInfo;
+using DependencySet = std::unordered_set<std::shared_ptr<AccessInfo>>;
+
+/* AccessInfo
+ *
+ * Represents a single bounded memory access to a buffer, for instance a Load or
+ * a Store. Holds information relating to the specific access and links to
+ * connected accesses in the dependency graph.
+ */
+class TORCH_API AccessInfo {
+ public:
+  AccessInfo(
+      size_t id,
+      AccessType type,
+      StmtPtr stmt,
+      VarPtr var,
+      IndexBounds bounds)
+      : id_(id),
+        type_(type),
+        stmt_(std::move(stmt)),
+        expr_(nullptr),
+        var_(std::move(var)),
+        bounds_(std::move(bounds)) {}
+
+  AccessInfo(
+      size_t id,
+      AccessType type,
+      ExprPtr expr,
+      StmtPtr stmt,
+      VarPtr var,
+      IndexBounds bounds)
+      : id_(id),
+        type_(type),
+        stmt_(std::move(stmt)),
+        expr_(std::move(expr)),
+        var_(std::move(var)),
+        bounds_(std::move(bounds)) {}
+
+  // Id is a unique int representing the order this access occurred in the
+  // graph.
+  size_t id() const {
+    return id_;
+  }
+
+  // The type of the access (Load, Store, etc).
+  AccessType type() const {
+    return type_;
+  }
+
+  // The enclosing Stmt this access represents. E.g. if this is a Store then
+  // Stmt is the Store itself, while if the access is caused by an Expr, this is
+  // the most immediate parent Stmt.
+  StmtPtr stmt() const {
+    return stmt_;
+  }
+
+  // If the access is represented by an Expr (such as Load or Call) then this is
+  // it, otherwise it's nullptr.
+  ExprPtr expr() const {
+    return expr_;
+  }
+
+  // The Var representing the underlying Buffer.
+  VarPtr var() const {
+    return var_;
+  }
+
+  // A vector of Bounds representing the start and end expression for each
+  // dimension.
+  IndexBounds& bounds() {
+    return bounds_;
+  }
+
+  // Each access that this depends upon,
+  // eg. if this is a Load, then it contains every Store that immediately
+  // contributes to a load of the bounds.
+  // or: if this is a Store, it contains all reads on the RHS of the Store.
+  const std::map<size_t, std::shared_ptr<AccessInfo>>& dependencies() const {
+    return dependencies_;
+  }
+
+  // Each access that depends on this one.
+  // ie. this access is present in the dependencies map of all accesses that are
+  // dependent.
+  std::map<size_t, std::shared_ptr<AccessInfo>> dependents() const {
+    std::map<size_t, std::shared_ptr<AccessInfo>> res;
+    for (const auto& kv : dependents_) {
+      res.emplace(kv.first, kv.second.lock());
+    }
+    return res;
+  }
+
+  // Returns the symbolic expression of the indices of this access.
+  std::vector<ExprPtr> getIndices() const;
+
+  // Establishes a dependency or dependent relationship with another access.
+  void addDependency(const std::shared_ptr<AccessInfo>& write);
+  void addDependent(const std::shared_ptr<AccessInfo>& read);
+
+  // helper for checking dependencies.
+  bool hasDependency(const std::shared_ptr<AccessInfo>& info) const;
+
+  // Returns the set of all nodes that are direct (immediate) dependencies of
+  // this access.
+  DependencySet getDirectDependencies();
+  // likewise, returns all nodes that directly depend on this one.
+  DependencySet getDirectDependents();
+
+  // Returns the full list of all nodes in the graph that this access depends
+  // on, and all nodes they depend on, and so forth, back to the inputs.
+  DependencySet getIndirectDependencies();
+  // likewise, returns the full list of all nodes that depend on this node, and
+  // all nodes that depend on those nodes and so on down to the outputs.
+  DependencySet getIndirectDependents();
+
+  // Does this access represent a read of memory (Load, ReduceOp, Call, etc).
+  bool isRead() const;
+  // Does this access represent a write of memory (Store, etc).
+  bool isWrite() const;
+
+  // Helpers for dumping accesses in various formats.
+  void print() const;
+  void dumpDOT(std::ostream& os) const;
+  const char* AccessTypeColour() const;
+
+ private:
+  size_t id_;
+  AccessType type_;
+  StmtPtr stmt_;
+  ExprPtr expr_;
+  VarPtr var_;
+  IndexBounds bounds_;
+
+  // Yes these should be sorted.
+  std::map<size_t, std::shared_ptr<AccessInfo>> dependencies_;
+  std::map<size_t, std::weak_ptr<AccessInfo>> dependents_;
+};
+
+using VarBoundMap = std::unordered_map<VarPtr, Bound>;
+
+/* MemDependencyChecker analyses a IR fragment and builds a dependency graph of
+ * accesses contained within.
+ *
+ * It's possible to retrieve the entire graph in node-object form, or can be
+ * used as an oracle for answering dependency questions. e.g:
+ *
+ *  analyzer.hasIndirectDependency(BufA, BufB); or,
+ *  analyzer.hasDirectDependency(LoadA, StoreB);
+ */
+class TORCH_API MemDependencyChecker : public IRVisitor {
+  struct Scope;
+
+ public:
+  MemDependencyChecker();
+  MemDependencyChecker(
+      const std::unordered_set<BufPtr>& inputs,
+      const std::unordered_set<BufPtr>& outputs);
+  MemDependencyChecker(
+      const std::vector<BufHandle>& inputs,
+      const std::vector<BufHandle>& outputs);
+
+  ~MemDependencyChecker() override = default;
+
+  // Whether or not to allow loop execution order to influence dependency
+  // calculation. If the loop may later be parallelized you don't want this.
+  bool allowLoopExecutionOrderAnalysis(bool allow = true);
+
+  // Dependency Checking API.
+  // The goal is to have enough overloads here so you don't really have to think
+  // about it.
+
+  // Returns true if any read in A has a direct dependence on a write in B.
+  bool dependsDirectly(const StmtPtr& A, const StmtPtr& B);
+  bool dependsDirectly(const ExprPtr& A, const StmtPtr& B);
+
+  // Returns true of the output depends directly on a write contained in B.
+  bool dependsDirectly(const BufPtr& output, const StmtPtr& B);
+
+  // Returns true if a read in A depends directly on the provided input.
+  bool dependsDirectly(const StmtPtr& A, const BufPtr& input);
+  bool dependsDirectly(const ExprPtr& A, const BufPtr& input);
+
+  // Outputs/inputs cannot depend directly.
+
+  // Returns true if the access A has B as an immediate dependency.
+  bool dependsDirectly(
+      const std::shared_ptr<AccessInfo>& A,
+      const std::shared_ptr<AccessInfo>& B);
+
+  // Returns true if any read in A has an ancestor write contained in B.
+  bool dependsIndirectly(const StmtPtr& A, const StmtPtr& B);
+  bool dependsIndirectly(const ExprPtr& A, const StmtPtr& B);
+
+  // Returns true of the output depends indirectly on a write contained in B.
+  bool dependsIndirectly(const BufPtr& output, const StmtPtr& B);
+
+  // Returns true if a read in A depends indirectly on the provided input.
+  bool dependsIndirectly(const StmtPtr& A, const BufPtr& input);
+  bool dependsIndirectly(const ExprPtr& A, const BufPtr& input);
+
+  // returns true if the output uses any load of the input.
+  bool dependsIndirectly(const BufPtr& output, const BufPtr& input);
+
+  // Returns true if the access A has a dependency chain to access B.
+  bool dependsIndirectly(
+      const std::shared_ptr<AccessInfo>& A,
+      const std::shared_ptr<AccessInfo>& B);
+
+  // Returns the AccessInfo
+  std::shared_ptr<AccessInfo> accessFor(const StmtPtr& A) const;
+  std::shared_ptr<AccessInfo> accessFor(const ExprPtr& A) const;
+
+  // Returns all AccessInfos.
+  std::unordered_set<std::shared_ptr<AccessInfo>> accessesWithin(
+      const StmtPtr& A) const;
+  // TODO: this will return only the AccessInfo for A. It's included for
+  // completeness but be aware it won't return accesses used in the computation
+  // of A.
+  std::unordered_set<std::shared_ptr<AccessInfo>> accessesWithin(
+      const ExprPtr& A) const;
+
+  // Accesses relating to input and output buffers.
+  std::shared_ptr<AccessInfo> input(const BufPtr& B) const;
+  std::shared_ptr<AccessInfo> output(const BufPtr& B) const;
+
+  // Returns the full history of reads and writes.
+  const std::vector<std::shared_ptr<AccessInfo>>& getHistory() const;
+
+  // Dumps the dependency graph in DOT format.
+  void dumpDAG(const std::string& filename) const;
+
+ private:
+  // Node visitors.
+  void visit(const StorePtr& v) override;
+  void visit(const LoadPtr& v) override;
+  void visit(const ForPtr& v) override;
+  void visit(const CondPtr& v) override;
+  void visit(const IfThenElsePtr& v) override;
+  void visit(const CompareSelectPtr& v) override;
+  void visit(const BlockPtr& v) override;
+  void visit(const LetPtr& v) override;
+  void visit(const AtomicAddPtr& v) override;
+  void visit(const AllocatePtr& v) override;
+  void visit(const FreePtr& v) override;
+
+  using BoundRelationship = std::pair<IndexBounds, std::shared_ptr<AccessInfo>>;
+
+  // An internal struct holding the accesses found within a scope Block.
+  struct Scope {
+    Scope(BlockPtr b, std::shared_ptr<Scope> p)
+        : block(std::move(b)), parent(std::move(p)) {}
+
+    BlockPtr block;
+    std::shared_ptr<Scope> parent;
+
+    std::unordered_map<VarPtr, Bound> shadowedVarBounds;
+    std::unordered_set<VarPtr> localVars;
+
+    std::vector<std::shared_ptr<AccessInfo>> accesses_;
+
+    std::unordered_map<VarPtr, std::list<BoundRelationship>> openWrites_;
+  };
+  std::shared_ptr<Scope> currentScope_;
+
+  bool allowExecutionOrderAnalysis_{false};
+
+  std::unordered_multimap<StmtPtr, std::shared_ptr<AccessInfo>> stmtToAccess_;
+  std::unordered_multimap<ExprPtr, std::shared_ptr<AccessInfo>> exprToAccess_;
+  std::unordered_map<StmtPtr, std::vector<std::shared_ptr<AccessInfo>>>
+      scopeToAccesses_;
+
+  VarBoundMap knownVarBounds_;
+
+  // Finds all accesses that are reads within the scope of v.
+  template <typename StmtOrExprPtr>
+  DependencySet getAllReadsWithin(const StmtOrExprPtr& v) {
+    DependencySet reads;
+    auto insertAllReads = [&](const auto& nodes) {
+      for (const auto& l : nodes) {
+        auto bound = exprToAccess_.equal_range(l);
+        for (auto it = bound.first; it != bound.second; ++it) {
+          if (it->second->isRead()) {
+            reads.insert(it->second);
+          }
+        }
+      }
+    };
+
+    // Look for and insert accesses belonging to all nodes that act like
+    // reads.
+    insertAllReads(NodeFinder<Load>::find(v));
+    insertAllReads(NodeFinder<ReduceOp>::find(v));
+
+    return reads;
+  }
+
+  // Finds all accesses that are writes within the scope of v.
+  // Writes cannot occur in Exprs, so this is a little simpler.
+  DependencySet getAllWritesWithin(const StmtPtr& v) {
+    DependencySet writes;
+
+    // writes just Store currently.
+    auto stores = NodeFinder<Store>::find(v);
+    for (const auto& s : stores) {
+      auto bound = stmtToAccess_.equal_range(s);
+      for (auto it = bound.first; it != bound.second; ++it) {
+        if (it->second->isWrite()) {
+          writes.insert(it->second);
+        }
+      }
+    }
+    return writes;
+  }
+
+  // Templated helpers to work on either Exprs or Stmts.
+  template <typename StmtOrExprPtr>
+  bool dependsDirectlyHelper(const StmtOrExprPtr& A, const StmtPtr& B) {
+    auto aReads = getAllReadsWithin(A);
+    auto bWrites = getAllWritesWithin(B);
+
+    for (auto& read : aReads) {
+      for (auto& depPair : read->dependencies()) {
+        if (bWrites.count(depPair.second) != 0) {
+          return true;
+        }
+      }
+    }
+
+    return false;
+  }
+
+  template <typename StmtOrExprPtr>
+  bool dependsIndirectlyHelper(StmtOrExprPtr A, const StmtPtr& B) {
+    auto aReads = getAllReadsWithin(A);
+    auto bWrites = getAllWritesWithin(B);
+
+    auto aDeps = getAllWriteDependencies(aReads);
+
+    for (auto& dependency : aDeps) {
+      if (bWrites.count(dependency) != 0) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  DependencySet getAllWriteDependencies(const DependencySet& products);
+
+  // Maps for inputs and outputs, since they aren't present directly in the IR.
+  std::unordered_map<BufPtr, std::shared_ptr<AccessInfo>> inputs_;
+  std::unordered_map<BufPtr, std::shared_ptr<AccessInfo>> outputs_;
+  std::unordered_map<VarPtr, std::shared_ptr<AccessInfo>> intermediates_;
+
+  // Inserts accesses for Buf's: specifically for inputs and outputs.
+  void insertBuffers(
+      std::unordered_map<BufPtr, std::shared_ptr<AccessInfo>>& bufs,
+      AccessType type);
+
+  // Update the write history with a new write, adding dependencies and closing
+  // any overlapped writes (if possible).
+  void updateWriteHistory(
+      std::list<BoundRelationship>& writeHistory,
+      const std::shared_ptr<AccessInfo>& info,
+      size_t latestAccessToClose,
+      bool closeOverlapped = true,
+      bool insert = true);
+
+  // Merge a child scope into a parent scope, adding dependencies for open
+  // writes in the parent to accesses in the child.
+  void mergeScope(
+      const std::shared_ptr<Scope>& child,
+      const std::shared_ptr<Scope>& parent,
+      bool closeOverlapped = true);
+
+  // Binds symbolic vars in indices with the low and high bound for those vars.
+  std::vector<Bound> getIndicesBounds(const std::vector<ExprPtr>& indices);
+
+  size_t nextAccess_{0};
+  StmtPtr lastStmt_{nullptr};
+};
+
+} // namespace torch::jit::tensorexpr::analysis
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/conv2d.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/conv2d.h
new file mode 100644
index 0000000000000000000000000000000000000000..ca63d7b9be1773cec3ada78416bd2c865427722d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/conv2d.h
@@ -0,0 +1,106 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/operators/misc.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+namespace torch::jit::tensorexpr {
+
+// An API to compute 2D depthwise convolutions with bias.
+TORCH_API Tensor conv2d_depthwise(
+    BufHandle input,
+    BufHandle weight,
+    BufHandle bias,
+    int stride,
+    int pad,
+    int groups);
+
+// An API to compute 2D depthwise convolutions without bias.
+TORCH_API Tensor conv2d_depthwise(
+    BufHandle input,
+    BufHandle weight,
+    int stride,
+    int pad,
+    int groups);
+
+TORCH_API Tensor conv2d_depthwise(
+    BufHandle input,
+    BufHandle weight,
+    BufHandle bias,
+    ExprHandle N,
+    ExprHandle C,
+    ExprHandle H,
+    ExprHandle W,
+    ExprHandle K,
+    ExprHandle CperG,
+    ExprHandle R,
+    ExprHandle S,
+    ExprHandle stride,
+    ExprHandle pad,
+    ExprHandle groups);
+
+TORCH_API Tensor conv2d_depthwise(
+    BufHandle input,
+    BufHandle weight,
+    ExprHandle N,
+    ExprHandle C,
+    ExprHandle H,
+    ExprHandle W,
+    ExprHandle K,
+    ExprHandle CperG,
+    ExprHandle R,
+    ExprHandle S,
+    ExprHandle stride,
+    ExprHandle pad,
+    ExprHandle groups);
+
+bool conv2dIsSupported(
+    const TensorInfo& input,
+    const TensorInfo& weight,
+    const TensorInfo& bias,
+    const std::vector<int64_t>& stride,
+    const std::vector<int64_t>& pad,
+    const std::vector<int64_t>& dilation,
+    int64_t groups);
+bool mkldnnPrepackedConvIsSupported(
+    const TensorInfo& input,
+    const TensorInfo& weight,
+    const std::vector<int64_t>& stride,
+    const std::vector<int64_t>& pad,
+    const std::vector<int64_t>& dilation,
+    int64_t groups);
+Tensor computeConv2d(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeConv1d(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computePrepackedConv2dClampRun(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computePrepackedLinearClampRun(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeMkldnnPrepackedConvRun(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/matmul.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/matmul.h
new file mode 100644
index 0000000000000000000000000000000000000000..1090455818d68d23348eb870b8930d1def34358b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/matmul.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/kernel.h>
+
+namespace torch::jit::tensorexpr {
+
+Tensor computeMatmul(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeAddMM(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/misc.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/misc.h
new file mode 100644
index 0000000000000000000000000000000000000000..6a1f61a2fb09e0342292b72c2c23d3a801df3a10
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/misc.h
@@ -0,0 +1,99 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+#include <torch/csrc/jit/tensorexpr/lowerings.h>
+#include <torch/csrc/jit/tensorexpr/tensor.h>
+
+namespace torch::jit::tensorexpr {
+
+struct TensorInfo {
+  std::vector<int64_t> dims;
+  c10::ScalarType dtype;
+};
+std::optional<TensorInfo> getTensorInfo(const BufHandle& b);
+
+int64_t normalizeAndCheckIndex(int64_t idx, int64_t list_size);
+
+// Convert boolean to integer, if needed.
+ExprHandle boolToInteger(const ExprHandle& x);
+ExprHandle promoteToDtype(ExprHandle e, ScalarType dt);
+void promoteInputs(
+    std::vector<ExprHandle>& inputs,
+    const int typeConstraints = kAllTypes);
+ExprHandle promoteIntegerToDefaultType(const ExprHandle& e);
+ExprHandle promoteHalfToFloat(const ExprHandle& e);
+ExprHandle demoteOutput(
+    const ExprHandle& e,
+    const std::optional<ScalarType> type);
+
+std::vector<ExprHandle> broadcastShapes(
+    std::vector<std::vector<ExprHandle>> shapes);
+std::vector<ExprHandle> broadcastShapes(
+    const std::vector<ExprHandle>& a,
+    const std::vector<ExprHandle>& b);
+
+std::vector<ExprHandle> valueShape(const ArgValue& v);
+ExprHandle tensorOrConstant(
+    const ArgValue& v,
+    const std::vector<ExprHandle>& axes);
+ExprHandle scalarOrConstant(const ArgValue& v);
+ExprHandle broadcast(const BufHandle& b, const std::vector<ExprHandle>& axes);
+ExprHandle constant(const ArgValue& v);
+
+ExprHandle clamp(
+    const ExprHandle& cmin,
+    const ExprHandle& cmax,
+    const ExprHandle& input);
+
+Tensor computeChunk(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeTranspose(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeExpand(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeReshape(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeFlatten(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeCatWoConditionals(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape);
+Tensor computeCat(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeEmbedding(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/norm.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/norm.h
new file mode 100644
index 0000000000000000000000000000000000000000..3d712cca1beae7547bd5f24f366f050d8783e5c5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/norm.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/kernel.h>
+
+namespace torch::jit::tensorexpr {
+
+Tensor computeBatchNorm(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/operators.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/operators.h
new file mode 100644
index 0000000000000000000000000000000000000000..8625cbf737729e3f33095526e2b712a3f35575c8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/operators.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/operators/conv2d.h>
+#include <torch/csrc/jit/tensorexpr/operators/matmul.h>
+#include <torch/csrc/jit/tensorexpr/operators/misc.h>
+#include <torch/csrc/jit/tensorexpr/operators/norm.h>
+#include <torch/csrc/jit/tensorexpr/operators/pointwise.h>
+#include <torch/csrc/jit/tensorexpr/operators/quantization.h>
+#include <torch/csrc/jit/tensorexpr/operators/reduction.h>
+#include <torch/csrc/jit/tensorexpr/operators/softmax.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/pointwise.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/pointwise.h
new file mode 100644
index 0000000000000000000000000000000000000000..cd8035fdaf773a5467a015e5113b8b5f4bd20fe8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/pointwise.h
@@ -0,0 +1,87 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/kernel.h>
+
+namespace torch::jit::tensorexpr {
+
+TORCH_API Tensor computeSign(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::optional<std::vector<ExprHandle>>& outputStrides = std::nullopt);
+
+Tensor computeOneOperand(
+    const std::string& name,
+    const std::vector<ArgValue>& inputValues,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    const std::function<ExprHandle(const ExprHandle&)>& innerExpr,
+    const int checkParamTypes = kAllTypes);
+Tensor computeTwoOperand(
+    const std::string& name,
+    const std::vector<ArgValue>& inputValues,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    const std::function<ExprHandle(const ExprHandle&, const ExprHandle&)>&
+        innerExpr);
+Tensor computeTwoOperandWithAlpha(
+    const std::string& name,
+    const std::vector<ArgValue>& inputValues,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    const std::function<ExprHandle(const ExprHandle&, const ExprHandle&)>&
+        innerExpr);
+Tensor computeConditionWithTwoOperand(
+    const std::string& name,
+    const std::vector<ArgValue>& inputValues,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    const std::function<
+        ExprHandle(const ExprHandle&, const ExprHandle&, const ExprHandle&)>&
+        innerExpr);
+Tensor computeThreeOperand(
+    const std::string& name,
+    const std::vector<ArgValue>& inputValues,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    const std::function<
+        ExprHandle(const ExprHandle&, const ExprHandle&, const ExprHandle&)>&
+        innerExpr,
+    bool promote_inputs = true);
+Tensor computeFourOperand(
+    const std::string& name,
+    const std::vector<ArgValue>& inputValues,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    const std::function<ExprHandle(
+        const ExprHandle&,
+        const ExprHandle&,
+        const ExprHandle&,
+        const ExprHandle&)>& innerExpr);
+Tensor computeNoop(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+Tensor computeScalar(
+    const std::string& name,
+    const std::vector<ArgValue>& inputValues,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    const std::function<ExprHandle(const ExprHandle&, const ExprHandle&)>&
+        innerExpr);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/quantization.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/quantization.h
new file mode 100644
index 0000000000000000000000000000000000000000..9ebd11bc633c053ad5a3ff6df24f56f0f95bdc35
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/quantization.h
@@ -0,0 +1,154 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/kernel.h>
+
+namespace torch::jit::tensorexpr {
+
+TORCH_API ExprHandle quantizePerTensorQParamFromArg(ArgValue arg);
+
+TORCH_API double immQScale(const BufHandle& qx);
+
+TORCH_API int64_t immQZero(const BufHandle& qx);
+
+TORCH_API ScalarType immQDType(const BufHandle& qx);
+
+TORCH_API bool isQuantized(const BufHandle& qx);
+
+TORCH_API Tensor computeQuantizePerTensor(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizePerTensorExternalCall(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedConv1d(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedConv2dPrepack(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedConv2d(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedConv2dRelu(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedLinear(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedLinearRelu(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedAdd(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+Tensor computeQuantizedAddExternalCall(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedMul(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedMulScalar(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedCat(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedRelu(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeDequantize(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeDequantizeExternalCall(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeUpsampleNearest2d(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeUpsampleNearest2dExternalCall(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+TORCH_API Tensor computeQuantizedSigmoidExternalCall(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device /*unused*/);
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/reduction.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/reduction.h
new file mode 100644
index 0000000000000000000000000000000000000000..c9e3cb67920a3984118cea1c7ddb1837d4080cec
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/reduction.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/kernel.h>
+
+namespace torch::jit::tensorexpr {
+
+TORCH_API Tensor computeSum(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+TORCH_API Tensor computeMean(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+TORCH_API Tensor computeAdaptiveAvgPool2d(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+Tensor computeMax(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    const std::optional<ScalarType>& outputType,
+    at::Device device);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/softmax.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/softmax.h
new file mode 100644
index 0000000000000000000000000000000000000000..675a6c0bc795913bc588be6084aad749eb4bf153
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/operators/softmax.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/kernel.h>
+
+namespace torch::jit::tensorexpr {
+
+Tensor computeSoftmax(
+    const std::vector<ArgValue>& inputs,
+    const std::vector<ExprHandle>& outputShape,
+    const std::vector<ExprHandle>& outputStrides,
+    bool log_softmax);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/reduction.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/reduction.h
new file mode 100644
index 0000000000000000000000000000000000000000..9faee10a839d2b3aa0ecd62724dd899e2f9f284a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/reduction.h
@@ -0,0 +1,311 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/tensorexpr/expr.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_printer.h>
+#include <torch/csrc/jit/tensorexpr/stmt.h>
+#include <torch/csrc/jit/tensorexpr/types.h>
+
+#include <functional>
+#include <utility>
+#include <vector>
+
+namespace torch::jit::tensorexpr {
+
+using ParameterList = const std::vector<VarHandle>;
+using ReduceInteraction = std::function<ExprHandle(ExprHandle, ExprHandle)>;
+
+// A Reducer is a user interface describing a particular reduction
+// operation. It has three components: An initialization value, a way of
+// interacting each value with the accumulation, and a method for obtaining the
+// current value to be reduced. It is materialized into a ReduceOp when loop
+// variables are known.
+class TORCH_API Reducer {
+ public:
+  Reducer(ExprHandle init, ReduceInteraction& interaction)
+      : init_(init.node()), interaction_(interaction) {}
+
+  template <typename RI>
+  Reducer(ExprHandle init, RI interaction)
+      : init_(init.node()), interaction_(std::move(interaction)) {}
+
+  ExprPtr initializer() const {
+    return init_;
+  }
+
+  ExprHandle operator()(
+      const BufHandle& result_buf,
+      ExprHandle body,
+      const std::vector<ExprHandle>& output,
+      const std::vector<VarHandle>& inner) const;
+
+  ReduceOpPtr operator()(
+      const BufPtr& result_buf,
+      ExprPtr body,
+      const std::vector<ExprPtr>& output,
+      const std::vector<VarPtr>& inner) const;
+
+  ExprHandle operator()(
+      const BufHandle& result_buf,
+      BufHandle acc_buf,
+      const ExprHandle& body,
+      const std::vector<ExprHandle>& output,
+      const std::vector<VarHandle>& inner) const;
+
+  // Polymorphic handling of Body functions with a variety of parameters.
+  static ExprHandle getReduceBody(
+      const std::function<ExprHandle(ParameterList&)>& func,
+      const std::vector<VarHandle>& vars) {
+    return func(vars);
+  }
+
+  static ExprHandle getReduceBody(
+      const std::function<ExprHandle(const VarHandle&)>& func,
+      const std::vector<VarHandle>& vars) {
+    if (vars.size() != 1) {
+      throw malformed_input("mismatch between reduce body and arg size (1)");
+    }
+
+    return func(vars[0]);
+  }
+
+  static ExprHandle getReduceBody(
+      const std::function<ExprHandle(const VarHandle&, const VarHandle&)>& func,
+      const std::vector<VarHandle>& vars) {
+    if (vars.size() != 2) {
+      throw malformed_input("mismatch between reduce body and arg size (2)");
+    }
+    return func(vars[0], vars[1]);
+  }
+
+  static ExprHandle getReduceBody(
+      const std::function<
+          ExprHandle(const VarHandle&, const VarHandle&, const VarHandle&)>&
+          func,
+      const std::vector<VarHandle>& vars) {
+    if (vars.size() != 3) {
+      throw malformed_input("mismatch between reduce body and arg size (3)");
+    }
+    return func(vars[0], vars[1], vars[2]);
+  }
+
+  static ExprHandle getReduceBody(
+      const std::function<ExprHandle(
+          const VarHandle&,
+          const VarHandle&,
+          const VarHandle&,
+          const VarHandle&)>& func,
+      const std::vector<VarHandle>& vars) {
+    if (vars.size() != 4) {
+      throw malformed_input("mismatch between reduce body and arg size (4)");
+    }
+    return func(vars[0], vars[1], vars[2], vars[3]);
+  }
+
+  // Completes the reduction operator by applying the interaction function to
+  // the accumulation and the body expression.
+  static ExprPtr complete(
+      const BufPtr& accumulator,
+      const ReduceInteraction& interaction,
+      ExprHandle body,
+      const std::vector<ExprPtr>& output_args,
+      const std::vector<VarPtr>& reduce_args) {
+    ExprHandle accum =
+        ExprHandle(alloc<Load>(body.dtype(), accumulator, output_args));
+    auto e = interaction(std::move(accum), std::move(body));
+    return e.node();
+  }
+  static ExprHandle complete(
+      const BufHandle& accumulator,
+      const ReduceInteraction& interaction,
+      ExprHandle body,
+      const std::vector<ExprHandle>& output_args,
+      const std::vector<VarHandle>& reduce_args) {
+    ExprHandle accum = Load::make(body.dtype(), accumulator, output_args);
+    auto e = interaction(std::move(accum), std::move(body));
+    return e;
+  }
+
+ private:
+  ExprPtr init_;
+  ReduceInteraction interaction_;
+};
+
+// An expression representing a Reduction operation (e.g. Sum, Max) broken into
+// it's component parts: initialization, accumulation var, acquisition of value
+// to be reduced and interaction.
+//
+// This is intended to be expanded in the loopnest and not make it to codegen.
+class TORCH_API ReduceOp : public ExprNode<ReduceOp> {
+ public:
+  ReduceOp(
+      const ExprPtr& body,
+      std::vector<VarPtr> reduce_args,
+      Reducer reducer)
+      : ExprNodeBase(body->dtype()),
+        body_(body),
+        reduce_args_(std::move(reduce_args)),
+        reducer_(std::move(reducer)) {
+    result_buf_ = nullptr;
+    acc_buf_ = nullptr;
+    ri_operand_ = nullptr;
+  }
+
+  ReduceOp(
+      const ExprPtr& body,
+      std::vector<VarPtr> reduce_args,
+      BufPtr result_buf,
+      BufPtr acc_buf,
+      ExprPtr ri_operand,
+      Reducer reducer)
+      : ExprNodeBase(body->dtype()),
+        body_(body),
+        reduce_args_(std::move(reduce_args)),
+        result_buf_(std::move(result_buf)),
+        acc_buf_(std::move(acc_buf)),
+        ri_operand_(std::move(ri_operand)),
+        reducer_(std::move(reducer)) {}
+
+  static ExprHandle make(
+      ExprHandle body,
+      const std::vector<VarHandle>& reduce_args,
+      const Reducer& reducer);
+
+  static ExprHandle make(
+      ExprHandle body,
+      const std::vector<VarHandle>& reduce_args,
+      BufHandle result_buf,
+      BufHandle acc_buf,
+      ExprHandle ri_operand,
+      const Reducer& reducer);
+
+  // return the body expression which obtains the value to be reduced.
+  ExprPtr body() const {
+    return body_;
+  }
+
+  // Returns the original Reducer factory that can create ReduceOps.
+  const Reducer& reducer() const {
+    return reducer_;
+  }
+
+  // returns variables associated with the axes of reduction.
+  const std::vector<VarPtr>& reduce_args() const {
+    return reduce_args_;
+  }
+
+  void setAccBuf(BufHandle acc_buf) {
+    acc_buf_ = acc_buf.node();
+  }
+  BufPtr getAccBuf() {
+    return acc_buf_;
+  }
+
+  void setResultBuf(BufHandle buf) {
+    result_buf_ = buf.node();
+  }
+  BufPtr getResultBuf() {
+    return result_buf_;
+  }
+
+  void setRiOperand(ExprHandle ri_operand) {
+    ri_operand_ = ri_operand.node();
+  }
+  ExprPtr getRiOperand() {
+    return ri_operand_;
+  }
+
+ private:
+  // body_ = reducer_->interaction_(result_buf_, ri_operand_)
+  ExprPtr body_;
+  std::vector<VarPtr> reduce_args_;
+
+  BufPtr result_buf_;
+  BufPtr acc_buf_;
+  ExprPtr ri_operand_;
+
+  const Reducer reducer_;
+};
+
+class Sum : public Reducer {
+ public:
+  Sum()
+      : Reducer(ExprHandle(0), [](const ExprHandle& a, const ExprHandle& b) {
+          return a + b;
+        }) {}
+};
+
+inline ExprHandle maximumVal(ScalarType type) {
+  switch (type) {
+#define MAX_BY_TYPE_CASE(Type, Name) \
+  case ScalarType::Name:             \
+    return ExprHandle(std::numeric_limits<Type>::max());
+    AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, MAX_BY_TYPE_CASE)
+#undef MAX_BY_TYPE_CASE
+    default:
+      throw unsupported_dtype();
+  }
+  return ExprHandle();
+}
+
+inline ExprHandle minimumVal(ScalarType type) {
+  switch (type) {
+#define MAX_BY_TYPE_CASE(Type, Name) \
+  case ScalarType::Name:             \
+    return ExprHandle(std::numeric_limits<Type>::min());
+    AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, MAX_BY_TYPE_CASE)
+#undef MAX_BY_TYPE_CASE
+    default:
+      throw unsupported_dtype();
+  }
+}
+
+class Maximum : public Reducer {
+ public:
+  // TODO possible to remove this arg by deferring the init value until we
+  // know the dtype of the body.
+  Maximum(Dtype dtype)
+      : Reducer(
+            minimumVal(dtype.scalar_type()),
+            [](const ExprHandle& a, const ExprHandle& b) {
+              return Max::make(a, b, true);
+            }) {}
+  Maximum(ExprHandle initializer)
+      : Reducer(
+            std::move(initializer),
+            [](const ExprHandle& a, const ExprHandle& b) {
+              return Max::make(a, b, true);
+            }) {}
+};
+
+class Minimum : public Reducer {
+ public:
+  Minimum(Dtype dtype)
+      : Reducer(
+            maximumVal(dtype.scalar_type()),
+            [](const ExprHandle& a, const ExprHandle& b) {
+              return Min::make(a, b, true);
+            }) {}
+  Minimum(const ExprHandle& initializer)
+      : Reducer(initializer, [](const ExprHandle& a, const ExprHandle& b) {
+          return Min::make(a, b, true);
+        }) {}
+};
+
+class ReductionExpander : public IRMutator {
+ public:
+  StmtPtr expand(const StmtPtr& s) {
+    return s->accept_mutator(this);
+  }
+
+  ExprPtr mutate(const ReduceOpPtr& v) override {
+    return v->body();
+  }
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/registerizer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/registerizer.h
new file mode 100644
index 0000000000000000000000000000000000000000..29236965a87abe2217ca1727d52b2133227e7b77
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/registerizer.h
@@ -0,0 +1,431 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/ScalarType.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/Export.h>
+
+#include <torch/csrc/jit/tensorexpr/hash_provider.h>
+#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
+#include <torch/csrc/jit/tensorexpr/ir_simplifier.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+
+#include <utility>
+#include <vector>
+
+namespace torch::jit::tensorexpr {
+namespace registerizer {
+
+/* The Registerizer performs scalar replacement by looking for common Stores and
+Loads to a single item in a buffer and replacing them with a local temporary
+scalar which is cheaper to write.
+
+For example it can replace:
+
+{
+  A[0] = 0;
+  for(const auto x : c10::irange(10)) {
+    A[0] = (A[0]) + x;
+  }
+}
+
+with:
+
+{
+  int A_ = 0;
+  for(const auto x : c10::irange(10)) {
+    A_ = x + A_;
+  }
+  A[0] = A_;
+}
+
+This is particularly useful on GPUs when parallelizing, since after replacing
+loops with metavars we have a lot of accesses like this. */
+
+class Scope;
+
+/*  Holds analysis information about accesses to a specific range of a
+ buffer, including the number of loads and stores and the lowest common parent
+ Block.
+ */
+class AccessInfo {
+ public:
+  AccessInfo() = default;
+  AccessInfo(
+      SimplifierHashType h,
+      BufPtr b,
+      std::vector<ExprPtr> i,
+      size_t accessOrder)
+      : hash_(h),
+        buf_(std::move(b)),
+        indices_(std::move(i)),
+        store_cost_(alloc<IntImm>(0)),
+        load_cost_(alloc<IntImm>(0)),
+        accessOrder_(accessOrder) {}
+
+  // Adds a Store to this access, which is in the provided scope.
+  void addStore(const StorePtr& store, const std::shared_ptr<Scope>& scope);
+
+  // Adds a Load to this access, which occurs in the usage Stmt in the provided
+  // scope.
+  void addLoad(
+      const LoadPtr& load,
+      const std::shared_ptr<Scope>& scope,
+      const StmtPtr& usage);
+
+  // Merge another AccessInfo into this one.
+  void merge(const std::shared_ptr<AccessInfo>& other);
+
+  // Returns true if the other AccessInfo's bounds may overlap this one.
+  bool overlaps(const std::shared_ptr<AccessInfo>& other);
+
+  // Returns true if the indices of this access depend on the provided Var.
+  bool dependsOnVar(const VarPtr& v);
+
+  // Clone this AccessInfo, and set this as the new accesses' hiddenAccess.
+  static std::shared_ptr<AccessInfo> cloneWithHiddenInfo(
+      const std::shared_ptr<AccessInfo>& orig);
+
+  // print for debugging.
+  void print() const;
+
+  SimplifierHashType hash() const {
+    return hash_;
+  }
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  const std::vector<ExprPtr>& indices() const {
+    return indices_;
+  }
+
+  BlockPtr block() const {
+    return block_;
+  }
+
+  void setEnclosingBlock(BlockPtr b) {
+    block_ = std::move(b);
+  }
+
+  StmtPtr first_usage() const {
+    return first_usage_;
+  }
+  StmtPtr last_usage() const {
+    return last_usage_;
+  }
+
+  void setUsageMarks(StmtPtr first, StmtPtr last) {
+    first_usage_ = std::move(first);
+    last_usage_ = std::move(last);
+  }
+
+  bool firstUsageOverlapped() const {
+    return firstUsageOverlapped_;
+  }
+
+  ExprPtr store_cost() const {
+    return store_cost_;
+  }
+
+  ExprPtr load_cost() const {
+    return load_cost_;
+  }
+
+  const std::vector<StorePtr>& stores() const {
+    return stores_;
+  }
+
+  const std::vector<LoadPtr>& loads() const {
+    return loads_;
+  }
+
+  void hoistCosts(const ExprPtr& extent) {
+    store_cost_ = IRSimplifier::simplify(alloc<Mul>(store_cost_, extent));
+    load_cost_ = IRSimplifier::simplify(alloc<Mul>(load_cost_, extent));
+  }
+
+  size_t conditionId() const {
+    return conditionId_;
+  }
+
+  void setConditionId(size_t c) {
+    conditionId_ = c;
+  }
+
+  size_t accessOrder() const {
+    return accessOrder_;
+  }
+
+  std::shared_ptr<AccessInfo> hiddenAccess() const {
+    return hiddenAccess_;
+  }
+
+  // Holds state relating to the scalar variable we will insert to replace some
+  // number of loads and stores.
+  struct ScalarReplacement {
+    VarPtr var{nullptr};
+    BufPtr var_wrapper{nullptr};
+    LetPtr initializer{nullptr};
+  };
+
+  ScalarReplacement& replacement() {
+    return replacement_;
+  }
+
+ private:
+  SimplifierHashType hash_;
+  BufPtr buf_;
+  std::vector<ExprPtr> indices_;
+  BlockPtr block_{nullptr};
+
+  StmtPtr first_usage_{nullptr};
+  StmtPtr last_usage_{nullptr};
+
+  // Whether or not this access is overlapped in the first Stmt it appears. This
+  // means we cannot use it's first Store as the initializer.
+  bool firstUsageOverlapped_{false};
+
+  // The cost in real ops that this access represents, to enable
+  // filtering accesses that won't save any loads or stores.
+  ExprPtr store_cost_;
+  ExprPtr load_cost_;
+
+  // The actual Stores and Loads which represent this access.
+  // Be careful with these, any mutator will invalidate these pointers.
+  std::vector<StorePtr> stores_;
+  std::vector<LoadPtr> loads_;
+
+  // An identifier representing the conditional block, if any, this access
+  // depends on.
+  size_t conditionId_{0};
+
+  // An identifier representing the order this access was first encountered, for
+  // sorting returned results.
+  size_t accessOrder_{0};
+
+  // Sometimes when traversing the tree we need to record what would happen if
+  // we hoisted an access, but sometimes it doesn't work out. This lets us
+  // "undo" some mutation and return to the internal hidden AccessInfo.
+  // It will be removed after any further additions to this AccessInfo.
+  std::shared_ptr<AccessInfo> hiddenAccess_;
+
+  ScalarReplacement replacement_;
+};
+
+using AccessHashMap =
+    std::unordered_map<SimplifierHashType, std::shared_ptr<AccessInfo>>;
+
+// Represents a scope block and holds all accesses contained within it.
+class Scope {
+ public:
+  Scope(BlockPtr b, std::shared_ptr<Scope> parent, size_t conditionId = 0)
+      : block_(std::move(b)),
+        parent_(std::move(parent)),
+        conditionId_(conditionId) {}
+
+  AccessHashMap& getAccessMapByBuf(const BufPtr& b);
+
+  std::unordered_map<BufPtr, AccessHashMap>& openAccesses() {
+    return openAccesses_;
+  }
+
+  std::vector<std::shared_ptr<AccessInfo>>& closedAccesses() {
+    return closedAccesses_;
+  }
+
+  BlockPtr block() const {
+    return block_;
+  }
+
+  std::shared_ptr<Scope> parent() const {
+    return parent_;
+  }
+
+  size_t conditionId() const {
+    return conditionId_;
+  }
+
+  const std::unordered_set<VarPtr>& localVars() const {
+    return localVars_;
+  }
+  void addLocalVar(VarPtr v) {
+    localVars_.insert(std::move(v));
+  }
+
+  void closeAccess(const std::shared_ptr<AccessInfo>& info);
+
+  void filterClosed();
+
+ private:
+  // Map of map to access, narrowing by Buf then by hash(Buf+Indices).
+  // This allows us to find a candidate access easily, and also check for
+  // overlap with other accesses to the same buf. Buf ->
+  //    Hash ->
+  //        Access
+  std::unordered_map<BufPtr, AccessHashMap> openAccesses_;
+  std::vector<std::shared_ptr<AccessInfo>> closedAccesses_;
+
+  // The Block object this scope represents.
+  BlockPtr block_;
+
+  // The enclosing scope object.
+  std::shared_ptr<Scope> parent_;
+
+  // An identifier representing the condition block this scope depends on.
+  size_t conditionId_;
+
+  // A set of variables local to this scope (e.g. loop vars).
+  std::unordered_set<VarPtr> localVars_;
+};
+
+/* Analyzes the graph and collects accesses to the same symbolic tensor element
+ * which can be replaced by a single local scalar.
+ *
+ * This works by recursively walking the tree in postfix order, building sets of
+ * accesses to the same symbolic element by scope and then merging lower scopes
+ * into their enclosing scope.
+ *
+ * It is safe to move two accesses of the same Tensor element to a local scalar
+ * Var if between all usages of the element there are no other Loads or Stores
+ * that may refer to it. In the comments I refer to this as overlapping the
+ * access, or "cutting" the existing AccessInfo. In the case where a candidate
+ * for registerization is cut, it may be possible to finalize the access early
+ * by writing it back to the Tensor and then create a new scalar variable after
+ * the overlapping access is complete. We will attempt to do this when it saves
+ * memory accesses.
+ *
+ * There are a few cases that make this more challenging:
+ *
+ *  - For: Loops change the number of real usages of a buffer by the loop
+ * extent, but only if we can pull the definition and finalization of the scalar
+ * variable out of the loop block.
+ *
+ * - Cond: Conditions complicate lifting scalars out of internal scopes.
+ * Generally we cannot lift an access outside of a conditional scope unless
+ * there is already a reference to that same access at the higher scope, since
+ * we don't know if the condition was guarding an array access not safe at the
+ * higher scope. In the comments I refer to this as the condition "hiding" the
+ * access, and the outer access "unhiding" it.
+ *
+ * - IfThenElse: Same situation as Cond, except since IfThenElse is an Expr
+ * rather than a Stmt we cannot insert the scalar definition or finalizer
+ * within the conditional scope. Accesses inside an IfThenElse can be safely
+ * combined with external accesses but cannot exist completely within.
+ *
+ * - Let: Accesses dependent on local variables via Let Stmts, or loop vars,
+ * cannot be raised outside of the scope of the dependent var.
+ */
+class TORCH_API RegisterizerAnalysis : public IRVisitor {
+ public:
+  RegisterizerAnalysis()
+      : currentScope_(std::make_shared<Scope>(nullptr, nullptr, 0)) {}
+  ~RegisterizerAnalysis() override = default;
+
+  void visit(const ForPtr& v) override;
+
+  void visit(const CondPtr& v) override;
+
+  void visit(const BlockPtr& v) override;
+
+  void visit(const StorePtr& v) override;
+
+  void visit(const LoadPtr& v) override;
+
+  void visit(const IfThenElsePtr& v) override;
+
+  void visit(const LetPtr& v) override;
+
+#define STMT_ON_STACK(Op)                 \
+  void visit(const Op##Ptr& v) override { \
+    stmtStack_.push_front(v);             \
+    IRVisitor::visit(v);                  \
+    stmtStack_.pop_front();               \
+  }
+
+  STMT_ON_STACK(AtomicAdd)
+  STMT_ON_STACK(Allocate)
+  STMT_ON_STACK(Free)
+
+#undef STMT_ON_STACK
+
+  std::vector<std::shared_ptr<AccessInfo>> getCandidates();
+
+ private:
+  void mergeCurrentScopeIntoParent();
+  void mergeHiddenScope(bool allowClosed);
+  void closeAccessIntoScope(
+      const std::shared_ptr<AccessInfo>& info,
+      const std::shared_ptr<Scope>& scope);
+
+  std::unordered_set<size_t> exprConditionals_;
+
+  // A stack of enclosing Stmts for tracking the usage Stmt of Loads.
+  std::deque<StmtPtr> stmtStack_;
+
+  // The current scope being analyzed.
+  std::shared_ptr<Scope> currentScope_;
+
+  HashProvider hasher_;
+
+  size_t conditionId_{0};
+  size_t accessOrder_{0};
+};
+
+/* Replaces each registerizable access with a Scalar variable, including
+ * definition, initializer and finalizer.
+ */
+class TORCH_API RegisterizerReplacer : public IRMutator {
+ public:
+  RegisterizerReplacer(std::vector<std::shared_ptr<AccessInfo>>& vec)
+      : infoSet_(vec) {
+    buildReplacements();
+  }
+
+  ExprPtr mutate(const LoadPtr& v) override;
+
+  StmtPtr mutate(const StorePtr& v) override;
+
+  StmtPtr mutate(const BlockPtr& v) override;
+
+ private:
+  struct ReplacerScope {
+    std::unordered_map<StmtPtr, std::deque<std::shared_ptr<AccessInfo>>>
+        initializerPoints_;
+    std::unordered_map<StmtPtr, std::deque<std::shared_ptr<AccessInfo>>>
+        finalizePoints_;
+  };
+
+  // Creates the various ReplacerScope objects and builds internal maps.
+  void buildReplacements();
+
+  // State relating to the accesses yet to be replaced.
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  std::vector<std::shared_ptr<AccessInfo>>& infoSet_;
+  std::unordered_map<StorePtr, std::shared_ptr<AccessInfo>> storeToAccess_;
+  std::unordered_map<LoadPtr, std::shared_ptr<AccessInfo>> loadToAccess_;
+  std::unordered_map<BlockPtr, ReplacerScope> parentToAccesses_;
+
+  // Holds the set of Stores that should be pulled into an initializer, so they
+  // can be eliminated.
+  std::set<StorePtr> eliminatedIntializers_;
+
+  // Tracks the number of times we've seen each buffer, so we can name the
+  // scalar Vars appropriately.
+  std::unordered_map<BufPtr, unsigned int> bufferAccessCounts_;
+  unsigned int getBufferAccessCount(const BufPtr& b) {
+    return ++bufferAccessCounts_[b];
+  }
+};
+} // namespace registerizer
+
+// Apply scalar replacement to all accesses in s.
+// To produce safe code, this must occur after handling parallelized axes and
+// atomics.
+TORCH_API StmtPtr registerize(StmtPtr s);
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/stmt.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/stmt.h
new file mode 100644
index 0000000000000000000000000000000000000000..331fc954825ce4d2416ec55877203891df1cee59
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/stmt.h
@@ -0,0 +1,1017 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <algorithm>
+#include <list>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include <torch/csrc/jit/tensorexpr/expr.h>
+
+namespace torch::jit::tensorexpr {
+
+// The common base between all statement node.
+class TORCH_API Stmt : public std::enable_shared_from_this<Stmt> {
+ public:
+  Stmt() = default;
+  virtual ~Stmt() = default;
+  virtual void accept(IRVisitor* visitor) = 0;
+  virtual StmtPtr accept_mutator(IRMutator* mutator) = 0;
+
+  StmtPtr get_parent() const {
+    return parent_ ? parent_->getptr() : nullptr;
+  }
+
+  /*
+   * Make a deep copy of the given statement.
+   *
+   * All statements and expressions used in children of the statement are
+   * cloned. Note that the variables are not deep-copied since they are
+   * immutable.
+   */
+  static StmtPtr clone(const StmtPtr& s);
+
+ protected:
+  static void set_parent(const StmtPtr& s, Stmt* new_parent) {
+    s->parent_ = new_parent;
+  }
+  std::shared_ptr<Stmt> getptr() {
+    return shared_from_this();
+  }
+
+ private:
+  Stmt* parent_ = nullptr;
+};
+
+template <class Op>
+class StmtNode : public Stmt {
+ public:
+  using StmtNodeBase = StmtNode<Op>;
+  void accept(IRVisitor* visitor) override {
+    visitor->visit(static_to<Op>(getptr()));
+  }
+  StmtPtr accept_mutator(IRMutator* mutator) override;
+  friend Op;
+
+ private:
+  StmtNode() = default;
+};
+
+template <class Op>
+StmtPtr StmtNode<Op>::accept_mutator(IRMutator* mutator) {
+  return mutator->mutate(static_to<Op>(getptr()));
+}
+
+// Concrete Stmt classes
+class TORCH_API Block : public StmtNode<Block> {
+ public:
+  static BlockPtr make(const std::vector<StmtPtr>& stmts) {
+    std::vector<StmtPtr> valid_stmts;
+    for (auto& stmt : stmts) {
+      if (!stmt) {
+        continue;
+      }
+      valid_stmts.push_back(stmt);
+    }
+    if (valid_stmts.empty()) {
+      return nullptr;
+    }
+    return alloc<Block>(valid_stmts);
+  }
+
+  size_t nstmts() const {
+    return stmts_.size();
+  }
+  bool empty() const {
+    return stmts_.empty();
+  }
+
+  void prepend_stmt(const StmtPtr& s) {
+    if (s->get_parent()) {
+      throw malformed_input("Block prepend Stmt with existing parent", s);
+    }
+
+    stmts_.push_front(s);
+    set_parent(s, this);
+  }
+  void append_stmt(const StmtPtr& s) {
+    if (s->get_parent()) {
+      throw malformed_input("Block append Stmt with existing parent", s);
+    }
+
+    stmts_.push_back(s);
+    set_parent(s, this);
+  }
+
+  void insert_stmt_before(const StmtPtr& s, const StmtPtr& before) {
+    if (s->get_parent()) {
+      throw malformed_input("Block append Stmt with existing parent", s);
+    }
+
+    auto pos = std::find(stmts_.begin(), stmts_.end(), before);
+    if (pos == stmts_.end()) {
+      throw malformed_input(
+          "Inserting after statement that is not in block", s);
+    }
+
+    stmts_.insert(pos, s);
+    set_parent(s, this);
+  }
+
+  void insert_stmt_after(const StmtPtr& s, const StmtPtr& after) {
+    if (s->get_parent()) {
+      throw malformed_input("Block append Stmt with existing parent", s);
+    }
+
+    auto pos = std::find(stmts_.begin(), stmts_.end(), after);
+    if (pos == stmts_.end()) {
+      throw malformed_input(
+          "Inserting after statement that is not in block", s);
+    }
+
+    ++pos;
+
+    stmts_.insert(pos, s);
+    set_parent(s, this);
+  }
+
+  bool replace_stmt(const StmtPtr& old_stmt, const StmtPtr& new_stmt) {
+    if (new_stmt->get_parent()) {
+      throw malformed_input(
+          "Block replace Stmt with existing parent", new_stmt);
+    }
+
+    auto pos = std::find(stmts_.begin(), stmts_.end(), old_stmt);
+    if (pos == stmts_.end()) {
+      return false;
+    }
+    stmts_.insert(pos, new_stmt);
+    stmts_.erase(pos);
+    set_parent(old_stmt, nullptr);
+    set_parent(new_stmt, this);
+    return true;
+  }
+
+  // Creates a new block by cloning `this` block and replacing the given
+  // statement with a new statement. Note that `old_stmt` refers to a statement
+  // in `this` block. If the `old_stmt` is not found, it will return `nullptr`.
+  BlockPtr clone_and_replace(const StmtPtr& old_stmt, const StmtPtr& new_stmt) {
+    if (new_stmt->get_parent()) {
+      throw malformed_input(
+          "Block replace Stmt with existing parent", new_stmt);
+    }
+
+    std::vector<StmtPtr> stmts(stmts_.begin(), stmts_.end());
+    std::vector<StmtPtr> cloned_stmts(stmts.size());
+    bool found = false;
+    for (int i = 0; i < static_cast<int>(stmts.size()); ++i) {
+      if (stmts[i] == old_stmt) {
+        found = true;
+        cloned_stmts[i] = new_stmt;
+      } else {
+        cloned_stmts[i] = Stmt::clone(stmts[i]);
+      }
+    }
+    if (!found) {
+      return nullptr;
+    }
+    return alloc<Block>(cloned_stmts);
+  }
+
+  bool remove_stmt(const StmtPtr& stmt) {
+    auto pos = std::find(stmts_.begin(), stmts_.end(), stmt);
+    if (pos == stmts_.end()) {
+      return false;
+    }
+
+    set_parent(stmt, nullptr);
+    stmts_.erase(pos);
+    return true;
+  }
+
+  std::list<StmtPtr> stmts() const {
+    return stmts_;
+  }
+
+  void clear() {
+    for (const auto& s : stmts_) {
+      set_parent(s, nullptr);
+    }
+    stmts_.clear();
+  }
+
+  void set_stmts(const std::vector<StmtPtr>& stmts) {
+    clear();
+    init(stmts);
+  }
+
+  explicit Block(const std::vector<StmtPtr>& stmts) {
+    init(stmts);
+  }
+
+  typedef std::list<StmtPtr>::iterator iterator;
+  typedef std::list<StmtPtr>::const_iterator const_iterator;
+
+  iterator begin() {
+    return stmts_.begin();
+  }
+
+  const_iterator begin() const {
+    return stmts_.begin();
+  }
+
+  iterator end() {
+    return stmts_.end();
+  }
+
+  const_iterator end() const {
+    return stmts_.end();
+  }
+
+  StmtPtr front() {
+    return stmts_.front();
+  }
+
+  StmtPtr front() const {
+    return stmts_.front();
+  }
+
+  StmtPtr back() {
+    return stmts_.back();
+  }
+
+  StmtPtr back() const {
+    return stmts_.back();
+  }
+
+  void splice(Block::iterator it, const BlockPtr& other) {
+    for (const StmtPtr& s : *other) {
+      set_parent(s, this);
+    }
+
+    stmts_.splice(it, other->stmts_);
+  }
+
+  static BlockPtr getSharedParent(StmtPtr p1, StmtPtr p2) {
+    std::unordered_set<BlockPtr> enclosing;
+
+    StmtPtr p1_p = std::move(p1);
+    while (p1_p) {
+      if (BlockPtr b = to<Block>(p1_p)) {
+        enclosing.insert(b);
+      }
+      p1_p = p1_p->get_parent();
+    }
+
+    StmtPtr p2_p = std::move(p2);
+    while (p2_p) {
+      if (BlockPtr b = to<Block>(p2_p)) {
+        if (enclosing.count(b) != 0) {
+          return b;
+        }
+      }
+      p2_p = p2_p->get_parent();
+    }
+
+    return nullptr;
+  }
+
+  // returns the immediate child containing statement s.
+  StmtPtr getEnclosedRoot(StmtPtr s) const {
+    while (s && s->get_parent().get() != this) {
+      s = s->get_parent();
+    }
+    return s;
+  }
+
+ private:
+  std::list<StmtPtr> stmts_;
+
+  void init(const std::vector<StmtPtr>& stmts) {
+    for (const StmtPtr& s : stmts) {
+      if (!s) {
+        continue;
+      }
+      if (!s->get_parent()) {
+        // If we get here, it's a bug, but we cannot throw an error from a
+        // constructor. But IR verifier would catch this.
+        set_parent(s, this);
+      }
+
+      stmts_.push_back(s);
+    }
+  }
+};
+
+class TORCH_API Store : public StmtNode<Store> {
+ public:
+  VarPtr base_handle() const {
+    return buf_->base_handle();
+  }
+  std::vector<ExprPtr> indices() const {
+    return indices_;
+  }
+  ExprPtr flat_index() const {
+    TORCH_CHECK(indices_.size() == 1, "Indices haven't been flattened.");
+    return indices_[0];
+  }
+  ExprPtr value() const {
+    return value_;
+  }
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  void set_buf(BufPtr buf) {
+    buf_ = std::move(buf);
+  }
+
+  void set_indices(std::vector<ExprPtr> indices) {
+    indices_ = std::move(indices);
+  }
+
+  void set_value(ExprPtr value) {
+    value_ = std::move(value);
+  }
+
+  static StorePtr make(
+      const BufHandle& buf,
+      const std::vector<ExprHandle>& indices,
+      const ExprHandle& value);
+
+  Store(BufPtr buf, std::vector<ExprPtr> indices, ExprPtr value);
+
+ private:
+  BufPtr buf_;
+  std::vector<ExprPtr> indices_;
+  ExprPtr value_;
+};
+
+// Allocate a buffer of given shapes and dtypes and bind it with the given
+// buffer var. The life span is at most through the current program, until it is
+// explicitly freed. An unfreed memory is likely considered an error.
+class TORCH_API Allocate : public StmtNode<Allocate> {
+ public:
+  static AllocatePtr make(const BufHandle& buf_handle) {
+    return alloc<Allocate>(buf_handle.node());
+  }
+
+  VarPtr buffer_var() const {
+    return buf_->base_handle();
+  }
+
+  Dtype dtype() const {
+    return buf_->dtype();
+  }
+
+  const std::vector<ExprPtr> dims() const {
+    return buf_->dims();
+  }
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  void set_buf(BufPtr buf) {
+    buf_ = std::move(buf);
+  }
+
+  explicit Allocate(BufPtr buf) : buf_(std::move(buf)) {}
+
+ private:
+  BufPtr buf_;
+  // TODO: add memory types.
+};
+
+// PlacementAllocate is a variation of the Allocate operator in NNC IR. It does
+// not allocate memory but reuse the memory of another buffer for the given
+// buffer.
+class TORCH_API PlacementAllocate : public StmtNode<PlacementAllocate> {
+ public:
+  static PlacementAllocatePtr make(
+      const BufHandle& buf_handle,
+      const BufHandle& buf_handle_to_reuse) {
+    return alloc<PlacementAllocate>(
+        buf_handle.node(), buf_handle_to_reuse.node());
+  }
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  BufPtr buf_to_reuse() const {
+    return buf_to_reuse_;
+  }
+
+  void set_buf(BufPtr buf) {
+    buf_ = std::move(buf);
+  }
+
+  void set_buf_to_reuse(BufPtr buf) {
+    buf_to_reuse_ = std::move(buf);
+  }
+
+  explicit PlacementAllocate(BufPtr buf, BufPtr buf_to_reuse)
+      : buf_(std::move(buf)), buf_to_reuse_(std::move(buf_to_reuse)) {}
+
+ private:
+  BufPtr buf_;
+  BufPtr buf_to_reuse_;
+};
+
+// Free the specific buffer. It is an error.
+class TORCH_API Free : public StmtNode<Free> {
+ public:
+  static FreePtr make(const BufHandle& buf_handle) {
+    return alloc<Free>(buf_handle.node());
+  }
+
+  VarPtr buffer_var() const {
+    return buf_->base_handle();
+  }
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  void set_buf(BufPtr buf) {
+    buf_ = std::move(buf);
+  }
+
+  explicit Free(BufPtr buf) : buf_(std::move(buf)) {}
+
+ private:
+  BufPtr buf_;
+};
+
+class TORCH_API FreeExt : public StmtNode<FreeExt> {
+ public:
+  static FreeExtPtr make(const std::vector<BufHandle>& bufs);
+
+  std::vector<BufPtr> bufs() const {
+    return bufs_;
+  }
+
+  void set_bufs(std::vector<BufPtr> bufs) {
+    bufs_ = std::move(bufs);
+  }
+
+  explicit FreeExt(std::vector<BufPtr> bufs) : bufs_(std::move(bufs)) {}
+
+ private:
+  std::vector<BufPtr> bufs_;
+};
+
+class TORCH_API Let : public StmtNode<Let> {
+ public:
+  static LetPtr make(const VarHandle& var, const ExprHandle& val) {
+    return alloc<Let>(var.node(), val.node());
+  }
+
+  Let(VarPtr var, ExprPtr val) : var_(std::move(var)), val_(std::move(val)) {}
+
+  VarPtr var() const {
+    return var_;
+  }
+
+  ExprPtr value() const {
+    return val_;
+  }
+
+  void set_var(VarPtr var) {
+    var_ = std::move(var);
+  }
+
+  void set_val(ExprPtr val) {
+    val_ = std::move(val);
+  }
+
+ private:
+  VarPtr var_;
+  ExprPtr val_;
+};
+
+class TORCH_API Cond : public StmtNode<Cond> {
+ public:
+  static CondPtr make(
+      const ExprHandle& condition,
+      const StmtPtr& true_stmt,
+      const StmtPtr& false_stmt) {
+    return alloc<Cond>(condition.node(), true_stmt, false_stmt);
+  }
+
+  ExprPtr condition() const {
+    return condition_;
+  }
+
+  BlockPtr true_stmt() const {
+    return true_stmt_;
+  }
+
+  BlockPtr false_stmt() const {
+    return false_stmt_;
+  }
+
+  void set_condition(ExprPtr condition) {
+    condition_ = std::move(condition);
+  }
+
+  void set_true_stmt(StmtPtr true_stmt) {
+    if (true_stmt) {
+      BlockPtr b = to<Block>(true_stmt);
+      if (!b) {
+        b = alloc<Block>(std::vector<StmtPtr>({std::move(true_stmt)}));
+      }
+      true_stmt_ = b;
+      set_parent(true_stmt_, this);
+    }
+  }
+
+  void set_false_stmt(StmtPtr false_stmt) {
+    if (false_stmt) {
+      BlockPtr b = to<Block>(false_stmt);
+      if (!b) {
+        b = alloc<Block>(std::vector<StmtPtr>({std::move(false_stmt)}));
+      }
+      false_stmt_ = b;
+      set_parent(false_stmt_, this);
+    }
+  }
+
+  Cond(ExprPtr condition, StmtPtr true_stmt, StmtPtr false_stmt)
+      : condition_(std::move(condition)) {
+    set_true_stmt(std::move(true_stmt));
+    set_false_stmt(std::move(false_stmt));
+  }
+
+  CondPtr cloneWithNewBodies(
+      const StmtPtr& true_stmt,
+      const StmtPtr& false_stmt) {
+    return alloc<Cond>(condition_, true_stmt, false_stmt);
+  }
+
+  CondPtr cloneWithNewBody(const StmtPtr& true_stmt) {
+    return alloc<Cond>(condition_, true_stmt, nullptr);
+  }
+
+ private:
+  ExprPtr condition_;
+  BlockPtr true_stmt_ = nullptr;
+  BlockPtr false_stmt_ = nullptr;
+};
+
+class TORCH_API LoopOptions {
+ public:
+  enum {
+    IDX_UNSET = -1,
+    IDX_X = 0,
+    IDX_Y = 1,
+    IDX_Z = 2,
+    IDX_W = 3,
+    IDX_MAX = IDX_W,
+  };
+  // GPU Block Index
+  bool is_gpu_block_index() const {
+    return gpu_block_index_ != IDX_UNSET;
+  }
+
+  int gpu_block_index() const {
+    return gpu_block_index_;
+  }
+
+  std::string gpu_block_index_str() const {
+    if (!is_gpu_block_index()) {
+      throw malformed_input("Has no GPU block index");
+    }
+
+    // NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
+    static constexpr const char* kBlockIndexNames[] = {
+        "blockIdx.x",
+        "blockIdx.y",
+        "blockIdx.z",
+        "blockIdx.w",
+    };
+
+    if (gpu_block_index_ < IDX_X || gpu_block_index_ > IDX_MAX) {
+      throw malformed_input("invalid GPU block index");
+    }
+
+    return kBlockIndexNames[gpu_block_index_];
+  }
+
+  void set_gpu_block_index(int index) {
+    if (index == IDX_UNSET) {
+      gpu_block_index_ = IDX_UNSET;
+    }
+
+    if (is_gpu_thread_index()) {
+      throw std::runtime_error("Cannot set both gpu block and thread index");
+    }
+    if (is_gpu_block_index() && gpu_block_index() != index) {
+      throw std::runtime_error("Cannot set a previously set block index");
+    }
+    gpu_block_index_ = index;
+  }
+
+  // GPU Thread Index
+  bool is_gpu_thread_index() const {
+    return gpu_thread_index() != IDX_UNSET;
+  }
+
+  int gpu_thread_index() const {
+    return gpu_thread_index_;
+  }
+
+  std::string gpu_thread_index_str() const {
+    if (!is_gpu_thread_index()) {
+      throw malformed_input("has no GPU thread index");
+    }
+
+    // NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
+    static constexpr const char* kThreadIndexNames[] = {
+        "threadIdx.x", "threadIdx.y", "threadIdx.z", "threadIdx.w"};
+
+    if (gpu_thread_index_ < IDX_X || gpu_thread_index_ > IDX_MAX) {
+      throw malformed_input("invalid GPU thread index");
+    }
+
+    return kThreadIndexNames[gpu_thread_index_];
+  }
+
+  void set_gpu_thread_index(int index) {
+    if (index == IDX_UNSET) {
+      gpu_thread_index_ = IDX_UNSET;
+    }
+
+    if (is_gpu_block_index()) {
+      throw std::runtime_error("Cannot set both gpu thread and block index");
+    }
+    if (is_gpu_thread_index() && gpu_thread_index() != index) {
+      throw std::runtime_error("Cannot set a previously set thread index");
+    }
+    gpu_thread_index_ = index;
+  }
+
+  void set_parallel() {
+    is_parallel_ = true;
+  }
+
+  bool is_parallel() const {
+    return is_parallel_;
+  }
+
+  std::string ToString() const {
+    if (is_gpu_block_index()) {
+      return gpu_block_index_str();
+    } else if (is_gpu_thread_index()) {
+      return gpu_thread_index_str();
+    } else if (is_parallel()) {
+      return "parallel";
+    }
+    return "";
+  }
+
+  bool isDefault() const {
+    return gpu_block_index_ == IDX_UNSET && gpu_thread_index_ == IDX_UNSET &&
+        !is_parallel_;
+  }
+
+  void set_buffer_mapping(const std::unordered_map<std::string, BufPtr>& map) {
+    map_input_to_tensor_bufs_ = map;
+  }
+
+  std::unordered_map<std::string, BufPtr> get_buffer_mapping() const {
+    return map_input_to_tensor_bufs_;
+  }
+
+ private:
+  int gpu_block_index_{IDX_UNSET};
+  int gpu_thread_index_{IDX_UNSET};
+  bool is_parallel_{false};
+  std::unordered_map<std::string, BufPtr> map_input_to_tensor_bufs_;
+};
+
+class TORCH_API For : public StmtNode<For> {
+ public:
+  VarPtr var() const {
+    return var_;
+  }
+  ExprPtr start() const {
+    return start_;
+  }
+  ExprPtr stop() const {
+    return stop_;
+  }
+  BlockPtr body() const {
+    return body_;
+  }
+  static ForPtr make(
+      const VarHandle& var,
+      const ExprHandle& start,
+      const ExprHandle& stop,
+      const StmtPtr& body) {
+    if (!body) {
+      return nullptr;
+    }
+    return alloc<For>(var.node(), start.node(), stop.node(), body);
+  }
+  static ForPtr make(
+      const VarHandle& var,
+      const ExprHandle& start,
+      const ExprHandle& stop,
+      const StmtPtr& body,
+      const LoopOptions& loop_options) {
+    if (!body) {
+      return nullptr;
+    }
+    return alloc<For>(
+        var.node(), start.node(), stop.node(), body, loop_options);
+  }
+  const LoopOptions loop_options() const {
+    return loop_options_;
+  }
+
+  For(VarPtr var, ExprPtr start, ExprPtr stop, StmtPtr body)
+      : var_(std::move(var)), start_(std::move(start)), stop_(std::move(stop)) {
+    BlockPtr b = to<Block>(body);
+    if (!b) {
+      b = alloc<Block>(std::vector<StmtPtr>({std::move(body)}));
+    }
+    body_ = b;
+    set_parent(body_, this);
+  }
+
+  For(VarPtr var,
+      ExprPtr start,
+      ExprPtr stop,
+      StmtPtr body,
+      LoopOptions loop_options)
+      : var_(std::move(var)),
+        start_(std::move(start)),
+        stop_(std::move(stop)),
+        loop_options_(std::move(loop_options)) {
+    if (!var_) {
+      throw malformed_input("invalid Var in For loop");
+    } else if (!start_) {
+      throw malformed_input("invalid Start in For loop");
+    } else if (!stop_) {
+      throw malformed_input("invalid Stop in For loop");
+    } else if (!body || body->get_parent()) {
+      throw malformed_input("invalid Body in For loop");
+    }
+
+    BlockPtr b = to<Block>(body);
+    if (!b) {
+      b = alloc<Block>(std::vector<StmtPtr>({std::move(body)}));
+    }
+    body_ = b;
+    set_parent(body_, this);
+  }
+
+  void set_gpu_block_index(int block_index) {
+    loop_options_.set_gpu_block_index(block_index);
+  }
+
+  void set_gpu_thread_index(int thread_index) {
+    loop_options_.set_gpu_thread_index(thread_index);
+  }
+
+  void set_parallel() {
+    loop_options_.set_parallel();
+  }
+
+  bool is_parallel() const {
+    return loop_options_.is_parallel();
+  }
+
+  void set_buffer_map(const std::unordered_map<std::string, BufPtr>& map) {
+    loop_options_.set_buffer_mapping(map);
+  }
+
+  ForPtr cloneWithNewBody(const StmtPtr& body) const {
+    return alloc<For>(var_, start_, stop_, body, loop_options_);
+  }
+
+  BlockPtr removeBody() {
+    auto res = body_;
+    set_parent(res, nullptr);
+    body_ = nullptr;
+    return res;
+  }
+
+  void set_body(StmtPtr body) {
+    BlockPtr b = to<Block>(body);
+    if (!b) {
+      b = alloc<Block>(std::vector<StmtPtr>({std::move(body)}));
+    }
+    body_ = b;
+    set_parent(body_, this);
+  }
+
+  void set_start(ExprPtr start) {
+    start_ = std::move(start);
+  }
+
+  void set_stop(ExprPtr stop) {
+    stop_ = std::move(stop);
+  }
+
+  void set_var(VarPtr var) {
+    var_ = std::move(var);
+  }
+
+ private:
+  VarPtr var_;
+  ExprPtr start_;
+  ExprPtr stop_;
+  BlockPtr body_;
+  LoopOptions loop_options_;
+};
+
+// A backend specific IR Node that implements atomic-add.
+// This node could only shows up as an internal with GPU backends.
+// TODO: move to this an internal IR.
+// TODO: make IR nodes extensible.
+class TORCH_API AtomicAdd : public StmtNode<AtomicAdd> {
+ public:
+  AtomicAdd(BufPtr buf, std::vector<ExprPtr> indices, ExprPtr value)
+      : buf_(std::move(buf)),
+        indices_(std::move(indices)),
+        value_(std::move(value)) {}
+
+  VarPtr base_handle() const {
+    return buf_->base_handle();
+  }
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  ExprPtr flat_index() const {
+    TORCH_CHECK(indices_.size() == 1, "Indices haven't been flattened.");
+    return indices_[0];
+  }
+
+  ExprPtr value() const {
+    return value_;
+  }
+
+  const std::vector<ExprPtr>& indices() const {
+    return indices_;
+  }
+
+  void set_buf(BufPtr buf) {
+    buf_ = std::move(buf);
+  }
+
+  void set_indices(std::vector<ExprPtr> indices) {
+    indices_ = std::move(indices);
+  }
+
+  void set_value(ExprPtr value) {
+    value_ = std::move(value);
+  }
+
+ private:
+  BufPtr buf_;
+  std::vector<ExprPtr> indices_;
+  ExprPtr value_;
+};
+
+class TORCH_API SyncThreads : public StmtNode<SyncThreads> {
+ public:
+  SyncThreads() = default;
+};
+
+/*
+ * ExternalCall statement represents a call to an external function that would
+ * compute the contents of the output buffer. An ExternalCall statement consists
+ * of:
+ *   1) output buffer - the buffer that'll be initialized by the call
+ *   2) external function name - a key from the NNC function registry to lookup
+ *      the actual function to call
+ *   3) buffer arguments - the input buffers used by the function
+ *   4) non-buffer arguments - scalar arguments to pass to the function
+ *
+ * An example:
+ *   A = nnc_conv2d(buf_args={Input, Weight, Bias}, args={1})
+ * Here 'A' is the output buffer, "nnc_conv2d" is the function name, the buffer
+ * arguments are 'Input', 'Weight', and 'Bias', and there is a single non-buffer
+ * argument - 1.
+ *
+ * The semantics of the scalar arguments is defined solely by the implementation
+ * of the external function.
+ */
+class TORCH_API ExternalCall : public StmtNode<ExternalCall> {
+ public:
+  static ExternalCallPtr make(
+      BufHandle buf,
+      const std::string& func_name,
+      const std::vector<BufHandle>& buf_args,
+      const std::vector<ExprHandle>& args);
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  std::string func_name() const {
+    return func_name_;
+  }
+
+  std::vector<BufPtr> buf_args() const {
+    return buf_args_;
+  }
+
+  std::vector<ExprPtr> args() const {
+    return args_;
+  }
+
+  void set_buf(BufPtr buf) {
+    buf_ = std::move(buf);
+  }
+
+  void set_buf_args(std::vector<BufPtr> buf_args) {
+    buf_args_ = std::move(buf_args);
+  }
+
+  void set_args(std::vector<ExprPtr> args) {
+    args_ = std::move(args);
+  }
+
+  ExternalCall(
+      BufPtr buf,
+      std::string func_name,
+      std::vector<BufPtr> buf_args,
+      std::vector<ExprPtr> args)
+      : buf_(std::move(buf)),
+        func_name_(std::move(func_name)),
+        buf_args_(std::move(buf_args)),
+        args_(std::move(args)) {}
+
+ private:
+  BufPtr buf_;
+  std::string func_name_;
+  std::vector<BufPtr> buf_args_;
+  std::vector<ExprPtr> args_;
+};
+
+class TORCH_API ExternalCallWithAlloc : public StmtNode<ExternalCallWithAlloc> {
+ public:
+  static ExternalCallWithAllocPtr make(
+      const std::string& func_name,
+      const std::vector<BufHandle>& buf_out_args,
+      const std::vector<BufHandle>& buf_args,
+      const std::vector<ExprHandle>& args);
+
+  std::vector<BufPtr> buf_out_args() const {
+    return buf_out_args_;
+  }
+
+  std::string func_name() const {
+    return func_name_;
+  }
+
+  std::vector<BufPtr> buf_args() const {
+    return buf_args_;
+  }
+
+  std::vector<ExprPtr> args() const {
+    return args_;
+  }
+
+  void set_buf_out_args(std::vector<BufPtr> buf_out_args) {
+    buf_out_args_ = std::move(buf_out_args);
+  }
+
+  void set_buf_args(std::vector<BufPtr> buf_args) {
+    buf_args_ = std::move(buf_args);
+  }
+
+  void set_args(std::vector<ExprPtr> args) {
+    args_ = std::move(args);
+  }
+
+  ExternalCallWithAlloc(
+      std::string func_name,
+      std::vector<BufPtr> buf_out_args,
+      std::vector<BufPtr> buf_args,
+      std::vector<ExprPtr> args)
+      : func_name_(std::move(func_name)),
+        buf_out_args_(std::move(buf_out_args)),
+        buf_args_(std::move(buf_args)),
+        args_(std::move(args)) {}
+
+ private:
+  std::string func_name_;
+  std::vector<BufPtr> buf_out_args_;
+  std::vector<BufPtr> buf_args_;
+  std::vector<ExprPtr> args_;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..54ff410d9de8bad9027a6e8fac7698f50722cc35
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/tensor.h
@@ -0,0 +1,326 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <functional>
+#include <utility>
+#include <vector>
+
+#include <torch/csrc/jit/tensorexpr/expr.h>
+#include <torch/csrc/jit/tensorexpr/reduction.h>
+
+namespace torch::jit::tensorexpr {
+
+class TORCH_API Tensor {
+ public:
+  Tensor(BufPtr buf, const std::vector<VarPtr>& args, const ExprPtr& body)
+      : buf_(std::move(buf)) {
+    stmt_ = constructStmt(args, body, {}, {});
+  }
+  Tensor(BufHandle buf, const std::vector<VarHandle>& args, ExprHandle body)
+      : Tensor(buf.node(), VarHandleVectorToVarVector(args), body.node()) {}
+
+  Tensor(
+      BufPtr buf,
+      const std::vector<VarPtr>& args,
+      const std::vector<ExprPtr>& reduce_dims,
+      const std::vector<VarPtr>& reduce_args,
+      const ExprPtr& body)
+      : buf_(std::move(buf)) {
+    stmt_ = constructStmt(args, body, reduce_dims, reduce_args);
+  }
+  Tensor(
+      BufHandle buf,
+      const std::vector<VarHandle>& args,
+      const std::vector<ExprHandle>& reduce_dims,
+      const std::vector<VarHandle>& reduce_args,
+      ExprHandle body)
+      : Tensor(
+            buf.node(),
+            VarHandleVectorToVarVector(args),
+            ExprHandleVectorToExprVector(reduce_dims),
+            VarHandleVectorToVarVector(reduce_args),
+            body.node()) {}
+
+  Tensor(BufPtr buf, StmtPtr stmt)
+      : buf_(std::move(buf)), stmt_(std::move(stmt)) {}
+
+  BufPtr buf() const {
+    return buf_;
+  }
+
+  StmtPtr stmt() const {
+    return stmt_;
+  }
+
+  template <typename T>
+  inline ExprHandle load(const std::vector<T>& args) const;
+  template <typename... Ts>
+  inline ExprHandle load(const Ts&... ts) const;
+
+ private:
+  StmtPtr constructStmt(
+      const std::vector<VarPtr>& args,
+      const ExprPtr& body,
+      const std::vector<ExprPtr>& reduce_dims,
+      const std::vector<VarPtr>& reduce_args) const;
+
+  BufPtr buf_;
+  StmtPtr stmt_;
+};
+
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const std::function<ExprHandle(const VarHandle&)>& body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::function<ExprHandle(const VarHandle&)>& body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const std::function<ExprHandle(const VarHandle&, const VarHandle&)>&
+        body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::function<ExprHandle(const VarHandle&, const VarHandle&)>&
+        body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const std::function<
+        ExprHandle(const VarHandle&, const VarHandle&, const VarHandle&)>&
+        body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::function<
+        ExprHandle(const VarHandle&, const VarHandle&, const VarHandle&)>&
+        body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const std::function<ExprHandle(
+        const VarHandle&,
+        const VarHandle&,
+        const VarHandle&,
+        const VarHandle&)>& body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::function<ExprHandle(
+        const VarHandle&,
+        const VarHandle&,
+        const VarHandle&,
+        const VarHandle&)>& body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const std::function<ExprHandle(const std::vector<VarHandle>&)>& body_func);
+TORCH_API Tensor Compute(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::function<ExprHandle(const std::vector<VarHandle>&)>& body_func);
+
+inline std::vector<VarHandle> create_index_vars(
+    const std::vector<ExprHandle>& dims) {
+  std::vector<VarHandle> vars;
+  vars.reserve(dims.size());
+  for (const ExprHandle& dim : dims) {
+    vars.emplace_back(alloc<Var>(
+        "i", dim.dtype().scalar_type() == ScalarType::Long ? kLong : kInt));
+  }
+  return vars;
+}
+
+// Handle reductions over a Reducer and a body_func which produces values.
+template <typename InitFunc, typename BodyFunc>
+Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const Reducer& reducer,
+    const InitFunc& init_func,
+    const BodyFunc& body_func,
+    const std::vector<ExprHandle>& reduce_dims) {
+  std::vector<VarHandle> vars = create_index_vars(dims);
+  std::vector<VarHandle> reduce_vars = create_index_vars(reduce_dims);
+
+  // If reduce_vars is empty, then it's not a reduction, but rather a simple
+  // copy
+  if (reduce_vars.empty()) {
+    ExprHandle body = Reducer::getReduceBody(body_func, vars);
+    BufHandle func_result =
+        Buf::make(func_name, dims, body.dtype(), std::nullopt, strides);
+    return Tensor(std::move(func_result), vars, std::move(body));
+  }
+
+  std::vector<VarHandle> all_vars;
+  all_vars.insert(all_vars.end(), vars.begin(), vars.end());
+  all_vars.insert(all_vars.end(), reduce_vars.begin(), reduce_vars.end());
+
+  ExprHandle body = Reducer::getReduceBody(body_func, all_vars);
+  std::vector<ExprHandle> output_args(vars.begin(), vars.end());
+  ExprHandle init_expr = Cast::make(body.dtype(), init_func(vars));
+  BufHandle func_result = Buf::make(func_name, dims, body.dtype(), init_expr);
+
+  ExprHandle reduce_op = reducer(func_result, body, output_args, reduce_vars);
+  if (body.dtype() == kBFloat16) {
+    ExprHandle init_expr_acc = Cast::make(kFloat, init_func(vars));
+    BufHandle func_result_acc =
+        Buf::make(func_name + "_acc", dims, kFloat, init_expr_acc);
+    reduce_op = reducer(
+        func_result,
+        std::move(func_result_acc),
+        body,
+        output_args,
+        reduce_vars);
+  }
+
+  Tensor t = Tensor(
+      std::move(func_result),
+      vars,
+      reduce_dims,
+      reduce_vars,
+      std::move(reduce_op));
+  return t;
+}
+template <typename InitFunc, typename BodyFunc>
+Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const Reducer& reducer,
+    const InitFunc& init_func,
+    const BodyFunc& body_func,
+    const std::vector<ExprHandle>& reduce_dims) {
+  return Reduce<InitFunc, BodyFunc>(
+      func_name,
+      dims,
+      std::nullopt,
+      reducer,
+      init_func,
+      body_func,
+      reduce_dims);
+}
+
+template <typename BodyFunc>
+Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const Reducer& reducer,
+    const BodyFunc& body_func,
+    const std::vector<ExprHandle>& reduce_dims) {
+  return Reduce(
+      func_name,
+      dims,
+      strides,
+      reducer,
+      [&](ParameterList& p [[maybe_unused]]) {
+        return ExprHandle(reducer.initializer());
+      },
+      body_func,
+      reduce_dims);
+}
+template <typename BodyFunc>
+Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const Reducer& reducer,
+    const BodyFunc& body_func,
+    const std::vector<ExprHandle>& reduce_dims) {
+  return Reduce<BodyFunc>(
+      func_name, dims, std::nullopt, reducer, body_func, reduce_dims);
+}
+
+// Overload which allows inline lambda functions for the body_func.
+template <typename BodyFunc>
+Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const Reducer& reducer,
+    const BodyFunc&& body_func,
+    const std::vector<ExprHandle>& reduce_dims) {
+  return Reduce(func_name, dims, strides, reducer, body_func, reduce_dims);
+}
+template <typename BodyFunc>
+Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const Reducer& reducer,
+    const BodyFunc&& body_func,
+    const std::vector<ExprHandle>& reduce_dims) {
+  return Reduce(func_name, dims, std::nullopt, reducer, body_func, reduce_dims);
+}
+
+TORCH_API Tensor Reduce(
+    const std::string& name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const Reducer& reducer,
+    const BufHandle& buffer,
+    const std::vector<ExprHandle>& reduce_dims);
+TORCH_API Tensor Reduce(
+    const std::string& name,
+    const std::vector<ExprHandle>& dims,
+    const Reducer& reducer,
+    const BufHandle& buffer,
+    const std::vector<ExprHandle>& reduce_dims);
+
+// Overload for the common case of all dimensions of a previously Computed
+// Tensor.
+TORCH_API Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const std::optional<std::vector<ExprHandle>>& strides,
+    const Reducer& reducer,
+    const Tensor& tensor,
+    const std::vector<ExprHandle>& reduce_dims);
+TORCH_API Tensor Reduce(
+    const std::string& func_name,
+    const std::vector<ExprHandle>& dims,
+    const Reducer& reducer,
+    const Tensor& tensor,
+    const std::vector<ExprHandle>& reduce_dims);
+
+template <typename... Ts>
+inline ExprHandle Tensor::load(const Ts&... ts) const {
+  std::vector<ExprHandle> params({ExprHandle(ts)...});
+  return Load::make(BufHandle(this->buf()), params);
+}
+
+template <typename T>
+inline ExprHandle Tensor::load(const std::vector<T>& args) const {
+  std::vector<ExprHandle> params(args.begin(), args.end());
+  return Load::make(BufHandle(this->buf()), params);
+}
+
+template <typename... Ts>
+inline ExprHandle BufHandle::load(const Ts&... ts) const {
+  std::vector<ExprHandle> params({ExprHandle(ts)...});
+  return ExprHandle(alloc<Load>(node(), ExprHandleVectorToExprVector(params)));
+}
+
+template <typename T>
+inline ExprHandle BufHandle::load(const std::vector<T>& args) const {
+  std::vector<ExprHandle> params(args.begin(), args.end());
+  return ExprHandle(alloc<Load>(node(), ExprHandleVectorToExprVector(params)));
+}
+
+inline ExprHandle BufHandle::load(const std::vector<ExprHandle>& args) const {
+  return this->template load<ExprHandle>(args);
+}
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/tensorexpr_init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/tensorexpr_init.h
new file mode 100644
index 0000000000000000000000000000000000000000..247b679de43aa244ab779527304da5adb74e9e7b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/tensorexpr_init.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/python/pybind.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::jit {
+// Initialize Python bindings for Tensor Expressions
+void initTensorExprBindings(PyObject* module);
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/types.h
new file mode 100644
index 0000000000000000000000000000000000000000..5615ba08951841b9e8086df7156d70d565ca155c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/types.h
@@ -0,0 +1,163 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstdint>
+#include <iosfwd>
+
+#include <c10/core/ScalarType.h>
+#include <c10/util/Logging.h>
+#include <torch/csrc/Export.h>
+
+#include <torch/csrc/jit/tensorexpr/exceptions.h>
+
+namespace torch::jit::tensorexpr {
+
+using int32 = std::int32_t;
+
+class Dtype;
+TORCH_API std::ostream& operator<<(std::ostream& stream, const Dtype& dtype);
+
+using ScalarType = c10::ScalarType;
+
+enum ElementType {
+  kAllTypes = 0,
+  kIntegralTypes = 1 << 0,
+  kFloatingPointTypes = 1 << 1,
+  kBoolType = 1 << 2,
+  kComplexTypes = 1 << 3,
+  kQintTypes = 1 << 4,
+  kNonComplexOrQintTypes = kIntegralTypes | kBoolType | kFloatingPointTypes,
+};
+
+// Data types for scalar and vector elements.
+class TORCH_API Dtype {
+ public:
+  explicit Dtype(int8_t type)
+      : scalar_type_(static_cast<ScalarType>(type)), lanes_(1) {}
+  explicit Dtype(ScalarType type) : scalar_type_(type), lanes_(1) {}
+  Dtype(int8_t type, int64_t lanes)
+      : scalar_type_(static_cast<ScalarType>(type)), lanes_(lanes) {}
+  Dtype(ScalarType type, int64_t lanes) : scalar_type_(type), lanes_(lanes) {}
+  Dtype(Dtype type, int64_t lanes)
+      : scalar_type_(type.scalar_type_), lanes_(lanes) {
+    if (type.lanes() != 1) {
+      throw malformed_input("dtype lanes dont match");
+    }
+  }
+  int64_t lanes() const {
+    return lanes_;
+  }
+  ScalarType scalar_type() const {
+    return scalar_type_;
+  }
+  Dtype scalar_dtype() const;
+  bool operator==(const Dtype& other) const {
+    return scalar_type_ == other.scalar_type_ && lanes_ == other.lanes_;
+  }
+  bool operator!=(const Dtype& other) const {
+    return !(*this == other);
+  }
+  int byte_size() const;
+  std::string ToCppString() const;
+
+  bool is_integral() const {
+    return c10::isIntegralType(scalar_type_, true);
+  }
+  bool is_floating_point() const {
+    return c10::isFloatingType(scalar_type_);
+  }
+  bool is_signed() const {
+    return c10::isSignedType(scalar_type_);
+  }
+
+  Dtype cloneWithScalarType(ScalarType nt) const {
+    return Dtype(nt, lanes_);
+  }
+
+ private:
+  friend TORCH_API std::ostream& operator<<(
+      std::ostream& stream,
+      const Dtype& dtype);
+  ScalarType scalar_type_;
+  int64_t lanes_; // the width of the element for a vector time
+};
+
+extern TORCH_API Dtype kHandle;
+
+#define NNC_DTYPE_DECLARATION(ctype, name) extern TORCH_API Dtype k##name;
+
+AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, NNC_DTYPE_DECLARATION)
+NNC_DTYPE_DECLARATION(c10::quint8, QUInt8)
+NNC_DTYPE_DECLARATION(c10::qint8, QInt8)
+#undef NNC_DTYPE_DECLARATION
+
+template <typename T>
+TORCH_API Dtype ToDtype();
+
+#define NNC_TODTYPE_DECLARATION(ctype, name) \
+  template <>                                \
+  inline Dtype ToDtype<ctype>() {            \
+    return k##name;                          \
+  }
+AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, NNC_TODTYPE_DECLARATION)
+NNC_TODTYPE_DECLARATION(c10::quint8, QUInt8)
+NNC_TODTYPE_DECLARATION(c10::qint8, QInt8)
+#undef NNC_TODTYPE_DECLARATION
+
+TORCH_API Dtype ToDtype(ScalarType type);
+
+inline Dtype promoteTypes(Dtype a, Dtype b) {
+  if (a.lanes() != b.lanes()) {
+    throw malformed_input("promoting types with different lanes");
+  }
+  return Dtype(
+      static_cast<ScalarType>(c10::promoteTypes(
+          static_cast<c10::ScalarType>(a.scalar_type()),
+          static_cast<c10::ScalarType>(b.scalar_type()))),
+      a.lanes());
+}
+
+inline Dtype BinaryOpDtype(
+    Dtype op1_dtype,
+    Dtype op2_dtype,
+    ScalarType ret_type = ScalarType::Undefined) {
+  if (op1_dtype == op2_dtype) {
+    if (ret_type == ScalarType::Undefined) {
+      return op1_dtype;
+    }
+
+    return ToDtype(ret_type);
+  }
+
+  if (op1_dtype.lanes() != op2_dtype.lanes()) {
+    throw malformed_input("lanes dont match");
+  }
+  int64_t lanes = op1_dtype.lanes();
+
+  Dtype resultType = promoteTypes(op1_dtype, op2_dtype);
+  if (resultType.scalar_type() == ScalarType::Undefined) {
+    throw malformed_input("scalar type doesn't match");
+  }
+
+  if (lanes == 1) {
+    // Use the fixed scalar Dtypes.
+    return ToDtype(resultType.scalar_type());
+  }
+
+  return resultType;
+}
+
+} // namespace torch::jit::tensorexpr
+
+namespace std {
+
+using torch::jit::tensorexpr::Dtype;
+std::string to_string(const Dtype& dtype);
+using torch::jit::tensorexpr::ScalarType;
+std::string to_string(const ScalarType& dtype);
+
+} // namespace std
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/unique_name_manager.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/unique_name_manager.h
new file mode 100644
index 0000000000000000000000000000000000000000..7ef8ec508cbffcf573d68d34f65636b4f04e11b4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/unique_name_manager.h
@@ -0,0 +1,38 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/tensorexpr/fwd_decls.h>
+
+namespace torch::jit::tensorexpr {
+
+class VarHandle;
+class Var;
+
+using VarNameMap = std::unordered_map<VarPtr, std::string>;
+
+// A manager to get unique names from vars.
+// It starts with the name hints of the var and append "_" + $counter until it
+// hits a unique name.
+class TORCH_API UniqueNameManager {
+ public:
+  const std::string& get_unique_name(const VarHandle& v);
+
+  const std::string& get_unique_name(const VarPtr& v);
+
+ private:
+  friend class ScopedVarName;
+  VarNameMap unique_name_mapping_;
+  std::unordered_map<std::string, int> unique_name_count_;
+  std::unordered_set<std::string> all_unique_names_;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/var_substitutor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/var_substitutor.h
new file mode 100644
index 0000000000000000000000000000000000000000..b74c53c9c9cb33318ccc10822058a0b0bea1448f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/tensorexpr/var_substitutor.h
@@ -0,0 +1,66 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <torch/csrc/jit/tensorexpr/analysis.h>
+#include <torch/csrc/jit/tensorexpr/ir.h>
+#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
+#include <torch/csrc/jit/tensorexpr/ir_visitor.h>
+#include <torch/csrc/jit/tensorexpr/reduction.h>
+
+namespace torch::jit::tensorexpr {
+
+using VarMapping = std::vector<std::pair<VarPtr, ExprPtr>>;
+
+class VarSubMutator : public IRMutator {
+ public:
+  VarSubMutator(const VarMapping& var_mapping) {
+    for (auto& entry : var_mapping) {
+      VarPtr key_var = entry.first;
+      ExprPtr value = entry.second;
+      if (!key_var) {
+        throw malformed_input("missing key in VarSubMutator");
+      }
+      var_mapping_[std::move(key_var)] = std::move(value);
+    }
+  }
+
+  ExprPtr mutate(const VarPtr& var) override {
+    auto iter = var_mapping_.find(var);
+    if (iter == var_mapping_.end()) {
+      return var;
+    }
+    return iter->second;
+  }
+
+  ExprPtr mutate(const ReduceOpPtr& var) override {
+    auto body = var->body()->accept_mutator(this);
+    std::vector<VarPtr> new_inner;
+
+    for (const auto& v : var->reduce_args()) {
+      ExprPtr e = v->accept_mutator(this);
+      if (VarPtr new_var = to<Var>(e)) {
+        new_inner.push_back(std::move(new_var));
+      } else {
+        VarFinder varFinder;
+        e->accept(&varFinder);
+        auto varlist = varFinder.vars();
+        new_inner.insert(new_inner.end(), varlist.begin(), varlist.end());
+      }
+    }
+
+    return alloc<ReduceOp>(body, new_inner, var->reducer());
+  }
+
+ private:
+  std::unordered_map<VarPtr, ExprPtr> var_mapping_;
+};
+
+} // namespace torch::jit::tensorexpr
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/catch_utils.hpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/catch_utils.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..a361da0933df1d61e9f4fea6a50120f8a26daffb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/catch_utils.hpp
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#define CATCH_CONFIG_PREFIX_ALL
+#include <catch.hpp>
+
+// CATCH_REQUIRE_THROWS is not defined identically to REQUIRE_THROWS and causes
+// warning; define our own version that doesn't warn.
+#define _CATCH_REQUIRE_THROWS(...) \
+  INTERNAL_CATCH_THROWS(           \
+      "CATCH_REQUIRE_THROWS", Catch::ResultDisposition::Normal, __VA_ARGS__)
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/file_check.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/file_check.h
new file mode 100644
index 0000000000000000000000000000000000000000..d21536d075419323681874e9cb02caa1fe713b5e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/file_check.h
@@ -0,0 +1,84 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <memory>
+#include <string>
+
+namespace torch::jit {
+
+struct Graph;
+
+namespace testing {
+
+struct FileCheckImpl;
+
+struct FileCheck {
+ public:
+  TORCH_API explicit FileCheck();
+  TORCH_API ~FileCheck();
+
+  // Run FileCheck against test string
+  TORCH_API void run(const std::string& test_string);
+
+  // Run FileCheck against dump of graph IR
+  TORCH_API void run(const Graph& graph);
+
+  // Parsing input checks string and run against test string / dump of graph IR
+  TORCH_API void run(
+      const std::string& input_checks_string,
+      const std::string& test_string);
+  TORCH_API void run(
+      const std::string& input_checks_string,
+      const Graph& graph);
+
+  // Checks that the string occurs, starting at the end of the most recent match
+  TORCH_API FileCheck* check(const std::string& str);
+
+  // Checks that the string does not occur between the previous match and next
+  // match. Consecutive check_nots test against the same previous match and next
+  // match
+  TORCH_API FileCheck* check_not(const std::string& str);
+
+  // Checks that the string occurs on the same line as the previous match
+  TORCH_API FileCheck* check_same(const std::string& str);
+
+  // Checks that the string occurs on the line immediately following the
+  // previous match
+  TORCH_API FileCheck* check_next(const std::string& str);
+
+  // Checks that the string occurs count number of times, starting at the end
+  // of the previous match. If exactly is true, checks that there are exactly
+  // count many matches
+  TORCH_API FileCheck* check_count(
+      const std::string& str,
+      size_t count,
+      bool exactly = false);
+
+  // A series of consecutive check_dags get turned into a group of checks
+  // which can appear in any order relative to each other. The checks begin
+  // at the end of the previous match, and the match for the check_dag group
+  // is the minimum match of all individual checks to the maximum match of all
+  // individual checks.
+  TORCH_API FileCheck* check_dag(const std::string& str);
+
+  // Checks that source token is highlighted in str (usually an error message).
+  TORCH_API FileCheck* check_source_highlighted(const std::string& str);
+
+  // Checks that the regex matched string occurs, starting at the end of the
+  // most recent match
+  TORCH_API FileCheck* check_regex(const std::string& str);
+
+  // reset checks
+  TORCH_API void reset();
+
+ private:
+  bool has_run = false;
+  std::unique_ptr<FileCheckImpl> fcImpl;
+};
+} // namespace testing
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/hooks_for_testing.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/hooks_for_testing.h
new file mode 100644
index 0000000000000000000000000000000000000000..5547574692493b6585214a68622523c962308322
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/jit/testing/hooks_for_testing.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/api/compilation_unit.h>
+#include <functional>
+#include <memory>
+
+namespace torch::jit {
+struct Module;
+
+using ModuleHook = std::function<void(Module module)>;
+using FunctionHook = std::function<void(StrongFunctionPtr function)>;
+
+TORCH_API void didFinishEmitModule(Module module);
+TORCH_API void didFinishEmitFunction(StrongFunctionPtr defined);
+TORCH_API void setEmitHooks(ModuleHook for_module, FunctionHook for_fn);
+
+TORCH_API std::pair<ModuleHook, FunctionHook> getEmitHooks();
+
+} // namespace torch::jit
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_data.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_data.h
new file mode 100644
index 0000000000000000000000000000000000000000..75f05b99b69aca3439c9b5be53685dd65622ebcc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_data.h
@@ -0,0 +1,64 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/backend/backend_device.h>
+#include <torch/csrc/lazy/core/shape.h>
+#include <cstring>
+
+namespace torch::lazy {
+
+class TORCH_API BackendData {
+ public:
+  struct Info {
+    /**
+     * Used by Lazy Graph Executor to tag info on BackendData objs
+     * */
+    virtual ~Info() = default;
+  };
+  /**
+   * Represents (Tensor) data stored on a backend device
+   * in its native format.
+   * */
+  using Handle = int64_t;
+
+  BackendData(BackendDevice device, Shape shape)
+      : device_(std::move(device)), shape_(std::move(shape)) {}
+
+  virtual ~BackendData() = default;
+
+  const BackendDevice& device() const {
+    return device_;
+  }
+
+  const Shape& shape() const {
+    return shape_;
+  }
+
+  Info* info() const {
+    return info_.get();
+  }
+
+  std::shared_ptr<Info> SetInfo(std::shared_ptr<Info> info) {
+    std::swap(info, info_);
+    return info;
+  }
+
+  virtual Handle GetHandle() = 0;
+
+  virtual void Assign(const BackendData& data) = 0;
+
+  virtual bool HasValue() const = 0;
+
+ private:
+  BackendDevice device_;
+  Shape shape_;
+  std::shared_ptr<Info> info_;
+};
+
+using BackendDataPtr = std::shared_ptr<BackendData>;
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_device.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_device.h
new file mode 100644
index 0000000000000000000000000000000000000000..99872dc5b1bb353e64761a4f603d0d8e52ebd11c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_device.h
@@ -0,0 +1,105 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <ostream>
+#include <string>
+
+#include <ATen/Tensor.h>
+#include <c10/macros/Export.h>
+#include <c10/util/Deprecated.h>
+#include <optional>
+
+namespace c10 {
+struct Device;
+}
+
+namespace torch::lazy {
+
+// Backend should extend it and define their own supported hardware types.
+struct TORCH_API BackendDeviceType {
+  int8_t type{(int8_t)at::kCPU};
+  // Note: previous default value was '0', which actually maps to at::kCPU, at
+  // least now it is explicit, we may want to make default/undefined semantics
+  // more clear though
+  BackendDeviceType() = default;
+  BackendDeviceType(int8_t type) : type(type) {}
+
+  virtual ~BackendDeviceType() = default;
+  virtual std::string toString() const {
+    return "Unknown";
+  }
+};
+
+class TORCH_API BackendDevice {
+ public:
+  // The default constructor will set both the device type and ordinal
+  // to backend specific defaults.
+  BackendDevice();
+  BackendDevice(std::shared_ptr<BackendDeviceType>&& type, int64_t ordinal);
+
+  int8_t type() const;
+  int64_t ordinal() const {
+    return ordinal_;
+  }
+
+  bool operator==(const BackendDevice& other) const {
+    return compare(other) == 0;
+  }
+  bool operator!=(const BackendDevice& other) const {
+    return compare(other) != 0;
+  }
+  bool operator<(const BackendDevice& rhs) const {
+    return compare(rhs) < 0;
+  }
+
+  std::string toString() const;
+
+ private:
+  int compare(const BackendDevice& rhs) const;
+
+  // Use shared_ptr instead of unique_ptr so that BackendDevice can be copied.
+  std::shared_ptr<BackendDeviceType> type_;
+  int64_t ordinal_;
+};
+
+TORCH_API std::ostream& operator<<(
+    std::ostream& os,
+    const BackendDevice& device);
+
+// Helpers for converting a c10::Device to BackendDevice and vice versa.
+TORCH_API BackendDevice atenDeviceToBackendDevice(const c10::Device& device);
+TORCH_API c10::Device backendDeviceToAtenDevice(const BackendDevice& device);
+
+// Tries to extract the backend device out of the lazy tensor. Returns nullopt
+// if the input is not a lazy tensor.
+TORCH_API std::optional<BackendDevice> GetBackendDevice(
+    const at::ITensorListRef tensors);
+TORCH_API std::optional<BackendDevice> GetBackendDevice(
+    const at::TensorList tensors);
+TORCH_API std::optional<BackendDevice> GetBackendDevice(
+    const at::Tensor& tensor);
+TORCH_API std::optional<BackendDevice> GetBackendDevice(
+    const std::optional<c10::Device>& device);
+
+// For variadic template.
+TORCH_API std::optional<BackendDevice> GetBackendDevice();
+
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Winfinite-recursion")
+template <typename T, typename... Args>
+std::optional<BackendDevice> GetBackendDevice(
+    const T& tensor,
+    const Args&... forward_tensors) {
+  auto optional_device = GetBackendDevice(tensor);
+  if (optional_device) {
+    return optional_device;
+  }
+  return GetBackendDevice(forward_tensors...);
+}
+C10_DIAGNOSTIC_POP()
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_interface.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_interface.h
new file mode 100644
index 0000000000000000000000000000000000000000..9f885d4d4c71618d6c824f197ef109b0c0c76dfc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/backend_interface.h
@@ -0,0 +1,160 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Tensor.h>
+#include <torch/csrc/lazy/backend/backend_data.h>
+#include <torch/csrc/lazy/backend/backend_device.h>
+#include <torch/csrc/lazy/backend/lowering_context.h>
+#include <torch/csrc/lazy/core/lazy_graph_executor.h>
+#include <torch/csrc/lazy/core/shape.h>
+#include <torch/csrc/lazy/core/tensor.h>
+
+namespace torch::lazy {
+
+struct IrBuilder;
+
+/**
+ * Work in progress- don't treat this as a stable interface yet!
+ */
+class TORCH_API BackendImplInterface {
+ public:
+  virtual ~BackendImplInterface() = default;
+
+  /**
+   * Initialization/Teardown
+   * */
+  // No-op by default. Allows custom functionality to be exposed through
+  // extension bindings.
+  virtual void InitializeAtenBindings() const {}
+
+  virtual void PrepareToExit() const = 0;
+
+  /**
+   * Configuration
+   * */
+
+  virtual void SetRngSeed(size_t seed) const = 0;
+
+  /**
+   * IR Tracing
+   * */
+
+  virtual const IrBuilder* GetIrBuilder() const = 0;
+
+  /**
+   * Data Transfer
+   * */
+
+  virtual BackendDataPtr MakeComputationDataFromTensor(
+      const at::Tensor& tensor,
+      const Shape& shape,
+      const BackendDevice& device) const = 0;
+  virtual BackendDataPtr MakeComputationDataFromScalar(
+      const at::Scalar& scalar,
+      const torch::lazy::BackendDevice& device) const = 0;
+  virtual BackendDataPtr CreateDataPlaceholder(
+      const BackendDevice& device,
+      const Shape& shape) const = 0;
+
+  // Gets backend data if the node is a device data node. Otherwise returns
+  // nullptr
+  virtual BackendDataPtr GetComputationDataFromNode(const Node*) const = 0;
+
+  virtual at::Tensor MakeTensorFromComputationData(
+      const BackendDataPtr data,
+      std::optional<at::ScalarType> logical_scalar_type) const = 0;
+
+  /**
+   * Lowering, Compilation, Execution
+   * */
+
+  virtual std::unique_ptr<LoweringContext> CreateLoweringContext(
+      const std::string& name,
+      BackendDevice device,
+      c10::ArrayRef<const torch::lazy::Node*> post_order,
+      Util::EmissionMap emit_status) const = 0;
+
+  virtual std::unique_ptr<LoweringContext> CreateLoweringContext(
+      const std::string& name,
+      BackendDevice device) const = 0;
+
+  // TODO(whc) need to keep this?
+  virtual std::vector<std::string> GetCompilationDevices(
+      const std::string& device,
+      c10::ArrayRef<std::string> devices) const = 0;
+
+  virtual std::vector<ComputationPtr> Compile(
+      std::vector<ComputationPtr> instances) const = 0;
+
+  virtual std::vector<BackendDataPtr> ExecuteComputation(
+      torch::lazy::ComputationPtr computation,
+      c10::ArrayRef<BackendDataPtr> arguments,
+      const BackendDevice& device) const = 0;
+
+  /**
+   * Device Configuration
+   * */
+
+  // Set or get the default device type.
+  // For backends used with virtual c10::Devices, this configures what real
+  // device type the backend should use, and matters if the backend supports
+  // more than one type of real device.
+  virtual std::shared_ptr<BackendDeviceType> GetDefaultDeviceType() const = 0;
+  virtual void SetDefaultDeviceType(int8_t type) = 0;
+
+  // Set or get the default device ordinal.
+  // For backends that supports multi-device, this configures what the
+  // default device the backend should use.
+  virtual int64_t GetDefaultDeviceOrdinal() const = 0;
+  virtual void SetDefaultDeviceOrdinal(int64_t) = 0;
+
+  // Specify which aten device should be used for eager fallback
+  // may change depending on current 'Default' DeviceType
+  virtual at::DeviceType EagerFallbackDeviceType() const = 0;
+
+  // Query all available backend devices
+  virtual std::vector<BackendDevice> GetBackendDevices() const = 0;
+
+  virtual std::string CreateMetricReport() const {
+    return "";
+  }
+
+  // Map a particular c10:: device to a concrete backend device
+  // Note:: c10:: devices may be virtual or concrete.  xla:: and lazy:: are
+  // virtual devices, meaning they may map to a gpu, tpu, etc. behind the
+  // scenes. In the future, non-virtual c10:: devices may also use lazy tensors
+  // through a mode, in which case these APIs should still work, but should be
+  // identity mappings.
+  virtual BackendDevice GetBackendDevice(c10::Device device) const = 0;
+
+  // TODO(whc)
+  // Additional APIs expected for supporting distributed training, to be
+  // designed
+
+  /**
+   * Debug/Metrics
+   * */
+
+  //   virtual std::map<std::string, Metric> GetMetrics() const = 0;
+
+  //   virtual MemoryInfo GetMemoryInfo(const std::string& device) = 0;
+
+  virtual std::string GetComputationBackendText(
+      const ComputationPtr computation) const = 0;
+};
+
+class TORCH_API BackendRegistrar {
+ public:
+  BackendRegistrar(const BackendImplInterface* backend_impl_interface);
+};
+
+TORCH_API bool hasBackend();
+TORCH_API const BackendImplInterface* getBackend();
+
+TORCH_API const IrBuilder* getIrBuilder();
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/lowering_context.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/lowering_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..8de72ec167fdd9d27412ba13b8ac143e16d2c0b1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/backend/lowering_context.h
@@ -0,0 +1,115 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <torch/csrc/lazy/backend/backend_data.h>
+#include <torch/csrc/lazy/backend/backend_device.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/ir_util.h>
+
+namespace torch::lazy {
+
+class TORCH_API Computation {
+ public:
+  virtual int parameters_size() const = 0;
+
+  virtual const std::vector<Shape>& parameter_shapes() const = 0;
+
+  virtual const std::vector<std::string>& parameter_names() const = 0;
+
+  virtual const Shape& result_shape() const = 0;
+
+  virtual const std::string to_string() const = 0;
+
+  virtual ~Computation() = default;
+
+  // Indicates whether this computation is being executed inside a mark step
+  // Assume false unless set otherwise
+  bool in_mark_step = false;
+};
+
+using ComputationPtr = std::shared_ptr<Computation>;
+
+// Keeps track of the code generation state.
+class TORCH_API LoweringContext {
+ public:
+  LoweringContext(const std::string& name, BackendDevice device);
+  LoweringContext(
+      const std::string& name,
+      BackendDevice device,
+      c10::ArrayRef<const torch::lazy::Node*> post_order,
+      Util::EmissionMap emit_status);
+
+  virtual ~LoweringContext() = default;
+
+  static std::unique_ptr<LoweringContext> Create(
+      const std::string& name,
+      BackendDevice device,
+      c10::ArrayRef<const torch::lazy::Node*> post_order,
+      Util::EmissionMap emit_status);
+
+  static std::unique_ptr<LoweringContext> Create(
+      const std::string& name,
+      BackendDevice device);
+
+  const BackendDevice& device() const {
+    return device_;
+  }
+
+  // Retrieves the vector holding all the tensors associated with the parameter
+  // instructions which have been created.
+  const std::vector<BackendDataPtr>& GetParametersData() const;
+
+  // Adds a new input/output alias.
+  virtual void SetUpAlias(
+      const std::vector<int64_t>& output_index,
+      int64_t param_number,
+      const std::vector<int64_t>& param_index,
+      bool must_alias = false) {
+    // Dummy default implementation to do nothing.
+  }
+
+  // Check if parameter shape matches result at index.
+  virtual bool CheckResultShape(
+      const BackendDataPtr& parameter_data,
+      size_t result_idx) {
+    // Dummy default implementation to do nothing.
+    return false;
+  }
+
+  // Adds the given output as a component of the result tuple and returns its
+  // assigned position within the tuple.
+  virtual size_t AddResult(const torch::lazy::Output& output) = 0;
+
+  // Associates the given output with the input parameter of the given index and
+  // shape. Only used for the operator-by-operator execution, mostly for
+  // debugging purposes.
+  virtual void AddParameter(
+      const torch::lazy::Output& output,
+      size_t index,
+      const Shape& shape,
+      const std::string& name) = 0;
+
+  // Build the computation capturing all the operations created with the
+  // embedded builder (returned by the builder() API).
+  virtual ComputationPtr Build() = 0;
+
+  size_t GetEmittedNodeCount() const {
+    return emit_status_.size();
+  }
+
+ protected:
+  BackendDevice device_;
+  std::vector<BackendDataPtr> parameters_;
+  std::vector<size_t> parameter_sequence_;
+  Util::EmissionMap emit_status_;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/cache.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/cache.h
new file mode 100644
index 0000000000000000000000000000000000000000..a34161654e64d8277e0af2508360dcaa9aa782b8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/cache.h
@@ -0,0 +1,148 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * Cache utils in this file is adapted from PyTorch/XLA
+ * https://github.com/pytorch/xla/blob/e0e5f937a0ba8d904f9608137dc8c51ba439df2d/third_party/xla_client/cache.h
+ */
+
+#pragma once
+
+#include <functional>
+#include <list>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <utility>
+
+namespace torch::lazy {
+
+// Generic key and object cache with LRU expiration policy. The objects of type
+// T will be stored as std::shared_ptr<T> and taken and returned as such, by the
+// cache API.
+template <
+    typename K,
+    typename T,
+    typename H = std::hash<K>,
+    typename E = std::equal_to<K>>
+class Cache {
+ public:
+  using TypePtr = std::shared_ptr<T>;
+  using Element = std::pair<K, TypePtr>;
+
+  explicit Cache(size_t max_size) : max_size_(max_size) {}
+
+  // Adds an object to the cache, unless it already exists. If the cache grows
+  // beyond the limit set during construction, the oldest used object will be
+  // removed from the cache.
+  TypePtr Add(K key, TypePtr object) {
+    if (!max_size_) {
+      return object;
+    }
+    std::lock_guard<std::mutex> slock(lock_);
+    element_list_.emplace_front(Element(std::move(key), std::move(object)));
+    auto it = element_list_.begin();
+    auto emplace_result = element_map_.emplace(&it->first, it);
+    if (!emplace_result.second) {
+      element_list_.erase(it);
+      DoLRU(emplace_result.first->second);
+    } else if (element_list_.size() > max_size_) {
+      Element* last = &element_list_.back();
+      element_map_.erase(&last->first);
+      element_list_.pop_back();
+    }
+    return emplace_result.first->second->second;
+  }
+
+  // Retrieves the existing object if it exists. If it does, its position in
+  // the LRU list gets moved to the head of the list.
+  // Returns nullptr if no object with the specified key is found within the
+  // cache.
+  TypePtr Get(const K& key) {
+    if (!max_size_) {
+      return nullptr;
+    }
+    std::lock_guard<std::mutex> slock(lock_);
+    auto it = element_map_.find(&key);
+    if (it == element_map_.end()) {
+      return nullptr;
+    }
+    DoLRU(it->second);
+    return it->second->second;
+  }
+
+  TypePtr GetLatest() {
+    std::lock_guard<std::mutex> g(lock_);
+    TORCH_CHECK(!element_list_.empty());
+    return element_list_.front().second;
+  }
+
+  bool Erase(const K& key) {
+    if (!max_size_) {
+      return false;
+    }
+    std::lock_guard<std::mutex> slock(lock_);
+    auto it = element_map_.find(&key);
+    if (it == element_map_.end()) {
+      return false;
+    }
+    auto lit = it->second;
+    element_map_.erase(it);
+    element_list_.erase(lit);
+    return true;
+  }
+
+  void Clear() {
+    if (!max_size_) {
+      return;
+    }
+    std::lock_guard<std::mutex> slock(lock_);
+    element_map_.clear();
+    element_list_.clear();
+  }
+
+  int Numel() const {
+    if (!max_size_) {
+      return 0;
+    }
+    std::lock_guard<std::mutex> g(lock_);
+    TORCH_CHECK(element_map_.size() == element_list_.size());
+    return element_map_.size();
+  }
+
+ private:
+  using ElementList = std::list<Element>;
+
+  struct Hasher {
+    size_t operator()(const K* key) const {
+      return hasher(*key);
+    }
+
+    H hasher;
+  };
+
+  struct Equaler {
+    bool operator()(const K* k1, const K* k2) const {
+      return equaler(*k1, *k2);
+    }
+
+    E equaler;
+  };
+
+  using ElementMap = std::
+      unordered_map<const K*, typename ElementList::iterator, Hasher, Equaler>;
+
+  void DoLRU(typename ElementList::iterator it) {
+    element_list_.splice(element_list_.begin(), element_list_, it);
+  }
+
+  mutable std::mutex lock_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const size_t max_size_ = 0;
+  ElementList element_list_;
+  ElementMap element_map_;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/config.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/config.h
new file mode 100644
index 0000000000000000000000000000000000000000..83ff42b145fa390c600e4f005c29c1993fdbbff9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/config.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/macros/Export.h>
+#include <c10/util/Flags.h>
+
+TORCH_DECLARE_bool(torch_lazy_ir_debug);
+TORCH_DECLARE_bool(torch_lazy_handle_special_scalars);
+TORCH_DECLARE_bool(torch_lazy_all_numbers_special_scalars);
+TORCH_DECLARE_bool(torch_lazy_param_aliasing);
+TORCH_DECLARE_bool(torch_lazy_reuse_ir);
+TORCH_DECLARE_bool(torch_lazy_use_thread_pool);
+TORCH_DECLARE_bool(torch_lazy_enable_device_data_cache);
+
+TORCH_DECLARE_int(torch_lazy_compilation_cache_size);
+TORCH_DECLARE_int(torch_lazy_device_data_cache_size);
+TORCH_DECLARE_int(torch_lazy_io_thread_pool_size);
+TORCH_DECLARE_int(torch_lazy_metrics_samples);
+TORCH_DECLARE_int(torch_lazy_trim_graph_check_frequency);
+TORCH_DECLARE_int(torch_lazy_trim_graph_size);
+
+TORCH_DECLARE_string(torch_lazy_metrics_percentiles);
+
+TORCH_DECLARE_int(torch_lazy_shape_cache_size);
+
+namespace torch::lazy {
+TORCH_API std::string& getLTCForceFallback();
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/debug_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/debug_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..7c3ce3171ce299d10a487dc5f85a9861d7d123c6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/debug_util.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include <torch/csrc/lazy/core/tensor.h>
+
+namespace torch::lazy {
+
+TORCH_API std::function<std::vector<SourceLocation>()>&
+GetPythonFramesFunction();
+
+TORCH_API std::string GetFirstUserFrameInPython();
+
+class TORCH_API DebugUtil {
+ public:
+  enum GraphFormat {
+    kText,
+    kDot,
+    kBackend,
+  };
+
+  static GraphFormat GetDefaultGraphFormat();
+
+  // Dumps the current Python frame and the IR Graph whose roots are the IR
+  // values held at the tensors. If indices is not nullptr, it selects the
+  // indices of the tensors whose graph will be emitted.
+  static std::string GetTensorsGraphInfo(
+      c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors,
+      const std::vector<size_t>* indices,
+      GraphFormat format = GetDefaultGraphFormat());
+
+  // If the environment variable LTC_SAVE_TENSORS_FILE is set to the proper
+  // output path, an instance of the report returned by GetTensorsGraphInfo() is
+  // saved.
+  static void SaveTensorsGraphInfo(
+      const char* name,
+      c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors,
+      const std::vector<size_t>* indices,
+      GraphFormat format = GetDefaultGraphFormat());
+
+  static bool ExperimentEnabled(const std::string& name);
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/dynamic_ir.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/dynamic_ir.h
new file mode 100644
index 0000000000000000000000000000000000000000..d21d6feb0ba58ad5e17f7e92b1e14c4334c3475e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/dynamic_ir.h
@@ -0,0 +1,53 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/symbol.h>
+
+#include <c10/core/ScalarType.h>
+#include <c10/util/Flags.h>
+#include <torch/csrc/lazy/core/hash.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/ir_metadata.h>
+#include <torch/csrc/lazy/ts_backend/ts_node.h>
+
+namespace torch::lazy {
+
+/**
+ * The goal of "dynamic" Nodes is to patch a hole in our tracing.
+ * Previously, if a user called `sizes` on a Tensor, it would leak out
+ * of our tracing system, as `sizes` returns a torch.Size or an int. To
+ * prevent this from happening, we introduce DimensionNode, a new type
+ * of Node that abstracts the operation of getting the dimensions of a
+ * Tensor.
+ *
+ * Consider the following example:
+ * ```
+ * numel = x.shape()[0] * x.shape()[1]
+ * ```
+ *
+ * Here, `x.shape()[i]` will be a SizeNode (subclass of DimensionNode),
+ * and the multiplication of the two SizeNodes will be represented by
+ * a SizeMul (also a subclass of DimensionNode). Through this, we can
+ * prevent `numel` from being represented as a Python int and thus
+ * burned into the Graph.
+ */
+
+class TORCH_API DimensionNode {
+ public:
+  virtual bool isSymbolic() const {
+    return false;
+  }
+  virtual int64_t getDynamicValue() const {
+    TORCH_CHECK(false, "NYI");
+  }
+  virtual int64_t getStaticValue() const {
+    TORCH_CHECK(false, "NYI");
+  }
+  virtual ~DimensionNode() = default;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/hash.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/hash.h
new file mode 100644
index 0000000000000000000000000000000000000000..f224aae7bf62257f83e4e5db89c843a9323777a1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/hash.h
@@ -0,0 +1,247 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * Hash utils in this file is adapted from PyTorch/XLA
+ * https://github.com/pytorch/xla/blob/e0e5f937a0ba8d904f9608137dc8c51ba439df2d/third_party/xla_client/util.h
+ */
+#pragma once
+
+#include <ATen/Tensor.h>
+#include <c10/core/Scalar.h>
+#include <c10/util/int128.h>
+#include <torch/csrc/Export.h>
+#include <cstring>
+#include <set>
+#include <string>
+#include <string_view>
+#include <vector>
+
+namespace torch::lazy {
+
+using size_t = std::size_t;
+
+class TORCH_API hash_t : public c10::uint128 {
+ public:
+  // Switch from typedef hash_t = uint128 to provide explicit casters
+  hash_t(int8_t val) : uint128(static_cast<uint32_t>(val)) {}
+  hash_t(int16_t val) : uint128(static_cast<uint32_t>(val)) {}
+  hash_t(int32_t val) : uint128(static_cast<uint32_t>(val)) {}
+  hash_t(int64_t val) : uint128(static_cast<uint64_t>(val)) {}
+  hash_t(uint32_t val) : uint128(val) {}
+  hash_t(uint64_t val) : uint128(val) {}
+  hash_t(uint128 val) : uint128(val) {}
+  hash_t(uint64_t top, uint64_t bottom) : uint128(top, bottom) {}
+  hash_t() = default;
+};
+
+// Std* functions use 64-bit hash
+size_t TORCH_API StdDataHash(const void* data, size_t size);
+
+size_t TORCH_API StdHashCombine(uintmax_t a, uintmax_t b);
+
+// Other functions are all 128-bit
+hash_t TORCH_API HashBlock(const void* data, size_t n, const hash_t& seed);
+
+hash_t TORCH_API DataHash(const void* data, size_t size);
+
+hash_t TORCH_API HashCombine(const hash_t& a, const hash_t& b);
+
+size_t TORCH_API HashReduce(const hash_t& a);
+
+// Returns a string representation of a hash
+std::string TORCH_API HashToString(const hash_t& a);
+
+struct HashReducer {
+  size_t operator()(const hash_t& value) const {
+    return HashReduce(value);
+  }
+};
+
+static inline hash_t StringHash(const char* data) {
+  return DataHash(data, std::strlen(data));
+}
+
+// Automatic templated implementation for 'arithmetic' types
+template <typename T, std::enable_if_t<std::is_arithmetic_v<T>>* = nullptr>
+hash_t Hash(const T& value) {
+  return DataHash(&value, sizeof(value));
+}
+
+// added because on macos builds the vector<bool> specialization
+// breaks falling through to the templated arithmetic types above
+hash_t TORCH_API Hash(const std::vector<bool>& value);
+
+// Specialized implementations for proprietary types
+static inline hash_t Hash(const c10::ScalarType& value) {
+  return DataHash(&value, sizeof(value));
+}
+
+static inline hash_t Hash(const c10::MemoryFormat& value) {
+  return DataHash(&value, sizeof(value));
+}
+
+static inline hash_t Hash(const c10::DeviceType& value) {
+  return DataHash(&value, sizeof(value));
+}
+
+static inline hash_t Hash(const c10::Device& value) {
+  return HashCombine(Hash(value.type()), Hash(value.index()));
+}
+
+static inline hash_t Hash(const c10::Layout& value) {
+  return DataHash(&value, sizeof(value));
+}
+
+static inline hash_t Hash(const c10::Scalar& value) {
+  switch (value.type()) {
+    case c10::ScalarType::ComplexDouble:
+      return Hash(value.toComplexDouble());
+    case c10::ScalarType::Double:
+      return Hash(value.toDouble());
+    case c10::ScalarType::Long:
+      return Hash(value.toLong());
+    case c10::ScalarType::Bool:
+      return Hash(value.toBool());
+    default:
+      TORCH_INTERNAL_ASSERT(false, "Unknown scalar type.", value.type());
+  }
+}
+
+static inline hash_t TensorHash(const at::Tensor& tensor) {
+  at::Tensor ctensor = tensor.contiguous();
+  int64_t size = ctensor.numel() * ctensor.element_size();
+  switch (ctensor.scalar_type()) {
+    case at::ScalarType::Bool:
+      return DataHash(ctensor.const_data_ptr<bool>(), size);
+    case at::ScalarType::Byte:
+      return DataHash(ctensor.const_data_ptr<uint8_t>(), size);
+    case at::ScalarType::Char:
+      return DataHash(ctensor.const_data_ptr<int8_t>(), size);
+    case at::ScalarType::Short:
+      return DataHash(ctensor.const_data_ptr<int16_t>(), size);
+    case at::ScalarType::Int:
+      return DataHash(ctensor.const_data_ptr<int32_t>(), size);
+    case at::ScalarType::Long:
+      return DataHash(ctensor.const_data_ptr<int64_t>(), size);
+    case at::ScalarType::Float:
+      return DataHash(ctensor.const_data_ptr<float>(), size);
+    case at::ScalarType::Double:
+      return DataHash(ctensor.const_data_ptr<double>(), size);
+    case at::ScalarType::BFloat16:
+      return DataHash(ctensor.const_data_ptr<at::BFloat16>(), size);
+    case at::ScalarType::Half:
+      return DataHash(ctensor.const_data_ptr<at::Half>(), size);
+    case at::ScalarType::ComplexFloat:
+      return DataHash(ctensor.const_data_ptr<c10::complex<float>>(), size);
+    case at::ScalarType::ComplexDouble:
+      return DataHash(ctensor.const_data_ptr<c10::complex<double>>(), size);
+    case at::ScalarType::UInt16:
+      return DataHash(ctensor.const_data_ptr<uint16_t>(), size);
+    case at::ScalarType::UInt32:
+      return DataHash(ctensor.const_data_ptr<uint32_t>(), size);
+    case at::ScalarType::UInt64:
+      return DataHash(ctensor.const_data_ptr<uint64_t>(), size);
+    default:
+      TORCH_INTERNAL_ASSERT(
+          false, "Unsupported scalar type:", ctensor.scalar_type());
+  }
+}
+
+static inline hash_t Hash(const std::string& value) {
+  return DataHash(value.data(), value.size());
+}
+
+static inline hash_t Hash(const std::string_view& value) {
+  return DataHash(value.data(), value.size());
+}
+
+static inline hash_t Hash(const at::Generator& value) {
+  return TensorHash(value.get_state());
+}
+
+// Taken from glibc's implementation of hashing optionals,
+// we want to include a contribution to the hash to distinguish
+// cases where one or another option was null, but we hope it doesn't
+// collide with an actually scalar value.
+//
+// Use an arbitrary randomly-selected 64-bit integer rather than a
+// small constant that we then hash at runtime so we don't have to
+// repeatedly hash a constant at runtime.
+// NOLINTNEXTLINE(*-narrowing-conversions)
+static const int64_t kNullOpt = 0x8655d738f3678dda;
+
+// Hashing for std::optional types contributes to hash
+// for optionals with null value, important to distinguish
+// between <nullopt, non-nullopt> and <non-nullopt, nullopt> cases
+template <typename T>
+hash_t Hash(const std::optional<T>& value) {
+  if (value.has_value()) {
+    return Hash(value.value());
+  } else {
+    return kNullOpt;
+  }
+}
+
+// Hashing of containers
+// Forward declare to allow hashes of vectors of vectors to work.
+template <typename T>
+hash_t ContainerHash(const T& values);
+
+template <typename T>
+hash_t Hash(const std::vector<T>& values) {
+  return ContainerHash(values);
+}
+
+// Need a special case for std::optional<container>?
+template <typename T>
+hash_t Hash(const std::optional<std::vector<T>>& value) {
+  if (value.has_value()) {
+    return ContainerHash(value.value());
+  } else {
+    return kNullOpt;
+  }
+}
+
+template <typename T>
+hash_t Hash(const std::set<T>& values) {
+  return ContainerHash(values);
+}
+
+template <typename T, typename S>
+hash_t Hash(const std::pair<T, S>& values) {
+  return HashCombine(Hash(values.first), Hash(values.second));
+}
+
+static inline hash_t Hash(const hash_t& value) {
+  return value;
+}
+
+template <typename T>
+hash_t Hash(c10::ArrayRef<T> values) {
+  return ContainerHash(values);
+}
+
+template <typename T>
+hash_t ContainerHash(const T& values) {
+  hash_t h(static_cast<uint64_t>(0x85ebca77c2b2ae63));
+  for (const auto& value : values) {
+    h = HashCombine(h, Hash(value));
+  }
+  return h;
+}
+
+// Varargs hashing
+template <typename T = void>
+hash_t MHash() {
+  return hash_t(static_cast<uint64_t>(0x165667b19e3779f9));
+}
+
+template <typename T, typename... Targs>
+hash_t MHash(T value, Targs... Fargs) {
+  return HashCombine(Hash(value), MHash(Fargs...));
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/helpers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/helpers.h
new file mode 100644
index 0000000000000000000000000000000000000000..440e880fcb76b7e60e7a383d6b820ab6a3c7ae34
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/helpers.h
@@ -0,0 +1,75 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/Scalar.h>
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <torch/csrc/lazy/core/permutation_util.h>
+#include <torch/csrc/lazy/core/shape.h>
+#include <torch/csrc/lazy/core/util.h>
+
+#include <complex>
+#include <functional>
+#include <optional>
+#include <tuple>
+#include <vector>
+
+// TODO: Consolidate this file with util.h
+
+namespace torch::lazy {
+
+// Converts an iterable container to a vector of int64's.
+template <typename S>
+static std::vector<int64_t> ToI64Vector(const S& input) {
+  return ToVector<int64_t>(input);
+}
+
+// Creates a set of dimension by dropping the drop_dims ones.
+TORCH_API std::vector<int64_t> DropDimensions(
+    c10::ArrayRef<int64_t> sizes,
+    c10::ArrayRef<int64_t> drop_dims);
+
+// Get the canonical dimension index in the [0, rank) interval. Negative
+// indices are interpreted as follows: -1 is rank-1, -2 is rank-2 etc.
+TORCH_API int64_t GetCanonicalDimensionIndex(int64_t dim, int64_t rank);
+
+// Same as above, for multiple dimensions.
+TORCH_API std::vector<int64_t> GetCanonicalDimensionIndices(
+    c10::ArrayRef<int64_t> dimensions,
+    int64_t rank);
+
+// Returns the canonical position in the dim dimension, handling negative
+// values for the position.
+TORCH_API int64_t GetCanonicalPosition(
+    c10::ArrayRef<int64_t> dimensions,
+    int64_t dim,
+    int64_t pos);
+
+// Creates a transposition from the given input and dimensions.
+TORCH_API std::vector<int64_t> MakeTransposePermutation(
+    int64_t dim0,
+    int64_t dim1,
+    int64_t rank);
+
+// Calculates the protomoted shape to which the input shapes should be
+// broadcasted for an elementwise operation. The size of the common dimensions
+// (2,3,4 for shape1, and 0,1,2 for shape2) must either match, or either one
+// of the two be 1.
+// Example:
+//   shape1       = [9, 7, 6, 1, 2]
+//   shape2       =       [6, 5, 2]
+//   result_shape = [9, 7, 6, 5, 2]
+TORCH_API std::vector<int64_t> GetPromotedShape(
+    c10::ArrayRef<int64_t> shape1_dims,
+    c10::ArrayRef<int64_t> shape2_dims);
+
+TORCH_API Shape
+GetPromotedBinaryOpShape(const Shape& shape1, const Shape& shape2);
+
+TORCH_API std::vector<std::string> StrSplit(std::string_view text, char delim);
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/internal_ops/ltc_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/internal_ops/ltc_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..29c0ebbbe8c4471051f45a7a00e3138c2e54c2f8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/internal_ops/ltc_ops.h
@@ -0,0 +1,54 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/core/ir.h>
+
+#include <c10/util/CallOnce.h>
+
+#include <string>
+
+namespace torch::lazy {
+
+class TORCH_API OpKindWrapper {
+ public:
+  explicit OpKindWrapper(const char* name) : name_(name) {}
+
+  const OpKind& operator*() const {
+    return get();
+  }
+
+  operator OpKind() const {
+    return get();
+  }
+
+ private:
+  const OpKind& get() const {
+    c10::call_once(once_, [this]() { op_kind_ = OpKind::Get(name_); });
+    return op_kind_;
+  }
+
+  const char* name_;
+  mutable OpKind op_kind_;
+  mutable c10::once_flag once_;
+};
+
+const OpKindWrapper ltc_all_to_all("lazy_tensors::all_to_all");
+const OpKindWrapper ltc_cast("lazy_tensors::cast");
+const OpKindWrapper ltc_collective_permute("lazy_tensors::collective_permute");
+const OpKindWrapper ltc_cross_replica_sum("lazy_tensors::cross_replica_sum");
+const OpKindWrapper ltc_device_data("lazy_tensors::device_data");
+const OpKindWrapper ltc_get_dimensions_size(
+    "lazy_tensors::ltc_get_dimensions_size");
+const OpKindWrapper ltc_moving_average("lazy_tensors::moving_average");
+const OpKindWrapper ltc_nms("lazy_tensors::nms");
+const OpKindWrapper ltc_not_supported("lazy_tensors::not_supported");
+const OpKindWrapper ltc_replication_pad("lazy_tensors::replication_pad");
+const OpKindWrapper ltc_replication_pad_backward(
+    "lazy_tensors::replication_pad_backward");
+const OpKindWrapper ltc_tensor_data("lazy_tensors::tensor_data");
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir.h
new file mode 100644
index 0000000000000000000000000000000000000000..79a4655a62810aa3bbe6fe843f42cd0408567b2a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir.h
@@ -0,0 +1,302 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/symbol.h>
+
+#include <functional>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include <c10/core/ScalarType.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/Flags.h>
+#include <torch/csrc/lazy/core/hash.h>
+#include <torch/csrc/lazy/core/ir_metadata.h>
+#include <torch/csrc/lazy/core/shape.h>
+
+TORCH_DECLARE_bool(ltc_enable_dynamic_shapes);
+
+namespace torch::lazy {
+
+static const hash_t kHashSeed(static_cast<uint32_t>(0x5a2d296e9));
+
+class Node;
+struct Output;
+struct Value;
+
+using NodePtr = std::shared_ptr<Node>;
+
+// The Kind of operation a Node can be associated to.
+struct TORCH_API OpKind {
+  OpKind() = default;
+  explicit OpKind(c10::Symbol op) : op(op) {}
+
+  bool operator==(const OpKind& rhs) const {
+    return op == rhs.op;
+  }
+  bool operator!=(const OpKind& rhs) const {
+    return !operator==(rhs);
+  }
+  bool operator<(const OpKind& rhs) const {
+    return c10::unique_t(op) < c10::unique_t(rhs.op);
+  }
+
+  hash_t hash() const;
+
+  std::string ToString() const {
+    return op.toQualString();
+  }
+
+  // Retrieves an existing operation object, or creates a new one. Operations
+  // that are specific to lazy tensors, should live within the 'lazy_tensors::'
+  // namespace.
+  static OpKind Get(const std::string& name);
+
+  c10::Symbol op;
+};
+
+inline std::ostream& operator<<(std::ostream& stream, const OpKind& op) {
+  stream << op.ToString();
+  return stream;
+}
+
+using OpList = c10::ArrayRef<Value>;
+
+hash_t OperandHashes(
+    const OpList& operands,
+    const hash_t& seed,
+    bool bakeInSizes);
+// A node in the graph. Nodes for operations which require extra data to be
+// stored for lowering should inherit from this class and add an operation
+// specific member there. For example, a constant might create a new
+// NodeConstant class (inheriting from Node) with an extra lazy_tensors::Literal
+// field, or a tensor value might create a new NodeTensor with a computation
+// client data handle in it.
+class TORCH_API Node {
+ public:
+  static bool enableDynamicShape();
+
+  // Creates a new node with the given op name. The op is a unique identifier
+  // for the operation. The num_outputs tells how many outputs a given operation
+  // generates.
+  //
+  // None leaf node's node_hash does not contains shape information always.
+  // So we pass in the hash value rather than a function.
+  Node(OpKind op, size_t num_outputs);
+
+  // Construct node with operands and shapes
+  Node(
+      OpKind op,
+      OpList operands,
+      std::vector<Shape>&& shapes,
+      size_t num_outputs = 1);
+
+  // Construct node with operands and no shape
+  Node(OpKind op, OpList operands, size_t num_outputs = 1);
+
+  // Construct node with shape and no operands
+  Node(OpKind op, Shape shape, size_t num_outputs = 1);
+
+  Node(const Node& rhs);
+
+  Node(Node&& rhs);
+
+  virtual ~Node();
+
+  Node& operator=(const Node& rhs);
+
+  Node& operator=(Node&& rhs);
+
+  const OpKind& op() const {
+    return op_;
+  }
+
+  size_t num_outputs() const {
+    return num_outputs_;
+  }
+
+  // Retrieves the full shape of the IR Node.
+  virtual c10::ArrayRef<Shape> shapes() const;
+
+  virtual const Shape& shape(size_t output_index = 0) const;
+
+  // Add the shape computed by the shape_fn
+  void addComputedShape(const std::function<Shape()>& shape_fn);
+
+  // Compute the shape using the provided shape_fn if not previously cached
+  Shape computeShape(const std::function<Shape()>& shape_fn);
+
+  virtual const std::vector<Output>& operands() const;
+
+  virtual const Output& operand(size_t i) const;
+
+  // Gets operand at index i if index is valid, or kNullOutput otherwise.
+  virtual const Output& nullable_operand(size_t i) const;
+
+  // Returns the hash of the dag used to look up the compiled graph
+  virtual hash_t hash() const = 0;
+
+  // Returns the hash of the dag used to for shape caching
+  virtual hash_t shapeHash() const = 0;
+
+  const MetaData& metadata() const {
+    return metadata_;
+  }
+
+  UserMetaData* user_metadata() const {
+    return user_metadata_.get();
+  }
+
+  std::shared_ptr<UserMetaData> SetUserMetadata(
+      std::shared_ptr<UserMetaData> user_meta) {
+    std::swap(user_metadata_, user_meta);
+    return user_meta;
+  }
+
+  virtual std::string ToString() const;
+
+ private:
+  // The ID of the operation captured by this node.
+  OpKind op_;
+  size_t num_outputs_ = 1;
+
+  // The IR specific metadata attached to the IR node.
+  MetaData metadata_;
+  // The IR framework user can attach a user defined metadata object deriving
+  // from UserMetaData.
+  std::shared_ptr<UserMetaData> user_metadata_;
+
+ protected:
+  // Adds node's index output number as operand.
+  void AddOperand(const NodePtr& node, size_t index = 0);
+
+  std::vector<Shape> shapes_;
+  // A node holds a real reference to its operands.
+  std::vector<NodePtr> operands_;
+  // Outputs do not hold references on the nodes, and neither do the uses, since
+  // otherwise we get into circular reference counting.
+  std::vector<Output> operands_as_outputs_;
+};
+
+inline std::ostream& operator<<(std::ostream& stream, const Node& node) {
+  stream << node.ToString();
+  return stream;
+}
+
+// Note: Keep this version of NodeCast for smooth PyTorch/XLA migration, and
+// clean up once the migration is done.
+template <typename T>
+const T* NodeCast(const Node* node, OpKind op) {
+  if (op != node->op()) {
+    return nullptr;
+  }
+#ifdef NDEBUG
+  return static_cast<const T*>(node);
+#else
+  return &dynamic_cast<const T&>(*node);
+#endif
+}
+
+template <typename T>
+const T* NodeCast(const Node* node) {
+  if (T::ClassOpKind() != node->op()) {
+    return nullptr;
+  }
+  // TODO: Some IR classes share the same opkind, such as Mean and MeanDim, so
+  // static_cast is not safe here. Unless we have opkind unique for each class,
+  // we have to use dynamic_cast here.
+  return dynamic_cast<const T*>(node);
+}
+
+// Represents a specific output produced by a node. Since the output of a node
+// can be composed by multiple outputs, the node+index coordinates fully qualify
+// each single output.
+struct TORCH_API Output {
+  struct Hasher {
+    size_t operator()(const Output& output) const;
+  };
+
+  Output() = default;
+  explicit Output(const Node* node, size_t index = 0)
+      : node(node), index(index) {}
+
+  hash_t hash() const;
+  hash_t shapeHash() const;
+
+  bool operator==(const Output& rhs) const {
+    return node == rhs.node && index == rhs.index;
+  }
+
+  // To compare the operands of to-be-constructed node and to-be-reused node
+  bool operator==(const Value& rhs) const;
+
+  bool operator!=(const Output& rhs) const {
+    return !operator==(rhs);
+  }
+
+  const Shape& shape() const {
+    return node->shape(index);
+  }
+
+  std::string ToString() const;
+
+  // The node providing the output.
+  const Node* node{nullptr};
+  // The index in the node's output this output refers to.
+  size_t index{0};
+};
+
+inline std::ostream& operator<<(std::ostream& stream, const Output& output) {
+  stream << output.ToString();
+  return stream;
+}
+
+template <typename T>
+using OutputMap = std::unordered_map<Output, T, Output::Hasher>;
+
+// Represents an input/operand for a Node object.
+struct TORCH_API Value {
+  Value() = default;
+  /* implicit */ Value(NodePtr&& node, size_t index = 0)
+      : node(std::move(node)), index(index) {}
+  /* implicit */ Value(const NodePtr& node, size_t index = 0)
+      : node(node), index(index) {}
+
+  hash_t hash() const;
+  hash_t shapeHash() const;
+
+  operator bool() const {
+    return node != nullptr;
+  }
+
+  operator Output() const {
+    return Output(node.get(), index);
+  }
+
+  const Shape& shape() const {
+    return node->shape(index);
+  }
+
+  Node* operator->() const {
+    return node.get();
+  }
+
+  NodePtr node;
+  size_t index = 0;
+};
+
+} // namespace torch::lazy
+
+namespace c10 {
+// Explicit template instantiation to make ArrayRef<Value> work
+template class at::ArrayRef<torch::lazy::Value>;
+} // namespace c10
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_builder.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_builder.h
new file mode 100644
index 0000000000000000000000000000000000000000..c1b1c2abc8f342f42fea0e315b2d6762cace2d33
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_builder.h
@@ -0,0 +1,153 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/lazy/backend/backend_interface.h>
+#include <torch/csrc/lazy/core/config.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/tensor.h>
+#include <torch/csrc/lazy/core/trie.h>
+#include <optional>
+#include <vector>
+
+// This file is part of the backend interface. So, ops shouldn't be added or
+// removed without due process The exception to this being the view ops which
+// will be removed soon pending functionalization
+
+namespace torch::lazy {
+
+template <typename T, typename... Args>
+NodePtr ReuseNode(Args&&... args) {
+  if (FLAGS_torch_lazy_reuse_ir) {
+    return LookupNodeFromTrieCache<T>(std::forward<Args>(args)...);
+  }
+  return nullptr;
+}
+
+// Caching an IR node into TrieCache
+static inline void CacheNode(NodePtr node) {
+  if (FLAGS_torch_lazy_reuse_ir) {
+    TrieCache::Get()->Insert(std::move(node));
+  }
+}
+
+template <typename T, typename... Args>
+NodePtr MakeNode(Args&&... args) {
+  return std::make_shared<T>(std::forward<Args>(args)...);
+}
+
+// op is passed in for a more efficient node casting, see the implementation of
+// NodeCast
+template <typename T, typename... Args>
+NodePtr ReuseOrMakeNode(Args&&... args) {
+  NodePtr node = ReuseNode<T>(std::forward<Args>(args)...);
+  if (!node) {
+    node = MakeNode<T>(std::forward<Args>(args)...);
+    CacheNode(node);
+  }
+  return node;
+}
+
+struct IrBuilder {
+  virtual NodePtr MakeDeviceData(
+      const std::shared_ptr<BackendData>& data) const = 0;
+  virtual NodePtr MakeScalar(
+      const at::Scalar& value,
+      const at::ScalarType& type) const = 0;
+  virtual NodePtr MakeExpand(
+      const Value& input0,
+      const std::vector<int64_t>& size,
+      const bool& is_scalar_expand) const = 0;
+  virtual NodePtr MakeCast(
+      const Value& input0,
+      const at::ScalarType& dtype,
+      const std::optional<at::ScalarType>& stype = std::nullopt) const = 0;
+  virtual NodePtr MakeTensorList(const OpList& inputs) const = 0;
+  virtual NodePtr MakeGeneric(
+      const OpKind& op,
+      const OpList& operands,
+      const Shape& shape,
+      const size_t& num_outputs = 1,
+      const hash_t& hash_seed = static_cast<uint32_t>(0x5a2d296e9)) const = 0;
+
+  // dynamic ir nodes
+  virtual NodePtr MakeSizeNode(const Value& input, size_t dim) const = 0;
+  virtual NodePtr MakeSizeAdd(const Value& a, const Value& b) const = 0;
+  virtual NodePtr MakeSizeMul(const Value& a, const Value& b) const = 0;
+  virtual NodePtr MakeSizeDiv(const Value& a, const Value& b) const = 0;
+
+  virtual ~IrBuilder() = default;
+};
+
+static inline NodePtr MakeDeviceData(const std::shared_ptr<BackendData>& data) {
+  return getIrBuilder()->MakeDeviceData(data);
+}
+static inline NodePtr MakeScalar(
+    const at::Scalar& value,
+    const at::ScalarType& type) {
+  return getIrBuilder()->MakeScalar(value, type);
+}
+static inline NodePtr MakeExpand(
+    const Value& input0,
+    const std::vector<int64_t>& size,
+    const bool& is_scalar_expand) {
+  return getIrBuilder()->MakeExpand(input0, size, is_scalar_expand);
+}
+static inline NodePtr MakeCast(
+    const Value& input0,
+    const at::ScalarType& dtype,
+    const std::optional<at::ScalarType>& stype = std::nullopt) {
+  return getIrBuilder()->MakeCast(input0, dtype, stype);
+}
+static inline NodePtr MakeTensorList(const OpList& inputs) {
+  return getIrBuilder()->MakeTensorList(inputs);
+}
+static inline NodePtr MakeGeneric(
+    const OpKind& op,
+    const OpList& operands,
+    const Shape& shape,
+    const size_t& num_outputs = 1,
+    const hash_t& hash_seed = static_cast<uint32_t>(0x5a2d296e9)) {
+  return getIrBuilder()->MakeGeneric(
+      op, operands, shape, num_outputs, hash_seed);
+}
+
+// dynamic ir nodes
+static inline NodePtr MakeSizeNode(const Value& input, size_t dim) {
+  return getIrBuilder()->MakeSizeNode(input, dim);
+}
+static inline NodePtr MakeSizeAdd(const Value& a, const Value& b) {
+  return getIrBuilder()->MakeSizeAdd(a, b);
+}
+static inline NodePtr MakeSizeMul(const Value& a, const Value& b) {
+  return getIrBuilder()->MakeSizeAdd(a, b);
+}
+static inline NodePtr MakeSizeDiv(const Value& a, const Value& b) {
+  return getIrBuilder()->MakeSizeDiv(a, b);
+}
+
+inline Value GetSymIntValue(const c10::SymInt& a) {
+  if (auto ma = a.maybe_as_int()) {
+    return Value(MakeScalar(*ma, at::kLong), 0);
+  } else {
+    return Value(
+        dynamic_cast<torch::lazy::SymNodeImpl*>(a.toSymNodeImplUnowned())
+            ->node_,
+        0);
+  }
+}
+
+// TODO: this should return Value
+inline std::vector<int64_t> GetSymIntArrayRefValue(c10::SymIntArrayRef arr) {
+  std::vector<int64_t> r;
+  for (const auto& a : arr) {
+    r.emplace_back(a.guard_int(__FILE__, __LINE__));
+  }
+  return r;
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_dump_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_dump_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..7a27bd8dbec82daf15acb6ac695a8aba86dac08d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_dump_util.h
@@ -0,0 +1,35 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/core/ir.h>
+
+#include <string>
+
+namespace torch::lazy {
+
+class BackendDevice;
+
+class TORCH_API DumpUtil {
+ public:
+  static std::string ToDot(c10::ArrayRef<const Node*> nodes);
+
+  static std::string PostOrderToDot(
+      c10::ArrayRef<const Node*> post_order,
+      c10::ArrayRef<const Node*> roots);
+
+  static std::string ToText(c10::ArrayRef<const Node*> nodes);
+
+  static std::string PostOrderToText(
+      c10::ArrayRef<const Node*> post_order,
+      c10::ArrayRef<const Node*> roots);
+
+  static std::string ToBackend(
+      c10::ArrayRef<Value> values,
+      const BackendDevice& device);
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_metadata.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_metadata.h
new file mode 100644
index 0000000000000000000000000000000000000000..8b913e2342810b322e4470fb01332a98a40d5116
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_metadata.h
@@ -0,0 +1,56 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Macros.h>
+
+#include <string>
+#include <vector>
+
+namespace torch::lazy {
+struct SourceLocation {
+  std::string file;
+  std::string function;
+  int line = -1;
+};
+
+TORCH_API void EmitShortFrameInfo(
+    std::ostream& stream,
+    const std::vector<SourceLocation>& frames);
+
+TORCH_API std::ostream& operator<<(
+    std::ostream& stream,
+    const std::vector<SourceLocation>& frames);
+
+// The base class for user defined metadata which is possible to attach to IR
+// nodes.
+struct TORCH_API UserMetaData {
+  virtual ~UserMetaData() = default;
+};
+
+struct TORCH_API MetaData {
+  std::string scope;
+  std::vector<SourceLocation> frame_info;
+};
+
+// TODO(whc) is this going to be used outside of in IR decompositions?
+// RAII data structure to be used a stack variable to enter a new IR scope. IR
+// scope names will appear in the IR and will help identifying the source of the
+// single IR nodes.
+struct TORCH_API ScopePusher {
+  explicit ScopePusher(const std::string& name);
+  ~ScopePusher();
+  ScopePusher(ScopePusher&& other) = delete;
+  ScopePusher(const ScopePusher&) = delete;
+  ScopePusher& operator=(const ScopePusher&) = delete;
+  ScopePusher& operator=(ScopePusher&&) = delete;
+
+  static void ResetScopes();
+};
+
+TORCH_API MetaData GetMetaDataIfDebugging();
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..bb2f2420a1028cd6aa1d5b58e456dcf767904e5e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ir_util.h
@@ -0,0 +1,50 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <unordered_map>
+#include <vector>
+
+#include <torch/csrc/lazy/core/ir.h>
+
+namespace torch::lazy {
+
+class TORCH_API Util {
+ public:
+  // Tracks the emission status of the nodes during the post-order generation.
+  // It helps tracking loops within the computation graphs.
+  enum EmitStatus {
+    kNotEmitted,
+    kEmitting,
+    kEmitted,
+  };
+
+  using EmissionMap = std::unordered_map<const Node*, EmitStatus>;
+
+  // Computes the post order from the given node, without using recursion. The
+  // emission map can be used as saved state, for multiple separate calls to
+  // this API. The returned post-order can be empty if the node has already been
+  // emitted inside the emission map. An error is generated if a loop is
+  // detected.
+  static std::vector<const Node*> ComputePostOrder(
+      const Node* node,
+      EmissionMap* emap);
+
+  static std::vector<const Node*> ComputePostOrder(
+      c10::ArrayRef<const Node*> nodes,
+      EmissionMap* emap);
+
+  // Same as above, but computes the post order on the set of nodes specified as
+  // argument.
+  static std::vector<const Node*> ComputePostOrder(
+      c10::ArrayRef<const Node*> nodes);
+
+  // Retrieves the number of nodes within the graph whose sink are passed in the
+  // nodes argument.
+  static size_t GetGraphSize(c10::ArrayRef<const Node*> nodes);
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/lazy_graph_executor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/lazy_graph_executor.h
new file mode 100644
index 0000000000000000000000000000000000000000..f7937602a8f3877b698ec06844bb5445d3e580ea
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/lazy_graph_executor.h
@@ -0,0 +1,434 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/ArrayRef.h>
+#include <torch/csrc/lazy/backend/lowering_context.h>
+#include <torch/csrc/lazy/core/cache.h>
+#include <torch/csrc/lazy/core/ir_util.h>
+#include <torch/csrc/lazy/core/multi_wait.h>
+#include <torch/csrc/lazy/core/tensor.h>
+#include <torch/csrc/lazy/core/util.h>
+
+namespace torch::lazy {
+
+class TORCH_API LazyGraphExecutor {
+ public:
+  struct DeviceDataInfo : public BackendData::Info {
+    DeviceDataInfo(int64_t tensor_id, bool read_only)
+        : tensor_id(tensor_id), read_only(read_only) {}
+
+    int64_t tensor_id = 0;
+    bool read_only = false;
+  };
+
+  // Register a lazy graph executor instance that can be retrieved using Get()
+  static void Register(LazyGraphExecutor* /*executor*/);
+  static LazyGraphExecutor* Get();
+
+  virtual ~LazyGraphExecutor() = default;
+
+  // Override these methods to perform custom tensor registration and
+  // unregistration Note: It is vital that the parent implementations are also
+  // called in order for the tensors to show up in the live tensor list
+  virtual void RegisterTensor(std::shared_ptr<LazyTensor::Data> data);
+  virtual void UnregisterTensor(LazyTensor::Data* data);
+
+  // Seed for random generator.
+  // Override to supply your own DeviceContextArena.
+  virtual Value GetRngSeed(const BackendDevice& device);
+  virtual uint64_t GetRunningSeed(const BackendDevice& device);
+  virtual void SetRngSeed(const BackendDevice& device, uint64_t seed);
+
+  void DeviceBarrier(const BackendDevice& device);
+
+  BackendDataPtr GetDeviceData(
+      const at::Tensor& tensor,
+      const BackendDevice& device);
+
+  BackendDataPtr GetDeviceData(
+      const at::Scalar& value,
+      at::ScalarType scalar_type,
+      const BackendDevice& device);
+
+  // Retrieves the set of lazy tensors which are currently live in the system,
+  // for the given device. If device is nullptr, the live tensors for all
+  // devices will be returned. Returned tensors are sorted by device as primary
+  // key, and by unique ID as secondary key.
+  std::vector<LazyTensorPtr> GetLiveTensors(const BackendDevice* device);
+
+  // Makes sure that any outstanding IR operation accumulated over live tensors,
+  // gets turned into device data. If wait is true, the sync operation will be
+  // run synchronously. The devices argument, if not empty, tells the devices
+  // which should be partecipating into the replicated computation.
+  virtual void SyncLiveTensorsGraph(
+      const BackendDevice* device,
+      c10::ArrayRef<std::string> devices,
+      bool wait);
+
+  // Applies all the pending IR operations queued over the input tensors. All
+  // the tensors must be on the same device. If wait is true, the sync operation
+  // will be run synchronously. The devices argument, if not empty, tells the
+  // devices which should be partecipating into the replicated computation.
+  void SyncTensorsGraph(
+      std::vector<LazyTensorPtr>* tensors,
+      c10::ArrayRef<std::string> devices,
+      bool wait,
+      bool sync_ltc_data);
+
+  // Marks an execution step, which allows the tensor framework to understand
+  // the computation boundaries.
+  // Override to supply your own DeviceContextArena.
+  virtual void MarkStep(const BackendDevice& device);
+
+  // Waits for all the outstanding operations on all the supplied devices.
+  // If devices is empty, the wait will happen for all local devices.
+  void WaitDeviceOps(c10::ArrayRef<BackendDevice> devices);
+
+  // Retrieves the PyTorch CPU tensors behind the lazy tensors IR operations.
+  // All the tensors must be on the same device.
+  std::vector<at::Tensor> GetTensors(std::vector<LazyTensorPtr>* tensors);
+
+  size_t IncTrimCounter() const;
+
+  // Dumps the backend specific text of the computation accumulated in the graph
+  // which is attached the tensors.
+  std::string DumpBackendComputation(const std::vector<LazyTensorPtr>& tensors);
+
+  Value GetDeviceDataIrValue(
+      const at::Scalar& value,
+      c10::ScalarType type,
+      const BackendDevice& device);
+  Value GetIrValueForScalar(
+      const at::Scalar& value,
+      c10::ScalarType type,
+      const BackendDevice& device);
+  Value GetIrValueForScalar(
+      const at::Scalar& value,
+      const BackendDevice& device);
+
+  // TODO: even though this API is currently used **only** in codegen to
+  // generate real scalar IR values vs scalar tensors, we would like to
+  // use it in other cases where `GetIrValueForXXXScalar` is used, as well
+  // In order to do that, we need to untangle the cases where we don't need
+  // `expand` and where we don't expect a scalar tensor
+  Value GetIrValueForScalarFromCodegen(
+      const at::Scalar& value,
+      const BackendDevice& device);
+  Value GetIrValueForExpandedScalar(
+      const at::Scalar& value,
+      const Shape& shape,
+      const BackendDevice& device);
+
+  struct CachedComputation {
+    explicit CachedComputation(ComputationPtr computation)
+        : computation(std::move(computation)) {}
+
+    ComputationPtr computation;
+  };
+
+  using ComputationCache = Cache<hash_t, CachedComputation, HashReducer>;
+
+  ComputationCache* GetComputationCache();
+
+  hash_t GetGraphHash(const std::vector<LazyTensorPtr>& tensors);
+
+  // Clear the computation cache.
+  void ClearComputationCache();
+  // Remove a specific computation cache entry from its hash.
+  void RemoveFromComputationCache(const hash_t& hash);
+
+ protected:
+  // TODO(alanwaketan): Revisit if all of them need to be accessible to
+  // derived classes.
+
+  struct SyncTensorsConfig {
+    // Whether we want to force data on the target tensors (hence trimming
+    // the IR graph above them).
+    bool force_ltc_data = true;
+    // Whether when setting the data, the other properties of the tensor
+    // state should be reset.
+    bool sync_ltc_data = true;
+  };
+
+  struct SyncTensorCollection {
+    SyncTensorCollection() : hash(0) {}
+
+    SyncTensorsConfig config;
+    std::vector<size_t> indices;
+    hash_t hash;
+    std::vector<ExceptionCleanup> unlocker;
+    BackendDevice device;
+  };
+
+  struct PostOrderData {
+    std::vector<const Node*> post_order;
+    Util::EmissionMap emission_map;
+    std::vector<BackendDataPtr> parameters_data;
+    std::vector<size_t> parameter_sequence;
+  };
+
+  // Locking:
+  // We perform two kinds of operations of tensors, synchronous and
+  // asynchronous. The ApplyPendingGraph() are synchronous, as we need the
+  // device data result immediately. Before the synchronous operations can
+  // start, they need to wait that the pending asynchronous operations have
+  // completed. Synchronous operations do not hold device locks, since they are
+  // strictly sequential, dictated by the PyTorch execution order. The
+  // SyncTensorsGraph() is asynchronous, and returns immediately after having
+  // scheduled the asynchronous operation. While executing, the asynchronous
+  // operations will hold locks on all the participating devices (in most common
+  // cases there will be only one device).
+  // Since asynchronous operations capture device locks, only one asynchronous
+  // operation can execute at the same time, on a given device. Tensor
+  // operations which send data to device do not need to hold any device locks
+  // while doing so. Only operations which _use_ device data (computations, and
+  // transfer from server) need to wait for asynchronous operations to complete
+  // (barrier).
+
+  class DeviceLocker {
+   public:
+    explicit DeviceLocker(BackendDevice device) : device_(std::move(device)) {}
+
+    const BackendDevice& device() const {
+      return device_;
+    }
+
+    void Lock();
+    void Unlock(std::exception_ptr exptr);
+    void Barrier();
+
+   private:
+    void CheckResetException();
+
+    BackendDevice device_;
+    std::mutex mutex_;
+    std::condition_variable cv_;
+    bool locked_ = false;
+    std::exception_ptr exptr_;
+  };
+
+  class DeviceLockerArena {
+   public:
+    static DeviceLockerArena* Get();
+
+    std::shared_ptr<DeviceLocker> GetLocker(const BackendDevice& device);
+
+    void DeviceBarrier(const BackendDevice& device);
+
+    // Use a set to impose an order on the device locking sequence (ABBA
+    // prevention).
+    std::vector<ExceptionCleanup> LockDevices(
+        const std::set<BackendDevice>& devices);
+
+   private:
+    ExceptionCleanup LockDevice(const BackendDevice& device);
+
+    std::mutex mutex_;
+    std::map<BackendDevice, std::shared_ptr<DeviceLocker>> lockers_;
+  };
+
+  class DataCacheArena {
+   public:
+    static DataCacheArena* Get();
+
+    BackendDataPtr GetDeviceData(
+        const at::Tensor& tensor,
+        const BackendDevice& device);
+
+    BackendDataPtr GetDeviceData(
+        const at::Scalar& value,
+        at::ScalarType scalar_type,
+        const BackendDevice& device);
+
+   private:
+    struct TensorHasher {
+      size_t operator()(const at::Tensor& tensor) const;
+    };
+    struct TensorComparer {
+      bool operator()(const at::Tensor& tensor1, const at::Tensor& tensor2)
+          const;
+    };
+
+    explicit DataCacheArena(size_t max_cache_size);
+
+    using DataCache =
+        Cache<at::Tensor, BackendData, TensorHasher, TensorComparer>;
+
+    DataCache* GetDataCache(const BackendDevice& device);
+
+    size_t max_cache_size_ = 0;
+    std::mutex mutex_;
+    std::map<BackendDevice, std::unique_ptr<DataCache>> device_caches_;
+  };
+
+  // The DeviceContextArena holds per device live information and statistics,
+  // among which the lazy tensors which are currently alive in the system. This
+  // is used to create computation "barriers" in order to flush pending
+  // operations and ensure the same computations are created during the training
+  // loops.
+  // TODO(alanwaketan): Add a registry such that we don't need to make all
+  // related methods virtual.
+  class DeviceContextArena {
+   protected:
+    struct DeviceContext {
+      std::mutex lock;
+      std::map<int64_t, std::weak_ptr<LazyTensor::Data>> tensors_data;
+      uint64_t seed = 101;
+      uint64_t running_seed = 101;
+      Value seed_ir_value;
+    };
+
+   public:
+    static DeviceContextArena* Get();
+    virtual ~DeviceContextArena() = default;
+
+    void RegisterTensor(std::shared_ptr<LazyTensor::Data> data);
+    void UnregisterTensor(LazyTensor::Data* data);
+
+    std::vector<LazyTensorPtr> GetLiveTensors(const BackendDevice* device);
+
+    // Overriding it allow derived class to use their own IRs for Value.
+    virtual Value GetRngSeed(const BackendDevice& device);
+    uint64_t GetRunningSeed(const BackendDevice& device);
+    void SetRngSeed(const BackendDevice& device, uint64_t seed);
+
+    void MarkStep(const BackendDevice& device);
+
+    std::vector<BackendDevice> GetActiveDevices();
+
+   protected:
+    DeviceContext* GetDeviceContext(const BackendDevice& device);
+
+    void ForAllDeviceContexts(
+        const std::function<void(DeviceContext*)>& fn,
+        const BackendDevice* device);
+
+    // Overriding it allow derived class to use their own conversions.
+    virtual Value IrValueFromScalar(
+        const at::Scalar& value,
+        at::ScalarType scalar_type,
+        const BackendDevice& device);
+
+   private:
+    std::vector<DeviceContext*> GetAllDeviceContexts();
+
+    std::mutex lock_;
+    std::map<BackendDevice, DeviceContext*> device_contexts_;
+  };
+
+  struct Async {
+    Async(
+        SyncTensorCollection* coll,
+        std::vector<BackendDataPtr> parameters_data,
+        std::vector<BackendDataPtr> tensors_data,
+        ComputationCache::TypePtr cached_computation);
+    virtual ~Async() = default;
+
+    void Wait();
+
+    MultiWait mwait;
+    std::vector<size_t> indices;
+    std::vector<ExceptionCleanup> unlocker;
+    std::vector<BackendDataPtr> parameters_data;
+    BackendDevice device;
+    ComputationCache::TypePtr cached_computation;
+    std::vector<BackendDataPtr> tensors_data;
+  };
+
+  void ResetTrimCounter() const;
+
+  // Waits for this SyncTensorCollection's device barrier and acquire the lock.
+  virtual void TensorCollectionBarrier(SyncTensorCollection* coll);
+
+  // One can override to insert your own profiler.
+  virtual PostOrderData RunPostOrder(
+      const std::vector<Value>& ir_values,
+      SyncTensorCollection* coll);
+
+ private:
+  struct CompilationResult {
+    BackendDevice device;
+    size_t emitted_nodes = 0;
+    ComputationPtr computation;
+    std::vector<BackendDataPtr> parameters_data;
+  };
+
+  virtual bool ShouldSyncTensor(const LazyTensorPtr& tensor) const;
+
+  SyncTensorCollection CollectSyncTensors(
+      const std::vector<LazyTensorPtr>& tensors,
+      const SyncTensorsConfig& config);
+
+  std::vector<Value> CollectRoots(
+      const std::vector<LazyTensorPtr>& tensors,
+      c10::ArrayRef<size_t> indices);
+
+  std::vector<BackendDataPtr> SetTensorData(
+      std::vector<LazyTensorPtr>* tensors,
+      const SyncTensorsConfig& config,
+      c10::ArrayRef<size_t> indices,
+      const std::vector<torch::lazy::BackendDataPtr>& tensor_data_vec);
+
+  void ExtractIRAndPrepareTensorData(
+      std::vector<LazyTensorPtr>* tensors,
+      const SyncTensorsConfig& config,
+      c10::ArrayRef<size_t> indices,
+      std::vector<Value>& ir_values,
+      std::vector<BackendDataPtr>& tensor_data_vec);
+
+  std::shared_ptr<Async> TryRunCachedSync(
+      std::vector<LazyTensorPtr>* tensors,
+      SyncTensorCollection* coll,
+      PostOrderData* po_data,
+      const std::vector<BackendDataPtr>& tensor_data_vec);
+
+  CompilationResult Compile(
+      const std::vector<LazyTensorPtr>& tensors,
+      c10::ArrayRef<std::string> devices,
+      const SyncTensorCollection& coll,
+      PostOrderData* po_data,
+      const std::vector<Value>& ir_values);
+
+  ComputationCache::TypePtr LookupCachedCompile(const hash_t& hash);
+
+  std::shared_ptr<Async> SyncTensorsGraphInternal(
+      std::vector<LazyTensorPtr>* tensors,
+      c10::ArrayRef<std::string> devices,
+      const SyncTensorsConfig& config);
+
+  // Schedules the execution of a sync tensors operation in background. The
+  // asynchronous operation will hold the device locks by capturing the ones
+  // present within the coll structure.
+  std::shared_ptr<Async> ScheduleSyncTensorsGraph(
+      SyncTensorCollection* coll,
+      std::vector<BackendDataPtr> parameters_data,
+      std::vector<BackendDataPtr> tensors_data,
+      ComputationCache::TypePtr cached_computation);
+
+  std::shared_ptr<Async> ScheduleSyncTensorsGraph(
+      std::vector<LazyTensorPtr>* tensors,
+      SyncTensorCollection* coll,
+      std::vector<BackendDataPtr> parameters_data,
+      ComputationCache::TypePtr cached_computation,
+      const std::vector<BackendDataPtr>& tensor_data_vec);
+
+  std::vector<at::Tensor> GetTensorsFused(std::vector<LazyTensorPtr>* tensors);
+
+  std::vector<at::Tensor> FetchTensors(
+      std::vector<LazyTensorPtr>* tensors,
+      c10::ArrayRef<BackendDataPtr> tensors_data,
+      const std::vector<size_t>* indices);
+
+  // Gathers the device data for all the input tensors, after an
+  // asynchronous operation.
+  std::vector<BackendDataPtr> GatherTensorsData(
+      const std::vector<LazyTensorPtr>& tensors,
+      c10::ArrayRef<size_t> indices,
+      c10::ArrayRef<BackendDataPtr> tensors_data);
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/metrics.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/metrics.h
new file mode 100644
index 0000000000000000000000000000000000000000..a175d9358ce87beb902226c1700403b5c0f70bc5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/metrics.h
@@ -0,0 +1,293 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * This file is adapted from PyTorch/XLA
+ * https://github.com/pytorch/xla/blob/e0e5f937a0ba8d904f9608137dc8c51ba439df2d/third_party/xla_client/metrics.h
+ */
+
+#pragma once
+
+#include <atomic>
+#include <functional>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <vector>
+
+#include <c10/macros/Export.h>
+
+namespace torch::lazy {
+
+struct TORCH_API Sample {
+  Sample() = default;
+  Sample(int64_t timestamp_ns, double value)
+      : timestamp_ns(timestamp_ns), value(value) {}
+
+  int64_t timestamp_ns = 0;
+  double value = 0;
+};
+
+using MetricReprFn = std::function<std::string(double)>;
+
+// Class used to collect time-stamped numeric samples. The samples are stored in
+// a circular buffer whose size can be configured at constructor time.
+class TORCH_API MetricData {
+ public:
+  // Creates a new MetricData object with the internal circular buffer storing
+  // max_samples samples. The repr_fn argument allow to specify a function which
+  // pretty-prints a sample value.
+  MetricData(MetricReprFn repr_fn, size_t max_samples);
+
+  // Returns the total values of all the samples being posted to this metric.
+  double Accumulator() const;
+
+  size_t TotalSamples() const;
+
+  void AddSample(int64_t timestamp_ns, double value);
+
+  // Returns a vector with all the current samples, from the oldest to the
+  // newer. If accumulator is not nullptr, it will receive the current value of
+  // the metrics' accumulator (the sum of all posted values). If total_samples
+  // is not nullptr, it will receive the count of the posted values.
+  std::vector<Sample> Samples(double* accumulator, size_t* total_samples) const;
+
+  std::string Repr(double value) const {
+    return repr_fn_(value);
+  }
+
+  void Reset();
+
+  bool IsValid() const {
+    return TotalSamples() > 0;
+  }
+
+ private:
+  mutable std::mutex lock_;
+  MetricReprFn repr_fn_;
+  size_t count_ = 0;
+  std::vector<Sample> samples_;
+  double accumulator_ = 0.0;
+};
+
+// Counters are a very lightweight form of metrics which do not need to track
+// sample time.
+class TORCH_API CounterData {
+ public:
+  CounterData() : value_(0) {}
+
+  void AddValue(int64_t value) {
+    value_ += value;
+  }
+
+  int64_t Value() const {
+    return value_;
+  }
+
+  void Reset() {
+    value_ = 0;
+  }
+
+  bool IsValid() const {
+    return value_ > 0;
+  }
+
+ private:
+  std::atomic<int64_t> value_;
+};
+
+class TORCH_API MetricsArena {
+ public:
+  static MetricsArena* Get();
+
+  void ResetCounters();
+  void ResetMetrics();
+
+  // Registers a new metric in the global arena.
+  void RegisterMetric(
+      const std::string& name,
+      MetricReprFn repr_fn,
+      size_t max_samples,
+      std::shared_ptr<MetricData>* data);
+
+  void RegisterCounter(
+      const std::string& name,
+      std::shared_ptr<CounterData>* data);
+
+  void ForEachMetric(
+      const std::function<void(const std::string&, MetricData*)>& metric_func);
+
+  void ForEachCounter(
+      const std::function<void(const std::string&, CounterData*)>&
+          counter_func);
+
+  std::vector<std::string> GetMetricNames();
+
+  MetricData* GetMetric(const std::string& name);
+
+  std::vector<std::string> GetCounterNames();
+
+  CounterData* GetCounter(const std::string& name);
+
+ private:
+  std::mutex lock_;
+  std::map<std::string, std::shared_ptr<MetricData>> metrics_;
+  std::map<std::string, std::shared_ptr<CounterData>> counters_;
+};
+
+// Emits the value in a to_string() conversion.
+TORCH_API std::string MetricFnValue(double value);
+// Emits the value in a humanized bytes representation.
+TORCH_API std::string MetricFnBytes(double value);
+// Emits the value in a humanized time representation. The value is expressed in
+// nanoseconds EPOCH time.
+TORCH_API std::string MetricFnTime(double value);
+
+// The typical use of a Metric is one in which it gets created either in a
+// global scope context:
+//   static Metric* metric = new Metric("RpcCount");
+// Or within a function scope:
+//   void MyFunction(...) {
+//     static Metric* metric = new Metric("RpcCount");
+//     ...
+//     metric->AddSample(ts_nanos, some_value);
+//   }
+class TORCH_API Metric {
+ public:
+  explicit Metric(
+      std::string name,
+      MetricReprFn repr_fn = MetricFnValue,
+      size_t max_samples = 0);
+
+  const std::string& Name() const {
+    return name_;
+  }
+
+  double Accumulator() const;
+
+  void AddSample(int64_t timestamp_ns, double value);
+
+  void AddSample(double value);
+
+  std::vector<Sample> Samples(double* accumulator, size_t* total_samples) const;
+
+  std::string Repr(double value) const;
+
+ private:
+  MetricData* GetData() const;
+
+  std::string name_;
+  MetricReprFn repr_fn_;
+  size_t max_samples_;
+  mutable std::shared_ptr<MetricData> data_ptr_;
+  mutable std::atomic<MetricData*> data_;
+};
+
+// A Counter is a lightweight form of metric which tracks an integer value which
+// can increase or decrease.
+// A typical use is as:
+//   static Counter* counter = new Counter("MyCounter");
+//   ...
+//   counter->AddValue(+1);
+class TORCH_API Counter {
+ public:
+  explicit Counter(std::string name);
+
+  void AddValue(int64_t value) {
+    GetData()->AddValue(value);
+  }
+
+  int64_t Value() const {
+    return GetData()->Value();
+  }
+
+ private:
+  CounterData* GetData() const;
+
+  std::string name_;
+  mutable std::shared_ptr<CounterData> data_ptr_;
+  mutable std::atomic<CounterData*> data_;
+};
+
+#define TORCH_LAZY_COUNTER(name, value)        \
+  do {                                         \
+    static ::torch::lazy::Counter* __counter = \
+        new ::torch::lazy::Counter(name);      \
+    __counter->AddValue(value);                \
+  } while (0)
+
+#define TORCH_LAZY_FN_COUNTER(ns) TORCH_LAZY_COUNTER(c10::str(ns, __func__), 1)
+
+#define TORCH_LAZY_VALUE_METRIC(name, value)                         \
+  do {                                                               \
+    static ::torch::lazy::Metric* __metric =                         \
+        new ::torch::lazy::Metric(name, torch::lazy::MetricFnValue); \
+    __metric->AddSample(value);                                      \
+  } while (0)
+
+// Creates a report with the current metrics statistics.
+TORCH_API std::string CreateMetricReport();
+
+// Creates a report with the selected metrics statistics.
+TORCH_API std::string CreateMetricReport(
+    const std::vector<std::string>& counter_names,
+    const std::vector<std::string>& metric_names);
+
+// Returns the currently registered metric names. Note that the list can grow
+// since metrics are usually function initialized (they are static function
+// variables).
+TORCH_API std::vector<std::string> GetMetricNames();
+
+// Retrieves the metric data of a given metric, or nullptr if such metric does
+// not exist.
+TORCH_API MetricData* GetMetric(const std::string& name);
+
+// Returns the currently registered counter names. Note that the list can grow
+// since counters are usually function initialized (they are static function
+// variables).
+TORCH_API std::vector<std::string> GetCounterNames();
+
+// Retrieves the counter data of a given counter, or nullptr if such counter
+// does not exist.
+TORCH_API CounterData* GetCounter(const std::string& name);
+
+// Retrieves the current EPOCH time in nanoseconds.
+TORCH_API int64_t NowNs();
+
+// Scope based utility class TORCH_API to measure the time the code takes within
+// a given C++ scope.
+class TORCH_API TimedSection {
+ public:
+  explicit TimedSection(Metric* metric) : metric_(metric), start_(NowNs()) {}
+
+  TimedSection(TimedSection&& other) = delete;
+  TimedSection(const TimedSection&) = delete;
+  TimedSection& operator=(const TimedSection&) = delete;
+  TimedSection& operator=(TimedSection&&) = delete;
+  ~TimedSection() {
+    int64_t now = NowNs();
+    metric_->AddSample(now, static_cast<double>(now - start_));
+  }
+
+  double Elapsed() const {
+    return 1e-9 * static_cast<double>(NowNs() - start_);
+  }
+
+ private:
+  Metric* metric_;
+  int64_t start_;
+};
+
+#define TORCH_LAZY_TIMED(name)                                  \
+  static torch::lazy::Metric* timed_metric =                    \
+      new torch::lazy::Metric(name, torch::lazy::MetricFnTime); \
+  torch::lazy::TimedSection timed_section(timed_metric)
+
+#define TORCH_LAZY_FN_COUNTER_TIMED_TRACING(ns) \
+  TORCH_LAZY_FN_COUNTER(ns);                    \
+  TORCH_LAZY_TIMED("LazyTracing")
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/multi_wait.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/multi_wait.h
new file mode 100644
index 0000000000000000000000000000000000000000..c808d3cc6dc6221ea61fee86f86e114e5fdee08c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/multi_wait.h
@@ -0,0 +1,65 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * This file is adapted from PyTorch/XLA
+ * https://github.com/pytorch/xla/blob/e0e5f937a0ba8d904f9608137dc8c51ba439df2d/third_party/xla_client/multi_wait.h
+ */
+
+#pragma once
+
+#include <condition_variable>
+#include <exception>
+#include <functional>
+#include <memory>
+#include <mutex>
+
+#include <c10/macros/Export.h>
+
+namespace torch::lazy {
+
+// Support waiting for a number of tasks to complete.
+class TORCH_API MultiWait {
+ public:
+  explicit MultiWait(size_t count) : count_(count) {}
+
+  // Signal the completion of a single task.
+  void Done();
+
+  // Waits until at least count (passed as constructor value) completions
+  // happened.
+  void Wait();
+
+  // Same as above, but waits up to wait_seconds.
+  void Wait(double wait_seconds);
+
+  // Resets the threshold counter for the MultiWait object. The completed count
+  // is also reset to zero.
+  void Reset(size_t count);
+
+  // Creates a completer functor which signals the mult wait object once func
+  // has completed. Handles exceptions by signaling the multi wait with the
+  // proper status value. This API returns a function which captures a MultiWait
+  // reference, so care must be taken such that the reference remains valid for
+  // the whole lifetime of the returned function.
+  std::function<void()> Completer(std::function<void()> func);
+
+  // Similar as the above API, but with explicit capture of the MultiWait shared
+  // pointer.
+  static std::function<void()> Completer(
+      std::shared_ptr<MultiWait> mwait,
+      std::function<void()> func);
+
+ private:
+  void Complete(const std::function<void()>& func);
+
+  std::mutex mutex_;
+  std::condition_variable cv_;
+  size_t count_ = 0;
+  size_t completed_count_ = 0;
+  std::exception_ptr exptr_;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ops/arithmetic_ir_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ops/arithmetic_ir_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..ff4fe2341d2bd27f1e13d716782ddaced8cb2b79
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ops/arithmetic_ir_ops.h
@@ -0,0 +1,17 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/core/ir.h>
+
+namespace torch::lazy {
+
+TORCH_API NodePtr operator+(const Value& node1, const Value& node2);
+TORCH_API NodePtr operator-(const Value& node1, const Value& node2);
+TORCH_API NodePtr operator*(const Value& node1, const Value& node2);
+TORCH_API NodePtr operator/(const Value& node1, const Value& node2);
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ops/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ops/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..f900b65fa2280f3cf08fc43942ba815ee3e6c45f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/ops/utils.h
@@ -0,0 +1,44 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <vector>
+
+#include <torch/csrc/lazy/core/tensor_util.h>
+#include <torch/csrc/lazy/core/util.h>
+
+namespace torch::lazy {
+
+TORCH_API bool StrideIsSupported(c10::ArrayRef<int64_t> stride);
+
+TORCH_API std::vector<int64_t> GetArrayStridePermutation(
+    c10::ArrayRef<int64_t> stride);
+
+TORCH_API Shape MakeDiagonalShape(
+    const Shape& shape,
+    int64_t offset,
+    int64_t dim1,
+    int64_t dim2);
+
+TORCH_API Shape
+MakePermuteShape(const Shape& source_shape, c10::ArrayRef<int64_t> permutation);
+
+TORCH_API Shape MakeSelectShape(
+    const Shape& shape,
+    int64_t dim,
+    int64_t start,
+    int64_t end,
+    int64_t stride);
+
+TORCH_API int64_t GetStride(int64_t start, int64_t end, int64_t stride);
+
+TORCH_API std::vector<int64_t> BuildSqueezedDimensions(
+    c10::ArrayRef<int64_t> dimensions,
+    int64_t squeeze_dim);
+
+TORCH_API std::vector<int64_t> BuildUnsqueezedDimensions(
+    c10::ArrayRef<int64_t> dimensions,
+    int64_t squeeze_dim);
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/permutation_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/permutation_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..6a3d5aaa7946fb920ac8d63c3b7524925a4b7b3c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/permutation_util.h
@@ -0,0 +1,46 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/ArrayRef.h>
+#include <c10/util/Exception.h>
+#include <c10/util/irange.h>
+
+#include <vector>
+
+namespace torch::lazy {
+
+TORCH_API std::vector<int64_t> InversePermutation(
+    c10::ArrayRef<int64_t> input_permutation);
+
+TORCH_API bool IsPermutation(c10::ArrayRef<int64_t> permutation);
+
+// Gathers the input using the order specified by the permutation. For each i,
+// output[i] = dimensions[permutation[i]]. The given permutation must be the
+// same size as the input.
+template <typename Container>
+std::vector<typename Container::value_type> PermuteDimensions(
+    c10::ArrayRef<int64_t> permutation,
+    const Container& dimensions) {
+  using T = typename Container::value_type;
+  TORCH_CHECK(
+      dimensions.size() == permutation.size(),
+      "Invalid permutation specified. dimensions.size() != permutation.size()  (",
+      dimensions.size(),
+      " vs. ",
+      permutation.size(),
+      ")");
+  TORCH_CHECK(
+      IsPermutation(permutation),
+      "Invalid permutation specified. Permutation is not permutation");
+  std::vector<T> output(dimensions.size());
+  for (const auto i : c10::irange(permutation.size())) {
+    output[i] = dimensions[permutation[i]];
+  }
+  return output;
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/shape.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/shape.h
new file mode 100644
index 0000000000000000000000000000000000000000..7d4fdb375aa018a79d2df149297705c7cc4bc3f3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/shape.h
@@ -0,0 +1,83 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ostream>
+#include <vector>
+
+#include <c10/core/Scalar.h>
+#include <torch/csrc/jit/passes/symbolic_shape_analysis.h>
+#include <torch/csrc/lazy/core/hash.h>
+
+TORCH_DECLARE_bool(ltc_enable_symbolic_shapes);
+
+namespace torch::lazy {
+
+class TORCH_API Shape {
+ public:
+  Shape() = default;
+
+  Shape(
+      at::ScalarType scalar_type,
+      c10::ArrayRef<int64_t> sizes,
+      std::optional<std::vector<bool>> is_symbolic = std::nullopt);
+
+  std::string to_string() const;
+
+  c10::ScalarType scalar_type() const {
+    return scalar_type_;
+  }
+  void set_scalar_type(at::ScalarType value) {
+    scalar_type_ = value;
+  }
+
+  int64_t dim() const {
+    return static_cast<int64_t>(sizes_.size());
+  }
+  c10::ArrayRef<int64_t> sizes() const {
+    return sizes_;
+  }
+  int64_t size(int64_t dim) const {
+    return sizes_.at(dim);
+  }
+  void set_size(int64_t dim, int64_t size) {
+    sizes_.at(dim) = size;
+  }
+
+  const std::optional<std::vector<bool>>& is_symbolic() const {
+    return is_symbolic_;
+  }
+
+  // Makes a copy with symbolic dims applied
+  Shape with_symbolic_dims(
+      std::optional<std::vector<bool>> symbolic_dims) const;
+
+  size_t numel() const;
+  hash_t hash(bool bakeInSizes) const;
+
+  bool operator==(const Shape& other) const;
+
+ private:
+  c10::ScalarType scalar_type_{c10::ScalarType::Undefined};
+
+  // Sizes are the upper bound sizes for a tensor, used by XLA.
+  std::vector<int64_t> sizes_;
+  // Stores which dimensions are symbolic
+  // If nullopt, either it hasn't been initialized or the symbolic
+  // dimensions are not calculable
+  std::optional<std::vector<bool>> is_symbolic_ = std::nullopt;
+};
+
+TORCH_API std::ostream& operator<<(std::ostream& out, const Shape& shape);
+
+TORCH_API bool symbolicShapeEnabled();
+// Calculate and applies symbolic shapes onto the
+// Shape objects passed to result_shapes
+TORCH_API void applySymbolicShapesOnLT(
+    const char* schema_str,
+    std::vector<c10::IValue> args,
+    std::vector<Shape>& result_shapes);
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/shape_inference.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/shape_inference.h
new file mode 100644
index 0000000000000000000000000000000000000000..9d1e9a1cd4c0087e9361f29e5c5319586cb05a3f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/shape_inference.h
@@ -0,0 +1,127 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Tensor.h>
+#include <c10/core/ScalarType.h>
+#include <c10/core/SymInt.h>
+#include <c10/core/SymIntArrayRef.h>
+#include <c10/core/SymNodeImpl.h>
+#include <c10/macros/Export.h>
+#include <torch/csrc/lazy/backend/backend_data.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/shape.h>
+#include <torch/csrc/lazy/core/tensor.h>
+#include <optional>
+#include <vector>
+
+namespace torch::lazy {
+// Turn clang-format off, as we rely on the whole signature being on one line
+// for codegen.
+// clang-format off
+TORCH_API std::vector<torch::lazy::Shape> compute_shape__adaptive_avg_pool2d(const at::Tensor & self, at::IntArrayRef output_size);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape__adaptive_avg_pool2d_backward(const at::Tensor & grad_output, const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape__adaptive_avg_pool3d(const at::Tensor & self, at::IntArrayRef output_size);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape__adaptive_avg_pool3d_backward(const at::Tensor & grad_output, const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_abs(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_arange_out(const at::Scalar & start, const at::Scalar & end, const at::Scalar & step, at::Tensor & out);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_bernoulli(const at::Tensor & self, ::std::optional<at::Generator> generator);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_bernoulli(const at::Tensor & self, double p, ::std::optional<at::Generator> generator);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_binary_cross_entropy(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_binary_cross_entropy_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_cat(at::TensorList tensors, int64_t dim);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_cholesky(const at::Tensor & self, bool upper);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_clamp_min(const at::Tensor & self, const at::Scalar & min);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_clone(const at::Tensor & self, ::std::optional<at::MemoryFormat> memory_format);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_constant_pad_nd(const at::Tensor & self, at::IntArrayRef pad, const at::Scalar & value);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_convolution(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_convolution_backward(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & weight, at::OptionalIntArrayRef bias_sizes, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups, ::std::array<bool,3> output_mask);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_embedding(const at::Tensor & weight, const at::Tensor & indices, int64_t padding_idx, bool scale_grad_by_freq, bool sparse);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_embedding_dense_backward(const at::Tensor & grad_output, const at::Tensor & indices, int64_t num_weights, int64_t padding_idx, bool scale_grad_by_freq);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_expand(const at::Tensor & self, at::IntArrayRef size, bool implicit);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_expand(const at::Tensor & self, c10::SymIntArrayRef size, bool implicit);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_flip(const at::Tensor & self, at::IntArrayRef dims);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_glu_backward(const at::Tensor & grad_output, const at::Tensor & self, int64_t dim);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_glu_jvp(const at::Tensor & glu, const at::Tensor & x, const at::Tensor & dx, int64_t dim);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_grid_sampler_2d(const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_grid_sampler_2d_backward(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners, ::std::array<bool,2> output_mask);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_index_select(const at::Tensor & self, int64_t dim, const at::Tensor & index);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_inverse(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_isnan(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_log_sigmoid_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & buffer);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_log_sigmoid_forward(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_logdet(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_logical_and(const at::Tensor & self, const at::Tensor & other);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_logical_not(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_logical_or(const at::Tensor & self, const at::Tensor & other);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_logical_xor(const at::Tensor & self, const at::Tensor & other);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_masked_fill(const at::Tensor & self, const at::Tensor & mask, const at::Scalar & value);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_masked_fill(const at::Tensor & self, const at::Tensor & mask, const at::Tensor & value);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_max(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_mean(const at::Tensor & self, ::std::optional<at::ScalarType> dtype);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_min(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_mv(const at::Tensor & self, const at::Tensor & vec);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_batch_norm(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, bool training, double momentum, double eps);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_batch_norm_backward(const at::Tensor & grad_out, const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, const ::std::optional<at::Tensor> & save_mean, const ::std::optional<at::Tensor> & save_invstd, bool train, double eps, ::std::array<bool,3> output_mask);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_dropout(const at::Tensor & input, double p, ::std::optional<bool> train);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_dropout_backward(const at::Tensor & grad_output, const at::Tensor & mask, double scale);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_layer_norm(const at::Tensor & input, at::IntArrayRef normalized_shape, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, double eps);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_layer_norm_backward(const at::Tensor & grad_out, const at::Tensor & input, at::IntArrayRef normalized_shape, const at::Tensor & mean, const at::Tensor & rstd, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, ::std::array<bool,3> output_mask);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_new_empty_strided(const at::Tensor & self, at::IntArrayRef size, at::IntArrayRef stride, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_nll_loss2d_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, int64_t ignore_index, const at::Tensor & total_weight);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_nll_loss2d_forward(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, int64_t ignore_index);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_nonzero(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_normal_functional(const at::Tensor & self, double mean, double std, ::std::optional<at::Generator> generator);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_random(const at::Tensor & self, ::std::optional<at::Generator> generator);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_random(const at::Tensor & self, int64_t to, ::std::optional<at::Generator> generator);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_random(const at::Tensor & self, int64_t from, ::std::optional<int64_t> to, ::std::optional<at::Generator> generator);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_relu(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_repeat(const at::Tensor & self, at::IntArrayRef repeats);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_slogdet(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_smooth_l1_loss_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction, double beta);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_sort(const at::Tensor & self, int64_t dim, bool descending);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_stack(at::TensorList tensors, int64_t dim);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_std(const at::Tensor & self, bool unbiased);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_std(const at::Tensor & self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_std(const at::Tensor & self, at::OptionalIntArrayRef dim, const ::std::optional<at::Scalar> & correction, bool keepdim);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_sum(const at::Tensor & self, ::std::optional<at::ScalarType> dtype);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape__to_copy(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, bool non_blocking, ::std::optional<at::MemoryFormat> memory_format);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_take(const at::Tensor & self, const at::Tensor & index);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_trace(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_zero(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_narrow_copy_symint(const at::Tensor & self, int64_t dim, int64_t start, c10::SymInt length);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_hardswish(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_hardswish_backward(const at::Tensor & grad_output, const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_selu(const at::Tensor & self);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_uniform(const at::Tensor & self, double from, double to, ::std::optional<at::Generator> generator);
+
+// Non-Native ops
+TORCH_API std::vector<Shape> compute_shape_scalar(const at::Scalar& value, const at::ScalarType& type);
+TORCH_API std::vector<Shape> compute_shape_expand(const Output& input0, const std::vector<int64_t>& size, const bool& is_scalar_expand);
+TORCH_API std::vector<Shape> compute_shape_view(const Output& input0, const std::vector<int64_t>& output_sizes);
+TORCH_API std::vector<Shape> compute_shape_cast(const Output& input0, const at::ScalarType& dtype, const ::std::optional<at::ScalarType>& stype);
+
+// View Ops
+// (Now that functionalization pass is used, we should kill these in a later PR)
+TORCH_API std::vector<Shape> compute_shape_as_strided_view_update(const Output& target, const Output& input, const std::vector<int64_t>& size, const std::vector<int64_t>& stride, const int64_t& storage_offset);
+TORCH_API std::vector<Shape> compute_shape_as_strided(const Output& input, const std::vector<int64_t>& size, const std::vector<int64_t>& stride, const int64_t& storage_offset);
+TORCH_API std::vector<Shape> compute_shape_diagonal_view_update(const Output& target, const Output& input, const int64_t& offset, const int64_t& dim1, const int64_t& dim2);
+TORCH_API std::vector<Shape> compute_shape_diagonal(const Output& input, const int64_t& offset, const int64_t& dim1, const int64_t& dim2);
+TORCH_API std::vector<Shape> compute_shape_narrow_view_update(const Output& input, const Output& source, const std::vector<int64_t>& base_indices);
+TORCH_API std::vector<Shape> compute_shape_narrow(const Output& input, const std::vector<int64_t>& base_indices, const std::vector<int64_t>& sizes);
+TORCH_API std::vector<Shape> compute_shape_permute(const Output& input, const std::vector<int64_t>& dims);
+TORCH_API std::vector<Shape> compute_shape_resize(const Output& input, const std::vector<int64_t>& size);
+TORCH_API std::vector<Shape> compute_shape_select_view_update(const Output& target, const Output& source, const int64_t& dim, const int64_t& start, const int64_t& end, const int64_t& stride);
+TORCH_API std::vector<Shape> compute_shape_select(const Output& input, const int64_t& dim, const int64_t& start, const int64_t& end, const int64_t& stride);
+TORCH_API std::vector<Shape> compute_shape_squeeze(const Output& input, const int& dim);
+TORCH_API std::vector<Shape> compute_shape_unsqueeze(const Output& input, const int& dim);
+
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_select_scatter(const at::Tensor & self, const at::Tensor & src, int64_t dim, int64_t index);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_diagonal_scatter(const at::Tensor & self, const at::Tensor & src, int64_t offset, int64_t dim1, int64_t dim2);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_slice_scatter_symint(const at::Tensor & self, const at::Tensor & src, int64_t dim, ::std::optional<c10::SymInt> start, ::std::optional<c10::SymInt> end, c10::SymInt step);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_as_strided_scatter_symint(const at::Tensor & self, const at::Tensor & src, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<c10::SymInt> storage_offset);
+// clang-format on
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..16a363079fa8fe784f82d8b0dc391264e6edd76f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor.h
@@ -0,0 +1,270 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/SymNodeImpl.h>
+#include <c10/util/intrusive_ptr.h>
+#include <torch/csrc/lazy/backend/backend_data.h>
+#include <torch/csrc/lazy/backend/backend_device.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/util.h>
+
+namespace torch::lazy {
+
+class TORCH_API SymNodeImpl : public c10::SymNodeImpl {
+ public:
+  SymNodeImpl(NodePtr ptr) : node_(std::move(ptr)) {}
+  NodePtr node_;
+};
+
+class LazyTensor;
+using LazyTensorPtr = c10::intrusive_ptr<LazyTensor>;
+
+class TORCH_API LazyTensor : public c10::intrusive_ptr_target {
+ public:
+  // This is the core lazy tensor data structure where all the tensor data is
+  // held. The lazy tensor is nothing more than a shared pointer to a Data
+  // object.
+  struct Data {
+    Data(BackendDataPtr handle, BackendDevice device)
+        : handle(std::move(handle)),
+          device(std::move(device)),
+          unique_id(GetNextTensorId()) {}
+    Data(Value ir_value, BackendDevice device)
+        : ir_value(std::move(ir_value)),
+          device(std::move(device)),
+          unique_id(GetNextTensorId()) {}
+    Data(at::Tensor tensor_data, BackendDevice device)
+        : tensor_data(std::move(tensor_data)),
+          device(std::move(device)),
+          unique_id(GetNextTensorId()) {}
+    // TODO(alanwaketan): Remove this ctor. This is a
+    // temporary ctor to ease XLA LTC migration. It depends on
+    // XLA's Functionalization integration.
+    Data(BackendDevice device)
+        : device(std::move(device)), unique_id(GetNextTensorId()) {}
+
+    Data(Data&& other) = delete;
+    Data(const Data&) = delete;
+    Data& operator=(const Data&) = delete;
+    Data& operator=(Data&&) = delete;
+    virtual ~Data();
+
+    BackendDataPtr handle;
+    Value ir_value;
+    std::optional<at::Tensor> tensor_data;
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+    const BackendDevice device;
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+    const int64_t unique_id = 0;
+    size_t generation = 1;
+  };
+
+  static LazyTensorPtr Create(
+      const at::Tensor& tensor,
+      const BackendDevice& device);
+  static LazyTensorPtr Create(Value ir_value, const BackendDevice& device);
+  static LazyTensorPtr Create(const BackendDataPtr& handle);
+  static LazyTensorPtr Create(std::shared_ptr<Data> data);
+
+  // The default ctor previously created a null LazyTensor (one with no 'data'
+  // obj). Creating a null LazyTensor is no longer possible, since the same can
+  // be achieved by creating a null LazyTensorPtr and it is way too confusing to
+  // have to check both lazy_tensor_ptr && *lazy_tensor_ptr, so everywhere that
+  // used to rely on a LazyTensor obj with a null Data can now rely on a null
+  // LazyTensorPtr instead.
+  LazyTensor() = delete;
+  LazyTensor(const LazyTensor&) = default;
+  LazyTensor(LazyTensor&&) noexcept = default;
+  LazyTensor& operator=(const LazyTensor&) = default;
+  LazyTensor& operator=(LazyTensor&&) noexcept = default;
+
+  ~LazyTensor() override = default;
+
+  size_t generation() const {
+    return data()->generation;
+  }
+
+  // Override it to use your own Shape.
+  virtual int64_t size(int64_t dim) const;
+
+  // Override it to use your own graph executor.
+  virtual at::Tensor ToTensor(bool detached);
+
+  void ShallowCopyTo(const LazyTensorPtr& dest) const;
+
+  // Assigns the tensor value to the lazy tensor.
+  void SetTensor(at::Tensor tensor);
+
+  void UpdateFromTensor(const at::Tensor& tensor, bool sync);
+  void UpdateFromTensorOut(const at::Tensor& tensor);
+  void UpdateFromTensorOut(const LazyTensorPtr& tensor);
+
+  const std::shared_ptr<Data>& data() const;
+
+  // Override it to use your own type conversion.
+  virtual at::ScalarType dtype() const;
+
+  MaybeRef<Shape> shape() const;
+
+  const BackendDevice& GetDevice() const;
+  int64_t GetUniqueId() const;
+
+  // Fetches the data behind the tensor. If the tensor has a graph defining
+  // its current value, executes the graph and fetches the data result.
+  BackendDataPtr GetDataHandle();
+
+  // Fetches the current value of the data, which can be missing (nullptr)
+  // in case the tensor has a graph defining its current value,
+  BackendDataPtr CurrentDataHandle() const;
+
+  void SetDataHandle(BackendDataPtr handle);
+  void SetDataHandle(BackendDataPtr handle, bool sync);
+
+  // Retrieves the current IR Node, or nullptr in case no active IR Node is
+  // available.
+  Value CurrentIrValue() const;
+
+  // Retrieves the IR Node representing this LazyTensor. One will be created if
+  // missing. Note that although this is a const API, it actually changes the
+  // internal state of the object.
+  Value GetIrValue() const;
+
+  void SetIrValue(Value ir_value);
+  void SetInPlaceIrValue(Value ir_value);
+
+  std::optional<at::Tensor> CurrentTensorData() const;
+
+  std::vector<LazyTensorPtr> MakeOutputTensors(const NodePtr& node) const;
+
+  LazyTensorPtr CopyTensorToDevice(const BackendDevice& device);
+
+  // Applies the queue of operations in preparation for using the data.
+  // Override it to use your own graph executor.
+  virtual void ApplyPendingGraph();
+
+  // Override it to set extra information.
+  virtual void AssignIrValue(Value ir_value) const;
+
+ protected:
+  explicit LazyTensor(std::shared_ptr<Data> data);
+
+  void SetTensorData(at::Tensor tensor_data);
+
+  // We build a graph accumulating operations, but at a given point we
+  // need to force a rendering, otherwise the graph can grow without control.
+  // Think:
+  //   for i in range(0, 100000):
+  //     a = a + b
+  void TryLimitGraphSize();
+
+  // Override it to instantiate your own data.
+  virtual Value GetIrValueForTensor(
+      const at::Tensor& tensor,
+      const BackendDevice& device) const;
+
+  Value CreateTensorNode(const BackendDataPtr& data, bool read_only) const;
+
+ private:
+  LazyTensor(const at::Tensor& tensor, const BackendDevice& device);
+  LazyTensor(Value ir_value, const BackendDevice& device);
+  explicit LazyTensor(const BackendDataPtr& handle);
+
+  static int64_t GetNextTensorId();
+
+  std::shared_ptr<Data> data_;
+};
+
+// Utils to convert at::Tensor to LazyTensor, and vice versa.
+
+// Section 0: c10::Tensorlist ==> lazy::TensorList
+// note: GetTensorList is not totally parallel to GetLtcTensor; A TensorList
+// skips
+//       the LazyTensor wrappers, assuming that the list of underlying IR nodes
+//       is actually more useful for downstream computations.  TBD.
+TORCH_API torch::lazy::Value GetTensorList(at::ITensorListRef tensors);
+
+// Section 1: at::Tensor => LazyTensor.
+// Extracts the LazyTensor out of an at::Tensor. Returns a null LazyTensor
+// if the tensor is not a lazy tensor.
+TORCH_API LazyTensorPtr TryGetLtcTensor(const at::Tensor& tensor);
+
+// Extracts the LazyTensor out of an at::Tensor. Throws an exception
+// if the tensor is not a lazy tensor.
+TORCH_API LazyTensorPtr GetLtcTensor(const at::Tensor& tensor);
+
+// Same as above, applied to a list of tensors.
+TORCH_API std::vector<LazyTensorPtr> GetLtcTensors(
+    c10::ArrayRef<at::Tensor> tensors);
+
+// If tensor is a lazy tensor type, returns the LazyTensor embedded within it,
+// otherwise creates a new lazy tensor type with tensor as data.
+TORCH_API LazyTensorPtr GetOrCreateLtcTensor(
+    const std::optional<at::Tensor>& tensor,
+    const BackendDevice& device);
+
+TORCH_API LazyTensorPtr GetLtcTensorOrCreateForWrappedNumber(
+    const at::Tensor& tensor,
+    const BackendDevice& device);
+
+// Section 2: LazyTensor => at::Tensor.
+// Creates an ATen tensor from an LazyTensor.
+TORCH_API at::Tensor CreateAtenFromLtcTensor(const LazyTensorPtr& ltc_tensor);
+TORCH_API at::Tensor CreateAtenFromLtcTensor(LazyTensor&& ltc_tensor);
+
+// Note [Lazy Tensor Functionalization]
+// The functionalization pass is implemented by wrapping all TensorImpl
+// objects in C++ with an extra FunctionalTensorWrapper object,
+// that knows how to perform functionalization
+//
+// Certain functions in the aten API serve as entry/exit points for
+// functionalization, where we need to perform the wrapping/unwrapping:
+// - aten::to.device
+// - aten::empty
+
+// Given a non-lazy tensor, this function creates a lazy tensor on the specified
+// (lazy) device. The functionalize_output determines whether or not we should
+// wrap the output in a "functional wrapper".
+//
+// How do you know whether to pass true/false for functionalize_output?
+//
+// Case 1: nonlazy -> lazy
+//   If you're implementing a function that takes in nonlazy tensors and returns
+//   lazy tensors, then you should think of that function as an "entrypoint" to
+//   functionalization, and use functionalize_output=true Examples include:
+//   - factory functions (the LTC kernel for at::empty)
+//   - CPU -> Lazy device conversions (the LTC kernel for at::to_device)
+//
+// Case 2: lazy -> lazy
+//   If you're implementing a function that takes in lazy tensors and returns
+//   lazy tensors,
+//   **but** requires creating lazy tensors internally,
+//   then you can assume that the current function is running inside of some
+//   outer context where functionalization is already running, that will take
+//   care of doing the wrapping for you, and use functionalize_output=true
+//   Examples include:
+//   - CPU fallback (takes in lazy tensors, converts to cpu, calls kernel,
+//   converts returns back to lazy tensors).
+TORCH_API at::Tensor to_lazy_tensor(
+    const at::Tensor& self,
+    const c10::TensorOptions& options,
+    at::Device device,
+    bool non_blocking,
+    bool functionalize_output);
+
+template <size_t... Indices>
+auto TupleAtenFromLtcTensorsImpl(
+    const std::vector<LazyTensorPtr>& tensors,
+    std::index_sequence<Indices...> /*unused*/) {
+  return std::make_tuple(CreateAtenFromLtcTensor(tensors[Indices])...);
+}
+
+template <size_t N>
+auto TupleAtenFromLtcTensors(const std::vector<LazyTensorPtr>& tensors) {
+  return TupleAtenFromLtcTensorsImpl(tensors, std::make_index_sequence<N>{});
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..161008918c8ccb22011e0f13ef2cae0e16f5b133
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor_impl.h
@@ -0,0 +1,66 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/Tensor.h>
+#include <c10/core/SymIntArrayRef.h>
+#include <c10/core/TensorImpl.h>
+
+#include <torch/csrc/lazy/core/tensor.h>
+
+namespace torch::lazy {
+
+// Tensor implementation class used to be fed to the at::Tensor.
+// Its scope is just to handle an LazyTensor.
+class TORCH_API LTCTensorImpl final : public c10::TensorImpl {
+ public:
+  explicit LTCTensorImpl(const LazyTensorPtr& tensor);
+  explicit LTCTensorImpl(const LazyTensor& tensor);
+  explicit LTCTensorImpl(LazyTensor&& tensor);
+
+  LazyTensorPtr tensor() {
+    return tensor_;
+  }
+
+  void set_tensor(const LazyTensorPtr& lazy_tensor);
+
+  void force_refresh_sizes() {
+    generation_ = 0;
+  }
+
+  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
+      const c10::VariableVersion& version_counter,
+      bool allow_tensor_metadata_change) const override;
+
+  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
+      c10::VariableVersion&& version_counter,
+      bool allow_tensor_metadata_change) const override;
+
+  void shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) override;
+
+  at::IntArrayRef sizes_custom() const override;
+  at::IntArrayRef strides_custom() const override;
+  int64_t numel_custom() const override;
+  int64_t storage_offset_custom() const override;
+  int64_t dim_custom() const override;
+  bool is_strides_like_custom(at::MemoryFormat memory_format) const override;
+  c10::SymBool sym_is_non_overlapping_and_dense_custom() const override;
+
+  c10::SymBool sym_is_contiguous_custom(
+      at::MemoryFormat memory_format) const override;
+  c10::SymIntArrayRef sym_sizes_custom() const override;
+  c10::SymIntArrayRef sym_strides_custom() const override;
+  c10::SymInt sym_numel_custom() const override;
+
+ private:
+  void setup_size_properties();
+
+  LazyTensorPtr tensor_;
+  mutable std::optional<std::vector<c10::SymInt>> sym_sizes_;
+  size_t generation_{0};
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..e47484f16265b9675645aa947f06d5cdf59d54cf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/tensor_util.h
@@ -0,0 +1,81 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/backend/backend_interface.h>
+#include <torch/csrc/lazy/core/shape.h>
+
+#include <ATen/FunctionalTensorWrapper.h>
+
+#include <string>
+#include <vector>
+
+namespace torch::lazy {
+
+TORCH_API std::vector<int64_t> ComputeArrayStrides(
+    c10::ArrayRef<int64_t> sizes);
+
+TORCH_API std::vector<at::Tensor> DataHandlesToTensors(
+    c10::ArrayRef<BackendDataPtr> data_handles,
+    at::ScalarType dest_element_type);
+
+// Uploads an ATEN tensor data to the device and fetches the corresponding
+// device data handle.
+TORCH_API BackendDataPtr
+TensorToDataHandle(const at::Tensor& tensor, const BackendDevice& device);
+
+// Retrieves the device data handles by parallel uploading data onto the
+// corresponding devices.
+TORCH_API std::vector<BackendDataPtr> CreateTensorsData(
+    const std::vector<at::Tensor>& tensors,
+    const std::vector<BackendDevice>& devices);
+
+// Makes a deep copy of an ATEN tensor.
+inline at::Tensor CopyTensor(const at::Tensor& ref) {
+  return ref.to(ref.options(), /*non_blocking=*/false, /*copy=*/true);
+}
+
+// Same as above, with an additional cast.
+inline at::Tensor CopyTensor(
+    const at::Tensor& ref,
+    at::ScalarType dest_type,
+    bool copy = true) {
+  return ref.to(ref.options().dtype(dest_type), /*non_blocking=*/false, copy);
+}
+
+template <typename T, typename S>
+T OptionalOr(const std::optional<S>& value, T defval) {
+  return value ? static_cast<T>(*value) : defval;
+}
+
+// Unwraps tensor to target dtype if it's a wrapped number.
+inline at::Tensor UnwrapNumber(const at::Tensor& tensor, at::ScalarType dtype) {
+  return tensor.unsafeGetTensorImpl()->is_wrapped_number() ? tensor.to(dtype)
+                                                           : tensor;
+}
+
+template <typename T>
+at::Scalar MakeIntScalar(T value) {
+  return at::Scalar(static_cast<int64_t>(value));
+}
+
+// Routing values to device data maximizes the changes for compilation cache
+// hits, but it can prevent the compiler to perform optimizations. So tensor
+// values which are within a given set, are routed to constant scalars if this
+// API returns true.
+TORCH_API bool IsSpecialScalar(const at::Scalar& value);
+
+// Note: returns a reference instead of a fresh tensor to avoid refcount bumps.
+inline const at::Tensor& maybe_unwrap_functional(const at::Tensor& tensor) {
+  if (at::functionalization::impl::isFunctionalTensor(tensor)) {
+    return at::functionalization::impl::unsafeGetFunctionalWrapper(tensor)
+        ->value();
+  } else {
+    return tensor;
+  }
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/thread_pool.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/thread_pool.h
new file mode 100644
index 0000000000000000000000000000000000000000..ac54c00f81a1bc0b7d47ce5d6402b98f1007f104
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/thread_pool.h
@@ -0,0 +1,41 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * This file is adapted from PyTorch/XLA
+ * https://github.com/pytorch/xla/blob/e0e5f937a0ba8d904f9608137dc8c51ba439df2d/third_party/xla_client/metrics.h
+ */
+
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <thread>
+
+#include <c10/macros/Export.h>
+
+namespace torch::lazy {
+
+// NOLINTNEXTLINE(cppcoreguidelines-special-member-functions)
+class TORCH_API Completion {
+ public:
+  class Data;
+
+  explicit Completion(std::shared_ptr<Data> data);
+
+  ~Completion();
+
+  void Wait();
+
+ private:
+  std::shared_ptr<Data> data_;
+};
+
+// Schedules a closure which might wait for IO or other events/conditions.
+TORCH_API void ScheduleIoClosure(std::function<void()> closure);
+TORCH_API Completion
+ScheduleIoClosureWithCompletion(std::function<void()> closure);
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/trie.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/trie.h
new file mode 100644
index 0000000000000000000000000000000000000000..ca5fc4645c2e69fd3894c382a7fcf47dc64ff398
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/trie.h
@@ -0,0 +1,82 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <atomic>
+#include <list>
+
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/metrics.h>
+
+namespace torch::lazy {
+
+struct TORCH_API TrieNode {
+  static size_t GetNextUniqueId() {
+    static thread_local size_t id_generator = 0;
+    return id_generator++;
+  }
+
+  size_t unique_id;
+  size_t hit_counter;
+  NodePtr ir_node;
+  std::list<std::shared_ptr<TrieNode>> successors;
+
+  TrieNode() : unique_id(GetNextUniqueId()), hit_counter(0), ir_node(nullptr) {}
+  explicit TrieNode(NodePtr node)
+      : unique_id(GetNextUniqueId()),
+        hit_counter(0),
+        ir_node(std::move(node)) {}
+};
+
+class TORCH_API TrieCache {
+ public:
+  static TrieCache* Get();
+
+  TrieNode* Current() const;
+  // Take an iterator as the input because we want to move the corresponding
+  // node in the successor list to achieve a LRU caching effect
+  void SetCurrent(std::list<std::shared_ptr<TrieNode>>::iterator& iter);
+  // Used in MarkStep to indicate the end of one tracing
+  void ResetCurrent();
+
+  // Create a new TrieNode for ir_node and insert into the TrieCache
+  void Insert(NodePtr ir_node);
+
+  // Clear all TrieCache nodes
+  // TODO: Because we don't expect user to explicitly call this function via
+  // a Python API, we may need to introduce a threshold on the size of the cache
+  // to avoid holding tensors for too long.
+  void Clear();
+
+  void DumpToDotFile(const std::string& file_name);
+
+ private:
+  TrieCache();
+
+  std::shared_ptr<TrieNode> root_;
+  TrieNode* current_;
+};
+
+template <typename T, typename... Args>
+NodePtr LookupNodeFromTrieCache(Args&&... args) {
+  auto& successors = TrieCache::Get()->Current()->successors;
+  for (auto it = successors.begin(); it != successors.end(); it++) {
+    NodePtr ir_node = (*it)->ir_node;
+    const T* concrete_node = NodeCast<T>(ir_node.get());
+    if (concrete_node &&
+        concrete_node->CanBeReused(std::forward<Args>(args)...)) {
+      TORCH_LAZY_COUNTER(
+          "IrNodeReused_" + c10::demangle((typeid(T).name())), 1);
+      (*it)->hit_counter++;
+      TrieCache::Get()->SetCurrent(it);
+      return ir_node;
+    }
+  }
+  return nullptr;
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/unique.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/unique.h
new file mode 100644
index 0000000000000000000000000000000000000000..718cac504751dc9b08fc72d69fcf91dbfbe77376
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/unique.h
@@ -0,0 +1,59 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * Unique in this file is adapted from PyTorch/XLA
+ * https://github.com/pytorch/xla/blob/e0e5f937a0ba8d904f9608137dc8c51ba439df2d/third_party/xla_client/unique.h
+ */
+
+#pragma once
+
+#include <optional>
+
+#include <functional>
+#include <set>
+
+namespace torch::lazy {
+
+// Helper class to allow tracking zero or more things, which should be forcibly
+// be one only thing.
+template <typename T, typename C = std::equal_to<T>>
+class Unique {
+ public:
+  std::pair<bool, const T&> set(const T& value) {
+    if (value_) {
+      TORCH_CHECK(C()(*value_, value), "'", *value_, "' vs '", value);
+      return std::pair<bool, const T&>(false, *value_);
+    }
+    value_ = value;
+    return std::pair<bool, const T&>(true, *value_);
+  }
+
+  operator bool() const {
+    return value_.has_value();
+  }
+  operator const T&() const {
+    return *value_;
+  }
+  const T& operator*() const {
+    return *value_;
+  }
+  const T* operator->() const {
+    return value_.operator->();
+  }
+
+  std::set<T> AsSet() const {
+    std::set<T> vset;
+    if (value_.has_value()) {
+      vset.insert(*value_);
+    }
+    return vset;
+  }
+
+ private:
+  std::optional<T> value_;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/util.h
new file mode 100644
index 0000000000000000000000000000000000000000..4324148de300340fa58151cc73ee6032f1f530ae
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/core/util.h
@@ -0,0 +1,130 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/**
+ * Most of the utils in this file is adapted from PyTorch/XLA
+ * https://github.com/pytorch/xla/blob/e0e5f937a0ba8d904f9608137dc8c51ba439df2d/third_party/xla_client/util.h
+ */
+
+#pragma once
+
+#include <exception>
+#include <functional>
+#include <vector>
+
+#include <c10/util/OptionalArrayRef.h>
+#include <optional>
+
+namespace torch::lazy {
+
+// Similar to c10::scope_exit but with a status.
+// TODO(alanwaketan): Consolidate it with c10::scope_exit.
+template <typename T>
+class Cleanup {
+ public:
+  using StatusType = T;
+
+  explicit Cleanup(std::function<void(StatusType&&)>&& func)
+      : func_(std::move(func)) {}
+  Cleanup(Cleanup&& ref) noexcept
+      : func_(std::move(ref.func_)), status_(std::move(ref.status_)) {}
+  Cleanup(const Cleanup&) = delete;
+
+  ~Cleanup() {
+    if (func_ != nullptr) {
+      func_(std::move(status_));
+    }
+  }
+
+  Cleanup& operator=(const Cleanup&) = delete;
+
+  Cleanup& operator=(Cleanup&& ref) noexcept {
+    if (this != &ref) {
+      func_ = std::move(ref.func_);
+      status_ = std::move(ref.status_);
+    }
+    return *this;
+  }
+
+  void Release() {
+    func_ = nullptr;
+  }
+
+  void SetStatus(StatusType&& status) {
+    status_ = std::move(status);
+  }
+
+  const StatusType& GetStatus() const {
+    return status_;
+  }
+
+ private:
+  std::function<void(StatusType&&)> func_;
+  StatusType status_;
+};
+
+using ExceptionCleanup = Cleanup<std::exception_ptr>;
+
+// Allows APIs which might return const references and values, to not be forced
+// to return values in the signature.
+// TODO(alanwaketan): This is clever, but is there really no std or c10
+// supports? Needs more investigations.
+template <typename T>
+class MaybeRef {
+ public:
+  /* implicit */ MaybeRef(const T& ref) : ref_(ref) {}
+  /* implicit */ MaybeRef(T&& value)
+      : storage_(std::move(value)), ref_(*storage_) {}
+
+  const T& Get() const {
+    return ref_;
+  }
+  const T& operator*() const {
+    return Get();
+  }
+  operator const T&() const {
+    return Get();
+  }
+
+  bool IsStored() const {
+    return storage_.has_value();
+  }
+
+ private:
+  std::optional<T> storage_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const T& ref_;
+};
+
+template <typename T>
+std::vector<T> Iota(size_t size, T init = 0, T incr = 1) {
+  std::vector<T> result(size);
+  T value = init;
+  for (size_t i = 0; i < size; ++i, value += incr) {
+    result[i] = value;
+  }
+  return result;
+}
+
+template <typename T, typename S>
+std::vector<T> ToVector(const S& input) {
+  return std::vector<T>(input.begin(), input.end());
+}
+
+template <typename T>
+std::optional<std::vector<T>> ToOptionalVector(
+    c10::OptionalArrayRef<T> arrayRef) {
+  if (arrayRef) {
+    return arrayRef->vec();
+  }
+  return std::nullopt;
+}
+
+template <typename T>
+std::underlying_type_t<T> GetEnumValue(T value) {
+  return static_cast<std::underlying_type_t<T>>(value);
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyIr.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyIr.h
new file mode 100644
index 0000000000000000000000000000000000000000..dfd6a881958c2d7eb0cbbe5006d4302669450d4d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyIr.h
@@ -0,0 +1,10312 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// This file contains autogenerated LazyTensor IR nodes
+#include <ATen/core/Formatting.h>
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/lazy/core/hash.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/shape.h>
+#include <optional>
+#include <vector>
+#include "torch/csrc/lazy/ts_backend/ts_node.h"
+
+namespace torch {
+namespace lazy {
+using at::operator<<;
+
+// kNullValue is used to contribute a static hash value any time
+// a node has an Optional<Value> input that is nullopt.  It is important
+// to differentiate between HASH(std::nullopt, something) and HASH(something, std::nullopt),
+// and using kNullValue in the hash function in the order of arguments
+// serves this purpose.
+static const torch::lazy::Value kNullValue = torch::lazy::Value();
+
+class AdaptiveAvgPool2d : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::_adaptive_avg_pool2d);
+  }
+
+  AdaptiveAvgPool2d(const torch::lazy::Value& self, const ::std::vector<int64_t>& output_size, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AdaptiveAvgPool2d::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(output_size)),
+        output_size(output_size)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", output_size=" << output_size;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& output_size) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->output_size == output_size);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("output_size", output_size);
+
+    torch::lazy::TSOpVector _adaptive_avg_pool2d_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_adaptive_avg_pool2d_out.size(), 1);
+
+    return _adaptive_avg_pool2d_out;
+
+  }
+
+
+  ::std::vector<int64_t> output_size;
+
+
+};
+
+class AdaptiveAvgPool2dBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::_adaptive_avg_pool2d_backward);
+  }
+
+  AdaptiveAvgPool2dBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AdaptiveAvgPool2dBackward::ClassOpKind(),
+              OpList{grad_output, self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector _adaptive_avg_pool2d_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_adaptive_avg_pool2d_backward_out.size(), 1);
+
+    return _adaptive_avg_pool2d_backward_out;
+
+  }
+
+
+
+
+
+};
+
+class LogSoftmax : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::_log_softmax);
+  }
+
+  LogSoftmax(const torch::lazy::Value& self, const int64_t& dim, const bool& half_to_float, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LogSoftmax::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, half_to_float)),
+        dim(dim),
+        half_to_float(half_to_float)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", half_to_float=" << half_to_float;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const bool& half_to_float) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim &&
+        this->half_to_float == half_to_float);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("half_to_float", half_to_float);
+
+    torch::lazy::TSOpVector _log_softmax_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_log_softmax_out.size(), 1);
+
+    return _log_softmax_out;
+
+  }
+
+
+  int64_t dim;
+  bool half_to_float;
+
+
+};
+
+class LogSoftmaxBackwardData : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::_log_softmax_backward_data);
+  }
+
+  LogSoftmaxBackwardData(const torch::lazy::Value& grad_output, const torch::lazy::Value& output, const int64_t& dim, const at::ScalarType& input_dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LogSoftmaxBackwardData::ClassOpKind(),
+              OpList{grad_output, output},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, input_dtype)),
+        dim(dim),
+        input_dtype(input_dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", input_dtype=" << input_dtype;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& output, const int64_t& dim, const at::ScalarType& input_dtype) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == output &&
+        this->dim == dim &&
+        this->input_dtype == input_dtype);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("input_dtype", input_dtype);
+
+    torch::lazy::TSOpVector _log_softmax_backward_data_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_log_softmax_backward_data_out.size(), 1);
+
+    return _log_softmax_backward_data_out;
+
+  }
+
+
+  int64_t dim;
+  at::ScalarType input_dtype;
+
+
+};
+
+class ReshapeAliasCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::_reshape_alias_copy);
+  }
+
+  ReshapeAliasCopy(const torch::lazy::Value& self, const ::std::vector<int64_t>& size, const ::std::vector<int64_t>& stride, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ReshapeAliasCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(size, stride)),
+        size(size),
+        stride(stride)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", size=" << size;
+    ss << ", stride=" << stride;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& size, const ::std::vector<int64_t>& stride) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->size == size &&
+        this->stride == stride);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("size", size);
+    arguments.emplace_back("stride", stride);
+
+    torch::lazy::TSOpVector _reshape_alias_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_reshape_alias_copy_out.size(), 1);
+
+    return _reshape_alias_copy_out;
+
+  }
+
+
+  ::std::vector<int64_t> size;
+  ::std::vector<int64_t> stride;
+
+
+};
+
+class Softmax : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::_softmax);
+  }
+
+  Softmax(const torch::lazy::Value& self, const int64_t& dim, const bool& half_to_float, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Softmax::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, half_to_float)),
+        dim(dim),
+        half_to_float(half_to_float)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", half_to_float=" << half_to_float;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const bool& half_to_float) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim &&
+        this->half_to_float == half_to_float);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("half_to_float", half_to_float);
+
+    torch::lazy::TSOpVector _softmax_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_softmax_out.size(), 1);
+
+    return _softmax_out;
+
+  }
+
+
+  int64_t dim;
+  bool half_to_float;
+
+
+};
+
+class SoftmaxBackwardData : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::_softmax_backward_data);
+  }
+
+  SoftmaxBackwardData(const torch::lazy::Value& grad_output, const torch::lazy::Value& output, const int64_t& dim, const at::ScalarType& input_dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SoftmaxBackwardData::ClassOpKind(),
+              OpList{grad_output, output},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, input_dtype)),
+        dim(dim),
+        input_dtype(input_dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", input_dtype=" << input_dtype;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& output, const int64_t& dim, const at::ScalarType& input_dtype) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == output &&
+        this->dim == dim &&
+        this->input_dtype == input_dtype);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("input_dtype", input_dtype);
+
+    torch::lazy::TSOpVector _softmax_backward_data_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_softmax_backward_data_out.size(), 1);
+
+    return _softmax_backward_data_out;
+
+  }
+
+
+  int64_t dim;
+  at::ScalarType input_dtype;
+
+
+};
+
+class Abs : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::abs);
+  }
+
+  Abs(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Abs::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector abs_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(abs_out.size(), 1);
+
+    return abs_out;
+
+  }
+
+
+
+
+
+};
+
+class AddTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::add);
+  }
+
+  AddTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, const torch::lazy::Value& alpha, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AddTensor::ClassOpKind(),
+              OpList{self, other, alpha},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other, const torch::lazy::Value& alpha) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other &&
+        operand(i++) == alpha);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("alpha", loctx->GetOutputOp(operand(i++)));
+    torch::lazy::TSOpVector add_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(add_out.size(), 1);
+
+    return add_out;
+
+  }
+
+
+
+
+
+};
+
+class Addcdiv : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::addcdiv);
+  }
+
+  Addcdiv(const torch::lazy::Value& self, const torch::lazy::Value& tensor1, const torch::lazy::Value& tensor2, const torch::lazy::Value& value, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Addcdiv::ClassOpKind(),
+              OpList{self, tensor1, tensor2, value},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& tensor1, const torch::lazy::Value& tensor2, const torch::lazy::Value& value) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == tensor1 &&
+        operand(i++) == tensor2 &&
+        operand(i++) == value);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("value", loctx->GetOutputOp(operand(i++)));
+    torch::lazy::TSOpVector addcdiv_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(addcdiv_out.size(), 1);
+
+    return addcdiv_out;
+
+  }
+
+
+
+
+
+};
+
+class Addcmul : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::addcmul);
+  }
+
+  Addcmul(const torch::lazy::Value& self, const torch::lazy::Value& tensor1, const torch::lazy::Value& tensor2, const torch::lazy::Value& value, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Addcmul::ClassOpKind(),
+              OpList{self, tensor1, tensor2, value},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& tensor1, const torch::lazy::Value& tensor2, const torch::lazy::Value& value) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == tensor1 &&
+        operand(i++) == tensor2 &&
+        operand(i++) == value);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("value", loctx->GetOutputOp(operand(i++)));
+    torch::lazy::TSOpVector addcmul_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(addcmul_out.size(), 1);
+
+    return addcmul_out;
+
+  }
+
+
+
+
+
+};
+
+class Addmm : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::addmm);
+  }
+
+  Addmm(const torch::lazy::Value& self, const torch::lazy::Value& mat1, const torch::lazy::Value& mat2, const torch::lazy::Value& beta, const torch::lazy::Value& alpha, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Addmm::ClassOpKind(),
+              OpList{self, mat1, mat2, beta, alpha},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& mat1, const torch::lazy::Value& mat2, const torch::lazy::Value& beta, const torch::lazy::Value& alpha) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == mat1 &&
+        operand(i++) == mat2 &&
+        operand(i++) == beta &&
+        operand(i++) == alpha);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(2);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("beta", loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("alpha", loctx->GetOutputOp(operand(i++)));
+    torch::lazy::TSOpVector addmm_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(addmm_out.size(), 1);
+
+    return addmm_out;
+
+  }
+
+
+
+
+
+};
+
+class AliasCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::alias_copy);
+  }
+
+  AliasCopy(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AliasCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector alias_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(alias_copy_out.size(), 1);
+
+    return alias_copy_out;
+
+  }
+
+
+
+
+
+};
+
+class All : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::all);
+  }
+
+  All(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              All::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector all_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(all_out.size(), 1);
+
+    return all_out;
+
+  }
+
+
+
+
+
+};
+
+class Any : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::any);
+  }
+
+  Any(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Any::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector any_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(any_out.size(), 1);
+
+    return any_out;
+
+  }
+
+
+
+
+
+};
+
+class ArangeStartOut : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::arange);
+  }
+
+  ArangeStartOut(const torch::lazy::Value& start, const torch::lazy::Value& end, const torch::lazy::Value& step, const torch::lazy::Value& out, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ArangeStartOut::ClassOpKind(),
+              OpList{start, end, step, out},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& start, const torch::lazy::Value& end, const torch::lazy::Value& step, const torch::lazy::Value& out) const {
+    size_t i = 0;
+    return (operand(i++) == start &&
+        operand(i++) == end &&
+        operand(i++) == step &&
+        operand(i++) == out);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("out", loctx->GetOutputOp(operand(i++)));
+    torch::lazy::TSOpVector arange_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(arange_out.size(), 1);
+
+    return arange_out;
+
+  }
+
+
+
+
+
+};
+
+class AsStridedCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::as_strided_copy);
+  }
+
+  AsStridedCopy(const torch::lazy::Value& self, const ::std::vector<int64_t>& size, const ::std::vector<int64_t>& stride, const ::std::optional<torch::lazy::Value>& storage_offset, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AsStridedCopy::ClassOpKind(),
+              OpList{self, storage_offset.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(size, stride)),
+        size(size),
+        stride(stride)
+  {
+    has_storage_offset = !!storage_offset;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", size=" << size;
+    ss << ", stride=" << stride;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& size, const ::std::vector<int64_t>& stride, const ::std::optional<torch::lazy::Value>& storage_offset) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        nullable_operand(i++) == storage_offset.value_or(kNullValue) &&
+        this->size == size &&
+        this->stride == stride);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("size", size);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back(has_storage_offset ? loctx->GetOutputOp(operand(i++)) : nullptr);
+
+    torch::lazy::TSOpVector as_strided_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(as_strided_copy_out.size(), 1);
+
+    return as_strided_copy_out;
+
+  }
+
+
+  ::std::vector<int64_t> size;
+  ::std::vector<int64_t> stride;
+  bool has_storage_offset: 1;
+
+};
+
+class AsStridedScatter : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::as_strided_scatter);
+  }
+
+  AsStridedScatter(const torch::lazy::Value& self, const torch::lazy::Value& src, const ::std::vector<int64_t>& size, const ::std::vector<int64_t>& stride, const ::std::optional<torch::lazy::Value>& storage_offset, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AsStridedScatter::ClassOpKind(),
+              OpList{self, src, storage_offset.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(size, stride)),
+        size(size),
+        stride(stride)
+  {
+    has_storage_offset = !!storage_offset;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", size=" << size;
+    ss << ", stride=" << stride;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& src, const ::std::vector<int64_t>& size, const ::std::vector<int64_t>& stride, const ::std::optional<torch::lazy::Value>& storage_offset) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == src &&
+        nullable_operand(i++) == storage_offset.value_or(kNullValue) &&
+        this->size == size &&
+        this->stride == stride);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("size", size);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back(has_storage_offset ? loctx->GetOutputOp(operand(i++)) : nullptr);
+
+    torch::lazy::TSOpVector as_strided_scatter_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(as_strided_scatter_out.size(), 1);
+
+    return as_strided_scatter_out;
+
+  }
+
+
+  ::std::vector<int64_t> size;
+  ::std::vector<int64_t> stride;
+  bool has_storage_offset: 1;
+
+};
+
+class AvgPool2d : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::avg_pool2d);
+  }
+
+  AvgPool2d(const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const bool& ceil_mode, const bool& count_include_pad, const ::std::optional<int64_t>& divisor_override, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AvgPool2d::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)),
+        kernel_size(kernel_size),
+        stride(stride),
+        padding(padding),
+        ceil_mode(ceil_mode),
+        count_include_pad(count_include_pad),
+        divisor_override(divisor_override)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", kernel_size=" << kernel_size;
+    ss << ", stride=" << stride;
+    ss << ", padding=" << padding;
+    ss << ", ceil_mode=" << ceil_mode;
+    ss << ", count_include_pad=" << count_include_pad;
+    if (divisor_override.has_value()) {
+      ss << ", divisor_override=" << divisor_override.value();
+    } else {
+      ss << ", divisor_override=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const bool& ceil_mode, const bool& count_include_pad, const ::std::optional<int64_t>& divisor_override) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->kernel_size == kernel_size &&
+        this->stride == stride &&
+        this->padding == padding &&
+        this->ceil_mode == ceil_mode &&
+        this->count_include_pad == count_include_pad &&
+        ((!this->divisor_override&&!divisor_override) || (this->divisor_override&&divisor_override && *(this->divisor_override) == *divisor_override)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(7);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("kernel_size", kernel_size);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back("padding", padding);
+    arguments.emplace_back("ceil_mode", ceil_mode);
+    arguments.emplace_back("count_include_pad", count_include_pad);
+    arguments.emplace_back("divisor_override", divisor_override);
+
+    torch::lazy::TSOpVector avg_pool2d_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(avg_pool2d_out.size(), 1);
+
+    return avg_pool2d_out;
+
+  }
+
+
+  ::std::vector<int64_t> kernel_size;
+  ::std::vector<int64_t> stride;
+  ::std::vector<int64_t> padding;
+  bool ceil_mode;
+  bool count_include_pad;
+  ::std::optional<int64_t> divisor_override;
+
+
+};
+
+class AvgPool2dBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::avg_pool2d_backward);
+  }
+
+  AvgPool2dBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const bool& ceil_mode, const bool& count_include_pad, const ::std::optional<int64_t>& divisor_override, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              AvgPool2dBackward::ClassOpKind(),
+              OpList{grad_output, self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)),
+        kernel_size(kernel_size),
+        stride(stride),
+        padding(padding),
+        ceil_mode(ceil_mode),
+        count_include_pad(count_include_pad),
+        divisor_override(divisor_override)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", kernel_size=" << kernel_size;
+    ss << ", stride=" << stride;
+    ss << ", padding=" << padding;
+    ss << ", ceil_mode=" << ceil_mode;
+    ss << ", count_include_pad=" << count_include_pad;
+    if (divisor_override.has_value()) {
+      ss << ", divisor_override=" << divisor_override.value();
+    } else {
+      ss << ", divisor_override=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const bool& ceil_mode, const bool& count_include_pad, const ::std::optional<int64_t>& divisor_override) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        this->kernel_size == kernel_size &&
+        this->stride == stride &&
+        this->padding == padding &&
+        this->ceil_mode == ceil_mode &&
+        this->count_include_pad == count_include_pad &&
+        ((!this->divisor_override&&!divisor_override) || (this->divisor_override&&divisor_override && *(this->divisor_override) == *divisor_override)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(8);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("kernel_size", kernel_size);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back("padding", padding);
+    arguments.emplace_back("ceil_mode", ceil_mode);
+    arguments.emplace_back("count_include_pad", count_include_pad);
+    arguments.emplace_back("divisor_override", divisor_override);
+
+    torch::lazy::TSOpVector avg_pool2d_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(avg_pool2d_backward_out.size(), 1);
+
+    return avg_pool2d_backward_out;
+
+  }
+
+
+  ::std::vector<int64_t> kernel_size;
+  ::std::vector<int64_t> stride;
+  ::std::vector<int64_t> padding;
+  bool ceil_mode;
+  bool count_include_pad;
+  ::std::optional<int64_t> divisor_override;
+
+
+};
+
+class Baddbmm : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::baddbmm);
+  }
+
+  Baddbmm(const torch::lazy::Value& self, const torch::lazy::Value& batch1, const torch::lazy::Value& batch2, const torch::lazy::Value& beta, const torch::lazy::Value& alpha, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Baddbmm::ClassOpKind(),
+              OpList{self, batch1, batch2, beta, alpha},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& batch1, const torch::lazy::Value& batch2, const torch::lazy::Value& beta, const torch::lazy::Value& alpha) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == batch1 &&
+        operand(i++) == batch2 &&
+        operand(i++) == beta &&
+        operand(i++) == alpha);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(2);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("beta", loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("alpha", loctx->GetOutputOp(operand(i++)));
+    torch::lazy::TSOpVector baddbmm_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(baddbmm_out.size(), 1);
+
+    return baddbmm_out;
+
+  }
+
+
+
+
+
+};
+
+class Bernoulli : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::bernoulli);
+  }
+
+  Bernoulli(const torch::lazy::Value& self, const ::std::optional<at::Generator>& generator, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Bernoulli::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(generator)),
+        generator(generator)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (generator.has_value()) {
+      ss << ", generator=" << "torch.Generator()";
+    } else {
+      ss << ", generator=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<at::Generator>& generator) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        ((!this->generator&&!generator) || (this->generator&&generator && *(this->generator) == *generator)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("generator", generator);
+    torch::lazy::TSOpVector bernoulli_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(bernoulli_out.size(), 1);
+
+    return bernoulli_out;
+
+  }
+
+
+  ::std::optional<at::Generator> generator;
+
+
+};
+
+class BernoulliP : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::bernoulli);
+  }
+
+  BernoulliP(const torch::lazy::Value& self, const double& p, const ::std::optional<at::Generator>& generator, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              BernoulliP::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(p, generator)),
+        p(p),
+        generator(generator)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", p=" << p;
+    if (generator.has_value()) {
+      ss << ", generator=" << "torch.Generator()";
+    } else {
+      ss << ", generator=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const double& p, const ::std::optional<at::Generator>& generator) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->p == p &&
+        ((!this->generator&&!generator) || (this->generator&&generator && *(this->generator) == *generator)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("p", p);
+    kwarguments.emplace_back("generator", generator);
+    torch::lazy::TSOpVector bernoulli_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(bernoulli_out.size(), 1);
+
+    return bernoulli_out;
+
+  }
+
+
+  double p;
+  ::std::optional<at::Generator> generator;
+
+
+};
+
+class BinaryCrossEntropy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::binary_cross_entropy);
+  }
+
+  BinaryCrossEntropy(const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              BinaryCrossEntropy::ClassOpKind(),
+              OpList{self, target, weight.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(reduction)),
+        reduction(reduction)
+  {
+    has_weight = !!weight;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == target &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        this->reduction == reduction);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("reduction", reduction);
+
+    torch::lazy::TSOpVector binary_cross_entropy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(binary_cross_entropy_out.size(), 1);
+
+    return binary_cross_entropy_out;
+
+  }
+
+
+  int64_t reduction;
+  bool has_weight: 1;
+
+};
+
+class BinaryCrossEntropyBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::binary_cross_entropy_backward);
+  }
+
+  BinaryCrossEntropyBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              BinaryCrossEntropyBackward::ClassOpKind(),
+              OpList{grad_output, self, target, weight.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(reduction)),
+        reduction(reduction)
+  {
+    has_weight = !!weight;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == target &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        this->reduction == reduction);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("reduction", reduction);
+
+    torch::lazy::TSOpVector binary_cross_entropy_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(binary_cross_entropy_backward_out.size(), 1);
+
+    return binary_cross_entropy_backward_out;
+
+  }
+
+
+  int64_t reduction;
+  bool has_weight: 1;
+
+};
+
+class BitwiseAndTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::bitwise_and);
+  }
+
+  BitwiseAndTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              BitwiseAndTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector bitwise_and_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(bitwise_and_out.size(), 1);
+
+    return bitwise_and_out;
+
+  }
+
+
+
+
+
+};
+
+class BitwiseOrTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::bitwise_or);
+  }
+
+  BitwiseOrTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              BitwiseOrTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector bitwise_or_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(bitwise_or_out.size(), 1);
+
+    return bitwise_or_out;
+
+  }
+
+
+
+
+
+};
+
+class Bmm : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::bmm);
+  }
+
+  Bmm(const torch::lazy::Value& self, const torch::lazy::Value& mat2, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Bmm::ClassOpKind(),
+              OpList{self, mat2},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& mat2) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == mat2);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector bmm_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(bmm_out.size(), 1);
+
+    return bmm_out;
+
+  }
+
+
+
+
+
+};
+
+class Cat : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::cat);
+  }
+
+  Cat(const torch::lazy::Value& tensors, const int64_t& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Cat::ClassOpKind(),
+              OpList{tensors},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& tensors, const int64_t& dim) const {
+    size_t i = 0;
+    return (operand(i++) == tensors &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector cat_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(cat_out.size(), 1);
+
+    return cat_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class Clamp : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::clamp);
+  }
+
+  Clamp(const torch::lazy::Value& self, const ::std::optional<torch::lazy::Value>& min, const ::std::optional<torch::lazy::Value>& max, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Clamp::ClassOpKind(),
+              OpList{self, min.value_or(kNullValue), max.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+    has_min = !!min;
+    has_max = !!max;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<torch::lazy::Value>& min, const ::std::optional<torch::lazy::Value>& max) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        nullable_operand(i++) == min.value_or(kNullValue) &&
+        nullable_operand(i++) == max.value_or(kNullValue));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_min ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_max ? loctx->GetOutputOp(operand(i++)) : nullptr);
+
+    torch::lazy::TSOpVector clamp_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(clamp_out.size(), 1);
+
+    return clamp_out;
+
+  }
+
+
+
+  bool has_min: 1;
+  bool has_max: 1;
+
+};
+
+class ClampMin : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::clamp_min);
+  }
+
+  ClampMin(const torch::lazy::Value& self, const torch::lazy::Value& min, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ClampMin::ClassOpKind(),
+              OpList{self, min},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& min) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == min);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector clamp_min_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(clamp_min_out.size(), 1);
+
+    return clamp_min_out;
+
+  }
+
+
+
+
+
+};
+
+class ConstantPadNd : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::constant_pad_nd);
+  }
+
+  ConstantPadNd(const torch::lazy::Value& self, const ::std::vector<int64_t>& pad, const torch::lazy::Value& value, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ConstantPadNd::ClassOpKind(),
+              OpList{self, value},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(pad)),
+        pad(pad)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", pad=" << pad;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& pad, const torch::lazy::Value& value) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == value &&
+        this->pad == pad);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("pad", pad);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector constant_pad_nd_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(constant_pad_nd_out.size(), 1);
+
+    return constant_pad_nd_out;
+
+  }
+
+
+  ::std::vector<int64_t> pad;
+
+
+};
+
+class Convolution : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::convolution);
+  }
+
+  Convolution(const torch::lazy::Value& input, const torch::lazy::Value& weight, const ::std::optional<torch::lazy::Value>& bias, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& transposed, const ::std::vector<int64_t>& output_padding, const int64_t& groups, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Convolution::ClassOpKind(),
+              OpList{input, weight, bias.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(stride, padding, dilation, transposed, output_padding, groups)),
+        stride(stride),
+        padding(padding),
+        dilation(dilation),
+        transposed(transposed),
+        output_padding(output_padding),
+        groups(groups)
+  {
+    has_bias = !!bias;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", stride=" << stride;
+    ss << ", padding=" << padding;
+    ss << ", dilation=" << dilation;
+    ss << ", transposed=" << transposed;
+    ss << ", output_padding=" << output_padding;
+    ss << ", groups=" << groups;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& input, const torch::lazy::Value& weight, const ::std::optional<torch::lazy::Value>& bias, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& transposed, const ::std::vector<int64_t>& output_padding, const int64_t& groups) const {
+    size_t i = 0;
+    return (operand(i++) == input &&
+        operand(i++) == weight &&
+        nullable_operand(i++) == bias.value_or(kNullValue) &&
+        this->stride == stride &&
+        this->padding == padding &&
+        this->dilation == dilation &&
+        this->transposed == transposed &&
+        this->output_padding == output_padding &&
+        this->groups == groups);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(9);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_bias ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back("padding", padding);
+    arguments.emplace_back("dilation", dilation);
+    arguments.emplace_back("transposed", transposed);
+    arguments.emplace_back("output_padding", output_padding);
+    arguments.emplace_back("groups", groups);
+
+    torch::lazy::TSOpVector convolution_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(convolution_out.size(), 1);
+
+    return convolution_out;
+
+  }
+
+
+  ::std::vector<int64_t> stride;
+  ::std::vector<int64_t> padding;
+  ::std::vector<int64_t> dilation;
+  bool transposed;
+  ::std::vector<int64_t> output_padding;
+  int64_t groups;
+  bool has_bias: 1;
+
+};
+
+class ConvolutionBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::convolution_backward);
+  }
+
+  ConvolutionBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& input, const torch::lazy::Value& weight, const ::std::optional<::std::vector<int64_t>>& bias_sizes, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& transposed, const ::std::vector<int64_t>& output_padding, const int64_t& groups, const ::std::vector<bool>& output_mask, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ConvolutionBackward::ClassOpKind(),
+              OpList{grad_output, input, weight},
+              std::move(shapes),
+              /* num_outputs */ 3,
+              torch::lazy::MHash(bias_sizes, stride, padding, dilation, transposed, output_padding, groups, output_mask)),
+        bias_sizes(bias_sizes),
+        stride(stride),
+        padding(padding),
+        dilation(dilation),
+        transposed(transposed),
+        output_padding(output_padding),
+        groups(groups),
+        output_mask(output_mask)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (bias_sizes.has_value()) {
+      ss << ", bias_sizes=" << bias_sizes.value();
+    } else {
+      ss << ", bias_sizes=null";
+    }
+    ss << ", stride=" << stride;
+    ss << ", padding=" << padding;
+    ss << ", dilation=" << dilation;
+    ss << ", transposed=" << transposed;
+    ss << ", output_padding=" << output_padding;
+    ss << ", groups=" << groups;
+    ss << ", output_mask=" << output_mask;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& input, const torch::lazy::Value& weight, const ::std::optional<::std::vector<int64_t>>& bias_sizes, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& transposed, const ::std::vector<int64_t>& output_padding, const int64_t& groups, const ::std::vector<bool>& output_mask) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == input &&
+        operand(i++) == weight &&
+        ((!this->bias_sizes&&!bias_sizes) || (this->bias_sizes&&bias_sizes && *(this->bias_sizes) == *bias_sizes)) &&
+        this->stride == stride &&
+        this->padding == padding &&
+        this->dilation == dilation &&
+        this->transposed == transposed &&
+        this->output_padding == output_padding &&
+        this->groups == groups &&
+        this->output_mask == output_mask);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(11);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("bias_sizes", bias_sizes);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back("padding", padding);
+    arguments.emplace_back("dilation", dilation);
+    arguments.emplace_back("transposed", transposed);
+    arguments.emplace_back("output_padding", output_padding);
+    arguments.emplace_back("groups", groups);
+    arguments.emplace_back("output_mask", output_mask);
+
+    torch::lazy::TSOpVector convolution_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(convolution_backward_out.size(), 3);
+
+    return convolution_backward_out;
+
+  }
+
+
+  ::std::optional<::std::vector<int64_t>> bias_sizes;
+  ::std::vector<int64_t> stride;
+  ::std::vector<int64_t> padding;
+  ::std::vector<int64_t> dilation;
+  bool transposed;
+  ::std::vector<int64_t> output_padding;
+  int64_t groups;
+  ::std::vector<bool> output_mask;
+
+
+};
+
+class Cos : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::cos);
+  }
+
+  Cos(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Cos::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector cos_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(cos_out.size(), 1);
+
+    return cos_out;
+
+  }
+
+
+
+
+
+};
+
+class Cumsum : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::cumsum);
+  }
+
+  Cumsum(const torch::lazy::Value& self, const int64_t& dim, const ::std::optional<at::ScalarType>& dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Cumsum::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, dtype)),
+        dim(dim),
+        dtype(dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    if (dtype.has_value()) {
+      ss << ", dtype=" << dtype.value();
+    } else {
+      ss << ", dtype=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const ::std::optional<at::ScalarType>& dtype) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim &&
+        ((!this->dtype&&!dtype) || (this->dtype&&dtype && *(this->dtype) == *dtype)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    kwarguments.emplace_back("dtype", dtype);
+    torch::lazy::TSOpVector cumsum_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(cumsum_out.size(), 1);
+
+    return cumsum_out;
+
+  }
+
+
+  int64_t dim;
+  ::std::optional<at::ScalarType> dtype;
+
+
+};
+
+class DetachCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::detach_copy);
+  }
+
+  DetachCopy(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              DetachCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector detach_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(detach_copy_out.size(), 1);
+
+    return detach_copy_out;
+
+  }
+
+
+
+
+
+};
+
+class DiagonalCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::diagonal_copy);
+  }
+
+  DiagonalCopy(const torch::lazy::Value& self, const int64_t& offset, const int64_t& dim1, const int64_t& dim2, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              DiagonalCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(offset, dim1, dim2)),
+        offset(offset),
+        dim1(dim1),
+        dim2(dim2)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", offset=" << offset;
+    ss << ", dim1=" << dim1;
+    ss << ", dim2=" << dim2;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& offset, const int64_t& dim1, const int64_t& dim2) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->offset == offset &&
+        this->dim1 == dim1 &&
+        this->dim2 == dim2);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("offset", offset);
+    arguments.emplace_back("dim1", dim1);
+    arguments.emplace_back("dim2", dim2);
+
+    torch::lazy::TSOpVector diagonal_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(diagonal_copy_out.size(), 1);
+
+    return diagonal_copy_out;
+
+  }
+
+
+  int64_t offset;
+  int64_t dim1;
+  int64_t dim2;
+
+
+};
+
+class DiagonalScatter : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::diagonal_scatter);
+  }
+
+  DiagonalScatter(const torch::lazy::Value& self, const torch::lazy::Value& src, const int64_t& offset, const int64_t& dim1, const int64_t& dim2, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              DiagonalScatter::ClassOpKind(),
+              OpList{self, src},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(offset, dim1, dim2)),
+        offset(offset),
+        dim1(dim1),
+        dim2(dim2)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", offset=" << offset;
+    ss << ", dim1=" << dim1;
+    ss << ", dim2=" << dim2;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& src, const int64_t& offset, const int64_t& dim1, const int64_t& dim2) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == src &&
+        this->offset == offset &&
+        this->dim1 == dim1 &&
+        this->dim2 == dim2);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("offset", offset);
+    arguments.emplace_back("dim1", dim1);
+    arguments.emplace_back("dim2", dim2);
+
+    torch::lazy::TSOpVector diagonal_scatter_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(diagonal_scatter_out.size(), 1);
+
+    return diagonal_scatter_out;
+
+  }
+
+
+  int64_t offset;
+  int64_t dim1;
+  int64_t dim2;
+
+
+};
+
+class DivTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::div);
+  }
+
+  DivTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              DivTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector div_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(div_out.size(), 1);
+
+    return div_out;
+
+  }
+
+
+
+
+
+};
+
+class DivTensorMode : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::div);
+  }
+
+  DivTensorMode(const torch::lazy::Value& self, const torch::lazy::Value& other, const ::std::optional<c10::string_view>& rounding_mode, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              DivTensorMode::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(rounding_mode)),
+        rounding_mode(rounding_mode.has_value() ? ::std::make_optional(std::string(*rounding_mode)) : ::std::nullopt)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (rounding_mode.has_value()) {
+      ss << ", rounding_mode=" << rounding_mode.value();
+    } else {
+      ss << ", rounding_mode=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other, const ::std::optional<c10::string_view>& rounding_mode) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other &&
+        ((!this->rounding_mode&&!rounding_mode) || (this->rounding_mode&&rounding_mode && *(this->rounding_mode) == *rounding_mode)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("rounding_mode", rounding_mode);
+    torch::lazy::TSOpVector div_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(div_out.size(), 1);
+
+    return div_out;
+
+  }
+
+
+  ::std::optional<std::string> rounding_mode;
+
+
+};
+
+class Elu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::elu);
+  }
+
+  Elu(const torch::lazy::Value& self, const torch::lazy::Value& alpha, const torch::lazy::Value& scale, const torch::lazy::Value& input_scale, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Elu::ClassOpKind(),
+              OpList{self, alpha, scale, input_scale},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& alpha, const torch::lazy::Value& scale, const torch::lazy::Value& input_scale) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == alpha &&
+        operand(i++) == scale &&
+        operand(i++) == input_scale);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector elu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(elu_out.size(), 1);
+
+    return elu_out;
+
+  }
+
+
+
+
+
+};
+
+class EluBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::elu_backward);
+  }
+
+  EluBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& alpha, const torch::lazy::Value& scale, const torch::lazy::Value& input_scale, const bool& is_result, const torch::lazy::Value& self_or_result, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              EluBackward::ClassOpKind(),
+              OpList{grad_output, alpha, scale, input_scale, self_or_result},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(is_result)),
+        is_result(is_result)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", is_result=" << is_result;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& alpha, const torch::lazy::Value& scale, const torch::lazy::Value& input_scale, const bool& is_result, const torch::lazy::Value& self_or_result) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == alpha &&
+        operand(i++) == scale &&
+        operand(i++) == input_scale &&
+        operand(i++) == self_or_result &&
+        this->is_result == is_result);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(6);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("is_result", is_result);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector elu_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(elu_backward_out.size(), 1);
+
+    return elu_backward_out;
+
+  }
+
+
+  bool is_result;
+
+
+};
+
+class Embedding : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::embedding);
+  }
+
+  Embedding(const torch::lazy::Value& weight, const torch::lazy::Value& indices, const int64_t& padding_idx, const bool& scale_grad_by_freq, const bool& sparse, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Embedding::ClassOpKind(),
+              OpList{weight, indices},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(padding_idx, scale_grad_by_freq, sparse)),
+        padding_idx(padding_idx),
+        scale_grad_by_freq(scale_grad_by_freq),
+        sparse(sparse)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", padding_idx=" << padding_idx;
+    ss << ", scale_grad_by_freq=" << scale_grad_by_freq;
+    ss << ", sparse=" << sparse;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& weight, const torch::lazy::Value& indices, const int64_t& padding_idx, const bool& scale_grad_by_freq, const bool& sparse) const {
+    size_t i = 0;
+    return (operand(i++) == weight &&
+        operand(i++) == indices &&
+        this->padding_idx == padding_idx &&
+        this->scale_grad_by_freq == scale_grad_by_freq &&
+        this->sparse == sparse);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("padding_idx", padding_idx);
+    arguments.emplace_back("scale_grad_by_freq", scale_grad_by_freq);
+    arguments.emplace_back("sparse", sparse);
+
+    torch::lazy::TSOpVector embedding_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(embedding_out.size(), 1);
+
+    return embedding_out;
+
+  }
+
+
+  int64_t padding_idx;
+  bool scale_grad_by_freq;
+  bool sparse;
+
+
+};
+
+class EmbeddingDenseBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::embedding_dense_backward);
+  }
+
+  EmbeddingDenseBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& indices, const int64_t& num_weights, const int64_t& padding_idx, const bool& scale_grad_by_freq, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              EmbeddingDenseBackward::ClassOpKind(),
+              OpList{grad_output, indices},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(num_weights, padding_idx, scale_grad_by_freq)),
+        num_weights(num_weights),
+        padding_idx(padding_idx),
+        scale_grad_by_freq(scale_grad_by_freq)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", num_weights=" << num_weights;
+    ss << ", padding_idx=" << padding_idx;
+    ss << ", scale_grad_by_freq=" << scale_grad_by_freq;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& indices, const int64_t& num_weights, const int64_t& padding_idx, const bool& scale_grad_by_freq) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == indices &&
+        this->num_weights == num_weights &&
+        this->padding_idx == padding_idx &&
+        this->scale_grad_by_freq == scale_grad_by_freq);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("num_weights", num_weights);
+    arguments.emplace_back("padding_idx", padding_idx);
+    arguments.emplace_back("scale_grad_by_freq", scale_grad_by_freq);
+
+    torch::lazy::TSOpVector embedding_dense_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(embedding_dense_backward_out.size(), 1);
+
+    return embedding_dense_backward_out;
+
+  }
+
+
+  int64_t num_weights;
+  int64_t padding_idx;
+  bool scale_grad_by_freq;
+
+
+};
+
+class EqScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::eq);
+  }
+
+  EqScalar(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              EqScalar::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector eq_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(eq_out.size(), 1);
+
+    return eq_out;
+
+  }
+
+
+
+
+
+};
+
+class EqTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::eq);
+  }
+
+  EqTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              EqTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector eq_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(eq_out.size(), 1);
+
+    return eq_out;
+
+  }
+
+
+
+
+
+};
+
+class Exp : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::exp);
+  }
+
+  Exp(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Exp::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector exp_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(exp_out.size(), 1);
+
+    return exp_out;
+
+  }
+
+
+
+
+
+};
+
+class ExpandCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::expand_copy);
+  }
+
+  ExpandCopy(const torch::lazy::Value& self, const ::std::vector<int64_t>& size, const bool& implicit, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ExpandCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(size, implicit)),
+        size(size),
+        implicit(implicit)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", size=" << size;
+    ss << ", implicit=" << implicit;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& size, const bool& implicit) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->size == size &&
+        this->implicit == implicit);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("size", size);
+    kwarguments.emplace_back("implicit", implicit);
+    torch::lazy::TSOpVector expand_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(expand_copy_out.size(), 1);
+
+    return expand_copy_out;
+
+  }
+
+
+  ::std::vector<int64_t> size;
+  bool implicit;
+
+
+};
+
+class Flip : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::flip);
+  }
+
+  Flip(const torch::lazy::Value& self, const ::std::vector<int64_t>& dims, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Flip::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dims)),
+        dims(dims)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dims=" << dims;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& dims) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dims == dims);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dims", dims);
+
+    torch::lazy::TSOpVector flip_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(flip_out.size(), 1);
+
+    return flip_out;
+
+  }
+
+
+  ::std::vector<int64_t> dims;
+
+
+};
+
+class Floor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::floor);
+  }
+
+  Floor(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Floor::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector floor_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(floor_out.size(), 1);
+
+    return floor_out;
+
+  }
+
+
+
+
+
+};
+
+class Frac : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::frac);
+  }
+
+  Frac(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Frac::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector frac_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(frac_out.size(), 1);
+
+    return frac_out;
+
+  }
+
+
+
+
+
+};
+
+class Gather : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::gather);
+  }
+
+  Gather(const torch::lazy::Value& self, const int64_t& dim, const torch::lazy::Value& index, const bool& sparse_grad, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Gather::ClassOpKind(),
+              OpList{self, index},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, sparse_grad)),
+        dim(dim),
+        sparse_grad(sparse_grad)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", sparse_grad=" << sparse_grad;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const torch::lazy::Value& index, const bool& sparse_grad) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == index &&
+        this->dim == dim &&
+        this->sparse_grad == sparse_grad);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("sparse_grad", sparse_grad);
+    torch::lazy::TSOpVector gather_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(gather_out.size(), 1);
+
+    return gather_out;
+
+  }
+
+
+  int64_t dim;
+  bool sparse_grad;
+
+
+};
+
+class GeScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::ge);
+  }
+
+  GeScalar(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GeScalar::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector ge_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(ge_out.size(), 1);
+
+    return ge_out;
+
+  }
+
+
+
+
+
+};
+
+class GeTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::ge);
+  }
+
+  GeTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GeTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector ge_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(ge_out.size(), 1);
+
+    return ge_out;
+
+  }
+
+
+
+
+
+};
+
+class Gelu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::gelu);
+  }
+
+  Gelu(const torch::lazy::Value& self, const c10::string_view& approximate, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Gelu::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(approximate)),
+        approximate(approximate)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", approximate=" << approximate;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const c10::string_view& approximate) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->approximate == approximate);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("approximate", approximate);
+    torch::lazy::TSOpVector gelu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(gelu_out.size(), 1);
+
+    return gelu_out;
+
+  }
+
+
+  std::string approximate;
+
+
+};
+
+class GeluBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::gelu_backward);
+  }
+
+  GeluBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const c10::string_view& approximate, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GeluBackward::ClassOpKind(),
+              OpList{grad_output, self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(approximate)),
+        approximate(approximate)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", approximate=" << approximate;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const c10::string_view& approximate) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        this->approximate == approximate);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("approximate", approximate);
+    torch::lazy::TSOpVector gelu_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(gelu_backward_out.size(), 1);
+
+    return gelu_backward_out;
+
+  }
+
+
+  std::string approximate;
+
+
+};
+
+class Glu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::glu);
+  }
+
+  Glu(const torch::lazy::Value& self, const int64_t& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Glu::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector glu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(glu_out.size(), 1);
+
+    return glu_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class GluBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::glu_backward);
+  }
+
+  GluBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const int64_t& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GluBackward::ClassOpKind(),
+              OpList{grad_output, self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const int64_t& dim) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector glu_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(glu_backward_out.size(), 1);
+
+    return glu_backward_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class GluJvp : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::glu_jvp);
+  }
+
+  GluJvp(const torch::lazy::Value& glu, const torch::lazy::Value& x, const torch::lazy::Value& dx, const int64_t& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GluJvp::ClassOpKind(),
+              OpList{glu, x, dx},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& glu, const torch::lazy::Value& x, const torch::lazy::Value& dx, const int64_t& dim) const {
+    size_t i = 0;
+    return (operand(i++) == glu &&
+        operand(i++) == x &&
+        operand(i++) == dx &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector glu_jvp_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(glu_jvp_out.size(), 1);
+
+    return glu_jvp_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class GridSampler2d : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::grid_sampler_2d);
+  }
+
+  GridSampler2d(const torch::lazy::Value& input, const torch::lazy::Value& grid, const int64_t& interpolation_mode, const int64_t& padding_mode, const bool& align_corners, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GridSampler2d::ClassOpKind(),
+              OpList{input, grid},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(interpolation_mode, padding_mode, align_corners)),
+        interpolation_mode(interpolation_mode),
+        padding_mode(padding_mode),
+        align_corners(align_corners)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", interpolation_mode=" << interpolation_mode;
+    ss << ", padding_mode=" << padding_mode;
+    ss << ", align_corners=" << align_corners;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& input, const torch::lazy::Value& grid, const int64_t& interpolation_mode, const int64_t& padding_mode, const bool& align_corners) const {
+    size_t i = 0;
+    return (operand(i++) == input &&
+        operand(i++) == grid &&
+        this->interpolation_mode == interpolation_mode &&
+        this->padding_mode == padding_mode &&
+        this->align_corners == align_corners);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("interpolation_mode", interpolation_mode);
+    arguments.emplace_back("padding_mode", padding_mode);
+    arguments.emplace_back("align_corners", align_corners);
+
+    torch::lazy::TSOpVector grid_sampler_2d_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(grid_sampler_2d_out.size(), 1);
+
+    return grid_sampler_2d_out;
+
+  }
+
+
+  int64_t interpolation_mode;
+  int64_t padding_mode;
+  bool align_corners;
+
+
+};
+
+class GridSampler2dBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::grid_sampler_2d_backward);
+  }
+
+  GridSampler2dBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& input, const torch::lazy::Value& grid, const int64_t& interpolation_mode, const int64_t& padding_mode, const bool& align_corners, const ::std::vector<bool>& output_mask, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GridSampler2dBackward::ClassOpKind(),
+              OpList{grad_output, input, grid},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(interpolation_mode, padding_mode, align_corners, output_mask)),
+        interpolation_mode(interpolation_mode),
+        padding_mode(padding_mode),
+        align_corners(align_corners),
+        output_mask(output_mask)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", interpolation_mode=" << interpolation_mode;
+    ss << ", padding_mode=" << padding_mode;
+    ss << ", align_corners=" << align_corners;
+    ss << ", output_mask=" << output_mask;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& input, const torch::lazy::Value& grid, const int64_t& interpolation_mode, const int64_t& padding_mode, const bool& align_corners, const ::std::vector<bool>& output_mask) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == input &&
+        operand(i++) == grid &&
+        this->interpolation_mode == interpolation_mode &&
+        this->padding_mode == padding_mode &&
+        this->align_corners == align_corners &&
+        this->output_mask == output_mask);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(7);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("interpolation_mode", interpolation_mode);
+    arguments.emplace_back("padding_mode", padding_mode);
+    arguments.emplace_back("align_corners", align_corners);
+    arguments.emplace_back("output_mask", output_mask);
+
+    torch::lazy::TSOpVector grid_sampler_2d_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(grid_sampler_2d_backward_out.size(), 2);
+
+    return grid_sampler_2d_backward_out;
+
+  }
+
+
+  int64_t interpolation_mode;
+  int64_t padding_mode;
+  bool align_corners;
+  ::std::vector<bool> output_mask;
+
+
+};
+
+class GtScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::gt);
+  }
+
+  GtScalar(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GtScalar::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector gt_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(gt_out.size(), 1);
+
+    return gt_out;
+
+  }
+
+
+
+
+
+};
+
+class GtTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::gt);
+  }
+
+  GtTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              GtTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector gt_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(gt_out.size(), 1);
+
+    return gt_out;
+
+  }
+
+
+
+
+
+};
+
+class Hardsigmoid : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::hardsigmoid);
+  }
+
+  Hardsigmoid(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Hardsigmoid::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector hardsigmoid_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(hardsigmoid_out.size(), 1);
+
+    return hardsigmoid_out;
+
+  }
+
+
+
+
+
+};
+
+class IndexSelect : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::index_select);
+  }
+
+  IndexSelect(const torch::lazy::Value& self, const int64_t& dim, const torch::lazy::Value& index, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              IndexSelect::ClassOpKind(),
+              OpList{self, index},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const torch::lazy::Value& index) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == index &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector index_select_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(index_select_out.size(), 1);
+
+    return index_select_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class LeScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::le);
+  }
+
+  LeScalar(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LeScalar::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector le_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(le_out.size(), 1);
+
+    return le_out;
+
+  }
+
+
+
+
+
+};
+
+class LeTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::le);
+  }
+
+  LeTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LeTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector le_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(le_out.size(), 1);
+
+    return le_out;
+
+  }
+
+
+
+
+
+};
+
+class LeakyRelu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::leaky_relu);
+  }
+
+  LeakyRelu(const torch::lazy::Value& self, const torch::lazy::Value& negative_slope, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LeakyRelu::ClassOpKind(),
+              OpList{self, negative_slope},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& negative_slope) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == negative_slope);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector leaky_relu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(leaky_relu_out.size(), 1);
+
+    return leaky_relu_out;
+
+  }
+
+
+
+
+
+};
+
+class LeakyReluBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::leaky_relu_backward);
+  }
+
+  LeakyReluBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& negative_slope, const bool& self_is_result, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LeakyReluBackward::ClassOpKind(),
+              OpList{grad_output, self, negative_slope},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(self_is_result)),
+        self_is_result(self_is_result)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", self_is_result=" << self_is_result;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& negative_slope, const bool& self_is_result) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == negative_slope &&
+        this->self_is_result == self_is_result);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("self_is_result", self_is_result);
+
+    torch::lazy::TSOpVector leaky_relu_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(leaky_relu_backward_out.size(), 1);
+
+    return leaky_relu_backward_out;
+
+  }
+
+
+  bool self_is_result;
+
+
+};
+
+class Log : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::log);
+  }
+
+  Log(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Log::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector log_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(log_out.size(), 1);
+
+    return log_out;
+
+  }
+
+
+
+
+
+};
+
+class Log2 : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::log2);
+  }
+
+  Log2(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Log2::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector log2_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(log2_out.size(), 1);
+
+    return log2_out;
+
+  }
+
+
+
+
+
+};
+
+class LogSigmoidBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::log_sigmoid_backward);
+  }
+
+  LogSigmoidBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& buffer, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LogSigmoidBackward::ClassOpKind(),
+              OpList{grad_output, self, buffer},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& buffer) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == buffer);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector log_sigmoid_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(log_sigmoid_backward_out.size(), 1);
+
+    return log_sigmoid_backward_out;
+
+  }
+
+
+
+
+
+};
+
+class LogSigmoidForward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::log_sigmoid_forward);
+  }
+
+  LogSigmoidForward(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LogSigmoidForward::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector log_sigmoid_forward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(log_sigmoid_forward_out.size(), 2);
+
+    return log_sigmoid_forward_out;
+
+  }
+
+
+
+
+
+};
+
+class Logdet : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::logdet);
+  }
+
+  Logdet(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Logdet::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector logdet_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(logdet_out.size(), 1);
+
+    return logdet_out;
+
+  }
+
+
+
+
+
+};
+
+class LtScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::lt);
+  }
+
+  LtScalar(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LtScalar::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector lt_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(lt_out.size(), 1);
+
+    return lt_out;
+
+  }
+
+
+
+
+
+};
+
+class LtTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::lt);
+  }
+
+  LtTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              LtTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector lt_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(lt_out.size(), 1);
+
+    return lt_out;
+
+  }
+
+
+
+
+
+};
+
+class MaskedFillScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::masked_fill);
+  }
+
+  MaskedFillScalar(const torch::lazy::Value& self, const torch::lazy::Value& mask, const torch::lazy::Value& value, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              MaskedFillScalar::ClassOpKind(),
+              OpList{self, mask, value},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& mask, const torch::lazy::Value& value) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == mask &&
+        operand(i++) == value);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector masked_fill_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(masked_fill_out.size(), 1);
+
+    return masked_fill_out;
+
+  }
+
+
+
+
+
+};
+
+class MaskedFillTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::masked_fill);
+  }
+
+  MaskedFillTensor(const torch::lazy::Value& self, const torch::lazy::Value& mask, const torch::lazy::Value& value, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              MaskedFillTensor::ClassOpKind(),
+              OpList{self, mask, value},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& mask, const torch::lazy::Value& value) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == mask &&
+        operand(i++) == value);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector masked_fill_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(masked_fill_out.size(), 1);
+
+    return masked_fill_out;
+
+  }
+
+
+
+
+
+};
+
+class MaxDim : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::max);
+  }
+
+  MaxDim(const torch::lazy::Value& self, const int64_t& dim, const bool& keepdim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              MaxDim::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(dim, keepdim)),
+        dim(dim),
+        keepdim(keepdim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", keepdim=" << keepdim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const bool& keepdim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim &&
+        this->keepdim == keepdim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("keepdim", keepdim);
+
+    torch::lazy::TSOpVector max_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(max_out.size(), 2);
+
+    return max_out;
+
+  }
+
+
+  int64_t dim;
+  bool keepdim;
+
+
+};
+
+class Max : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::max);
+  }
+
+  Max(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Max::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector max_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(max_out.size(), 1);
+
+    return max_out;
+
+  }
+
+
+
+
+
+};
+
+class MaxPool2dWithIndices : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::max_pool2d_with_indices);
+  }
+
+  MaxPool2dWithIndices(const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& ceil_mode, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              MaxPool2dWithIndices::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(kernel_size, stride, padding, dilation, ceil_mode)),
+        kernel_size(kernel_size),
+        stride(stride),
+        padding(padding),
+        dilation(dilation),
+        ceil_mode(ceil_mode)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", kernel_size=" << kernel_size;
+    ss << ", stride=" << stride;
+    ss << ", padding=" << padding;
+    ss << ", dilation=" << dilation;
+    ss << ", ceil_mode=" << ceil_mode;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& ceil_mode) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->kernel_size == kernel_size &&
+        this->stride == stride &&
+        this->padding == padding &&
+        this->dilation == dilation &&
+        this->ceil_mode == ceil_mode);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(6);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("kernel_size", kernel_size);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back("padding", padding);
+    arguments.emplace_back("dilation", dilation);
+    arguments.emplace_back("ceil_mode", ceil_mode);
+
+    torch::lazy::TSOpVector max_pool2d_with_indices_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(max_pool2d_with_indices_out.size(), 2);
+
+    return max_pool2d_with_indices_out;
+
+  }
+
+
+  ::std::vector<int64_t> kernel_size;
+  ::std::vector<int64_t> stride;
+  ::std::vector<int64_t> padding;
+  ::std::vector<int64_t> dilation;
+  bool ceil_mode;
+
+
+};
+
+class MaxPool2dWithIndicesBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::max_pool2d_with_indices_backward);
+  }
+
+  MaxPool2dWithIndicesBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& ceil_mode, const torch::lazy::Value& indices, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              MaxPool2dWithIndicesBackward::ClassOpKind(),
+              OpList{grad_output, self, indices},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(kernel_size, stride, padding, dilation, ceil_mode)),
+        kernel_size(kernel_size),
+        stride(stride),
+        padding(padding),
+        dilation(dilation),
+        ceil_mode(ceil_mode)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", kernel_size=" << kernel_size;
+    ss << ", stride=" << stride;
+    ss << ", padding=" << padding;
+    ss << ", dilation=" << dilation;
+    ss << ", ceil_mode=" << ceil_mode;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const ::std::vector<int64_t>& kernel_size, const ::std::vector<int64_t>& stride, const ::std::vector<int64_t>& padding, const ::std::vector<int64_t>& dilation, const bool& ceil_mode, const torch::lazy::Value& indices) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == indices &&
+        this->kernel_size == kernel_size &&
+        this->stride == stride &&
+        this->padding == padding &&
+        this->dilation == dilation &&
+        this->ceil_mode == ceil_mode);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(8);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("kernel_size", kernel_size);
+    arguments.emplace_back("stride", stride);
+    arguments.emplace_back("padding", padding);
+    arguments.emplace_back("dilation", dilation);
+    arguments.emplace_back("ceil_mode", ceil_mode);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector max_pool2d_with_indices_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(max_pool2d_with_indices_backward_out.size(), 1);
+
+    return max_pool2d_with_indices_backward_out;
+
+  }
+
+
+  ::std::vector<int64_t> kernel_size;
+  ::std::vector<int64_t> stride;
+  ::std::vector<int64_t> padding;
+  ::std::vector<int64_t> dilation;
+  bool ceil_mode;
+
+
+};
+
+class Maximum : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::maximum);
+  }
+
+  Maximum(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Maximum::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector maximum_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(maximum_out.size(), 1);
+
+    return maximum_out;
+
+  }
+
+
+
+
+
+};
+
+class Mean : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::mean);
+  }
+
+  Mean(const torch::lazy::Value& self, const ::std::optional<at::ScalarType>& dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Mean::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dtype)),
+        dtype(dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (dtype.has_value()) {
+      ss << ", dtype=" << dtype.value();
+    } else {
+      ss << ", dtype=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<at::ScalarType>& dtype) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        ((!this->dtype&&!dtype) || (this->dtype&&dtype && *(this->dtype) == *dtype)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("dtype", dtype);
+    torch::lazy::TSOpVector mean_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(mean_out.size(), 1);
+
+    return mean_out;
+
+  }
+
+
+  ::std::optional<at::ScalarType> dtype;
+
+
+};
+
+class MeanDim : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::mean);
+  }
+
+  MeanDim(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const bool& keepdim, const ::std::optional<at::ScalarType>& dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              MeanDim::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, keepdim, dtype)),
+        dim(dim),
+        keepdim(keepdim),
+        dtype(dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (dim.has_value()) {
+      ss << ", dim=" << dim.value();
+    } else {
+      ss << ", dim=null";
+    }
+    ss << ", keepdim=" << keepdim;
+    if (dtype.has_value()) {
+      ss << ", dtype=" << dtype.value();
+    } else {
+      ss << ", dtype=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const bool& keepdim, const ::std::optional<at::ScalarType>& dtype) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        ((!this->dim&&!dim) || (this->dim&&dim && *(this->dim) == *dim)) &&
+        this->keepdim == keepdim &&
+        ((!this->dtype&&!dtype) || (this->dtype&&dtype && *(this->dtype) == *dtype)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("keepdim", keepdim);
+    kwarguments.emplace_back("dtype", dtype);
+    torch::lazy::TSOpVector mean_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(mean_out.size(), 1);
+
+    return mean_out;
+
+  }
+
+
+  ::std::optional<::std::vector<int64_t>> dim;
+  bool keepdim;
+  ::std::optional<at::ScalarType> dtype;
+
+
+};
+
+class Min : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::min);
+  }
+
+  Min(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Min::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector min_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(min_out.size(), 1);
+
+    return min_out;
+
+  }
+
+
+
+
+
+};
+
+class Minimum : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::minimum);
+  }
+
+  Minimum(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Minimum::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector minimum_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(minimum_out.size(), 1);
+
+    return minimum_out;
+
+  }
+
+
+
+
+
+};
+
+class Mm : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::mm);
+  }
+
+  Mm(const torch::lazy::Value& self, const torch::lazy::Value& mat2, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Mm::ClassOpKind(),
+              OpList{self, mat2},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& mat2) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == mat2);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector mm_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(mm_out.size(), 1);
+
+    return mm_out;
+
+  }
+
+
+
+
+
+};
+
+class MulTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::mul);
+  }
+
+  MulTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              MulTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector mul_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(mul_out.size(), 1);
+
+    return mul_out;
+
+  }
+
+
+
+
+
+};
+
+class Mv : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::mv);
+  }
+
+  Mv(const torch::lazy::Value& self, const torch::lazy::Value& vec, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Mv::ClassOpKind(),
+              OpList{self, vec},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& vec) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == vec);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector mv_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(mv_out.size(), 1);
+
+    return mv_out;
+
+  }
+
+
+
+
+
+};
+
+class NativeBatchNorm : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::native_batch_norm);
+  }
+
+  NativeBatchNorm(const torch::lazy::Value& input, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& bias, const ::std::optional<torch::lazy::Value>& running_mean, const ::std::optional<torch::lazy::Value>& running_var, const bool& training, const double& momentum, const double& eps, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NativeBatchNorm::ClassOpKind(),
+              OpList{input, weight.value_or(kNullValue), bias.value_or(kNullValue), running_mean.value_or(kNullValue), running_var.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 3,
+              torch::lazy::MHash(training, momentum, eps)),
+        training(training),
+        momentum(momentum),
+        eps(eps)
+  {
+    has_weight = !!weight;
+    has_bias = !!bias;
+    has_running_mean = !!running_mean;
+    has_running_var = !!running_var;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", training=" << training;
+    ss << ", momentum=" << momentum;
+    ss << ", eps=" << eps;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& input, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& bias, const ::std::optional<torch::lazy::Value>& running_mean, const ::std::optional<torch::lazy::Value>& running_var, const bool& training, const double& momentum, const double& eps) const {
+    size_t i = 0;
+    return (operand(i++) == input &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        nullable_operand(i++) == bias.value_or(kNullValue) &&
+        nullable_operand(i++) == running_mean.value_or(kNullValue) &&
+        nullable_operand(i++) == running_var.value_or(kNullValue) &&
+        this->training == training &&
+        this->momentum == momentum &&
+        this->eps == eps);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(8);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_bias ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_running_mean ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_running_var ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("training", training);
+    arguments.emplace_back("momentum", momentum);
+    arguments.emplace_back("eps", eps);
+
+    torch::lazy::TSOpVector native_batch_norm_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(native_batch_norm_out.size(), 3);
+
+    return native_batch_norm_out;
+
+  }
+
+
+  bool training;
+  double momentum;
+  double eps;
+  bool has_weight: 1;
+  bool has_bias: 1;
+  bool has_running_mean: 1;
+  bool has_running_var: 1;
+
+};
+
+class NativeBatchNormBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::native_batch_norm_backward);
+  }
+
+  NativeBatchNormBackward(const torch::lazy::Value& grad_out, const torch::lazy::Value& input, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& running_mean, const ::std::optional<torch::lazy::Value>& running_var, const ::std::optional<torch::lazy::Value>& save_mean, const ::std::optional<torch::lazy::Value>& save_invstd, const bool& train, const double& eps, const ::std::vector<bool>& output_mask, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NativeBatchNormBackward::ClassOpKind(),
+              OpList{grad_out, input, weight.value_or(kNullValue), running_mean.value_or(kNullValue), running_var.value_or(kNullValue), save_mean.value_or(kNullValue), save_invstd.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 3,
+              torch::lazy::MHash(train, eps, output_mask)),
+        train(train),
+        eps(eps),
+        output_mask(output_mask)
+  {
+    has_weight = !!weight;
+    has_running_mean = !!running_mean;
+    has_running_var = !!running_var;
+    has_save_mean = !!save_mean;
+    has_save_invstd = !!save_invstd;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", train=" << train;
+    ss << ", eps=" << eps;
+    ss << ", output_mask=" << output_mask;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_out, const torch::lazy::Value& input, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& running_mean, const ::std::optional<torch::lazy::Value>& running_var, const ::std::optional<torch::lazy::Value>& save_mean, const ::std::optional<torch::lazy::Value>& save_invstd, const bool& train, const double& eps, const ::std::vector<bool>& output_mask) const {
+    size_t i = 0;
+    return (operand(i++) == grad_out &&
+        operand(i++) == input &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        nullable_operand(i++) == running_mean.value_or(kNullValue) &&
+        nullable_operand(i++) == running_var.value_or(kNullValue) &&
+        nullable_operand(i++) == save_mean.value_or(kNullValue) &&
+        nullable_operand(i++) == save_invstd.value_or(kNullValue) &&
+        this->train == train &&
+        this->eps == eps &&
+        this->output_mask == output_mask);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(10);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_running_mean ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_running_var ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_save_mean ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_save_invstd ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("train", train);
+    arguments.emplace_back("eps", eps);
+    arguments.emplace_back("output_mask", output_mask);
+
+    torch::lazy::TSOpVector native_batch_norm_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(native_batch_norm_backward_out.size(), 3);
+
+    return native_batch_norm_backward_out;
+
+  }
+
+
+  bool train;
+  double eps;
+  ::std::vector<bool> output_mask;
+  bool has_weight: 1;
+  bool has_running_mean: 1;
+  bool has_running_var: 1;
+  bool has_save_mean: 1;
+  bool has_save_invstd: 1;
+
+};
+
+class NativeDropout : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::native_dropout);
+  }
+
+  NativeDropout(const torch::lazy::Value& input, const double& p, const ::std::optional<bool>& train, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NativeDropout::ClassOpKind(),
+              OpList{input},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(p, train)),
+        p(p),
+        train(train)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", p=" << p;
+    if (train.has_value()) {
+      ss << ", train=" << train.value();
+    } else {
+      ss << ", train=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& input, const double& p, const ::std::optional<bool>& train) const {
+    size_t i = 0;
+    return (operand(i++) == input &&
+        this->p == p &&
+        ((!this->train&&!train) || (this->train&&train && *(this->train) == *train)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("p", p);
+    arguments.emplace_back("train", train);
+
+    torch::lazy::TSOpVector native_dropout_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(native_dropout_out.size(), 2);
+
+    return native_dropout_out;
+
+  }
+
+
+  double p;
+  ::std::optional<bool> train;
+
+
+};
+
+class NativeDropoutBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::native_dropout_backward);
+  }
+
+  NativeDropoutBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& mask, const double& scale, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NativeDropoutBackward::ClassOpKind(),
+              OpList{grad_output, mask},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(scale)),
+        scale(scale)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", scale=" << scale;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& mask, const double& scale) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == mask &&
+        this->scale == scale);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("scale", scale);
+
+    torch::lazy::TSOpVector native_dropout_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(native_dropout_backward_out.size(), 1);
+
+    return native_dropout_backward_out;
+
+  }
+
+
+  double scale;
+
+
+};
+
+class NativeLayerNorm : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::native_layer_norm);
+  }
+
+  NativeLayerNorm(const torch::lazy::Value& input, const ::std::vector<int64_t>& normalized_shape, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& bias, const double& eps, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NativeLayerNorm::ClassOpKind(),
+              OpList{input, weight.value_or(kNullValue), bias.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 3,
+              torch::lazy::MHash(normalized_shape, eps)),
+        normalized_shape(normalized_shape),
+        eps(eps)
+  {
+    has_weight = !!weight;
+    has_bias = !!bias;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", normalized_shape=" << normalized_shape;
+    ss << ", eps=" << eps;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& input, const ::std::vector<int64_t>& normalized_shape, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& bias, const double& eps) const {
+    size_t i = 0;
+    return (operand(i++) == input &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        nullable_operand(i++) == bias.value_or(kNullValue) &&
+        this->normalized_shape == normalized_shape &&
+        this->eps == eps);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("normalized_shape", normalized_shape);
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_bias ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("eps", eps);
+
+    torch::lazy::TSOpVector native_layer_norm_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(native_layer_norm_out.size(), 3);
+
+    return native_layer_norm_out;
+
+  }
+
+
+  ::std::vector<int64_t> normalized_shape;
+  double eps;
+  bool has_weight: 1;
+  bool has_bias: 1;
+
+};
+
+class NativeLayerNormBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::native_layer_norm_backward);
+  }
+
+  NativeLayerNormBackward(const torch::lazy::Value& grad_out, const torch::lazy::Value& input, const ::std::vector<int64_t>& normalized_shape, const torch::lazy::Value& mean, const torch::lazy::Value& rstd, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& bias, const ::std::vector<bool>& output_mask, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NativeLayerNormBackward::ClassOpKind(),
+              OpList{grad_out, input, mean, rstd, weight.value_or(kNullValue), bias.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 3,
+              torch::lazy::MHash(normalized_shape, output_mask)),
+        normalized_shape(normalized_shape),
+        output_mask(output_mask)
+  {
+    has_weight = !!weight;
+    has_bias = !!bias;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", normalized_shape=" << normalized_shape;
+    ss << ", output_mask=" << output_mask;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_out, const torch::lazy::Value& input, const ::std::vector<int64_t>& normalized_shape, const torch::lazy::Value& mean, const torch::lazy::Value& rstd, const ::std::optional<torch::lazy::Value>& weight, const ::std::optional<torch::lazy::Value>& bias, const ::std::vector<bool>& output_mask) const {
+    size_t i = 0;
+    return (operand(i++) == grad_out &&
+        operand(i++) == input &&
+        operand(i++) == mean &&
+        operand(i++) == rstd &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        nullable_operand(i++) == bias.value_or(kNullValue) &&
+        this->normalized_shape == normalized_shape &&
+        this->output_mask == output_mask);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(8);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("normalized_shape", normalized_shape);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_bias ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("output_mask", output_mask);
+
+    torch::lazy::TSOpVector native_layer_norm_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(native_layer_norm_backward_out.size(), 3);
+
+    return native_layer_norm_backward_out;
+
+  }
+
+
+  ::std::vector<int64_t> normalized_shape;
+  ::std::vector<bool> output_mask;
+  bool has_weight: 1;
+  bool has_bias: 1;
+
+};
+
+class NeScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::ne);
+  }
+
+  NeScalar(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NeScalar::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector ne_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(ne_out.size(), 1);
+
+    return ne_out;
+
+  }
+
+
+
+
+
+};
+
+class NeTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::ne);
+  }
+
+  NeTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NeTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector ne_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(ne_out.size(), 1);
+
+    return ne_out;
+
+  }
+
+
+
+
+
+};
+
+class Neg : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::neg);
+  }
+
+  Neg(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Neg::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector neg_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(neg_out.size(), 1);
+
+    return neg_out;
+
+  }
+
+
+
+
+
+};
+
+class NllLoss2dBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::nll_loss2d_backward);
+  }
+
+  NllLoss2dBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index, const torch::lazy::Value& total_weight, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NllLoss2dBackward::ClassOpKind(),
+              OpList{grad_output, self, target, weight.value_or(kNullValue), total_weight},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(reduction, ignore_index)),
+        reduction(reduction),
+        ignore_index(ignore_index)
+  {
+    has_weight = !!weight;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    ss << ", ignore_index=" << ignore_index;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index, const torch::lazy::Value& total_weight) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == target &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        operand(i++) == total_weight &&
+        this->reduction == reduction &&
+        this->ignore_index == ignore_index);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(7);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("reduction", reduction);
+    arguments.emplace_back("ignore_index", ignore_index);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector nll_loss2d_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(nll_loss2d_backward_out.size(), 1);
+
+    return nll_loss2d_backward_out;
+
+  }
+
+
+  int64_t reduction;
+  int64_t ignore_index;
+  bool has_weight: 1;
+
+};
+
+class NllLoss2dForward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::nll_loss2d_forward);
+  }
+
+  NllLoss2dForward(const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NllLoss2dForward::ClassOpKind(),
+              OpList{self, target, weight.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(reduction, ignore_index)),
+        reduction(reduction),
+        ignore_index(ignore_index)
+  {
+    has_weight = !!weight;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    ss << ", ignore_index=" << ignore_index;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == target &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        this->reduction == reduction &&
+        this->ignore_index == ignore_index);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("reduction", reduction);
+    arguments.emplace_back("ignore_index", ignore_index);
+
+    torch::lazy::TSOpVector nll_loss2d_forward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(nll_loss2d_forward_out.size(), 2);
+
+    return nll_loss2d_forward_out;
+
+  }
+
+
+  int64_t reduction;
+  int64_t ignore_index;
+  bool has_weight: 1;
+
+};
+
+class NllLossBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::nll_loss_backward);
+  }
+
+  NllLossBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index, const torch::lazy::Value& total_weight, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NllLossBackward::ClassOpKind(),
+              OpList{grad_output, self, target, weight.value_or(kNullValue), total_weight},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(reduction, ignore_index)),
+        reduction(reduction),
+        ignore_index(ignore_index)
+  {
+    has_weight = !!weight;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    ss << ", ignore_index=" << ignore_index;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index, const torch::lazy::Value& total_weight) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == target &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        operand(i++) == total_weight &&
+        this->reduction == reduction &&
+        this->ignore_index == ignore_index);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(7);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("reduction", reduction);
+    arguments.emplace_back("ignore_index", ignore_index);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector nll_loss_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(nll_loss_backward_out.size(), 1);
+
+    return nll_loss_backward_out;
+
+  }
+
+
+  int64_t reduction;
+  int64_t ignore_index;
+  bool has_weight: 1;
+
+};
+
+class NllLossForward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::nll_loss_forward);
+  }
+
+  NllLossForward(const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NllLossForward::ClassOpKind(),
+              OpList{self, target, weight.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(reduction, ignore_index)),
+        reduction(reduction),
+        ignore_index(ignore_index)
+  {
+    has_weight = !!weight;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    ss << ", ignore_index=" << ignore_index;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& target, const ::std::optional<torch::lazy::Value>& weight, const int64_t& reduction, const int64_t& ignore_index) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == target &&
+        nullable_operand(i++) == weight.value_or(kNullValue) &&
+        this->reduction == reduction &&
+        this->ignore_index == ignore_index);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_weight ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("reduction", reduction);
+    arguments.emplace_back("ignore_index", ignore_index);
+
+    torch::lazy::TSOpVector nll_loss_forward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(nll_loss_forward_out.size(), 2);
+
+    return nll_loss_forward_out;
+
+  }
+
+
+  int64_t reduction;
+  int64_t ignore_index;
+  bool has_weight: 1;
+
+};
+
+class Nonzero : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::nonzero);
+  }
+
+  Nonzero(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Nonzero::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector nonzero_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(nonzero_out.size(), 1);
+
+    return nonzero_out;
+
+  }
+
+
+
+
+
+};
+
+class NormScalaroptDim : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::norm);
+  }
+
+  NormScalaroptDim(const torch::lazy::Value& self, const ::std::optional<torch::lazy::Value>& p, const ::std::vector<int64_t>& dim, const bool& keepdim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NormScalaroptDim::ClassOpKind(),
+              OpList{self, p.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, keepdim)),
+        dim(dim),
+        keepdim(keepdim)
+  {
+    has_p = !!p;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", keepdim=" << keepdim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<torch::lazy::Value>& p, const ::std::vector<int64_t>& dim, const bool& keepdim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        nullable_operand(i++) == p.value_or(kNullValue) &&
+        this->dim == dim &&
+        this->keepdim == keepdim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(has_p ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("keepdim", keepdim);
+
+    torch::lazy::TSOpVector norm_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(norm_out.size(), 1);
+
+    return norm_out;
+
+  }
+
+
+  ::std::vector<int64_t> dim;
+  bool keepdim;
+  bool has_p: 1;
+
+};
+
+class NormalFunctional : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::normal_functional);
+  }
+
+  NormalFunctional(const torch::lazy::Value& self, const double& mean, const double& std, const ::std::optional<at::Generator>& generator, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              NormalFunctional::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(mean, std, generator)),
+        mean(mean),
+        std(std),
+        generator(generator)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", mean=" << mean;
+    ss << ", std=" << std;
+    if (generator.has_value()) {
+      ss << ", generator=" << "torch.Generator()";
+    } else {
+      ss << ", generator=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const double& mean, const double& std, const ::std::optional<at::Generator>& generator) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->mean == mean &&
+        this->std == std &&
+        ((!this->generator&&!generator) || (this->generator&&generator && *(this->generator) == *generator)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("mean", mean);
+    arguments.emplace_back("std", std);
+    kwarguments.emplace_back("generator", generator);
+    torch::lazy::TSOpVector normal_functional_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(normal_functional_out.size(), 1);
+
+    return normal_functional_out;
+
+  }
+
+
+  double mean;
+  double std;
+  ::std::optional<at::Generator> generator;
+
+
+};
+
+class PermuteCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::permute_copy);
+  }
+
+  PermuteCopy(const torch::lazy::Value& self, const ::std::vector<int64_t>& dims, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              PermuteCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dims)),
+        dims(dims)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dims=" << dims;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& dims) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dims == dims);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dims", dims);
+
+    torch::lazy::TSOpVector permute_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(permute_copy_out.size(), 1);
+
+    return permute_copy_out;
+
+  }
+
+
+  ::std::vector<int64_t> dims;
+
+
+};
+
+class PowTensorTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::pow);
+  }
+
+  PowTensorTensor(const torch::lazy::Value& self, const torch::lazy::Value& exponent, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              PowTensorTensor::ClassOpKind(),
+              OpList{self, exponent},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& exponent) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == exponent);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector pow_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(pow_out.size(), 1);
+
+    return pow_out;
+
+  }
+
+
+
+
+
+};
+
+class PowTensorScalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::pow);
+  }
+
+  PowTensorScalar(const torch::lazy::Value& self, const torch::lazy::Value& exponent, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              PowTensorScalar::ClassOpKind(),
+              OpList{self, exponent},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& exponent) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == exponent);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector pow_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(pow_out.size(), 1);
+
+    return pow_out;
+
+  }
+
+
+
+
+
+};
+
+class RandomFrom : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::random);
+  }
+
+  RandomFrom(const torch::lazy::Value& self, const int64_t& from, const ::std::optional<int64_t>& to, const ::std::optional<at::Generator>& generator, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              RandomFrom::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(from, to, generator)),
+        from(from),
+        to(to),
+        generator(generator)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", from=" << from;
+    if (to.has_value()) {
+      ss << ", to=" << to.value();
+    } else {
+      ss << ", to=null";
+    }
+    if (generator.has_value()) {
+      ss << ", generator=" << "torch.Generator()";
+    } else {
+      ss << ", generator=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& from, const ::std::optional<int64_t>& to, const ::std::optional<at::Generator>& generator) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->from == from &&
+        ((!this->to&&!to) || (this->to&&to && *(this->to) == *to)) &&
+        ((!this->generator&&!generator) || (this->generator&&generator && *(this->generator) == *generator)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("from", from);
+    arguments.emplace_back("to", to);
+    kwarguments.emplace_back("generator", generator);
+    torch::lazy::TSOpVector random_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(random_out.size(), 1);
+
+    return random_out;
+
+  }
+
+
+  int64_t from;
+  ::std::optional<int64_t> to;
+  ::std::optional<at::Generator> generator;
+
+
+};
+
+class RandomTo : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::random);
+  }
+
+  RandomTo(const torch::lazy::Value& self, const int64_t& to, const ::std::optional<at::Generator>& generator, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              RandomTo::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(to, generator)),
+        to(to),
+        generator(generator)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", to=" << to;
+    if (generator.has_value()) {
+      ss << ", generator=" << "torch.Generator()";
+    } else {
+      ss << ", generator=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& to, const ::std::optional<at::Generator>& generator) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->to == to &&
+        ((!this->generator&&!generator) || (this->generator&&generator && *(this->generator) == *generator)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("to", to);
+    kwarguments.emplace_back("generator", generator);
+    torch::lazy::TSOpVector random_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(random_out.size(), 1);
+
+    return random_out;
+
+  }
+
+
+  int64_t to;
+  ::std::optional<at::Generator> generator;
+
+
+};
+
+class Random : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::random);
+  }
+
+  Random(const torch::lazy::Value& self, const ::std::optional<at::Generator>& generator, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Random::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(generator)),
+        generator(generator)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (generator.has_value()) {
+      ss << ", generator=" << "torch.Generator()";
+    } else {
+      ss << ", generator=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<at::Generator>& generator) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        ((!this->generator&&!generator) || (this->generator&&generator && *(this->generator) == *generator)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("generator", generator);
+    torch::lazy::TSOpVector random_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(random_out.size(), 1);
+
+    return random_out;
+
+  }
+
+
+  ::std::optional<at::Generator> generator;
+
+
+};
+
+class Reciprocal : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::reciprocal);
+  }
+
+  Reciprocal(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Reciprocal::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector reciprocal_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(reciprocal_out.size(), 1);
+
+    return reciprocal_out;
+
+  }
+
+
+
+
+
+};
+
+class Relu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::relu);
+  }
+
+  Relu(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Relu::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector relu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(relu_out.size(), 1);
+
+    return relu_out;
+
+  }
+
+
+
+
+
+};
+
+class RemainderTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::remainder);
+  }
+
+  RemainderTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              RemainderTensor::ClassOpKind(),
+              OpList{self, other},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector remainder_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(remainder_out.size(), 1);
+
+    return remainder_out;
+
+  }
+
+
+
+
+
+};
+
+class Repeat : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::repeat);
+  }
+
+  Repeat(const torch::lazy::Value& self, const ::std::vector<int64_t>& repeats, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Repeat::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(repeats)),
+        repeats(repeats)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", repeats=" << repeats;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& repeats) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->repeats == repeats);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("repeats", repeats);
+
+    torch::lazy::TSOpVector repeat_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(repeat_out.size(), 1);
+
+    return repeat_out;
+
+  }
+
+
+  ::std::vector<int64_t> repeats;
+
+
+};
+
+class Rsqrt : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::rsqrt);
+  }
+
+  Rsqrt(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Rsqrt::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector rsqrt_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(rsqrt_out.size(), 1);
+
+    return rsqrt_out;
+
+  }
+
+
+
+
+
+};
+
+class ScatterAdd : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::scatter_add);
+  }
+
+  ScatterAdd(const torch::lazy::Value& self, const int64_t& dim, const torch::lazy::Value& index, const torch::lazy::Value& src, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ScatterAdd::ClassOpKind(),
+              OpList{self, index, src},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const torch::lazy::Value& index, const torch::lazy::Value& src) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == index &&
+        operand(i++) == src &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector scatter_add_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(scatter_add_out.size(), 1);
+
+    return scatter_add_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class SelectCopyInt : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::select_copy);
+  }
+
+  SelectCopyInt(const torch::lazy::Value& self, const int64_t& dim, const int64_t& index, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SelectCopyInt::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, index)),
+        dim(dim),
+        index(index)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", index=" << index;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const int64_t& index) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim &&
+        this->index == index);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("index", index);
+
+    torch::lazy::TSOpVector select_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(select_copy_out.size(), 1);
+
+    return select_copy_out;
+
+  }
+
+
+  int64_t dim;
+  int64_t index;
+
+
+};
+
+class SelectScatter : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::select_scatter);
+  }
+
+  SelectScatter(const torch::lazy::Value& self, const torch::lazy::Value& src, const int64_t& dim, const int64_t& index, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SelectScatter::ClassOpKind(),
+              OpList{self, src},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, index)),
+        dim(dim),
+        index(index)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", index=" << index;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& src, const int64_t& dim, const int64_t& index) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == src &&
+        this->dim == dim &&
+        this->index == index);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("index", index);
+
+    torch::lazy::TSOpVector select_scatter_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(select_scatter_out.size(), 1);
+
+    return select_scatter_out;
+
+  }
+
+
+  int64_t dim;
+  int64_t index;
+
+
+};
+
+class Selu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::selu);
+  }
+
+  Selu(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Selu::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector selu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(selu_out.size(), 1);
+
+    return selu_out;
+
+  }
+
+
+
+
+
+};
+
+class Sgn : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::sgn);
+  }
+
+  Sgn(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Sgn::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector sgn_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sgn_out.size(), 1);
+
+    return sgn_out;
+
+  }
+
+
+
+
+
+};
+
+class Sigmoid : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::sigmoid);
+  }
+
+  Sigmoid(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Sigmoid::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector sigmoid_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sigmoid_out.size(), 1);
+
+    return sigmoid_out;
+
+  }
+
+
+
+
+
+};
+
+class SigmoidBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(c10::Symbol::fromQualString("aten::sigmoid_backward"));
+  }
+
+  SigmoidBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& output, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SigmoidBackward::ClassOpKind(),
+              OpList{grad_output, output},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& output) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == output);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector sigmoid_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sigmoid_backward_out.size(), 1);
+
+    return sigmoid_backward_out;
+
+  }
+
+
+
+
+
+};
+
+class Silu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::silu);
+  }
+
+  Silu(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Silu::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector silu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(silu_out.size(), 1);
+
+    return silu_out;
+
+  }
+
+
+
+
+
+};
+
+class SliceCopyTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::slice_copy);
+  }
+
+  SliceCopyTensor(const torch::lazy::Value& self, const int64_t& dim, const ::std::optional<torch::lazy::Value>& start, const ::std::optional<torch::lazy::Value>& end, const torch::lazy::Value& step, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SliceCopyTensor::ClassOpKind(),
+              OpList{self, start.value_or(kNullValue), end.value_or(kNullValue), step},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+    has_start = !!start;
+    has_end = !!end;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const ::std::optional<torch::lazy::Value>& start, const ::std::optional<torch::lazy::Value>& end, const torch::lazy::Value& step) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        nullable_operand(i++) == start.value_or(kNullValue) &&
+        nullable_operand(i++) == end.value_or(kNullValue) &&
+        operand(i++) == step &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back(has_start ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_end ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector slice_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(slice_copy_out.size(), 1);
+
+    return slice_copy_out;
+
+  }
+
+
+  int64_t dim;
+  bool has_start: 1;
+  bool has_end: 1;
+
+};
+
+class SliceScatter : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::slice_scatter);
+  }
+
+  SliceScatter(const torch::lazy::Value& self, const torch::lazy::Value& src, const int64_t& dim, const ::std::optional<torch::lazy::Value>& start, const ::std::optional<torch::lazy::Value>& end, const torch::lazy::Value& step, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SliceScatter::ClassOpKind(),
+              OpList{self, src, start.value_or(kNullValue), end.value_or(kNullValue), step},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+    has_start = !!start;
+    has_end = !!end;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& src, const int64_t& dim, const ::std::optional<torch::lazy::Value>& start, const ::std::optional<torch::lazy::Value>& end, const torch::lazy::Value& step) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == src &&
+        nullable_operand(i++) == start.value_or(kNullValue) &&
+        nullable_operand(i++) == end.value_or(kNullValue) &&
+        operand(i++) == step &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(6);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back(has_start ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(has_end ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector slice_scatter_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(slice_scatter_out.size(), 1);
+
+    return slice_scatter_out;
+
+  }
+
+
+  int64_t dim;
+  bool has_start: 1;
+  bool has_end: 1;
+
+};
+
+class SmoothL1Loss : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::smooth_l1_loss);
+  }
+
+  SmoothL1Loss(const torch::lazy::Value& self, const torch::lazy::Value& target, const int64_t& reduction, const double& beta, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SmoothL1Loss::ClassOpKind(),
+              OpList{self, target},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(reduction, beta)),
+        reduction(reduction),
+        beta(beta)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    ss << ", beta=" << beta;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& target, const int64_t& reduction, const double& beta) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == target &&
+        this->reduction == reduction &&
+        this->beta == beta);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("reduction", reduction);
+    arguments.emplace_back("beta", beta);
+
+    torch::lazy::TSOpVector smooth_l1_loss_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(smooth_l1_loss_out.size(), 1);
+
+    return smooth_l1_loss_out;
+
+  }
+
+
+  int64_t reduction;
+  double beta;
+
+
+};
+
+class SmoothL1LossBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::smooth_l1_loss_backward);
+  }
+
+  SmoothL1LossBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const int64_t& reduction, const double& beta, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SmoothL1LossBackward::ClassOpKind(),
+              OpList{grad_output, self, target},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(reduction, beta)),
+        reduction(reduction),
+        beta(beta)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", reduction=" << reduction;
+    ss << ", beta=" << beta;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& target, const int64_t& reduction, const double& beta) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == target &&
+        this->reduction == reduction &&
+        this->beta == beta);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("reduction", reduction);
+    arguments.emplace_back("beta", beta);
+
+    torch::lazy::TSOpVector smooth_l1_loss_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(smooth_l1_loss_backward_out.size(), 1);
+
+    return smooth_l1_loss_backward_out;
+
+  }
+
+
+  int64_t reduction;
+  double beta;
+
+
+};
+
+class Softplus : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::softplus);
+  }
+
+  Softplus(const torch::lazy::Value& self, const torch::lazy::Value& beta, const torch::lazy::Value& threshold, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Softplus::ClassOpKind(),
+              OpList{self, beta, threshold},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& beta, const torch::lazy::Value& threshold) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == beta &&
+        operand(i++) == threshold);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector softplus_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(softplus_out.size(), 1);
+
+    return softplus_out;
+
+  }
+
+
+
+
+
+};
+
+class SoftplusBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::softplus_backward);
+  }
+
+  SoftplusBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& beta, const torch::lazy::Value& threshold, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SoftplusBackward::ClassOpKind(),
+              OpList{grad_output, self, beta, threshold},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& beta, const torch::lazy::Value& threshold) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == beta &&
+        operand(i++) == threshold);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector softplus_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(softplus_backward_out.size(), 1);
+
+    return softplus_backward_out;
+
+  }
+
+
+
+
+
+};
+
+class Sort : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::sort);
+  }
+
+  Sort(const torch::lazy::Value& self, const int64_t& dim, const bool& descending, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Sort::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(dim, descending)),
+        dim(dim),
+        descending(descending)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    ss << ", descending=" << descending;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim, const bool& descending) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim &&
+        this->descending == descending);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("descending", descending);
+
+    torch::lazy::TSOpVector sort_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sort_out.size(), 2);
+
+    return sort_out;
+
+  }
+
+
+  int64_t dim;
+  bool descending;
+
+
+};
+
+class Sqrt : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::sqrt);
+  }
+
+  Sqrt(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Sqrt::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector sqrt_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sqrt_out.size(), 1);
+
+    return sqrt_out;
+
+  }
+
+
+
+
+
+};
+
+class SqueezeCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::squeeze_copy);
+  }
+
+  SqueezeCopy(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SqueezeCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector squeeze_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(squeeze_copy_out.size(), 1);
+
+    return squeeze_copy_out;
+
+  }
+
+
+
+
+
+};
+
+class SqueezeCopyDim : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::squeeze_copy);
+  }
+
+  SqueezeCopyDim(const torch::lazy::Value& self, const int64_t& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SqueezeCopyDim::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector squeeze_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(squeeze_copy_out.size(), 1);
+
+    return squeeze_copy_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class SqueezeCopyDims : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::squeeze_copy);
+  }
+
+  SqueezeCopyDims(const torch::lazy::Value& self, const ::std::vector<int64_t>& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SqueezeCopyDims::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& dim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector squeeze_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(squeeze_copy_out.size(), 1);
+
+    return squeeze_copy_out;
+
+  }
+
+
+  ::std::vector<int64_t> dim;
+
+
+};
+
+class Stack : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::stack);
+  }
+
+  Stack(const torch::lazy::Value& tensors, const int64_t& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Stack::ClassOpKind(),
+              OpList{tensors},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& tensors, const int64_t& dim) const {
+    size_t i = 0;
+    return (operand(i++) == tensors &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector stack_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(stack_out.size(), 1);
+
+    return stack_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class Std : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::std);
+  }
+
+  Std(const torch::lazy::Value& self, const bool& unbiased, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Std::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(unbiased)),
+        unbiased(unbiased)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", unbiased=" << unbiased;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const bool& unbiased) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->unbiased == unbiased);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("unbiased", unbiased);
+
+    torch::lazy::TSOpVector std_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(std_out.size(), 1);
+
+    return std_out;
+
+  }
+
+
+  bool unbiased;
+
+
+};
+
+class StdDim : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::std);
+  }
+
+  StdDim(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const bool& unbiased, const bool& keepdim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              StdDim::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, unbiased, keepdim)),
+        dim(dim),
+        unbiased(unbiased),
+        keepdim(keepdim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (dim.has_value()) {
+      ss << ", dim=" << dim.value();
+    } else {
+      ss << ", dim=null";
+    }
+    ss << ", unbiased=" << unbiased;
+    ss << ", keepdim=" << keepdim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const bool& unbiased, const bool& keepdim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        ((!this->dim&&!dim) || (this->dim&&dim && *(this->dim) == *dim)) &&
+        this->unbiased == unbiased &&
+        this->keepdim == keepdim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("unbiased", unbiased);
+    arguments.emplace_back("keepdim", keepdim);
+
+    torch::lazy::TSOpVector std_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(std_out.size(), 1);
+
+    return std_out;
+
+  }
+
+
+  ::std::optional<::std::vector<int64_t>> dim;
+  bool unbiased;
+  bool keepdim;
+
+
+};
+
+class StdCorrection : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::std);
+  }
+
+  StdCorrection(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const ::std::optional<torch::lazy::Value>& correction, const bool& keepdim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              StdCorrection::ClassOpKind(),
+              OpList{self, correction.value_or(kNullValue)},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, keepdim)),
+        dim(dim),
+        keepdim(keepdim)
+  {
+    has_correction = !!correction;
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (dim.has_value()) {
+      ss << ", dim=" << dim.value();
+    } else {
+      ss << ", dim=null";
+    }
+    ss << ", keepdim=" << keepdim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const ::std::optional<torch::lazy::Value>& correction, const bool& keepdim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        nullable_operand(i++) == correction.value_or(kNullValue) &&
+        ((!this->dim&&!dim) || (this->dim&&dim && *(this->dim) == *dim)) &&
+        this->keepdim == keepdim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(2);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    kwarguments.emplace_back("correction", has_correction ? loctx->GetOutputOp(operand(i++)) : nullptr);
+    kwarguments.emplace_back("keepdim", keepdim);
+    torch::lazy::TSOpVector std_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(std_out.size(), 1);
+
+    return std_out;
+
+  }
+
+
+  ::std::optional<::std::vector<int64_t>> dim;
+  bool keepdim;
+  bool has_correction: 1;
+
+};
+
+class SubTensor : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::sub);
+  }
+
+  SubTensor(const torch::lazy::Value& self, const torch::lazy::Value& other, const torch::lazy::Value& alpha, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SubTensor::ClassOpKind(),
+              OpList{self, other, alpha},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other, const torch::lazy::Value& alpha) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == other &&
+        operand(i++) == alpha);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("alpha", loctx->GetOutputOp(operand(i++)));
+    torch::lazy::TSOpVector sub_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sub_out.size(), 1);
+
+    return sub_out;
+
+  }
+
+
+
+
+
+};
+
+class Sum : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::sum);
+  }
+
+  Sum(const torch::lazy::Value& self, const ::std::optional<at::ScalarType>& dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Sum::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dtype)),
+        dtype(dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (dtype.has_value()) {
+      ss << ", dtype=" << dtype.value();
+    } else {
+      ss << ", dtype=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<at::ScalarType>& dtype) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        ((!this->dtype&&!dtype) || (this->dtype&&dtype && *(this->dtype) == *dtype)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("dtype", dtype);
+    torch::lazy::TSOpVector sum_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sum_out.size(), 1);
+
+    return sum_out;
+
+  }
+
+
+  ::std::optional<at::ScalarType> dtype;
+
+
+};
+
+class SumDimIntlist : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::sum);
+  }
+
+  SumDimIntlist(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const bool& keepdim, const ::std::optional<at::ScalarType>& dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              SumDimIntlist::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim, keepdim, dtype)),
+        dim(dim),
+        keepdim(keepdim),
+        dtype(dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    if (dim.has_value()) {
+      ss << ", dim=" << dim.value();
+    } else {
+      ss << ", dim=null";
+    }
+    ss << ", keepdim=" << keepdim;
+    if (dtype.has_value()) {
+      ss << ", dtype=" << dtype.value();
+    } else {
+      ss << ", dtype=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::optional<::std::vector<int64_t>>& dim, const bool& keepdim, const ::std::optional<at::ScalarType>& dtype) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        ((!this->dim&&!dim) || (this->dim&&dim && *(this->dim) == *dim)) &&
+        this->keepdim == keepdim &&
+        ((!this->dtype&&!dtype) || (this->dtype&&dtype && *(this->dtype) == *dtype)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("keepdim", keepdim);
+    kwarguments.emplace_back("dtype", dtype);
+    torch::lazy::TSOpVector sum_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(sum_out.size(), 1);
+
+    return sum_out;
+
+  }
+
+
+  ::std::optional<::std::vector<int64_t>> dim;
+  bool keepdim;
+  ::std::optional<at::ScalarType> dtype;
+
+
+};
+
+class TCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::t_copy);
+  }
+
+  TCopy(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              TCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector t_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(t_copy_out.size(), 1);
+
+    return t_copy_out;
+
+  }
+
+
+
+
+
+};
+
+class Tanh : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::tanh);
+  }
+
+  Tanh(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Tanh::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector tanh_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(tanh_out.size(), 1);
+
+    return tanh_out;
+
+  }
+
+
+
+
+
+};
+
+class TanhBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::tanh_backward);
+  }
+
+  TanhBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& output, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              TanhBackward::ClassOpKind(),
+              OpList{grad_output, output},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& output) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == output);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector tanh_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(tanh_backward_out.size(), 1);
+
+    return tanh_backward_out;
+
+  }
+
+
+
+
+
+};
+
+class Threshold : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::threshold);
+  }
+
+  Threshold(const torch::lazy::Value& self, const torch::lazy::Value& threshold, const torch::lazy::Value& value, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Threshold::ClassOpKind(),
+              OpList{self, threshold, value},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& threshold, const torch::lazy::Value& value) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        operand(i++) == threshold &&
+        operand(i++) == value);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector threshold_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(threshold_out.size(), 1);
+
+    return threshold_out;
+
+  }
+
+
+
+
+
+};
+
+class ThresholdBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::threshold_backward);
+  }
+
+  ThresholdBackward(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& threshold, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ThresholdBackward::ClassOpKind(),
+              OpList{grad_output, self, threshold},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const torch::lazy::Value& self, const torch::lazy::Value& threshold) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        operand(i++) == self &&
+        operand(i++) == threshold);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector threshold_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(threshold_backward_out.size(), 1);
+
+    return threshold_backward_out;
+
+  }
+
+
+
+
+
+};
+
+class Topk : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::topk);
+  }
+
+  Topk(const torch::lazy::Value& self, const int64_t& k, const int64_t& dim, const bool& largest, const bool& sorted, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Topk::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 2,
+              torch::lazy::MHash(k, dim, largest, sorted)),
+        k(k),
+        dim(dim),
+        largest(largest),
+        sorted(sorted)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", k=" << k;
+    ss << ", dim=" << dim;
+    ss << ", largest=" << largest;
+    ss << ", sorted=" << sorted;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& k, const int64_t& dim, const bool& largest, const bool& sorted) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->k == k &&
+        this->dim == dim &&
+        this->largest == largest &&
+        this->sorted == sorted);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("k", k);
+    arguments.emplace_back("dim", dim);
+    arguments.emplace_back("largest", largest);
+    arguments.emplace_back("sorted", sorted);
+
+    torch::lazy::TSOpVector topk_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(topk_out.size(), 2);
+
+    return topk_out;
+
+  }
+
+
+  int64_t k;
+  int64_t dim;
+  bool largest;
+  bool sorted;
+
+
+};
+
+class Trace : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::trace);
+  }
+
+  Trace(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Trace::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector trace_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(trace_out.size(), 1);
+
+    return trace_out;
+
+  }
+
+
+
+
+
+};
+
+class TransposeCopyInt : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::transpose_copy);
+  }
+
+  TransposeCopyInt(const torch::lazy::Value& self, const int64_t& dim0, const int64_t& dim1, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              TransposeCopyInt::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim0, dim1)),
+        dim0(dim0),
+        dim1(dim1)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim0=" << dim0;
+    ss << ", dim1=" << dim1;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim0, const int64_t& dim1) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim0 == dim0 &&
+        this->dim1 == dim1);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim0", dim0);
+    arguments.emplace_back("dim1", dim1);
+
+    torch::lazy::TSOpVector transpose_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(transpose_copy_out.size(), 1);
+
+    return transpose_copy_out;
+
+  }
+
+
+  int64_t dim0;
+  int64_t dim1;
+
+
+};
+
+class Tril : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::tril);
+  }
+
+  Tril(const torch::lazy::Value& self, const int64_t& diagonal, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Tril::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(diagonal)),
+        diagonal(diagonal)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", diagonal=" << diagonal;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& diagonal) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->diagonal == diagonal);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("diagonal", diagonal);
+
+    torch::lazy::TSOpVector tril_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(tril_out.size(), 1);
+
+    return tril_out;
+
+  }
+
+
+  int64_t diagonal;
+
+
+};
+
+class Triu : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::triu);
+  }
+
+  Triu(const torch::lazy::Value& self, const int64_t& diagonal, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Triu::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(diagonal)),
+        diagonal(diagonal)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", diagonal=" << diagonal;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& diagonal) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->diagonal == diagonal);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("diagonal", diagonal);
+
+    torch::lazy::TSOpVector triu_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(triu_out.size(), 1);
+
+    return triu_out;
+
+  }
+
+
+  int64_t diagonal;
+
+
+};
+
+class Trunc : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::trunc);
+  }
+
+  Trunc(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Trunc::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector trunc_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(trunc_out.size(), 1);
+
+    return trunc_out;
+
+  }
+
+
+
+
+
+};
+
+class UnfoldCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::unfold_copy);
+  }
+
+  UnfoldCopy(const torch::lazy::Value& self, const int64_t& dimension, const int64_t& size, const int64_t& step, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              UnfoldCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dimension, size, step)),
+        dimension(dimension),
+        size(size),
+        step(step)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dimension=" << dimension;
+    ss << ", size=" << size;
+    ss << ", step=" << step;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dimension, const int64_t& size, const int64_t& step) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dimension == dimension &&
+        this->size == size &&
+        this->step == step);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dimension", dimension);
+    arguments.emplace_back("size", size);
+    arguments.emplace_back("step", step);
+
+    torch::lazy::TSOpVector unfold_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(unfold_copy_out.size(), 1);
+
+    return unfold_copy_out;
+
+  }
+
+
+  int64_t dimension;
+  int64_t size;
+  int64_t step;
+
+
+};
+
+class Uniform : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::uniform);
+  }
+
+  Uniform(const torch::lazy::Value& self, const double& from, const double& to, const ::std::optional<at::Generator>& generator, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Uniform::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(from, to, generator)),
+        from(from),
+        to(to),
+        generator(generator)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", from=" << from;
+    ss << ", to=" << to;
+    if (generator.has_value()) {
+      ss << ", generator=" << "torch.Generator()";
+    } else {
+      ss << ", generator=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const double& from, const double& to, const ::std::optional<at::Generator>& generator) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->from == from &&
+        this->to == to &&
+        ((!this->generator&&!generator) || (this->generator&&generator && *(this->generator) == *generator)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(3);
+    kwarguments.reserve(1);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("from", from);
+    arguments.emplace_back("to", to);
+    kwarguments.emplace_back("generator", generator);
+    torch::lazy::TSOpVector uniform_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(uniform_out.size(), 1);
+
+    return uniform_out;
+
+  }
+
+
+  double from;
+  double to;
+  ::std::optional<at::Generator> generator;
+
+
+};
+
+class UnsqueezeCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::unsqueeze_copy);
+  }
+
+  UnsqueezeCopy(const torch::lazy::Value& self, const int64_t& dim, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              UnsqueezeCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dim)),
+        dim(dim)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dim=" << dim;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const int64_t& dim) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dim == dim);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dim", dim);
+
+    torch::lazy::TSOpVector unsqueeze_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(unsqueeze_copy_out.size(), 1);
+
+    return unsqueeze_copy_out;
+
+  }
+
+
+  int64_t dim;
+
+
+};
+
+class UpsampleBilinear2d : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::upsample_bilinear2d);
+  }
+
+  UpsampleBilinear2d(const torch::lazy::Value& self, const ::std::vector<int64_t>& output_size, const bool& align_corners, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              UpsampleBilinear2d::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(output_size, align_corners, scales_h, scales_w)),
+        output_size(output_size),
+        align_corners(align_corners),
+        scales_h(scales_h),
+        scales_w(scales_w)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", output_size=" << output_size;
+    ss << ", align_corners=" << align_corners;
+    if (scales_h.has_value()) {
+      ss << ", scales_h=" << scales_h.value();
+    } else {
+      ss << ", scales_h=null";
+    }
+    if (scales_w.has_value()) {
+      ss << ", scales_w=" << scales_w.value();
+    } else {
+      ss << ", scales_w=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& output_size, const bool& align_corners, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->output_size == output_size &&
+        this->align_corners == align_corners &&
+        ((!this->scales_h&&!scales_h) || (this->scales_h&&scales_h && *(this->scales_h) == *scales_h)) &&
+        ((!this->scales_w&&!scales_w) || (this->scales_w&&scales_w && *(this->scales_w) == *scales_w)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("output_size", output_size);
+    arguments.emplace_back("align_corners", align_corners);
+    arguments.emplace_back("scales_h", scales_h);
+    arguments.emplace_back("scales_w", scales_w);
+
+    torch::lazy::TSOpVector upsample_bilinear2d_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(upsample_bilinear2d_out.size(), 1);
+
+    return upsample_bilinear2d_out;
+
+  }
+
+
+  ::std::vector<int64_t> output_size;
+  bool align_corners;
+  ::std::optional<double> scales_h;
+  ::std::optional<double> scales_w;
+
+
+};
+
+class UpsampleBilinear2dBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::upsample_bilinear2d_backward);
+  }
+
+  UpsampleBilinear2dBackward(const torch::lazy::Value& grad_output, const ::std::vector<int64_t>& output_size, const ::std::vector<int64_t>& input_size, const bool& align_corners, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              UpsampleBilinear2dBackward::ClassOpKind(),
+              OpList{grad_output},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(output_size, input_size, align_corners, scales_h, scales_w)),
+        output_size(output_size),
+        input_size(input_size),
+        align_corners(align_corners),
+        scales_h(scales_h),
+        scales_w(scales_w)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", output_size=" << output_size;
+    ss << ", input_size=" << input_size;
+    ss << ", align_corners=" << align_corners;
+    if (scales_h.has_value()) {
+      ss << ", scales_h=" << scales_h.value();
+    } else {
+      ss << ", scales_h=null";
+    }
+    if (scales_w.has_value()) {
+      ss << ", scales_w=" << scales_w.value();
+    } else {
+      ss << ", scales_w=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const ::std::vector<int64_t>& output_size, const ::std::vector<int64_t>& input_size, const bool& align_corners, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        this->output_size == output_size &&
+        this->input_size == input_size &&
+        this->align_corners == align_corners &&
+        ((!this->scales_h&&!scales_h) || (this->scales_h&&scales_h && *(this->scales_h) == *scales_h)) &&
+        ((!this->scales_w&&!scales_w) || (this->scales_w&&scales_w && *(this->scales_w) == *scales_w)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(6);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("output_size", output_size);
+    arguments.emplace_back("input_size", input_size);
+    arguments.emplace_back("align_corners", align_corners);
+    arguments.emplace_back("scales_h", scales_h);
+    arguments.emplace_back("scales_w", scales_w);
+
+    torch::lazy::TSOpVector upsample_bilinear2d_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(upsample_bilinear2d_backward_out.size(), 1);
+
+    return upsample_bilinear2d_backward_out;
+
+  }
+
+
+  ::std::vector<int64_t> output_size;
+  ::std::vector<int64_t> input_size;
+  bool align_corners;
+  ::std::optional<double> scales_h;
+  ::std::optional<double> scales_w;
+
+
+};
+
+class UpsampleNearest2d : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::upsample_nearest2d);
+  }
+
+  UpsampleNearest2d(const torch::lazy::Value& self, const ::std::vector<int64_t>& output_size, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              UpsampleNearest2d::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(output_size, scales_h, scales_w)),
+        output_size(output_size),
+        scales_h(scales_h),
+        scales_w(scales_w)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", output_size=" << output_size;
+    if (scales_h.has_value()) {
+      ss << ", scales_h=" << scales_h.value();
+    } else {
+      ss << ", scales_h=null";
+    }
+    if (scales_w.has_value()) {
+      ss << ", scales_w=" << scales_w.value();
+    } else {
+      ss << ", scales_w=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& output_size, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->output_size == output_size &&
+        ((!this->scales_h&&!scales_h) || (this->scales_h&&scales_h && *(this->scales_h) == *scales_h)) &&
+        ((!this->scales_w&&!scales_w) || (this->scales_w&&scales_w && *(this->scales_w) == *scales_w)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(4);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("output_size", output_size);
+    arguments.emplace_back("scales_h", scales_h);
+    arguments.emplace_back("scales_w", scales_w);
+
+    torch::lazy::TSOpVector upsample_nearest2d_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(upsample_nearest2d_out.size(), 1);
+
+    return upsample_nearest2d_out;
+
+  }
+
+
+  ::std::vector<int64_t> output_size;
+  ::std::optional<double> scales_h;
+  ::std::optional<double> scales_w;
+
+
+};
+
+class UpsampleNearest2dBackward : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::upsample_nearest2d_backward);
+  }
+
+  UpsampleNearest2dBackward(const torch::lazy::Value& grad_output, const ::std::vector<int64_t>& output_size, const ::std::vector<int64_t>& input_size, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              UpsampleNearest2dBackward::ClassOpKind(),
+              OpList{grad_output},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(output_size, input_size, scales_h, scales_w)),
+        output_size(output_size),
+        input_size(input_size),
+        scales_h(scales_h),
+        scales_w(scales_w)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", output_size=" << output_size;
+    ss << ", input_size=" << input_size;
+    if (scales_h.has_value()) {
+      ss << ", scales_h=" << scales_h.value();
+    } else {
+      ss << ", scales_h=null";
+    }
+    if (scales_w.has_value()) {
+      ss << ", scales_w=" << scales_w.value();
+    } else {
+      ss << ", scales_w=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& grad_output, const ::std::vector<int64_t>& output_size, const ::std::vector<int64_t>& input_size, const ::std::optional<double>& scales_h, const ::std::optional<double>& scales_w) const {
+    size_t i = 0;
+    return (operand(i++) == grad_output &&
+        this->output_size == output_size &&
+        this->input_size == input_size &&
+        ((!this->scales_h&&!scales_h) || (this->scales_h&&scales_h && *(this->scales_h) == *scales_h)) &&
+        ((!this->scales_w&&!scales_w) || (this->scales_w&&scales_w && *(this->scales_w) == *scales_w)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(5);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("output_size", output_size);
+    arguments.emplace_back("input_size", input_size);
+    arguments.emplace_back("scales_h", scales_h);
+    arguments.emplace_back("scales_w", scales_w);
+
+    torch::lazy::TSOpVector upsample_nearest2d_backward_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(upsample_nearest2d_backward_out.size(), 1);
+
+    return upsample_nearest2d_backward_out;
+
+  }
+
+
+  ::std::vector<int64_t> output_size;
+  ::std::vector<int64_t> input_size;
+  ::std::optional<double> scales_h;
+  ::std::optional<double> scales_w;
+
+
+};
+
+class ViewCopy : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::view_copy);
+  }
+
+  ViewCopy(const torch::lazy::Value& self, const ::std::vector<int64_t>& size, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ViewCopy::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(size)),
+        size(size)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", size=" << size;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const ::std::vector<int64_t>& size) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->size == size);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("size", size);
+
+    torch::lazy::TSOpVector view_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(view_copy_out.size(), 1);
+
+    return view_copy_out;
+
+  }
+
+
+  ::std::vector<int64_t> size;
+
+
+};
+
+class ViewCopyDtype : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::view_copy);
+  }
+
+  ViewCopyDtype(const torch::lazy::Value& self, const at::ScalarType& dtype, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              ViewCopyDtype::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dtype)),
+        dtype(dtype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dtype=" << dtype;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self, const at::ScalarType& dtype) const {
+    size_t i = 0;
+    return (operand(i++) == self &&
+        this->dtype == dtype);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(2);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    arguments.emplace_back("dtype", dtype);
+
+    torch::lazy::TSOpVector view_copy_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(view_copy_out.size(), 1);
+
+    return view_copy_out;
+
+  }
+
+
+  at::ScalarType dtype;
+
+
+};
+
+class Zero : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::zero);
+  }
+
+  Zero(const torch::lazy::Value& self, std::vector<torch::lazy::Shape>&& shapes)
+      : TsNode(
+              Zero::ClassOpKind(),
+              OpList{self},
+              std::move(shapes),
+              /* num_outputs */ 1,
+              torch::lazy::MHash())
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& self) const {
+    size_t i = 0;
+    return (operand(i++) == self);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+        std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(0);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+
+    torch::lazy::TSOpVector zero_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(zero_out.size(), 1);
+
+    return zero_out;
+
+  }
+
+
+
+
+
+};
+
+} // namespace lazy
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyNativeFunctions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyNativeFunctions.h
new file mode 100644
index 0000000000000000000000000000000000000000..82f63eb0acc9037329361360459e52f8eda6d55c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyNativeFunctions.h
@@ -0,0 +1,216 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// an external backend might generate file within its code tree
+// and check all the source files within the tree with clang-format.
+// so, disable it since the backend might have a different config.
+// clang-format off
+
+// Autogenerated file by gen_backend_stubs.py. Do not edit directly!
+
+#include <ATen/Tensor.h>
+
+namespace torch {
+namespace lazy {
+
+struct LazyNativeFunctions {
+
+static ::std::tuple<at::Tensor,at::Tensor,at::Tensor> convolution_backward(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & weight, at::OptionalIntArrayRef bias_sizes, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups, ::std::array<bool,3> output_mask);
+static ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_batch_norm(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, bool training, double momentum, double eps);
+static ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_batch_norm_backward(const at::Tensor & grad_out, const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, const ::std::optional<at::Tensor> & save_mean, const ::std::optional<at::Tensor> & save_invstd, bool train, double eps, ::std::array<bool,3> output_mask);
+static ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_layer_norm(const at::Tensor & input, at::IntArrayRef normalized_shape, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, double eps);
+static ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_layer_norm_backward(const at::Tensor & grad_out, const at::Tensor & input, at::IntArrayRef normalized_shape, const at::Tensor & mean, const at::Tensor & rstd, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, ::std::array<bool,3> output_mask);
+static ::std::tuple<at::Tensor,at::Tensor,at::Tensor> svd(const at::Tensor & self, bool some, bool compute_uv);
+static ::std::tuple<at::Tensor,at::Tensor> grid_sampler_2d_backward(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners, ::std::array<bool,2> output_mask);
+static ::std::tuple<at::Tensor,at::Tensor> log_sigmoid_forward(const at::Tensor & self);
+static ::std::tuple<at::Tensor,at::Tensor> max(const at::Tensor & self, int64_t dim, bool keepdim);
+static ::std::tuple<at::Tensor,at::Tensor> max_pool2d_with_indices(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode);
+static ::std::tuple<at::Tensor,at::Tensor> max_pool3d_with_indices(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode);
+static ::std::tuple<at::Tensor,at::Tensor> native_dropout(const at::Tensor & input, double p, ::std::optional<bool> train);
+static ::std::tuple<at::Tensor,at::Tensor> nll_loss2d_forward(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, int64_t ignore_index);
+static ::std::tuple<at::Tensor,at::Tensor> nll_loss_forward(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, int64_t ignore_index);
+static ::std::tuple<at::Tensor,at::Tensor> sort(const at::Tensor & self, int64_t dim, bool descending);
+static ::std::tuple<at::Tensor,at::Tensor> topk(const at::Tensor & self, int64_t k, int64_t dim, bool largest, bool sorted);
+static at::Tensor & arange_out(const at::Scalar & start, const at::Scalar & end, const at::Scalar & step, at::Tensor & out);
+static at::Tensor & fill_(at::Tensor & self, const at::Scalar & value);
+static at::Tensor & logsumexp_out(const at::Tensor & self, at::IntArrayRef dim, bool keepdim, at::Tensor & out);
+static at::Tensor _adaptive_avg_pool2d(const at::Tensor & self, at::IntArrayRef output_size);
+static at::Tensor _adaptive_avg_pool2d_backward(const at::Tensor & grad_output, const at::Tensor & self);
+static at::Tensor _copy_from(const at::Tensor & self, const at::Tensor & dst, bool non_blocking);
+static at::Tensor _copy_from_and_resize(const at::Tensor & self, const at::Tensor & dst);
+static at::Tensor _log_softmax(const at::Tensor & self, int64_t dim, bool half_to_float);
+static at::Tensor _log_softmax_backward_data(const at::Tensor & grad_output, const at::Tensor & output, int64_t dim, at::ScalarType input_dtype);
+static at::Tensor _reshape_alias_copy_symint(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride);
+static at::Tensor _softmax(const at::Tensor & self, int64_t dim, bool half_to_float);
+static at::Tensor _softmax_backward_data(const at::Tensor & grad_output, const at::Tensor & output, int64_t dim, at::ScalarType input_dtype);
+static at::Tensor _to_copy(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, bool non_blocking, ::std::optional<at::MemoryFormat> memory_format);
+static at::Tensor _trilinear(const at::Tensor & i1, const at::Tensor & i2, const at::Tensor & i3, at::IntArrayRef expand1, at::IntArrayRef expand2, at::IntArrayRef expand3, at::IntArrayRef sumdim, int64_t unroll_dim);
+static at::Tensor _unsafe_view(const at::Tensor & self, at::IntArrayRef size);
+static at::Tensor abs(const at::Tensor & self);
+static at::Tensor add(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha);
+static at::Tensor addcdiv(const at::Tensor & self, const at::Tensor & tensor1, const at::Tensor & tensor2, const at::Scalar & value);
+static at::Tensor addcmul(const at::Tensor & self, const at::Tensor & tensor1, const at::Tensor & tensor2, const at::Scalar & value);
+static at::Tensor addmm(const at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, const at::Scalar & beta, const at::Scalar & alpha);
+static at::Tensor alias_copy(const at::Tensor & self);
+static at::Tensor all(const at::Tensor & self);
+static at::Tensor any(const at::Tensor & self);
+static at::Tensor as_strided_copy_symint(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<c10::SymInt> storage_offset);
+static at::Tensor as_strided_scatter_symint(const at::Tensor & self, const at::Tensor & src, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<c10::SymInt> storage_offset);
+static at::Tensor avg_pool2d(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, bool ceil_mode, bool count_include_pad, ::std::optional<int64_t> divisor_override);
+static at::Tensor avg_pool2d_backward(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, bool ceil_mode, bool count_include_pad, ::std::optional<int64_t> divisor_override);
+static at::Tensor baddbmm(const at::Tensor & self, const at::Tensor & batch1, const at::Tensor & batch2, const at::Scalar & beta, const at::Scalar & alpha);
+static at::Tensor bernoulli(const at::Tensor & self, ::std::optional<at::Generator> generator);
+static at::Tensor bernoulli(const at::Tensor & self, double p, ::std::optional<at::Generator> generator);
+static at::Tensor binary_cross_entropy(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction);
+static at::Tensor binary_cross_entropy_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction);
+static at::Tensor bitwise_and(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor bitwise_or(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor block_diag(at::TensorList tensors);
+static at::Tensor bmm(const at::Tensor & self, const at::Tensor & mat2);
+static at::Tensor cat(const at::ITensorListRef & tensors, int64_t dim);
+static at::Tensor clamp(const at::Tensor & self, const ::std::optional<at::Scalar> & min, const ::std::optional<at::Scalar> & max);
+static at::Tensor clamp_min(const at::Tensor & self, const at::Scalar & min);
+static at::Tensor clone(const at::Tensor & self, ::std::optional<at::MemoryFormat> memory_format);
+static at::Tensor constant_pad_nd(const at::Tensor & self, at::IntArrayRef pad, const at::Scalar & value);
+static at::Tensor convolution(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups);
+static at::Tensor cos(const at::Tensor & self);
+static at::Tensor cumsum(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype);
+static at::Tensor detach_copy(const at::Tensor & self);
+static at::Tensor diag_embed(const at::Tensor & self, int64_t offset, int64_t dim1, int64_t dim2);
+static at::Tensor diagonal_backward_symint(const at::Tensor & grad_output, c10::SymIntArrayRef input_sizes, int64_t offset, int64_t dim1, int64_t dim2);
+static at::Tensor diagonal_copy(const at::Tensor & self, int64_t offset, int64_t dim1, int64_t dim2);
+static at::Tensor diagonal_scatter(const at::Tensor & self, const at::Tensor & src, int64_t offset, int64_t dim1, int64_t dim2);
+static at::Tensor div(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor div(const at::Tensor & self, const at::Tensor & other, ::std::optional<c10::string_view> rounding_mode);
+static at::Tensor elu(const at::Tensor & self, const at::Scalar & alpha, const at::Scalar & scale, const at::Scalar & input_scale);
+static at::Tensor elu_backward(const at::Tensor & grad_output, const at::Scalar & alpha, const at::Scalar & scale, const at::Scalar & input_scale, bool is_result, const at::Tensor & self_or_result);
+static at::Tensor embedding(const at::Tensor & weight, const at::Tensor & indices, int64_t padding_idx, bool scale_grad_by_freq, bool sparse);
+static at::Tensor embedding_dense_backward(const at::Tensor & grad_output, const at::Tensor & indices, int64_t num_weights, int64_t padding_idx, bool scale_grad_by_freq);
+static at::Tensor empty_strided_symint(c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory);
+static at::Tensor empty_symint(c10::SymIntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format);
+static at::Tensor eq(const at::Tensor & self, const at::Scalar & other);
+static at::Tensor eq(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor exp(const at::Tensor & self);
+static at::Tensor expand_copy_symint(const at::Tensor & self, c10::SymIntArrayRef size, bool implicit);
+static at::Tensor flip(const at::Tensor & self, at::IntArrayRef dims);
+static at::Tensor floor(const at::Tensor & self);
+static at::Tensor frac(const at::Tensor & self);
+static at::Tensor gather(const at::Tensor & self, int64_t dim, const at::Tensor & index, bool sparse_grad);
+static at::Tensor ge(const at::Tensor & self, const at::Scalar & other);
+static at::Tensor ge(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor gelu(const at::Tensor & self, c10::string_view approximate);
+static at::Tensor gelu_backward(const at::Tensor & grad_output, const at::Tensor & self, c10::string_view approximate);
+static at::Tensor glu(const at::Tensor & self, int64_t dim);
+static at::Tensor glu_backward(const at::Tensor & grad_output, const at::Tensor & self, int64_t dim);
+static at::Tensor glu_jvp(const at::Tensor & glu, const at::Tensor & x, const at::Tensor & dx, int64_t dim);
+static at::Tensor grid_sampler_2d(const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners);
+static at::Tensor gt(const at::Tensor & self, const at::Scalar & other);
+static at::Tensor gt(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor hardsigmoid(const at::Tensor & self);
+static at::Tensor index_select(const at::Tensor & self, int64_t dim, const at::Tensor & index);
+static at::Tensor le(const at::Tensor & self, const at::Scalar & other);
+static at::Tensor le(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor leaky_relu(const at::Tensor & self, const at::Scalar & negative_slope);
+static at::Tensor leaky_relu_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & negative_slope, bool self_is_result);
+static at::Tensor lift(const at::Tensor & self);
+static at::Tensor lift_fresh(const at::Tensor & self);
+static at::Tensor linalg_pinv(const at::Tensor & self, const ::std::optional<at::Tensor> & atol, const ::std::optional<at::Tensor> & rtol, bool hermitian);
+static at::Tensor log(const at::Tensor & self);
+static at::Tensor log2(const at::Tensor & self);
+static at::Tensor log_sigmoid_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & buffer);
+static at::Tensor logdet(const at::Tensor & self);
+static at::Tensor lt(const at::Tensor & self, const at::Scalar & other);
+static at::Tensor lt(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor masked_fill(const at::Tensor & self, const at::Tensor & mask, const at::Scalar & value);
+static at::Tensor masked_fill(const at::Tensor & self, const at::Tensor & mask, const at::Tensor & value);
+static at::Tensor max(const at::Tensor & self);
+static at::Tensor max_pool2d_with_indices_backward(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode, const at::Tensor & indices);
+static at::Tensor max_pool3d_with_indices_backward(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode, const at::Tensor & indices);
+static at::Tensor maximum(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor mean(const at::Tensor & self, ::std::optional<at::ScalarType> dtype);
+static at::Tensor mean(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype);
+static at::Tensor min(const at::Tensor & self);
+static at::Tensor minimum(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor mm(const at::Tensor & self, const at::Tensor & mat2);
+static at::Tensor mul(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor mv(const at::Tensor & self, const at::Tensor & vec);
+static at::Tensor narrow_copy_symint(const at::Tensor & self, int64_t dim, c10::SymInt start, c10::SymInt length);
+static at::Tensor native_dropout_backward(const at::Tensor & grad_output, const at::Tensor & mask, double scale);
+static at::Tensor ne(const at::Tensor & self, const at::Scalar & other);
+static at::Tensor ne(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor neg(const at::Tensor & self);
+static at::Tensor new_empty_strided_symint(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory);
+static at::Tensor nll_loss2d_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, int64_t ignore_index, const at::Tensor & total_weight);
+static at::Tensor nll_loss_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, int64_t ignore_index, const at::Tensor & total_weight);
+static at::Tensor nonzero(const at::Tensor & self);
+static at::Tensor norm(const at::Tensor & self, const ::std::optional<at::Scalar> & p, at::IntArrayRef dim, bool keepdim);
+static at::Tensor normal_functional(const at::Tensor & self, double mean, double std, ::std::optional<at::Generator> generator);
+static at::Tensor permute_copy(const at::Tensor & self, at::IntArrayRef dims);
+static at::Tensor pixel_shuffle(const at::Tensor & self, int64_t upscale_factor);
+static at::Tensor pixel_unshuffle(const at::Tensor & self, int64_t downscale_factor);
+static at::Tensor pow(const at::Tensor & self, const at::Scalar & exponent);
+static at::Tensor pow(const at::Tensor & self, const at::Tensor & exponent);
+static at::Tensor random(const at::Tensor & self, ::std::optional<at::Generator> generator);
+static at::Tensor random(const at::Tensor & self, int64_t from, ::std::optional<int64_t> to, ::std::optional<at::Generator> generator);
+static at::Tensor random(const at::Tensor & self, int64_t to, ::std::optional<at::Generator> generator);
+static at::Tensor reciprocal(const at::Tensor & self);
+static at::Tensor relu(const at::Tensor & self);
+static at::Tensor remainder(const at::Tensor & self, const at::Tensor & other);
+static at::Tensor repeat(const at::Tensor & self, at::IntArrayRef repeats);
+static at::Tensor rsqrt(const at::Tensor & self);
+static at::Tensor scatter_add(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src);
+static at::Tensor select_backward_symint(const at::Tensor & grad_output, c10::SymIntArrayRef input_sizes, int64_t dim, c10::SymInt index);
+static at::Tensor select_copy(const at::Tensor & self, int64_t dim, int64_t index);
+static at::Tensor select_scatter(const at::Tensor & self, const at::Tensor & src, int64_t dim, int64_t index);
+static at::Tensor sgn(const at::Tensor & self);
+static at::Tensor sigmoid(const at::Tensor & self);
+static at::Tensor sigmoid_backward(const at::Tensor & grad_output, const at::Tensor & output);
+static at::Tensor silu(const at::Tensor & self);
+static at::Tensor slice_backward_symint(const at::Tensor & grad_output, c10::SymIntArrayRef input_sizes, int64_t dim, c10::SymInt start, c10::SymInt end, c10::SymInt step);
+static at::Tensor slice_copy_symint(const at::Tensor & self, int64_t dim, ::std::optional<c10::SymInt> start, ::std::optional<c10::SymInt> end, c10::SymInt step);
+static at::Tensor slice_scatter_symint(const at::Tensor & self, const at::Tensor & src, int64_t dim, ::std::optional<c10::SymInt> start, ::std::optional<c10::SymInt> end, c10::SymInt step);
+static at::Tensor smooth_l1_loss(const at::Tensor & self, const at::Tensor & target, int64_t reduction, double beta);
+static at::Tensor smooth_l1_loss_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction, double beta);
+static at::Tensor softplus(const at::Tensor & self, const at::Scalar & beta, const at::Scalar & threshold);
+static at::Tensor softplus_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & beta, const at::Scalar & threshold);
+static at::Tensor sqrt(const at::Tensor & self);
+static at::Tensor squeeze_copy(const at::Tensor & self);
+static at::Tensor squeeze_copy(const at::Tensor & self, at::IntArrayRef dim);
+static at::Tensor squeeze_copy(const at::Tensor & self, int64_t dim);
+static at::Tensor stack(at::TensorList tensors, int64_t dim);
+static at::Tensor std(const at::Tensor & self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim);
+static at::Tensor std(const at::Tensor & self, at::OptionalIntArrayRef dim, const ::std::optional<at::Scalar> & correction, bool keepdim);
+static at::Tensor std(const at::Tensor & self, bool unbiased);
+static at::Tensor sub(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha);
+static at::Tensor sum(const at::Tensor & self, ::std::optional<at::ScalarType> dtype);
+static at::Tensor sum(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype);
+static at::Tensor t_copy(const at::Tensor & self);
+static at::Tensor tanh(const at::Tensor & self);
+static at::Tensor tanh_backward(const at::Tensor & grad_output, const at::Tensor & output);
+static at::Tensor threshold(const at::Tensor & self, const at::Scalar & threshold, const at::Scalar & value);
+static at::Tensor threshold_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & threshold);
+static at::Tensor trace(const at::Tensor & self);
+static at::Tensor transpose_copy(const at::Tensor & self, int64_t dim0, int64_t dim1);
+static at::Tensor tril(const at::Tensor & self, int64_t diagonal);
+static at::Tensor triu(const at::Tensor & self, int64_t diagonal);
+static at::Tensor trunc(const at::Tensor & self);
+static at::Tensor unfold_copy(const at::Tensor & self, int64_t dimension, int64_t size, int64_t step);
+static at::Tensor uniform(const at::Tensor & self, double from, double to, ::std::optional<at::Generator> generator);
+static at::Tensor unsqueeze_copy(const at::Tensor & self, int64_t dim);
+static at::Tensor upsample_bilinear2d(const at::Tensor & self, at::IntArrayRef output_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w);
+static at::Tensor upsample_bilinear2d_backward(const at::Tensor & grad_output, at::IntArrayRef output_size, at::IntArrayRef input_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w);
+static at::Tensor upsample_nearest2d(const at::Tensor & self, at::IntArrayRef output_size, ::std::optional<double> scales_h, ::std::optional<double> scales_w);
+static at::Tensor upsample_nearest2d_backward(const at::Tensor & grad_output, at::IntArrayRef output_size, at::IntArrayRef input_size, ::std::optional<double> scales_h, ::std::optional<double> scales_w);
+static at::Tensor view_copy(const at::Tensor & self, at::ScalarType dtype);
+static at::Tensor view_copy_symint(const at::Tensor & self, c10::SymIntArrayRef size);
+static at::Tensor zero(const at::Tensor & self);
+static ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_group_norm(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, int64_t N, int64_t C, int64_t HxW, int64_t group, double eps);
+static at::Tensor max_pool3d(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode);
+
+};
+} // namespace lazy
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyNonNativeIr.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyNonNativeIr.h
new file mode 100644
index 0000000000000000000000000000000000000000..4cea78933ef5fb0fe8fb8dea61f56825200d92bc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/generated/LazyNonNativeIr.h
@@ -0,0 +1,160 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/ir_builder.h>
+#include <torch/csrc/lazy/core/internal_ops/ltc_ops.h>
+#include <torch/csrc/lazy/core/shape_inference.h>
+#include <torch/csrc/lazy/ts_backend/ts_node.h>
+
+// This file contains autogenerated LazyTensor Non Native IR nodes
+
+namespace torch {
+namespace lazy {
+
+class Scalar : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::prim::Constant);
+  }
+
+  Scalar(const at::Scalar& value, const at::ScalarType& type)
+      : TsNode(
+              Scalar::ClassOpKind(),
+              OpList{},
+              compute_shape_scalar(value, type),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(value, type)),
+        value(value),
+        type(type)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", value=" << value;
+    ss << ", type=" << type;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const at::Scalar& value, const at::ScalarType& type) const;
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override;
+
+  at::Scalar value;
+  at::ScalarType type;
+
+
+};
+
+class Expand : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(at::aten::expand);
+  }
+
+  Expand(const torch::lazy::Value& input, const ::std::vector<int64_t>& size, const bool& is_scalar_expand)
+      : TsNode(
+              Expand::ClassOpKind(),
+              OpList{input},
+              [&](){ return compute_shape_expand(operand(0), size, is_scalar_expand)[0]; },
+              /* num_outputs */ 1,
+              torch::lazy::MHash(size, is_scalar_expand)),
+        size(size),
+        is_scalar_expand(is_scalar_expand)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", size=" << size;
+    ss << ", is_scalar_expand=" << is_scalar_expand;
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& input, const ::std::vector<int64_t>& size, const bool& is_scalar_expand) const {
+    size_t i = 0;
+    return (operand(i++) == input &&
+        this->size == size &&
+        this->is_scalar_expand == is_scalar_expand);
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override;
+
+  ::std::vector<int64_t> size;
+  bool is_scalar_expand;
+
+
+};
+
+class Cast : public TsNode {
+ public:
+  static torch::lazy::OpKind ClassOpKind() {
+    return torch::lazy::OpKind(ltc_cast);
+  }
+
+  Cast(const torch::lazy::Value& input, const at::ScalarType& dtype, const ::std::optional<at::ScalarType>& stype)
+      : TsNode(
+              Cast::ClassOpKind(),
+              OpList{input},
+              compute_shape_cast(input, dtype, stype),
+              /* num_outputs */ 1,
+              torch::lazy::MHash(dtype, stype)),
+        dtype(dtype),
+        stype(stype)
+  {
+
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << TsNode::ToString();
+    ss << ", dtype=" << dtype;
+    if (stype.has_value()) {
+      ss << ", stype=" << stype.value();
+    } else {
+      ss << ", stype=null";
+    }
+    return ss.str();
+  }
+
+
+
+  bool CanBeReused(const torch::lazy::Value& input, const at::ScalarType& dtype, const ::std::optional<at::ScalarType>& stype) const {
+    size_t i = 0;
+    return (operand(i++) == input &&
+        this->dtype == dtype &&
+        ((!this->stype&&!stype) || (this->stype&&stype && *(this->stype) == *stype)));
+  }
+
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override;
+
+  at::ScalarType dtype;
+  ::std::optional<at::ScalarType> stype;
+
+
+};
+
+} // namespace lazy
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/python/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/python/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..56e5f624fcff53187e799126003008a0d2874429
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/python/init.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <pybind11/pybind11.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::lazy {
+
+TORCH_PYTHON_API void initLazyBindings(PyObject* module);
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/python/python_util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/python/python_util.h
new file mode 100644
index 0000000000000000000000000000000000000000..dc5777bb0f45af1f31d11fc08adb7f873f7bfe46
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/python/python_util.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/Export.h>
+#include <torch/csrc/lazy/core/ir_metadata.h>
+#include <optional>
+#include <vector>
+
+namespace torch::lazy {
+
+std::optional<SourceLocation> TORCH_PYTHON_API GetPythonFrameTop();
+
+std::vector<SourceLocation> TORCH_PYTHON_API GetPythonFrames();
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/config.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/config.h
new file mode 100644
index 0000000000000000000000000000000000000000..0157b30fc7817ed9a74cca3ceb769db3a31b320b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/config.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Flags.h>
+
+// TODO(whc) unclear if this is useful, has only been tested as true
+TORCH_DECLARE_bool(torch_lazy_ts_tensor_update_sync);
+
+TORCH_DECLARE_bool(torch_lazy_ts_cuda);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/dynamic_ir.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/dynamic_ir.h
new file mode 100644
index 0000000000000000000000000000000000000000..4c42b0831100df2c145d7d5ddc0933f7c2085460
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/dynamic_ir.h
@@ -0,0 +1,82 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/symbol.h>
+
+#include <memory>
+#include <string>
+
+#include <c10/core/ScalarType.h>
+#include <c10/util/Flags.h>
+#include <torch/csrc/lazy/core/dynamic_ir.h>
+#include <torch/csrc/lazy/core/hash.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/ir_metadata.h>
+#include <torch/csrc/lazy/ts_backend/ts_node.h>
+
+TORCH_DECLARE_bool(ltc_enable_dynamic_shapes);
+
+namespace torch::lazy {
+
+/**
+ * The goal of "dynamic" Nodes is to patch a hole in our tracing.
+ * Previously, if a user called `sizes` on a Tensor, it would leak out
+ * of our tracing system, as `sizes` returns a torch.Size or an int. To
+ * prevent this from happening, we introduce DimensionNode, a new type
+ * of Node that abstracts the operation of getting the dimensions of a
+ * Tensor.
+ *
+ * Consider the following example:
+ * ```
+ * numel = x.shape()[0] * x.shape()[1]
+ * ```
+ *
+ * Here, `x.shape()[i]` will be a SizeNode (subclass of DimensionNode),
+ * and the multiplication of the two SizeNodes will be represented by
+ * a SizeMul (also a subclass of DimensionNode). Through this, we can
+ * prevent `numel` from being represented as a Python int and thus
+ * burned into the Graph.
+ */
+
+// Represents the result of calling `size` on a Tensor
+class TORCH_API SizeNode : public TsNode, public DimensionNode {
+ public:
+  SizeNode(Value input, size_t dim);
+  int64_t getStaticValue() const override;
+  bool isSymbolic() const override;
+  std::string ToString() const override;
+  size_t dim_ = 0;
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      TSLoweringContext* loctx) const override;
+};
+
+class TORCH_API SizeAdd : public TsNode, public DimensionNode {
+ public:
+  SizeAdd(Value a, Value b);
+  int64_t getStaticValue() const override;
+  bool isSymbolic() const override;
+  std::string ToString() const override;
+};
+
+class TORCH_API SizeMul : public TsNode, public DimensionNode {
+ public:
+  SizeMul(Value a, Value b);
+  int64_t getStaticValue() const override;
+  bool isSymbolic() const override;
+  std::string ToString() const override;
+};
+
+class TORCH_API SizeDiv : public TsNode, public DimensionNode {
+ public:
+  SizeDiv(Value a, Value b);
+  int64_t getStaticValue() const override;
+  bool isSymbolic() const override;
+  std::string ToString() const override;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ir_builder.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ir_builder.h
new file mode 100644
index 0000000000000000000000000000000000000000..7d8dd8c804cc68910334fc737043c58e90cc4ea0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ir_builder.h
@@ -0,0 +1,74 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/core/internal_ops/ltc_ops.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/ir_builder.h>
+#include <torch/csrc/lazy/core/shape_inference.h>
+#include <torch/csrc/lazy/generated/LazyNonNativeIr.h>
+#include <torch/csrc/lazy/ts_backend/dynamic_ir.h>
+#include <torch/csrc/lazy/ts_backend/ops/device_data.h>
+#include <torch/csrc/lazy/ts_backend/ops/generic.h>
+#include <torch/csrc/lazy/ts_backend/ts_node.h>
+
+namespace torch::lazy {
+
+struct TorchScriptIrBuilder : IrBuilder {
+  NodePtr MakeDeviceData(
+      const std::shared_ptr<BackendData>& data) const override {
+    return DeviceData::Create(data);
+  }
+  // TODO: Scalar node is not currently used by ts_backend. Enable reusing
+  // Scalar node later if needed.
+  NodePtr MakeScalar(const at::Scalar& value, const at::ScalarType& type)
+      const override {
+    return MakeNode<Scalar>(value, type);
+  }
+  NodePtr MakeExpand(
+      const Value& input0,
+      const std::vector<int64_t>& size,
+      const bool& is_scalar_expand) const override {
+    return ReuseOrMakeNode<Expand>(input0, size, is_scalar_expand);
+  }
+  NodePtr MakeCast(
+      const Value& input0,
+      const at::ScalarType& dtype,
+      const std::optional<at::ScalarType>& stype =
+          std::nullopt) const override {
+    return ReuseOrMakeNode<Cast>(input0, dtype, stype);
+  }
+  NodePtr MakeTensorList(const OpList& inputs) const override {
+    return ReuseOrMakeNode<TensorList>(inputs);
+  }
+  // Generic needs cleanup
+  NodePtr MakeGeneric(
+      const OpKind& op,
+      const OpList& operands,
+      const Shape& shape,
+      const size_t& num_outputs = 1,
+      const hash_t& hash_seed =
+          static_cast<uint32_t>(0x5a2d296e9)) const override {
+    return MakeNode<Generic>(op, operands, shape, num_outputs, hash_seed);
+  }
+
+  // dynamic ir nodes
+  // TODO: verify if IR node reusing works for Dynamic shape ops
+  NodePtr MakeSizeNode(const Value& input, size_t dim) const override {
+    return MakeNode<SizeNode>(input, dim);
+  }
+  NodePtr MakeSizeAdd(const Value& a, const Value& b) const override {
+    return MakeNode<SizeAdd>(a, b);
+  }
+  NodePtr MakeSizeMul(const Value& a, const Value& b) const override {
+    return MakeNode<SizeMul>(a, b);
+  }
+  NodePtr MakeSizeDiv(const Value& a, const Value& b) const override {
+    return MakeNode<SizeDiv>(a, b);
+  }
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/device_data.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/device_data.h
new file mode 100644
index 0000000000000000000000000000000000000000..f258cf99c580b129f640070fd14839230a6f881f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/device_data.h
@@ -0,0 +1,55 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/backend/backend_data.h>
+#include <torch/csrc/lazy/core/internal_ops/ltc_ops.h>
+#include <torch/csrc/lazy/ts_backend/ts_node.h>
+
+#include <utility>
+
+namespace torch::lazy {
+
+class TORCH_API DeviceData : public TsNode {
+ public:
+  static OpKind ClassOpKind() {
+    return ltc_device_data;
+  }
+
+  explicit DeviceData(std::shared_ptr<BackendData> data);
+
+  // A DeviceData node can be reused if the shape matches,
+  // but we will substitute the actual data_ pointer under
+  // the hood.
+  bool CanBeReused(const std::shared_ptr<BackendData>& data) const {
+    return data_->shape() == data->shape();
+  }
+
+  std::string ToString() const override;
+
+  const std::shared_ptr<BackendData>& data() const {
+    return data_;
+  }
+
+  void SetData(std::shared_ptr<BackendData> data) {
+    data_ = std::move(data);
+  }
+
+  static const DeviceData* Cast(const Node* node);
+
+  // To reuse IR nodes, use this method to create DeviceData nodes
+  // instead of calling the constructor directconst ly.
+  static NodePtr Create(const std::shared_ptr<BackendData>& data);
+
+  TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      TSLoweringContext* loctx) const override;
+
+ private:
+  std::shared_ptr<BackendData> data_;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/generic.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/generic.h
new file mode 100644
index 0000000000000000000000000000000000000000..8334391a593aa80e966f057b94eb47b43834c03e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/generic.h
@@ -0,0 +1,57 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/ts_backend/ts_node.h>
+
+#include <torch/csrc/lazy/core/ir_builder.h>
+
+namespace torch::lazy {
+
+// Generic IR Node implementation for nodes which can simply be described by a
+// specific OpKind and a lowering function. IR nodes carrying
+// metadata should not be using this class TORCH_API (and have the metadata
+// captured by the LowerFn), but they should instead create a dedicated IR node.
+// Doing the former would limit IR introspection.
+class TORCH_API Generic : public TsNode {
+ public:
+  Generic(
+      OpKind op,
+      OpList operands,
+      Shape shape,
+      size_t num_outputs = 1,
+      hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9));
+
+  Generic(
+      OpKind op,
+      OpList operands,
+      const std::function<Shape()>& shape_fn,
+      size_t num_outputs = 1,
+      hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9));
+
+  Generic(
+      OpKind op,
+      OpList operands,
+      size_t num_outputs = 1,
+      hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9));
+
+  Generic(OpKind op, Shape shape, size_t num_outputs, hash_t hash_seed);
+
+ private:
+  hash_t hash_seed_;
+};
+
+inline NodePtr GenericOp(
+    OpKind op,
+    OpList operands,
+    Shape shape,
+    size_t num_outputs = 1,
+    hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9)) {
+  return MakeNode<Generic>(
+      op, operands, std::move(shape), num_outputs, hash_seed);
+}
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/to_copy.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/to_copy.h
new file mode 100644
index 0000000000000000000000000000000000000000..121cd0ffcbc99e5680c16e312013059d6dd5e74c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ops/to_copy.h
@@ -0,0 +1,130 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/ts_backend/ts_node.h>
+
+namespace torch::lazy {
+
+// This IR was copied from code-generated output, but the entire _to_copy
+// operator cannot be trivially code generated since it is only desirable to
+// capture IR for certain permutations of _to_copy (e.g. dtype), and for the
+// others it is difficult to even invoke the aten/eager fallback necessitating
+// directly implementing the right to(device) behavior
+class ToCopy : public torch::lazy::TsNode {
+ public:
+  static OpKind ClassOpKind() {
+    return OpKind(at::aten::_to_copy);
+  }
+
+  ToCopy(
+      const torch::lazy::Value& self,
+      const std::optional<at::ScalarType>& dtype,
+      const std::optional<at::Layout>& layout,
+      const std::optional<at::Device>& device,
+      const std::optional<bool>& pin_memory,
+      const bool& non_blocking,
+      const std::optional<at::MemoryFormat>& memory_format,
+      std::vector<torch::lazy::Shape>&& shapes)
+      : torch::lazy::TsNode(
+            ClassOpKind(),
+            {self},
+            std::move(shapes),
+            /* num_outputs */ 1,
+            torch::lazy::MHash(
+                dtype,
+                layout,
+                device,
+                pin_memory,
+                non_blocking,
+                memory_format)),
+
+        dtype(dtype),
+        layout(layout),
+        device(device),
+        pin_memory(pin_memory),
+        non_blocking(non_blocking),
+        memory_format(memory_format) {}
+
+  bool CanBeReused(
+      const torch::lazy::Value& self,
+      const std::optional<at::ScalarType>& dtype,
+      const std::optional<at::Layout>& layout,
+      const std::optional<at::Device>& device,
+      const std::optional<bool>& pin_memory,
+      const bool& non_blocking,
+      const std::optional<at::MemoryFormat>& memory_format) const {
+    size_t i = 0;
+    return (
+        operand(i++) == self && this->dtype == dtype &&
+        this->layout == layout && this->device == device &&
+        this->pin_memory == pin_memory && this->non_blocking == non_blocking &&
+        this->memory_format == memory_format);
+  }
+
+  std::string ToString() const override {
+    std::stringstream ss;
+    ss << torch::lazy::TsNode::ToString();
+    if (dtype.has_value()) {
+      ss << ", dtype=" << dtype.value();
+    } else {
+      ss << ", dtype=null";
+    }
+    if (layout.has_value()) {
+      ss << ", layout=" << layout.value();
+    } else {
+      ss << ", layout=null";
+    }
+    if (device.has_value()) {
+      ss << ", device=" << device.value();
+    } else {
+      ss << ", device=null";
+    }
+    if (pin_memory.has_value()) {
+      ss << ", pin_memory=" << pin_memory.value();
+    } else {
+      ss << ", pin_memory=null";
+    }
+    ss << ", non_blocking=" << non_blocking;
+    if (memory_format.has_value()) {
+      ss << ", memory_format=" << memory_format.value();
+    } else {
+      ss << ", memory_format=null";
+    }
+    return ss.str();
+  }
+
+  torch::lazy::TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      torch::lazy::TSLoweringContext* loctx) const override {
+    std::vector<torch::jit::NamedValue> arguments;
+    std::vector<torch::jit::NamedValue> kwarguments;
+    arguments.reserve(1);
+    kwarguments.reserve(6);
+    size_t i = 0;
+    arguments.emplace_back(loctx->GetOutputOp(operand(i++)));
+    kwarguments.emplace_back("dtype", dtype);
+    kwarguments.emplace_back("layout", layout);
+    kwarguments.emplace_back("device", device);
+    kwarguments.emplace_back("pin_memory", pin_memory);
+    kwarguments.emplace_back("non_blocking", non_blocking);
+    kwarguments.emplace_back("memory_format", memory_format);
+    torch::lazy::TSOpVector _to_copy_out =
+        torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
+    TORCH_CHECK_EQ(_to_copy_out.size(), 1);
+
+    return _to_copy_out;
+  }
+
+  std::optional<at::ScalarType> dtype;
+  std::optional<at::Layout> layout;
+  std::optional<at::Device> device;
+  std::optional<bool> pin_memory;
+  bool non_blocking;
+  std::optional<at::MemoryFormat> memory_format;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/tensor_aten_ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/tensor_aten_ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..b6d42fb70a08f6942e1a9c89ea387b0221878cb2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/tensor_aten_ops.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/core/tensor.h>
+
+namespace torch::lazy {
+
+//////////////////////////////////////////////////////////////////////////////
+// ATEN operators follows here, listed in alphabetical order.
+//////////////////////////////////////////////////////////////////////////////
+
+void copy_(torch::lazy::LazyTensorPtr& input, torch::lazy::LazyTensorPtr& src);
+// Fills the input with the given value.
+void fill_(torch::lazy::LazyTensorPtr& input, const at::Scalar& value);
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_autograd_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_autograd_functions.h
new file mode 100644
index 0000000000000000000000000000000000000000..3d7ba8436e4923dac4fe138a1a90cb62b084b546
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_autograd_functions.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/autograd/custom_function.h>
+
+namespace torch::lazy {
+
+struct MaxPool3dAutogradFunctionTS
+    : public torch::autograd::Function<MaxPool3dAutogradFunctionTS> {
+  static at::Tensor forward(
+      torch::autograd::AutogradContext* ctx,
+      const at::Tensor& self,
+      at::IntArrayRef kernel_size,
+      at::IntArrayRef stride,
+      at::IntArrayRef padding,
+      at::IntArrayRef dilation,
+      bool ceil_mode);
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output);
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_backend_impl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_backend_impl.h
new file mode 100644
index 0000000000000000000000000000000000000000..d00d8f1812545994e00b2137df9cdc11c1cf20e8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_backend_impl.h
@@ -0,0 +1,57 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/lazy/backend/backend_interface.h>
+
+#include <utility>
+
+namespace torch::lazy {
+
+class TORCH_API TSData : public torch::lazy::BackendData {
+ public:
+  TSData(const at::Scalar& scalar, const torch::lazy::BackendDevice& device)
+      : torch::lazy::BackendData(device, torch::lazy::Shape(scalar.type(), {})),
+        scalar(scalar) {}
+
+  TSData(
+      at::Tensor data,
+      const torch::lazy::Shape& shape,
+      const torch::lazy::BackendDevice& device)
+      : torch::lazy::BackendData(device, shape), data_(std::move(data)) {}
+
+  TSData(
+      const torch::lazy::Shape& shape,
+      const torch::lazy::BackendDevice& device)
+      : torch::lazy::BackendData(device, shape) {}
+
+  Handle GetHandle() override {
+    return reinterpret_cast<int64_t>(this);
+  }
+
+  void Assign(const torch::lazy::BackendData& data) override {
+    data_ = static_cast<const TSData&>(data).data_;
+  }
+
+  bool HasValue() const override {
+    return data_.defined();
+  }
+
+  at::Tensor data() {
+    return data_;
+  }
+
+  std::optional<at::Scalar> scalar;
+
+ private:
+  at::Tensor data_;
+};
+
+TORCH_API torch::lazy::BackendImplInterface* GetTSBackendImpl();
+
+TORCH_PYTHON_API void InitTorchScriptBackend();
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_eager_fallback.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_eager_fallback.h
new file mode 100644
index 0000000000000000000000000000000000000000..3cbf6f8a37d864acfe1f2be569a1b7cc436801fc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_eager_fallback.h
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/stack.h>
+#include <functional>
+
+namespace torch::lazy {
+
+bool force_eager_fallback(c10::Symbol op);
+void ltc_eager_fallback(
+    const c10::OperatorHandle& op,
+    torch::jit::Stack* stack);
+
+void ts_eager_fallback(
+    const c10::OperatorHandle& op,
+    torch::jit::Stack* stack,
+    c10::DeviceType device_type);
+
+// The TorchScript backend does not register itself with pytorch dispatcher
+// until it is explicitly initialized.  This function should only be called
+// by the main Torchscript backend init function.
+void register_ts_ltc_eager_fallback();
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_lowering_context.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_lowering_context.h
new file mode 100644
index 0000000000000000000000000000000000000000..3ab1b3191135cd0ef213962515cc264459f9f28b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_lowering_context.h
@@ -0,0 +1,156 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <sstream>
+
+#include <c10/util/Exception.h>
+#include <torch/csrc/api/include/torch/jit.h>
+#include <torch/csrc/jit/runtime/graph_executor.h>
+#include <torch/csrc/lazy/backend/lowering_context.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/ts_backend/ts_node_lowering.h>
+
+namespace torch::lazy {
+
+using TSOpVector = std::vector<torch::jit::Value*>;
+
+class TORCH_API TSComputation : public Computation {
+ public:
+  TSComputation(const std::shared_ptr<torch::jit::Graph>& graph)
+      : graph_(graph), graph_executor_(graph, "") {
+    for (torch::jit::Value* input : graph_->inputs()) {
+      parameter_names_.push_back(input->debugName());
+    }
+  }
+
+  int parameters_size() const override {
+    return static_cast<int>(parameter_names_.size());
+  }
+
+  const std::vector<Shape>& parameter_shapes() const override {
+    TORCH_CHECK(
+        false, "TODO(whc) implement TS computation shapes or change interface");
+    return parameter_shapes_;
+  }
+
+  const std::vector<std::string>& parameter_names() const override {
+    return parameter_names_;
+  }
+
+  const Shape& result_shape() const override {
+    TORCH_CHECK(
+        false, "TODO(whc) implement TS computation shapes or change interface");
+    return result_shape_;
+  }
+
+  const std::string to_string() const override {
+    std::ostringstream oss;
+    oss << *graph_;
+    return oss.str();
+  }
+
+  std::shared_ptr<torch::jit::Graph> graph() const {
+    return graph_;
+  }
+
+  torch::jit::GraphExecutor& graph_executor() {
+    return graph_executor_;
+  }
+
+ private:
+  std::shared_ptr<torch::jit::Graph> graph_;
+  torch::jit::GraphExecutor graph_executor_;
+  std::vector<std::string> parameter_names_;
+  std::vector<Shape> parameter_shapes_;
+  Shape result_shape_;
+};
+
+class TORCH_API TSLoweringContext : public LoweringContext {
+ public:
+  TSLoweringContext(const std::string& name, const BackendDevice device);
+
+  TSLoweringContext(
+      const std::string& name,
+      BackendDevice device,
+      c10::ArrayRef<const Node*> post_order,
+      Util::EmissionMap emit_status);
+
+  size_t AddResult(const Output& output) override {
+    return AddResult(GetOutputOp(output));
+  }
+
+  void AddParameter(
+      const torch::lazy::Output& output,
+      size_t index,
+      const Shape& shape,
+      const std::string& name) override {
+    TORCH_INTERNAL_ASSERT(false, "not implemented");
+  }
+
+  void Lower(const Node* node);
+
+  ComputationPtr Build() override {
+    for (torch::jit::Value* output : root_tuple_) {
+      graph_->block()->registerOutput(output);
+    }
+    return std::make_shared<TSComputation>(graph_);
+  }
+
+  // Retrieves the lowered operation for an output. If the requested output is
+  // not available yet, the graph behind the output's Node is lowered, and the
+  // corresponding TS operation returned.
+  torch::jit::Value* GetOutputOp(const Output& output) {
+    auto it = emitted_outputs_.find(output);
+    if (it == emitted_outputs_.end()) {
+      auto post_order = Util::ComputePostOrder(output.node, &emit_status_);
+      for (auto node : post_order) {
+        Lower(node);
+      }
+      // At this point the output better be present, otherwise there is an issue
+      // with the lowering code.
+      it = emitted_outputs_.find(output);
+      TORCH_CHECK(
+          it != emitted_outputs_.end(),
+          "No TS operation emitted for output: ",
+          output.ToString());
+    }
+    return it->second;
+  }
+
+  // Assigns the given TS operation to the specified output. As outputs are
+  // lowered in a post-order fashion, later nodes should always find their
+  // operands among the emitted outputs.
+  void AssignOutputOp(const Output& output, torch::jit::Value* op);
+
+  // If a parameter associated with data has already been declared, it will be
+  // returned. Otherwise a new one will be created, associated with the tensor
+  // held in data.
+  torch::jit::Value* GetParameter(const BackendDataPtr& data);
+
+  std::shared_ptr<torch::jit::Graph> graph() const {
+    return graph_;
+  }
+
+ private:
+  struct Parameter {
+    torch::jit::Value* param{nullptr};
+    size_t index = 0;
+  };
+
+  size_t AddResult(torch::jit::Value* op) {
+    root_tuple_.push_back(op);
+    return root_tuple_.size() - 1;
+  }
+
+  std::shared_ptr<torch::jit::Graph> graph_;
+  std::shared_ptr<torch::jit::GraphFunction> function_;
+  std::unordered_map<BackendData::Handle, Parameter> parameters_map_;
+  std::vector<torch::jit::Value*> root_tuple_;
+  OutputMap<torch::jit::Value*> emitted_outputs_;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_node.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_node.h
new file mode 100644
index 0000000000000000000000000000000000000000..5efd7eed90acd7f260b9f9a64fd86819cc9fb6c3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_node.h
@@ -0,0 +1,109 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/ArrayRef.h>
+#include <torch/csrc/jit/api/function_impl.h>
+#include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/lazy/backend/lowering_context.h>
+#include <torch/csrc/lazy/core/ir.h>
+#include <torch/csrc/lazy/core/shape.h>
+#include <torch/csrc/lazy/ts_backend/ts_lowering_context.h>
+
+namespace torch::lazy {
+
+using TSOpVector = std::vector<torch::jit::Value*>;
+
+class TORCH_API TsNode : public lazy::Node {
+ public:
+  TsNode(
+      OpKind op,
+      OpList operands,
+      std::vector<Shape>&& shapes,
+      size_t num_outputs,
+      hash_t hash_seed = kHashSeed);
+
+  TsNode(
+      OpKind op,
+      OpList operands,
+      const std::function<Shape()>& shape_fn,
+      size_t num_outputs,
+      hash_t hash_seed = kHashSeed);
+
+  TsNode(
+      OpKind op,
+      OpList operands,
+      size_t num_outputs,
+      hash_t hash_seed = kHashSeed);
+
+  TsNode(
+      OpKind op,
+      Shape shape,
+      size_t num_outputs,
+      hash_t hash_seed = kHashSeed);
+
+  ~TsNode() override = default;
+
+  hash_t hash() const override;
+
+  hash_t shapeHash() const override;
+
+  const std::string getPythonStacktrace() const;
+
+  // Lower is a backend-specific method since it returns a backend specific
+  // type. hence, it is convenient to define it differently per-backend rather
+  // than at Node API
+  virtual TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      TSLoweringContext* loctx) const;
+
+ private:
+  // The hash of the dag WITH size info. Used for shape caching
+  hash_t shape_hash_;
+  // The hash of the dag used to look up the compiled graph by a hash
+  // in this case, we will use the dag hash WITHOUT size info if dynamic shape
+  // is enabled and use the dag hash WITH size info otherwise.
+  hash_t dag_hash_;
+};
+
+// Note: this OpKind is separate from ltc_ops.h since it would be a circular
+// import otherwise, I like leaving TensorList in this file, and I think most of
+// ltc_ops special cases will be deleted anyway
+const OpKind tensor_list_opkind = OpKind::Get("lazy_tensors::tensor_list");
+
+// TensorList represents an at::TensorList which is a vector[Tensor] but is also
+// a first-class IValue and can be fed as a single input to a TS program.  It is
+// much easier to handle TensorLists in Lazy Tensor code if they are represented
+// as a single Node so there can be more than one TensorList and more than one
+// Tensor side-by-side as operands to an op.
+//
+// Note: shape is undefined for TensorList.  We assert in some places that
+// #shapes matches #outputs and this stems from
+//       the fact that currently all IR nodes represent tensors (there is no
+//       type system for this IR).  Because of this, TensorList is a bit of a
+//       hack.
+//
+// TODO(whc) once Shape() API is moved to Node base, also make it virtual, and
+// then implement it as NotImplemented for TensorList, also fixing the assertion
+// that would fail.
+struct TORCH_API TensorList : public TsNode {
+  static OpKind ClassOpKind() {
+    return tensor_list_opkind;
+  }
+
+  TensorList() = delete;
+  TensorList(OpList values);
+
+  bool CanBeReused(OpList values) const {
+    return operands() == std::vector<Output>(values.begin(), values.end());
+  }
+
+  TSOpVector Lower(
+      std::shared_ptr<torch::jit::GraphFunction> function,
+      TSLoweringContext* loctx) const override;
+};
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_node_lowering.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_node_lowering.h
new file mode 100644
index 0000000000000000000000000000000000000000..37a11e964bb5e8e4eeaff374b8a5b7792c3502b0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/lazy/ts_backend/ts_node_lowering.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/api/include/torch/jit.h>
+#include <torch/csrc/lazy/backend/lowering_context.h>
+
+namespace torch::lazy {
+using TSOpVector = std::vector<torch::jit::Value*>;
+
+TORCH_API TSOpVector LowerTSBuiltin(
+    const std::shared_ptr<torch::jit::GraphFunction>& function,
+    c10::Symbol sym,
+    const std::vector<torch::jit::NamedValue>& arguments,
+    const std::vector<torch::jit::NamedValue>& kwarguments = {});
+
+} // namespace torch::lazy
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/counters.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/counters.h
new file mode 100644
index 0000000000000000000000000000000000000000..137c149ed6737e3484f6536a6dba0462c47f33ee
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/counters.h
@@ -0,0 +1,285 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <bitset>
+#include <mutex>
+#include <sstream>
+#include <unordered_map>
+#include <vector>
+
+#include <c10/macros/Macros.h>
+
+#include <torch/csrc/monitor/events.h>
+
+namespace torch::monitor {
+
+constexpr int NUM_AGGREGATIONS = 7;
+
+// Aggregation is the list of possible aggregations for Stats.
+// These use bitwise flags so they can be efficiently stored.
+enum class C10_API_ENUM Aggregation {
+  // NONE means no aggregations are set.
+  NONE = 0,
+  // VALUE exports the most recently set value.
+  VALUE = 1,
+  // MEAN computes the mean of the set values within the window. Zero if no
+  // values.
+  MEAN = 2,
+  // COUNT tracks the number of times a value is set within the window.
+  COUNT = 3,
+  // SUM computes the sum of the values set within the window.
+  SUM = 4,
+  // MIN computes the minimum of the values set within the window. Zero if no
+  // values.
+  MAX = 5,
+  // MAX computes the maximum of the values set within the window. Zero if no
+  // values.
+  MIN = 6,
+};
+
+struct TORCH_API AggregationHash{template <typename T> std::size_t operator()(
+    T t) const {return static_cast<std::size_t>(t);
+} // namespace torch::monitor
+}
+;
+
+// aggregationName returns the human readable name corresponding to the
+// aggregation.
+TORCH_API const char* aggregationName(Aggregation agg);
+
+template <typename T>
+class Stat;
+
+namespace {
+template <typename T>
+inline std::bitset<NUM_AGGREGATIONS> merge(T& list) {
+  std::bitset<NUM_AGGREGATIONS> a;
+  for (Aggregation b : list) {
+    a.set(static_cast<int>(b));
+  }
+  return a;
+}
+} // namespace
+
+namespace detail {
+void TORCH_API registerStat(Stat<double>* stat);
+void TORCH_API registerStat(Stat<int64_t>* stat);
+void TORCH_API unregisterStat(Stat<double>* stat);
+void TORCH_API unregisterStat(Stat<int64_t>* stat);
+} // namespace detail
+
+// Stat is used to compute summary statistics in a performant way over fixed
+// intervals. Stat logs the statistics as an Event once every `windowSize`
+// duration. When the window closes the stats are logged via the event handlers
+// as a `torch.monitor.Stat` event.
+//
+// `windowSize` should be set to something relatively high to avoid a huge
+// number of events being logged. Ex: 60s. Stat uses millisecond precision.
+//
+// If maxSamples is set, the stat will cap the number of samples per window by
+// discarding `add` calls once `maxSamples` adds have occurred. If it's not set,
+// all `add` calls during the window will be included.
+// This is an optional field to make aggregations more directly comparable
+// across windows when the number of samples might vary.
+//
+// Stats support double and int64_t data types depending on what needs to be
+// logged and needs to be templatized with one of them.
+//
+// When the Stat is destructed it will log any remaining data even if the window
+// hasn't elapsed.
+template <typename T>
+class Stat {
+ private:
+  struct Values {
+    T value{0};
+    T sum{0};
+    T min{0};
+    T max{0};
+    int64_t count{0};
+  };
+
+ public:
+  Stat(
+      std::string name,
+      std::initializer_list<Aggregation> aggregations,
+      std::chrono::milliseconds windowSize,
+      int64_t maxSamples = std::numeric_limits<int64_t>::max())
+      : name_(std::move(name)),
+        aggregations_(merge(aggregations)),
+        windowSize_(windowSize),
+        maxSamples_(maxSamples) {
+    detail::registerStat(this);
+  }
+
+  Stat(
+      std::string name,
+      std::vector<Aggregation> aggregations,
+      std::chrono::milliseconds windowSize,
+      int64_t maxSamples = std::numeric_limits<int64_t>::max())
+      : name_(std::move(name)),
+        aggregations_(merge(aggregations)),
+        windowSize_(windowSize),
+        maxSamples_(maxSamples) {
+    detail::registerStat(this);
+  }
+  Stat(const Stat&) = delete;
+  Stat(Stat&&) = delete;
+  Stat& operator=(const Stat&) = delete;
+  Stat& operator=(Stat&&) = delete;
+
+  virtual ~Stat() {
+    {
+      // on destruction log if there's unlogged data
+      std::lock_guard<std::mutex> guard(mu_);
+      logLocked();
+    }
+    detail::unregisterStat(this);
+  }
+
+  // add adds the value v to the current window.
+  void add(T v) {
+    std::lock_guard<std::mutex> guard(mu_);
+    maybeLogLocked();
+
+    if (alreadyLogged()) {
+      return;
+    }
+
+    if (aggregations_.test(static_cast<int>(Aggregation::VALUE))) {
+      current_.value = v;
+    }
+    if (aggregations_.test(static_cast<int>(Aggregation::MEAN)) ||
+        aggregations_.test(static_cast<int>(Aggregation::SUM))) {
+      current_.sum += v;
+    }
+
+    if (aggregations_.test(static_cast<int>(Aggregation::MAX))) {
+      if (current_.max < v || current_.count == 0) {
+        current_.max = v;
+      }
+    }
+    if (aggregations_.test(static_cast<int>(Aggregation::MIN))) {
+      if (current_.min > v || current_.count == 0) {
+        current_.min = v;
+      }
+    }
+
+    current_.count += 1;
+    maybeLogLocked();
+  }
+
+  const std::string& name() const noexcept {
+    return name_;
+  }
+
+  // count returns the number of items in the current open window.
+  int64_t count() noexcept {
+    std::lock_guard<std::mutex> guard(mu_);
+
+    return current_.count;
+  }
+
+  std::unordered_map<Aggregation, T, AggregationHash> get() noexcept {
+    std::lock_guard<std::mutex> guard(mu_);
+    return getLocked();
+  }
+
+ protected:
+  virtual uint64_t currentWindowId() const {
+    std::chrono::milliseconds now =
+        std::chrono::duration_cast<std::chrono::milliseconds>(
+            std::chrono::steady_clock::now().time_since_epoch());
+
+    // always returns a currentWindowId of at least 1 to avoid 0 window issues
+    return (now / windowSize_) + 1;
+  }
+
+ private:
+  bool alreadyLogged() {
+    return lastLoggedWindowId_ == currentWindowId();
+  }
+
+  void maybeLogLocked() {
+    auto windowId = currentWindowId();
+    bool shouldLog = windowId_ != windowId || current_.count >= maxSamples_;
+    if (shouldLog && !alreadyLogged()) {
+      logLocked();
+      lastLoggedWindowId_ = windowId_;
+      windowId_ = windowId;
+    }
+  }
+
+  void logLocked() {
+    prev_ = current_;
+    current_ = Values();
+
+    // don't log event if there's no data
+    if (prev_.count == 0) {
+      return;
+    }
+
+    Event e;
+    e.name = "torch.monitor.Stat";
+    e.timestamp = std::chrono::system_clock::now();
+
+    auto stats = getLocked();
+    e.data.reserve(stats.size());
+    for (auto& kv : stats) {
+      std::stringstream key;
+      key << name_;
+      key << '.';
+      key << aggregationName(kv.first);
+      e.data[key.str()] = kv.second;
+    }
+
+    logEvent(e);
+  }
+
+  std::unordered_map<Aggregation, T, AggregationHash> getLocked()
+      const noexcept {
+    std::unordered_map<Aggregation, T, AggregationHash> out;
+    out.reserve(aggregations_.count());
+
+    if (aggregations_.test(static_cast<int>(Aggregation::VALUE))) {
+      out.emplace(Aggregation::VALUE, prev_.value);
+    }
+    if (aggregations_.test(static_cast<int>(Aggregation::MEAN))) {
+      if (prev_.count == 0) {
+        out.emplace(Aggregation::MEAN, 0);
+      } else {
+        out.emplace(Aggregation::MEAN, prev_.sum / prev_.count);
+      }
+    }
+    if (aggregations_.test(static_cast<int>(Aggregation::COUNT))) {
+      out.emplace(Aggregation::COUNT, prev_.count);
+    }
+    if (aggregations_.test(static_cast<int>(Aggregation::SUM))) {
+      out.emplace(Aggregation::SUM, prev_.sum);
+    }
+    if (aggregations_.test(static_cast<int>(Aggregation::MAX))) {
+      out.emplace(Aggregation::MAX, prev_.max);
+    }
+    if (aggregations_.test(static_cast<int>(Aggregation::MIN))) {
+      out.emplace(Aggregation::MIN, prev_.min);
+    }
+
+    return out;
+  }
+
+  const std::string name_;
+  const std::bitset<NUM_AGGREGATIONS> aggregations_;
+
+  std::mutex mu_;
+  Values current_;
+  Values prev_;
+
+  uint64_t windowId_{0};
+  uint64_t lastLoggedWindowId_{0};
+  const std::chrono::milliseconds windowSize_;
+  const int64_t maxSamples_;
+};
+} // namespace torch::monitor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/events.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/events.h
new file mode 100644
index 0000000000000000000000000000000000000000..320d63457f91c46c501bfc9c62d1e3aeb8624676
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/events.h
@@ -0,0 +1,76 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <chrono>
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+#include <c10/macros/Macros.h>
+#include <variant>
+
+namespace torch::monitor {
+
+// data_value_t is the type for Event data values.
+using data_value_t = std::variant<std::string, double, int64_t, bool>;
+
+// Event represents a single event that can be logged out to an external
+// tracker. This does acquire a lock on logging so should be used relatively
+// infrequently to avoid performance issues.
+struct TORCH_API Event {
+  // name is the name of the event. This is a static string that's used to
+  // differentiate between event types for programmatic access. The type should
+  // be in the format of a fully qualified Python-style class name.
+  // Ex: torch.monitor.MonitorEvent
+  std::string name;
+
+  // timestamp is a timestamp relative to the Unix epoch time.
+  std::chrono::system_clock::time_point timestamp;
+
+  // data contains rich information about the event. The contents are event
+  // specific so you should check the type to ensure it's what you expect before
+  // accessing the data.
+  //
+  // NOTE: these events are not versioned and it's up to the consumer of the
+  // events to check the fields to ensure backwards compatibility.
+  std::unordered_map<std::string, data_value_t> data;
+};
+
+inline bool operator==(const Event& lhs, const Event& rhs) {
+  return lhs.name == rhs.name && lhs.timestamp == rhs.timestamp &&
+      lhs.data == rhs.data;
+}
+
+// EventHandler represents an abstract event handler that can be registered to
+// capture events. Every time an event is logged every handler will be called
+// with the events contents.
+//
+// NOTE: The handlers should avoid any IO, blocking calls or heavy computation
+// as this may block the main thread and cause performance issues.
+class TORCH_API EventHandler {
+ public:
+  virtual ~EventHandler() = default;
+
+  // handle needs to be implemented to handle the events. This may be called
+  // from multiple threads so needs to be thread safe.
+  virtual void handle(const Event& e) = 0;
+};
+
+// logEvent calls each registered event handler with the event. This method can
+// be called from concurrently from multiple threads.
+TORCH_API void logEvent(const Event& e);
+
+// registerEventHandler registers an EventHandler so it receives any logged
+// events. Typically an EventHandler will be registered during program
+// setup and unregistered at the end.
+TORCH_API void registerEventHandler(std::shared_ptr<EventHandler> p);
+
+// unregisterEventHandler unregisters the event handler pointed to by the
+// shared_ptr.
+TORCH_API void unregisterEventHandler(const std::shared_ptr<EventHandler>& p);
+
+} // namespace torch::monitor
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/python_init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/python_init.h
new file mode 100644
index 0000000000000000000000000000000000000000..bfe66eacb1ec44f7734a1ac71bfba64614e1d6a8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/monitor/python_init.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::monitor {
+
+void initMonitorBindings(PyObject* module);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mps/Module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mps/Module.h
new file mode 100644
index 0000000000000000000000000000000000000000..f636e580b673f35d4e0fbbc7308c79f020d72bc5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mps/Module.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::mps {
+
+PyMethodDef* python_functions();
+void initModule(PyObject* module);
+
+} // namespace torch::mps
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mtia/Module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mtia/Module.h
new file mode 100644
index 0000000000000000000000000000000000000000..84259c7651ebceace7f29d7429e5e3217eba57a4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mtia/Module.h
@@ -0,0 +1,15 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::mtia {
+
+// PyMethodDef* python_functions();
+void initModule(PyObject* module);
+
+} // namespace torch::mtia
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mtia/profiler/MTIAMemoryProfiler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mtia/profiler/MTIAMemoryProfiler.h
new file mode 100644
index 0000000000000000000000000000000000000000..b50d495618218335dd644f3da308eb59f28a9eac
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/mtia/profiler/MTIAMemoryProfiler.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/profiler/orchestration/python_tracer.h>
+
+namespace torch::mtia {
+using namespace torch::profiler::impl::python_tracer;
+
+void initMemoryProfiler();
+
+std::unique_ptr<PythonMemoryTracerBase> getMemoryTracer();
+
+class MTIAMemoryProfiler final : public PythonMemoryTracerBase {
+ public:
+  explicit MTIAMemoryProfiler() = default;
+  ~MTIAMemoryProfiler() override = default;
+  void start() override;
+  void stop() override;
+  void export_memory_history(const std::string& path) override;
+};
+
+} // namespace torch::mtia
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/multiprocessing/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/multiprocessing/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..1800ffa2ce1dd605a0daae7e707516920537060f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/multiprocessing/init.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::multiprocessing {
+
+const PyMethodDef* python_functions();
+
+} // namespace torch::multiprocessing
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/back_compat.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/back_compat.h
new file mode 100644
index 0000000000000000000000000000000000000000..5132fe262ac095863a64efb85bf0358dbde14e35
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/back_compat.h
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <onnx/onnx_pb.h>
+
+namespace torch::onnx {
+
+// The following constants are defined here to avoid breaking Meta's internal
+// usage of ONNX which pre-dates ONNX 1.14 and thus does not support FLOAT8:
+// cf. https://github.com/pytorch/pytorch/pull/106379#issuecomment-1675189340
+// -abock, 2023-08-25
+//
+// ::ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN
+constexpr auto TensorProto_DataType_FLOAT8E4M3FN =
+    static_cast<::ONNX_NAMESPACE::TensorProto_DataType>(17);
+// ::ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FNUZ
+constexpr auto TensorProto_DataType_FLOAT8E4M3FNUZ =
+    static_cast<::ONNX_NAMESPACE::TensorProto_DataType>(18);
+// ::ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E5M2
+constexpr auto TensorProto_DataType_FLOAT8E5M2 =
+    static_cast<::ONNX_NAMESPACE::TensorProto_DataType>(19);
+// ::ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E5M2FNUZ
+constexpr auto TensorProto_DataType_FLOAT8E5M2FNUZ =
+    static_cast<::ONNX_NAMESPACE::TensorProto_DataType>(20);
+
+} // namespace torch::onnx
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..4c451df8c9e5f9af3b5ece63f82e9c93cdd785ee
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/init.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::onnx {
+
+void initONNXBindings(PyObject* module);
+
+} // namespace torch::onnx
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/onnx.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/onnx.h
new file mode 100644
index 0000000000000000000000000000000000000000..6a5eb189134d0eaac3e056b939ef8cf539edb8d9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/onnx/onnx.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::onnx {
+
+enum class OperatorExportTypes {
+  ONNX, // Strict ONNX export
+  ONNX_ATEN, // ONNX With ATen op everywhere
+  ONNX_ATEN_FALLBACK, // ONNX export with ATen fallback
+  ONNX_FALLTHROUGH, // Export supported ONNX ops. Pass through unsupported ops.
+};
+
+enum class TrainingMode {
+  EVAL, // Inference mode
+  PRESERVE, // Preserve model state (eval/training)
+  TRAINING, // Training mode
+};
+
+constexpr auto kOnnxNodeNameAttribute = "onnx_name";
+
+} // namespace torch::onnx
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/api.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/api.h
new file mode 100644
index 0000000000000000000000000000000000000000..02514ee197c4cacafa6b324da402be8ed4504787
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/api.h
@@ -0,0 +1,19 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/profiler/orchestration/observer.h>
+
+// There are some components which use these symbols. Until we migrate them
+// we have to mirror them in the old autograd namespace.
+
+namespace torch::autograd::profiler {
+using torch::profiler::impl::ActivityType;
+using torch::profiler::impl::getProfilerConfig;
+using torch::profiler::impl::ProfilerConfig;
+using torch::profiler::impl::profilerEnabled;
+using torch::profiler::impl::ProfilerState;
+} // namespace torch::autograd::profiler
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/collection.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/collection.h
new file mode 100644
index 0000000000000000000000000000000000000000..98fbf924244663974e55609f445efae537ea9a62
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/collection.h
@@ -0,0 +1,715 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <type_traits>
+#include <utility>
+#include <variant>
+
+#include <ATen/Context.h>
+#include <c10/core/Device.h>
+#include <c10/core/TensorImpl.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/ApproximateClock.h>
+#include <c10/util/flat_hash_map.h>
+#include <c10/util/strong_type.h>
+#include <torch/csrc/profiler/containers.h>
+#include <torch/csrc/profiler/data_flow.h>
+#include <torch/csrc/profiler/events.h>
+#include <torch/csrc/profiler/kineto_shim.h>
+#include <torch/csrc/profiler/orchestration/python_tracer.h>
+#include <torch/csrc/profiler/perf.h>
+#include <torch/csrc/profiler/stubs/base.h>
+#include <torch/csrc/profiler/util.h>
+#include <torch/csrc/utils/python_stub.h>
+
+namespace torch::profiler::impl {
+
+enum class EventType : uint8_t {
+  TorchOp = 0,
+  Backend,
+  Vulkan,
+  Allocation,
+  OutOfMemory,
+  PyCall,
+  PyCCall,
+  Kineto,
+  PythonGC
+};
+
+// ============================================================================
+// == Value (Tensor, Scalar) summary ==========================================
+// ============================================================================
+struct TORCH_API RawTensorMetadataBase {
+  RawTensorMetadataBase() = default;
+  explicit RawTensorMetadataBase(const at::Tensor& t);
+
+  StorageImplData data_;
+  c10::ScalarType dtype_{c10::ScalarType::Undefined};
+  c10::Layout layout_{c10::Layout::Strided};
+  uint32_t size_dim_{0};
+};
+
+// Collected during profiling.
+struct TORCH_API RawTensorMetadata : RawTensorMetadataBase {
+  RawTensorMetadata() = default;
+  RawTensorMetadata(const RawTensorMetadata&) = default;
+  RawTensorMetadata(RawTensorMetadata&&) noexcept = default;
+  RawTensorMetadata& operator=(const RawTensorMetadata&) = default;
+  RawTensorMetadata& operator=(RawTensorMetadata&&) noexcept = default;
+  ~RawTensorMetadata() = default;
+  explicit RawTensorMetadata(const at::Tensor& t);
+
+  // Wrap `weak_self_` in `std::optional` and split device into components to
+  // keep struct default constructable. (which the std::array initializer needs)
+  std::optional<WeakTensor> weak_self_;
+  c10::DeviceType device_type_{c10::DeviceType::CPU};
+  c10::DeviceIndex device_index_{-1};
+};
+
+// Used during post processing.
+struct TORCH_API TensorMetadata : public RawTensorMetadataBase {
+  TensorMetadata(
+      const RawTensorMetadata& r,
+      std::vector<int64_t> sizes,
+      std::vector<int64_t> strides);
+
+  TensorImplAddress impl() const {
+    return weak_self_.get();
+  }
+
+  WeakTensor weak_self_;
+  c10::Device device_;
+  std::vector<int64_t> sizes_;
+  std::vector<int64_t> strides_;
+
+  // Set during `calculateUniqueTensorIDs`.
+  std::optional<TensorID> id_;
+  std::optional<AllocationID> allocation_id_;
+};
+
+// Used during post processing.
+struct TORCH_API ProfilerStepInfo {
+  int64_t start_time_ns; // start time of the profiler step
+  int64_t end_time_ns; // end time of the profiler step
+  uint64_t out_idx; // index of the profiler step in the profiler "out" var in
+                    // getRecords
+
+  ProfilerStepInfo(int64_t start, int64_t end, uint64_t out_idx)
+      : start_time_ns(start), end_time_ns(end), out_idx(out_idx) {}
+};
+
+using op_input_t = std::variant<
+    TensorMetadata,
+    std::vector<TensorMetadata>,
+    c10::IValue,
+    std::nullopt_t>;
+
+// ============================================================================
+// == ExtraFields =============================================================
+// ============================================================================
+template <EventType>
+struct ExtraFields;
+
+struct TorchOpBasicFields {
+  int64_t sequence_number_{0};
+  uint64_t forward_tid_{0};
+  at::RecordScope scope_{};
+  bool is_async_{false};
+  uint64_t record_function_id_{0};
+  int64_t debug_handle_{0};
+  std::string name_;
+  std::string overload_name_;
+
+  // Set in the exit callback.
+  uint64_t end_tid_{0};
+};
+
+using jit_stack_t = std::vector<std::string>;
+using jit_modules_t = std::vector<std::string>;
+using extra_args_t = std::unordered_map<std::string, c10::IValue>;
+using extra_meta_t = std::unordered_map<std::string, std::string>;
+using kwinputs_t = std::unordered_map<std::string, c10::IValue>;
+
+// Mirrors `libkineto::GenericTraceActivity::Flow`. Used during post processing
+// to embed Kineto events into the broader profiler tree structure.
+struct Flow {
+  uint32_t id{0};
+  uint32_t type{0};
+  uint32_t start{0};
+};
+
+struct FallbackPair {
+  ProfilerVoidEventStub device_event_start_ = nullptr;
+  ProfilerVoidEventStub device_event_end_ = nullptr;
+};
+
+template <>
+struct ExtraFields<EventType::TorchOp> : TorchOpBasicFields {
+  ExtraFields(
+      TorchOpBasicFields&& f,
+      uint64_t correlation_id,
+      c10::time_t end_time_ns,
+      std::vector<op_input_t>&& inputs,
+      std::vector<op_input_t>&& concrete_inputs,
+      jit_stack_t&& jit_stack,
+      jit_modules_t&& jit_modules,
+      extra_args_t&& extra_args,
+      extra_meta_t&& extra_meta,
+      kwinputs_t&& kwinputs,
+      FallbackPair&& device_fallback,
+      bool allow_tf32_cublas,
+      std::unique_ptr<perf_counters_t>&& perf_event_counters)
+      : TorchOpBasicFields(std::move(f)),
+        correlation_id_{correlation_id},
+        end_time_ns_{end_time_ns},
+        inputs_{std::move(inputs)},
+        concrete_inputs_{std::move(concrete_inputs)},
+        jit_stack_{std::move(jit_stack)},
+        jit_modules_{std::move(jit_modules)},
+        extra_args_{std::move(extra_args)},
+        extra_meta_{std::move(extra_meta)},
+        kwinputs_{std::move(kwinputs)},
+        device_fallback_{std::move(device_fallback)},
+        allow_tf32_cublas_{allow_tf32_cublas},
+        perf_event_counters_{std::move(perf_event_counters)} {}
+  uint64_t correlation_id_;
+  c10::time_t end_time_ns_;
+  std::vector<op_input_t> inputs_;
+  std::vector<op_input_t> concrete_inputs_;
+  jit_stack_t jit_stack_;
+  jit_modules_t jit_modules_;
+  extra_args_t extra_args_;
+  extra_meta_t extra_meta_;
+  kwinputs_t kwinputs_;
+  FallbackPair device_fallback_;
+  bool allow_tf32_cublas_;
+  std::unique_ptr<perf_counters_t> perf_event_counters_;
+  std::string metadata_json_;
+  Flow flow;
+};
+
+template <>
+struct ExtraFields<EventType::Backend> {
+  int64_t start_time_us_;
+  int64_t end_time_us_;
+  int64_t debug_handle_;
+  at::RecordScope scope_;
+  std::string name_;
+  std::string backend_;
+  jit_stack_t jit_stack_;
+  jit_modules_t jit_modules_;
+};
+
+template <>
+struct ExtraFields<EventType::PythonGC> {
+  std::string phase;
+  int64_t duration_ns_;
+};
+
+template <>
+struct ExtraFields<EventType::Vulkan> {
+  using raw_event_t = std::pair<c10::approx_time_t, vulkan_id_t>;
+  std::string name_;
+  int64_t duration_ns_{0};
+  // While building the event tree, we want to report a vulkan event's duration
+  // as 0 so that its end time doesn't exceed that of its parent cpu op
+  bool in_tree_building_{false};
+};
+
+struct RawAllocation {
+  c10::approx_time_t start_time_;
+  void* ptr_;
+  int64_t alloc_size_;
+  size_t total_allocated_;
+  size_t total_reserved_;
+  c10::DeviceType device_type_;
+  c10::DeviceIndex device_index_;
+};
+
+// For performance.
+static_assert(
+    std::is_trivial_v<RawAllocation>,
+    "Non-Trivial member of RawAllocation.");
+
+template <>
+struct ExtraFields<EventType::Allocation> : RawAllocation {
+  ExtraFields(const RawAllocation& allocation) : RawAllocation(allocation) {}
+
+  c10::Device device() const {
+    return {device_type_, device_index_};
+  }
+
+  std::optional<TensorID> id_;
+  std::optional<AllocationID> allocation_id_;
+};
+
+template <>
+struct ExtraFields<EventType::OutOfMemory> {
+  c10::approx_time_t start_time_;
+  int64_t alloc_size_;
+  size_t total_allocated_;
+  size_t total_reserved_;
+  c10::DeviceType device_type_;
+  c10::DeviceIndex device_index_;
+};
+
+// For performance.
+static_assert(
+    std::is_trivial_v<ExtraFields<EventType::OutOfMemory>>,
+    "Non-Trivial member of ExtraFields<EventType::OutOfMemory>.");
+
+struct PyFrameState {
+  int line_no_;
+  at::StringView filename_;
+  at::StringView funcname_;
+};
+
+template <typename T, typename Tag>
+using strong_t = strong::
+    type<T, Tag, strong::regular, strong::convertible_to<T>, strong::hashable>;
+
+using PyModuleSelf = strong_t<PyObject*, struct PyModuleSelf_>;
+using PyModuleCls = strong_t<PyObject*, struct PyModuleCls_>;
+using PyMethod = strong_t</*PyMethodDef*/ void*, struct PyMethod_>;
+using PyOptimizerSelf = strong_t<PyObject*, struct PyOptSelf_>;
+using PyOptimizerCls = strong_t<PyObject*, struct PyOptimizer_>;
+
+struct NNModuleInfo {
+  struct ParameterInfo {
+    std::string name_;
+    TensorMetadata metadata_;
+    std::optional<TensorMetadata> grad_metadata_;
+  };
+
+  PyModuleSelf self_;
+  PyModuleCls cls_;
+  at::StringView cls_name_;
+
+  std::vector<ParameterInfo> parameters_;
+  // Indicates that `self_` is the kth instance of `cls_` observed.
+  size_t id_{std::numeric_limits<size_t>::max()};
+};
+
+struct OptimizerInfo {
+  struct ParameterInfo {
+    TensorMetadata metadata_;
+    std::optional<TensorMetadata> grad_metadata_;
+    std::vector<std::pair<std::string, TensorMetadata>> state_;
+  };
+
+  PyOptimizerSelf self_;
+  PyOptimizerCls cls_;
+  at::StringView cls_name_;
+
+  std::vector<ParameterInfo> parameters_;
+};
+
+struct PyExtraFieldsBase {
+  PyExtraFieldsBase(
+      c10::time_t end_time_ns,
+      size_t python_tid,
+      PyFrameState caller)
+      : end_time_ns_{end_time_ns},
+        python_tid_{python_tid},
+        caller_{std::move(caller)} {}
+
+  c10::time_t end_time_ns_;
+  size_t python_tid_;
+  PyFrameState caller_;
+
+  // kth python event observed. (Used by TensorBoard)
+  size_t id_{std::numeric_limits<size_t>::max()};
+};
+
+template <>
+struct ExtraFields<EventType::PyCall> : public PyExtraFieldsBase {
+  struct args_t {
+    PyFrameState frame_state_;
+    std::optional<NNModuleInfo> module_info_;
+    std::optional<OptimizerInfo> optimizer_info_;
+  };
+
+  ExtraFields(
+      c10::time_t end_time_ns,
+      size_t python_tid,
+      PyFrameState caller,
+      args_t args)
+      : PyExtraFieldsBase(end_time_ns, python_tid, std::move(caller)),
+        callsite_{std::move(args.frame_state_)},
+        module_{std::move(args.module_info_)},
+        optimizer_{std::move(args.optimizer_info_)} {}
+
+  PyFrameState callsite_;
+  std::optional<NNModuleInfo> module_;
+  std::optional<OptimizerInfo> optimizer_;
+};
+
+template <>
+struct ExtraFields<EventType::PyCCall> : public PyExtraFieldsBase {
+  using args_t = at::StringView;
+
+  ExtraFields(
+      c10::time_t end_time_ns,
+      size_t python_tid,
+      PyFrameState caller,
+      args_t args)
+      : PyExtraFieldsBase(end_time_ns, python_tid, std::move(caller)),
+        function_name_{std::move(args)} {}
+
+  at::StringView function_name_;
+};
+
+template <>
+struct ExtraFields<EventType::Kineto> {
+  std::string name_;
+  int64_t duration_ns_{0};
+  uint64_t correlation_id_{0};
+  libkineto::ActivityType activity_type_;
+  Flow flow;
+  std::weak_ptr<Result> linked_activity_;
+  std::string metadata_json_;
+  extra_meta_t extra_meta_;
+};
+
+struct TORCH_API Result : public std::enable_shared_from_this<Result> {
+  template <typename... Args>
+  [[nodiscard]] static std::shared_ptr<Result> create(Args... args) {
+    return std::shared_ptr<Result>(new Result(std::forward<Args>(args)...));
+  }
+
+  template <typename T>
+  auto visit(T&& visitor) {
+    return std::visit(std::forward<T>(visitor), extra_fields_);
+  }
+
+  template <typename T>
+  auto visit(T&& visitor) const {
+    return std::visit(std::forward<T>(visitor), extra_fields_);
+  }
+
+  template <typename T, typename Fn>
+  void visit_if_base(const Fn& fn) const {
+    visit([&](const auto& extra_fields) {
+      using extra_fields_t = typename std::remove_cv_t<
+          typename std::remove_reference_t<decltype(extra_fields)>>;
+
+      if constexpr (std::is_base_of_v<T, extra_fields_t>) {
+        fn(extra_fields);
+      }
+    });
+  }
+
+  EventType tag() const {
+    return visit([](const auto& i) { return deduceTag(i); });
+  }
+
+  std::string name() const;
+  std::string overload_name() const;
+  libkineto::ActivityType kinetoType() const;
+  uint64_t correlationID() const;
+  int64_t endTimeNS() const;
+  uint64_t endTID() const;
+  c10::DeviceType deviceType() const;
+
+  int64_t start_time_ns_;
+  uint64_t start_tid_;
+  kineto::DeviceAndResource kineto_info_;
+  std::variant<
+      ExtraFields<EventType::TorchOp>,
+      ExtraFields<EventType::Backend>,
+      ExtraFields<EventType::Vulkan>,
+      ExtraFields<EventType::Allocation>,
+      ExtraFields<EventType::OutOfMemory>,
+      ExtraFields<EventType::PyCall>,
+      ExtraFields<EventType::PyCCall>,
+      ExtraFields<EventType::Kineto>,
+      ExtraFields<EventType::PythonGC>>
+      extra_fields_;
+
+  std::weak_ptr<Result> parent_;
+  std::vector<std::shared_ptr<Result>> children_;
+  bool finished_{false};
+  bool hidden_{false};
+  const torch::profiler::impl::kineto::activity_t* kineto_activity_{nullptr};
+
+ private:
+  template <EventType E>
+  Result(
+      int64_t start_time_ns,
+      uint64_t start_tid,
+      kineto::DeviceAndResource kineto_info,
+      ExtraFields<E>&& extra_fields)
+      : start_time_ns_{start_time_ns},
+        start_tid_{start_tid},
+        kineto_info_{kineto_info},
+        extra_fields_{std::move(extra_fields)} {}
+
+  template <EventType E>
+  static EventType deduceTag(const ExtraFields<E>& /*unused*/) {
+    return E;
+  }
+};
+
+struct KinetoObserverContext : public at::ObserverContext {
+  struct Event {
+    TorchOpBasicFields basic_fields_;
+    c10::approx_time_t start_time_;
+
+    // Set in the exit callback.
+    c10::approx_time_t end_time_{
+        std::numeric_limits<c10::approx_time_t>::min()};
+
+    bool allow_tf32_cublas_;
+    std::unique_ptr<perf_counters_t> counters_;
+    extra_meta_t* extra_nccl_meta_{};
+  };
+
+  explicit KinetoObserverContext(Event* event) : event_{event} {}
+
+  Event* event_;
+  FallbackPair* fallback_{nullptr};
+};
+
+constexpr int IO_ENCODER_DEFAULT_BLOCK_SIZE = 1024;
+
+constexpr int SCALAR_LIST_LENGTH_LIMIT = 30;
+
+// InputOutputEncoder
+// Stores each op_events' shapes and dtypes, and concrete values into a
+// contiguous AppendOnlyList so that we no longer create vectors for shapes
+// and dtypes on every op. Those vectors can be created during
+// post-processing.
+// It splits the data into two categories: input shapes and concrete inputs.
+class InputOutputEncoder final {
+ public:
+  void push(c10::ArrayRef<const c10::IValue> values);
+
+  // Used during post-processing to unpack the encoded data.
+  // Each method returns a "supplier" lambda which takes no arguments;
+  // invoking the lambda once will return a list of args that represent
+  // the inputs for one op.
+  // The data is split into two streams: "input shapes" and "concrete inputs".
+  // Note: "auto" only works because these are only used in collection.cpp,
+  // where they are implemented.
+  auto getInputShapeGenerator();
+  auto getConcreteInputGenerator();
+
+  bool isSupportedScalarList(const c10::IValue& list_candidate);
+
+  void clear();
+
+  enum class Tag {
+    Tensor = 0,
+    UndefinedTensor,
+    TensorListBegin, // TODO: generalize to other lists.
+    ScalarList,
+    Scalar,
+    Other,
+    TERMINATOR
+  };
+
+  enum class IOType { Shapes, ConcreteInputs, None };
+
+ private:
+  void push(const at::Tensor& t);
+
+  // Implementation detail for getInputShapeGenerator and
+  // getConcreteInputGenerator
+  auto getIValueGenerator(const IOType& io_type);
+
+  AppendOnlyList<Tag, IO_ENCODER_DEFAULT_BLOCK_SIZE> tags_;
+  AppendOnlyList<RawTensorMetadata, IO_ENCODER_DEFAULT_BLOCK_SIZE>
+      tensor_metadata_;
+  AppendOnlyList<int64_t, IO_ENCODER_DEFAULT_BLOCK_SIZE> tensor_sizes_strides_;
+  AppendOnlyList<c10::IValue, IO_ENCODER_DEFAULT_BLOCK_SIZE> ivalues_;
+};
+
+using perf_profiler_t = torch::profiler::impl::linux_perf::PerfProfiler;
+
+class TORCH_API ThreadLocalSubqueue {
+ public:
+  ThreadLocalSubqueue(const uint64_t tid, ProfilerConfig config);
+
+  std::unique_ptr<KinetoObserverContext> begin_op(const at::RecordFunction& fn);
+
+  template <class... Args>
+  void emplace_backend_event(Args&&... args) {
+    backend_events_.emplace_back(std::forward<Args>(args)...);
+  }
+
+  template <class... Args>
+  void emplace_vulkan_event(Args&&... args) {
+    vulkan_events_.emplace_back(std::forward<Args>(args)...);
+  }
+
+  template <class... Args>
+  void emplace_allocation_event(Args&&... args) {
+    allocations_.emplace_back(std::forward<Args>(args)...);
+  }
+
+  template <class... Args>
+  void emplace_ooms_event(Args&&... args) {
+    ooms_.emplace_back(std::forward<Args>(args)...);
+  }
+
+  template <class... Args>
+  void emplace_py_call(Args&&... args) {
+    py_calls_.emplace_back(std::forward<Args>(args)...);
+  }
+
+  template <class... Args>
+  void emplace_gc_call(Args&&... args) {
+    pythongc_.emplace_back(std::forward<Args>(args)...);
+  }
+
+  uint64_t tid() const {
+    return tid_;
+  }
+
+  const kineto::DeviceAndResource& kineto_info() const {
+    return kineto_info_;
+  }
+
+  inline void disable_perf_profiler(perf_counters_t& counters) const {
+    perf_profiler_->Disable(counters);
+  }
+
+ private:
+  uint64_t tid_;
+  ProfilerConfig config_;
+  kineto::DeviceAndResource kineto_info_;
+  std::unique_ptr<perf_profiler_t> perf_profiler_;
+
+  friend class RecordQueue;
+  // See `containers.h` for block size benchmarks.
+  static constexpr size_t BlockSize = 512;
+
+  struct TorchOpStorage {
+    // NB: This is a destructive operation.
+    void materialize(
+        std::vector<std::shared_ptr<Result>>& out,
+        std::vector<ProfilerStepInfo>& step_info,
+        const std::function<c10::time_t(c10::approx_time_t)>& time_converter,
+        const uint64_t tid,
+        const kineto::DeviceAndResource& kineto_info);
+
+    template <typename T, size_t ChunkSize>
+    class EventBlock : public std::array<T, ChunkSize> {
+     public:
+      EventBlock();
+      uint64_t correlation_id(const T* ptr) const;
+
+     private:
+      uint64_t id_start_;
+    };
+
+    using event_t = KinetoObserverContext::Event;
+    class OpList : public AppendOnlyList<event_t, BlockSize, EventBlock> {
+     public:
+      template <class... Args>
+      std::pair<event_t*, uint64_t> emplace_back(Args&&... args);
+      static uint64_t correlationID(const OpList::Iterator& e);
+    } op_events_;
+
+    // report_input_shapes
+    InputOutputEncoder inputs_outputs_;
+
+    // with_stack (JIT)
+    AppendOnlyList<jit_stack_t, BlockSize> jit_stack_;
+
+    // with_modules
+    AppendOnlyList<jit_modules_t, BlockSize> jit_modules_;
+
+    // with_flops
+    AppendOnlyList<extra_args_t, BlockSize> extra_args_;
+
+    // report extra metadata, i.e. collective communication meta
+    AppendOnlyList<extra_meta_t, BlockSize> extra_meta_;
+
+    // report kwinputs
+    AppendOnlyList<kwinputs_t, BlockSize> kwinputs_;
+
+    // ProfilerState::KINETO_GPU_FALLBACK or
+    // ProfilerState::KINETO_PRIVATEUSE1_FALLBACK
+    AppendOnlyList<FallbackPair, BlockSize> device_fallback_;
+  } torch_ops_;
+
+  // reportBackendEventToActiveKinetoProfiler
+  AppendOnlyList<ExtraFields<EventType::Backend>, BlockSize> backend_events_;
+
+  // _reportVulkanEventToProfiler
+  AppendOnlyList<ExtraFields<EventType::Vulkan>::raw_event_t, BlockSize>
+      vulkan_events_;
+
+  // reportMemoryUsage
+  AppendOnlyList<RawAllocation, BlockSize> allocations_;
+
+  // reportOOMs
+  AppendOnlyList<ExtraFields<EventType::OutOfMemory>, BlockSize> ooms_;
+
+  // with_stack (Python)
+  AppendOnlyList<
+      std::pair<python_tracer::TraceKey, c10::approx_time_t>,
+      BlockSize>
+      py_calls_;
+  // gc with_stack (Python)
+  AppendOnlyList<std::pair<std::string, c10::approx_time_t>, BlockSize>
+      pythongc_;
+};
+
+class TORCH_API RecordQueue {
+ public:
+  RecordQueue(ProfilerConfig config, std::set<ActivityType> activities);
+
+  bool tracePython() const;
+  bool getPythonGcEvents() const;
+  ThreadLocalSubqueue* getSubqueue();
+  void stop();
+  void restart();
+
+  // NB: This is a destructive operation.
+  std::pair<
+      std::vector<std::shared_ptr<Result>>,
+      std::unique_ptr<torch::profiler::impl::kineto::ActivityTraceWrapper>>
+  getRecords(
+      std::function<c10::time_t(c10::approx_time_t)> time_converter,
+      uint64_t start_time_ns,
+      uint64_t end_time_ns);
+
+ private:
+  uint32_t id_;
+  ProfilerConfig config_;
+  std::set<ActivityType> activities_;
+  ska::flat_hash_map<uint64_t, std::unique_ptr<ThreadLocalSubqueue>>
+      sub_queues_;
+  std::mutex sub_queue_mutex_;
+  std::unique_ptr<python_tracer::PythonTracerBase> python_tracer_;
+};
+
+TORCH_API bool get_record_concrete_inputs_enabled();
+TORCH_API void set_record_concrete_inputs_enabled_fn(
+    std::function<bool()> /*fn*/);
+TORCH_API void set_record_concrete_inputs_enabled_val(bool /*val*/);
+
+TORCH_API bool get_fwd_bwd_enabled();
+TORCH_API void set_fwd_bwd_enabled_fn(std::function<bool()> /*fn*/);
+TORCH_API void set_fwd_bwd_enabled_val(bool /*val*/);
+
+TORCH_API bool get_cuda_sync_enabled();
+TORCH_API void set_cuda_sync_enabled_fn(std::function<bool()> /*fn*/);
+TORCH_API void set_cuda_sync_enabled_val(bool /*val*/);
+
+// Comms related RecordFunctions will record information about tensor storage
+// locations.
+TORCH_API bool get_record_tensor_addrs_enabled();
+TORCH_API void set_record_tensor_addrs_enabled_fn(std::function<bool()> /*fn*/);
+TORCH_API void set_record_tensor_addrs_enabled_val(bool /*val*/);
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/combined_traceback.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/combined_traceback.h
new file mode 100644
index 0000000000000000000000000000000000000000..1553cdc2727203cdaf7d4e16f4731b336dbd6a2d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/combined_traceback.h
@@ -0,0 +1,119 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <torch/csrc/profiler/unwind/unwind.h>
+
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace torch {
+
+// struct that holds the result of symbolizing multiple tracebacks
+// each traceback is a list of indices into all_frames
+// (lots of Frames get duplicated across traces)
+struct TORCH_API SymbolizedTracebacks {
+  std::vector<unwind::Frame> all_frames;
+  // index into all_frames, so that
+  // it is possible to dedupe frame objects in
+  // construction of python objects
+  std::vector<std::vector<uint64_t>> tracebacks;
+};
+
+struct TORCH_API CapturedTraceback : public c10::GatheredContext {
+  struct PyFrame {
+    void* code; // PyCodeObject*, but python headers not present
+    int lasti;
+  };
+
+  static std::shared_ptr<CapturedTraceback> gather(
+      bool python,
+      bool script,
+      bool cpp);
+  CapturedTraceback() = default;
+  CapturedTraceback(const CapturedTraceback&) = delete;
+  CapturedTraceback& operator=(const CapturedTraceback&) = delete;
+  CapturedTraceback(CapturedTraceback&&) noexcept = default;
+  CapturedTraceback& operator=(CapturedTraceback&&) noexcept = delete;
+  ~CapturedTraceback() override;
+
+  using visitproc = int (*)(void* self, void* arg);
+
+  struct Python {
+    // Check if it's safe to gather Python frames from the current thread.
+    // Returns false for pure C++ threads that cannot acquire the GIL.
+    virtual bool canGather() = 0;
+    virtual std::vector<PyFrame> gather() = 0;
+    virtual void release(std::vector<PyFrame>& frames) = 0;
+    virtual void appendSymbolized(
+        const std::vector<PyFrame>& to_symbolize,
+        SymbolizedTracebacks& st) = 0;
+    // tp_traverse/tp_clear implementations
+    virtual int traverse(
+        std::vector<PyFrame>& frames,
+        visitproc visit,
+        void* arg) = 0;
+    virtual int clear(std::vector<PyFrame>& frames) = 0;
+    // Gather forward traceback from the current autograd node's anomaly
+    // metadata. Returns a vector of strings representing the forward stack
+    // trace, or empty if not available.
+    virtual std::vector<std::string> gatherForwardTraceback() {
+      return {};
+    }
+    virtual ~Python() = default;
+    Python* next_ = nullptr;
+  };
+  // called once by each python interpreter to
+  // register python stack recording functionality
+  // p cannot be deleted once added.
+  static void addPythonUnwinder(Python* p);
+
+  int traversePython(visitproc visit, void* arg);
+  int clearPython();
+
+ private:
+  std::vector<PyFrame> frames_;
+  std::vector<void*> cpp_frames_;
+  std::vector<jit::StackEntry> script_frames_;
+  friend TORCH_API SymbolizedTracebacks
+  symbolize(const std::vector<CapturedTraceback*>& to_symbolize);
+
+  // non-owning reference to one of the immortal Python* objects
+  // registered above.
+  Python* python_ = nullptr;
+
+  // Optional forward traceback from anomaly mode.
+  // This is a list of Python strings representing the forward stack trace
+  // when the autograd Node was created. Used to correlate backward allocations
+  // with forward operations.
+  std::optional<std::vector<std::string>> forward_traceback_;
+
+ public:
+  // Set the forward traceback from anomaly mode metadata
+  void set_forward_traceback(std::vector<std::string> traceback) {
+    forward_traceback_ = std::move(traceback);
+  }
+
+  // Get the forward traceback if available
+  const std::optional<std::vector<std::string>>& forward_traceback() const {
+    return forward_traceback_;
+  }
+};
+
+TORCH_API SymbolizedTracebacks
+symbolize(const std::vector<CapturedTraceback*>& to_symbolize);
+
+inline CapturedTraceback* getCapturedTracebackFromContext(
+    const std::shared_ptr<c10::GatheredContext>& x) {
+  auto* traceback = dynamic_cast<CapturedTraceback*>(x.get());
+  TORCH_CHECK(
+      traceback, "attempting to gather stack context from the wrong type.");
+  return traceback;
+}
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/containers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/containers.h
new file mode 100644
index 0000000000000000000000000000000000000000..49c872babfc71e7edd119d15c359877962304717
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/containers.h
@@ -0,0 +1,208 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <algorithm>
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <forward_list>
+#include <utility>
+
+#include <c10/macros/Macros.h>
+#include <c10/util/ArrayRef.h>
+#include <c10/util/Exception.h>
+
+namespace torch::profiler::impl {
+
+// ============================================================================
+// == AppendOnlyList ==========================================================
+// ============================================================================
+//   During profiling, we have a very predictable access pattern: we only
+// append to the end of the container. We can specialize and outperform both
+// std::vector (which must realloc) and std::deque (which performs a double
+// indirection), and this class of operation is sufficiently important to the
+// profiling hot path to warrant specializing:
+//   https://godbolt.org/z/rTjozf1c4
+//   https://quick-bench.com/q/mmfuu71ogwaiULDCJyHdKnHZms4    (Prototype #1,
+//   int) https://quick-bench.com/q/5vWDW6jjdXVdoffev2zst8D09no    (Prototype
+//   #1, int pair) https://quick-bench.com/q/IfEkfAQMeJSNBA52xtMP6Agcl-Q
+//   (Prototype #2, int pair)
+//   https://quick-bench.com/q/wJV2lKmuXL4XyGJzcI5hs4gEHFg    (Prototype #3, int
+//   pair) https://quick-bench.com/q/xiO8ZaBEkYRYUA9dFrMuPLlW9fo    (Full impl,
+//   int pair)
+// AppendOnlyList has 2x lower emplace overhead compared to more generic STL
+// containers.
+//
+//   The optimal value of `ChunkSize` will vary by use case, but testing shows
+// that a value of 1024 does a good job amortizing the `malloc` cost of growth.
+// Performance drops off for larger values, so testing on a case-by-case basis
+// is recommended if performance is absolutely critical.
+
+template <
+    typename T,
+    size_t ChunkSize,
+    template <typename U, size_t N> class block_t = std::array>
+class AppendOnlyList {
+ public:
+  using array_t = block_t<T, ChunkSize>;
+  static_assert(
+      std::is_base_of_v<std::array<T, ChunkSize>, array_t>,
+      "AppendOnlyList expects raw low level pointer storage.");
+  static_assert(ChunkSize > 0, "Block cannot be empty.");
+
+  AppendOnlyList() : buffer_last_{buffer_.before_begin()} {}
+  AppendOnlyList(const AppendOnlyList&) = delete;
+  AppendOnlyList(AppendOnlyList&&) = delete;
+  AppendOnlyList& operator=(const AppendOnlyList&) = delete;
+  AppendOnlyList& operator=(AppendOnlyList&&) = delete;
+  ~AppendOnlyList() = default;
+
+  size_t size() const {
+    return n_blocks_ * ChunkSize - (size_t)(end_ - next_);
+  }
+
+  template <class... Args>
+  T* emplace_back(Args&&... args) {
+    maybe_grow();
+    if constexpr (
+        std::is_trivially_destructible_v<T> &&
+        std::is_trivially_destructible_v<array_t>) {
+      ::new ((void*)next_) T{std::forward<Args>(args)...};
+    } else {
+      *next_ = T{std::forward<Args>(args)...};
+    }
+    return next_++;
+  }
+
+  template <typename T0>
+  std::enable_if_t<std::is_same_v<T0, T> && std::is_trivially_copyable_v<T>>
+  copy(c10::ArrayRef<T0> src) {
+    size_t n = src.size();
+    if (C10_UNLIKELY(n == 0)) {
+      return;
+    }
+    maybe_grow();
+    if (C10_LIKELY(next_ && (next_ + n <= end_))) {
+      std::memcpy((void*)next_, (void*)src.begin(), n * sizeof(T0));
+      next_ += n;
+    } else {
+      // We could chunk this into several `memcpy`s, but because we expect this
+      // fallback to be infrequent (n << ChunkSize) the performance impact is
+      // negligible.
+      for (auto i : src) {
+        emplace_back(i);
+      }
+    }
+  }
+
+  void clear() {
+    buffer_.clear();
+    buffer_last_ = buffer_.before_begin();
+    n_blocks_ = 0;
+    next_ = nullptr;
+    end_ = nullptr;
+  }
+
+  struct Iterator {
+    using iterator_category = std::forward_iterator_tag;
+    using difference_type = std::ptrdiff_t;
+    using value_type = T;
+    using pointer = T*;
+    using reference = T&;
+
+    Iterator(std::forward_list<array_t>& buffer, const size_t size)
+        : block_{buffer.begin()}, size_{size} {}
+
+    // End iterator.
+    Iterator() = default;
+
+    bool exhausted() const {
+      return current_ >= size_;
+    }
+
+    reference operator*() const {
+      return *current_ptr(/*checked=*/true);
+    }
+    pointer operator->() {
+      return current_ptr(/*checked=*/true);
+    }
+
+    // Prefix increment
+    Iterator& operator++() {
+      if (!(++current_ % ChunkSize)) {
+        block_++;
+      }
+      return *this;
+    }
+
+    // Postfix increment
+    Iterator operator++(int) {
+      Iterator tmp = *this;
+      ++(*this);
+      return tmp;
+    }
+
+    friend bool operator==(const Iterator& a, const Iterator& b) {
+      return a.current_ptr() == b.current_ptr();
+    }
+    friend bool operator!=(const Iterator& a, const Iterator& b) {
+      return a.current_ptr() != b.current_ptr();
+    }
+
+    std::pair<array_t*, size_t> address() const {
+      if (current_ >= size_) {
+        return {nullptr, 0};
+      }
+      return {&(*block_), current_ % ChunkSize};
+    }
+
+   private:
+    T* current_ptr(bool checked = false) const {
+      auto a = address();
+      if (a.first == nullptr) {
+        TORCH_INTERNAL_ASSERT(!checked, "Invalid access on AppendOnlyList.");
+        return nullptr;
+      }
+      return a.first->data() + a.second;
+    }
+
+    typename std::forward_list<array_t>::iterator block_;
+    size_t current_{0};
+    size_t size_{0};
+  };
+
+  Iterator begin() {
+    return Iterator(buffer_, size());
+  }
+  Iterator end() {
+    return Iterator();
+  }
+  // TODO: cbegin and cend()
+
+ private:
+  void maybe_grow() {
+    if (C10_UNLIKELY(next_ == end_)) {
+      buffer_last_ = buffer_.emplace_after(buffer_last_);
+      n_blocks_++;
+      next_ = buffer_last_->data();
+      end_ = next_ + ChunkSize;
+    }
+  }
+
+  std::forward_list<array_t> buffer_;
+
+  // We maintain a pointer to the last element of `buffer_` so that we can
+  // insert at the end in O(1) time.
+  size_t n_blocks_{0};
+  T* next_{nullptr};
+  T* end_{nullptr};
+
+ protected:
+  typename std::forward_list<array_t>::iterator buffer_last_;
+};
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/data_flow.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/data_flow.h
new file mode 100644
index 0000000000000000000000000000000000000000..115bd72209f71641b412d124643e5a9242b1b6ab
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/data_flow.h
@@ -0,0 +1,95 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+
+#include <ATen/core/TensorBody.h>
+#include <c10/core/TensorImpl.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/strong_type.h>
+
+namespace torch::profiler::impl {
+
+// Identity is a complex concept in PyTorch. A Tensor might not have a
+// an associated storage, multiple Tensors might share the same underlying
+// storage, the storage of a Tensor might change over time, etc.
+//
+// For the purpose of profiling we're mostly interested in data flow
+// analysis. As a result, we can take an expansive view of identity:
+// Tensors share an ID if they share a TensorImpl or storage data.
+//
+// This identity equality is transitive; If Tensors T0 and T1 share a storage
+// S0 and T1 later points to a different storage S1 then all Tensors which
+// point to either S0 or S1 are considered to have the same identity. (Since
+// profiler cannot reason beyond that.)
+//
+// The profiler will handle lifetime analysis to ensure that identities do
+// not run afoul of the ABA problem. This does, however, mean that identities
+// can only be assigned when memory profiling is enabled.
+using TensorID = strong::type<size_t, struct TensorID_, strong::regular>;
+
+// Uniquely identifies an allocation. (Generally a StorageImpl's data ptr.)
+using AllocationID = strong::type<
+    size_t,
+    struct StorageID_,
+    strong::ordered,
+    strong::regular,
+    strong::hashable>;
+
+// We use a Tensor's TensorImpl address and StorageImpl data start to build the
+// data flow graph. We do not hold an owning reference so we wrap them in strong
+// types to prevent direct access.
+using TensorImplAddress = strong::type<
+    const c10::TensorImpl*,
+    struct TensorImplAddress_,
+    strong::regular,
+    strong::hashable,
+    strong::boolean>;
+
+using StorageImplData = strong::type<
+    const void*,
+    struct StorageImplData_,
+    strong::regular,
+    strong::hashable,
+    strong::boolean>;
+
+// ============================================================================
+// == weak_intrusive_ptr and the ABA problem for TensorImpl* ==================
+// ============================================================================
+// Tracking `TensorImpl`s is an important part of identity tracking, because
+// a Tensor might change storage; however when it does we want to retain the
+// fact that the old and new storage belong to the same logical Tensor. We
+// cannot take an owning reference to the Tensor because that would change
+// program semantics by extending the lifetime of the Tensor. However if we
+// store a raw TensorImpl* pointer the TensorImpl might be deleted and a new
+// TensorImpl might be created that reuses the address. (ABA problem)
+//
+// Fortunately, there is a feature of `c10::intrusive_ptr` that we can use to
+// prevent address reuse for the duration of profiling: the weak intrusive ptr.
+// When a Tensor's refcount reaches zero but there are outstanding weak
+// references (`weakcount_ > 0`) it will free the underlying managed resources
+// by calling `target_->release_resources()`, but it will not call `delete`.
+// (Instead, `delete` is called when the last weak reference is destroyed.)
+// This means that we can safely use address identity to track `TensorImpls`.
+class WeakTensor {
+ public:
+  explicit WeakTensor(const at::Tensor& t) : weak_self_(t.getIntrusivePtr()) {}
+
+  auto get() const {
+    return TensorImplAddress{weak_self_._unsafe_get_target()};
+  }
+
+ private:
+  c10::weak_intrusive_ptr<c10::TensorImpl> weak_self_;
+};
+
+struct Result;
+
+void calculateUniqueTensorIDs(
+    std::vector<std::shared_ptr<Result>>& sorted_results);
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/events.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/events.h
new file mode 100644
index 0000000000000000000000000000000000000000..83cc186e15f19dc61fb6d0123593631c5d5961d0
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/events.h
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <cstdint>
+#include <cstring>
+#include <vector>
+
+namespace torch::profiler {
+
+/* A vector type to hold a list of performance counters */
+using perf_counters_t = std::vector<uint64_t>;
+
+/* Standard list of performance events independent of hardware or backend */
+constexpr std::array<const char*, 2> ProfilerPerfEvents = {
+    /*
+     * Number of Processing Element (PE) cycles between two points of interest
+     * in time. This should correlate positively with wall-time. Measured in
+     * uint64_t. PE can be non cpu. TBD reporting behavior for multiple PEs
+     * participating (i.e. threadpool).
+     */
+    "cycles",
+
+    /* Number of PE instructions between two points of interest in time. This
+     * should correlate positively with wall time and the amount of computation
+     * (i.e. work). Across repeat executions, the number of instructions should
+     * be more or less invariant. Measured in uint64_t. PE can be non cpu.
+     */
+    "instructions"};
+} // namespace torch::profiler
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/kineto_client_interface.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/kineto_client_interface.h
new file mode 100644
index 0000000000000000000000000000000000000000..12b936c44d1fd363c8c53942f3d2e99a7bbb8f04
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/kineto_client_interface.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/runtime/interpreter.h>
+#include <torch/csrc/profiler/unwind/unwind.h>
+
+namespace torch {
+
+// declare global_kineto_init for libtorch_cpu.so to call
+TORCH_API void global_kineto_init();
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/kineto_shim.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/kineto_shim.h
new file mode 100644
index 0000000000000000000000000000000000000000..2349701c533b6d0a09039bf018d3a780b05836ad
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/kineto_shim.h
@@ -0,0 +1,157 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <memory>
+#include <string>
+
+// Skip Kineto dependency on mobile unless explicitly asked for.
+// When is it explicitly asked for?
+//   KinetoEdgeCPUProfiler uses KinetoProfiler for cpu
+//   event profiling. This has a dependency on cpu only libkineto
+#if defined(USE_KINETO) && defined(C10_MOBILE) && \
+    !defined(EDGE_PROFILER_USE_KINETO)
+#undef USE_KINETO
+#endif
+
+#include <ActivityType.h>
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/profiler/api.h>
+
+#ifdef USE_KINETO
+// Forward declarations so we don't have to include `libkineto.h` in a header.
+namespace libkineto {
+class GenericTraceActivity;
+struct CpuTraceBuffer;
+class ActivityTraceInterface;
+} // namespace libkineto
+#endif
+
+namespace torch {
+namespace profiler {
+
+#ifdef USE_KINETO
+constexpr bool kKinetoAvailable{true};
+#else
+constexpr bool kKinetoAvailable{false};
+#endif
+
+namespace impl::kineto {
+
+// ----------------------------------------------------------------------------
+// -- Interface (Does not require Kineto) -------------------------------------
+// ----------------------------------------------------------------------------
+struct DeviceAndResource {
+  int32_t device;
+  int32_t resource;
+};
+const DeviceAndResource kineto_ids();
+
+#ifdef USE_KINETO
+using trace_t = libkineto::CpuTraceBuffer;
+using interface_trace_t = libkineto::ActivityTraceInterface;
+using activity_t = libkineto::GenericTraceActivity;
+#else
+struct DummyTraceBuffer {};
+struct DummyTraceInterface {};
+
+using trace_t = DummyTraceBuffer;
+using interface_trace_t = DummyTraceBuffer;
+struct activity_t;
+#endif // USE_KINETO
+
+void addMetadata(
+    activity_t* activity,
+    const std::string& key,
+    const std::string& value);
+
+// Wraps: libkineto::CpuTraceBuffer
+struct TraceWrapper {
+  TraceWrapper(const int64_t start_time, const std::string& name);
+
+  // The caller is expected to hold a mutex when calling `addCPUActivity`.
+  activity_t* addCPUActivity(
+      const std::string& name,
+      const libkineto::ActivityType type,
+      const DeviceAndResource device_and_resource,
+      const uint64_t correlation_id,
+      const int64_t start_time,
+      const int64_t end_time);
+
+  void transferCpuTrace(int64_t end_time);
+
+  explicit operator bool() const;
+
+  std::unique_ptr<trace_t>& get() {
+    return cpu_trace_;
+  }
+
+ private:
+  std::unique_ptr<trace_t> cpu_trace_;
+};
+
+// Wraps libkineto::ActivityTraceInterface
+struct ActivityTraceWrapper {
+  explicit ActivityTraceWrapper(std::unique_ptr<interface_trace_t>&& trace);
+  ActivityTraceWrapper() = default;
+  explicit operator bool() const;
+  void save(const std::string& path);
+
+  const std::unique_ptr<interface_trace_t>& get() {
+    return trace_;
+  }
+
+ private:
+  std::unique_ptr<interface_trace_t> trace_;
+#ifdef USE_KINETO
+  bool saved_ = false; // Kineto's save is destructive
+#endif
+};
+
+using ActivitySet = std::set<torch::autograd::profiler::ActivityType>;
+using ActivityFilter = std::unordered_map<
+    torch::autograd::profiler::ActivityType,
+    std::unordered_set<std::string>>;
+void prepareTrace(
+    const bool cpuOnly,
+    const ActivitySet& activities,
+    const torch::profiler::impl::ExperimentalConfig& config,
+    const std::string& trace_id = "",
+    const ActivityFilter& activity_filter = {});
+
+void toggleCollectionDynamic(const bool enable);
+void startTrace();
+ActivityTraceWrapper stopTrace();
+void pushCorrelationId(uint64_t correlation_id);
+void pushUserCorrelationId(uint64_t correlation_id);
+void popCorrelationId();
+void popUserCorrelationId();
+void recordThreadInfo();
+bool collectivesProfilerExists();
+
+void logInvariantViolation(
+    const std::string& assertion,
+    const std::string& error,
+    const std::string& profile_id,
+    const std::string& group_profile_id);
+
+} // namespace impl::kineto
+
+} // namespace profiler
+
+namespace autograd::profiler {
+c10::DeviceType deviceTypeFromActivity(libkineto::ActivityType activity_type);
+
+TORCH_API void addMetadataJson(
+    const std::string& key,
+    const std::string& value);
+
+TORCH_API void profilerStep();
+
+} // namespace autograd::profiler
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/observer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/observer.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ed82f18a2e5b8f545163f4e504b9e3291a40861
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/observer.h
@@ -0,0 +1,223 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/record_function.h>
+#include <torch/csrc/Export.h>
+
+#include <utility>
+
+namespace torch::profiler::impl {
+
+// ----------------------------------------------------------------------------
+// -- Profiler Config ---------------------------------------------------------
+// ----------------------------------------------------------------------------
+enum class C10_API_ENUM ActivityType {
+  CPU = 0,
+  XPU, // XPU kernels, runtime
+  CUDA, // CUDA kernels, runtime
+  HPU, // HPU kernels, runtime
+  MTIA, // MTIA kernels, runtime
+  PrivateUse1, // PrivateUse1 kernels, runtime
+  NUM_KINETO_ACTIVITIES, // must be the last one
+};
+
+inline std::string actToString(ActivityType t) {
+  const std::array<
+      std::string,
+      static_cast<size_t>(ActivityType::NUM_KINETO_ACTIVITIES)>
+      ActivityTypeNames = {"CPU", "XPU", "CUDA", "MTIA", "PrivateUse1"};
+  return ActivityTypeNames[static_cast<int>(t)];
+}
+
+enum class C10_API_ENUM ProfilerState {
+  Disabled = 0,
+  CPU, // CPU-only profiling
+  CUDA, // CPU + CUDA events
+  NVTX, // only emit NVTX markers
+  ITT, // only emit ITT markers
+  PRIVATEUSE1, // only emit PRIVATEUSE1 markers
+  KINETO, // use libkineto
+  KINETO_GPU_FALLBACK, // use CUDA events when CUPTI is not available
+  KINETO_PRIVATEUSE1_FALLBACK, // use PrivateUse1 events
+  KINETO_PRIVATEUSE1, // use Kineto with registered IActivityProfiler
+  KINETO_ONDEMAND, // run the profiler in on-demand mode
+  NUM_PROFILER_STATES, // must be the last one
+};
+
+enum class C10_API_ENUM ActiveProfilerType {
+  NONE = 0,
+  LEGACY,
+  KINETO,
+  NVTX,
+  ITT,
+  PRIVATEUSE1
+};
+
+struct TORCH_API ExperimentalConfig {
+  ExperimentalConfig(
+      std::vector<std::string> profiler_metrics = {},
+      bool profiler_measure_per_kernel = false,
+      bool verbose = false,
+      std::vector<std::string> performance_events = {},
+      bool enable_cuda_sync_events = false,
+      bool adjust_profiler_step = false,
+      bool disable_external_correlation = false,
+      bool profile_all_threads = false,
+      bool capture_overload_names = false,
+      bool record_python_gc_info = false,
+      bool expose_kineto_event_metadata = false,
+      std::string custom_profiler_config = "",
+      bool adjust_timestamps = false);
+  explicit operator bool() const;
+
+  std::vector<std::string> profiler_metrics;
+  bool profiler_measure_per_kernel;
+  bool verbose;
+  /*
+   * List of performance events to be profiled.
+   * An empty list will disable performance event based profiling altogether.
+   */
+  std::vector<std::string> performance_events;
+  /*
+   * For CUDA profiling mode, enable adding CUDA synchronization events
+   * that expose CUDA device, stream and event synchronization activities.
+   * This feature is new and currently disabled by default.
+   */
+  bool enable_cuda_sync_events;
+  /*
+   * Controls whether or not timestamp adjustment for ProfilerStep and parent
+   * Python events occurs after profiling. This occurs at an O(n) cost and
+   * affects only the start of profiler step events.
+   */
+  bool adjust_profiler_step;
+  /*
+   * Controls whether or not external correlation is disabled. This is used to
+   * lower the amount of events received by CUPTI as correlation events are
+   * paired with runtime/gpu events for each kind of correlation
+   */
+  bool disable_external_correlation;
+
+  /* controls whether profiler records cpu events on threads
+   * that are not spawned from the main thread on which the
+   * profiler was enabled, similar to on_demand mode */
+  bool profile_all_threads;
+
+  /* controls whether overload names are queried from an ATen
+   * function schema and stored in the profile  */
+  bool capture_overload_names;
+
+  /*
+   * Controls whether or not python gc info is recorded. This is used to
+   * determine if gc collect is slowing down your profile.
+   */
+  bool record_python_gc_info;
+
+  /* controls whether KinetoEvent metadata is exposed to FunctionEvent
+   * in the PyTorch Profiler as a JSON string */
+  bool expose_kineto_event_metadata;
+
+  /*
+   * A custom_profiler_config option is introduced to allow custom backends
+   * to apply custom configurations as needed.
+   */
+  std::string custom_profiler_config;
+
+  /*
+   * Controls whether or not timestamp adjustment occurs after profiling.
+   * The purpose of this is to adjust Vulkan event timelines to align with those
+   * of their parent CPU events.
+   * This sometimes requires increasing CPU event durations (to fully contain
+   * their child events) and delaying CPU event start times (to
+   * prevent overlaps), so this should not be used unless Vulkan events are
+   * being profiled and it is ok to use this modified timestamp/duration
+   * information instead of the original information.
+   */
+  bool adjust_timestamps;
+};
+
+struct TORCH_API ProfilerConfig {
+  explicit ProfilerConfig(
+      ProfilerState state,
+      bool report_input_shapes = false,
+      bool profile_memory = false,
+      bool with_stack = false,
+      bool with_flops = false,
+      bool with_modules = false,
+      ExperimentalConfig experimental_config = ExperimentalConfig(),
+      std::string trace_id = "");
+
+  bool disabled() const;
+  bool global() const;
+  bool pushGlobalCallbacks() const;
+
+  ProfilerState state;
+  ExperimentalConfig experimental_config;
+  bool report_input_shapes;
+  bool profile_memory;
+  bool with_stack;
+  bool with_flops;
+  bool with_modules;
+  std::string trace_id;
+
+  // For serialization
+  at::IValue toIValue() const;
+  static ProfilerConfig fromIValue(const at::IValue& profilerConfigIValue);
+};
+
+// ----------------------------------------------------------------------------
+// -- Profiler base class -----------------------------------------------------
+// ----------------------------------------------------------------------------
+struct TORCH_API ProfilerStateBase : public c10::MemoryReportingInfoBase {
+  explicit ProfilerStateBase(ProfilerConfig config);
+  ProfilerStateBase(const ProfilerStateBase&) = delete;
+  ProfilerStateBase(ProfilerStateBase&&) = delete;
+  ProfilerStateBase& operator=(const ProfilerStateBase&) = delete;
+  ProfilerStateBase& operator=(ProfilerStateBase&&) = delete;
+  ~ProfilerStateBase() override;
+
+  static ProfilerStateBase* get(bool global);
+  static ProfilerStateBase* get() {
+    auto* out = get(/*global=*/true);
+    return out ? out : get(/*global=*/false);
+  }
+
+  static void push(std::shared_ptr<ProfilerStateBase>&& state);
+
+  static std::shared_ptr<ProfilerStateBase> pop(bool global);
+  static std::shared_ptr<ProfilerStateBase> pop() {
+    auto out = pop(/*global=*/true);
+    return out ? std::move(out) : pop(/*global=*/false);
+  }
+
+  const ProfilerConfig& config() const {
+    return config_;
+  }
+
+  void setCallbackHandle(at::CallbackHandle handle);
+  void removeCallback();
+
+  bool memoryProfilingEnabled() const override {
+    return config_.profile_memory;
+  }
+
+  virtual ActiveProfilerType profilerType() = 0;
+
+ protected:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  std::mutex state_mutex_;
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  ProfilerConfig config_ = ProfilerConfig(ProfilerState::Disabled);
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  at::CallbackHandle handle_ = 0;
+};
+
+// Note: The following are only for the active *thread local* profiler.
+TORCH_API bool profilerEnabled();
+TORCH_API ActiveProfilerType profilerType();
+TORCH_API ProfilerConfig getProfilerConfig();
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/python_tracer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/python_tracer.h
new file mode 100644
index 0000000000000000000000000000000000000000..3cfd6d54b173ddb17adf295d33aab4636c3638a2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/python_tracer.h
@@ -0,0 +1,82 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include <c10/util/ApproximateClock.h>
+#include <c10/util/strong_type.h>
+
+#include <torch/csrc/profiler/kineto_shim.h>
+#include <torch/csrc/profiler/util.h>
+
+namespace torch::profiler::impl {
+
+class RecordQueue;
+struct Result;
+namespace python_tracer {
+
+using TraceKey = strong::type<
+    uint64_t,
+    struct TraceKey_,
+    strong::regular,
+    strong::hashable,
+    strong::ostreamable>;
+
+struct CompressedEvent {
+  TraceKey key_;
+  uint64_t system_tid_{};
+  kineto::DeviceAndResource kineto_info_{};
+  c10::time_t enter_t_{};
+};
+
+/*
+Libtorch does not depend on Python (e.g. cannot #include <Python.h>); however
+when we call the profiler from libtorch_python we need the profiler to be able
+to ingest the data that we collect from the Python tracer. (`PyEval_SetProfile`)
+
+In order to solve this dependency issue we define a virtual base and a function
+to register a getter. The python tracer then implements these functions and
+exposes itself by calling `registerTracer` from `torch/csrc/autograd/init.cpp`.
+This pattern of registration for faux python dependencies in libtorch is common
+in the PyTorch codebase.
+*/
+struct TORCH_API PythonTracerBase {
+  static std::unique_ptr<PythonTracerBase> make(RecordQueue* queue);
+  virtual ~PythonTracerBase() = default;
+
+  virtual void stop() = 0;
+  virtual void restart() = 0;
+  virtual void register_gc_callback() = 0;
+  virtual std::vector<std::shared_ptr<Result>> getEvents(
+      std::function<c10::time_t(c10::approx_time_t)> time_converter,
+      std::vector<CompressedEvent>& enters,
+      c10::time_t end_time_ns) = 0;
+};
+
+using MakeFn = std::unique_ptr<PythonTracerBase> (*)(RecordQueue*);
+TORCH_API void registerTracer(MakeFn make_tracer);
+
+/**
+ * Memory Tracer Implementation
+ */
+struct TORCH_API PythonMemoryTracerBase {
+  static std::unique_ptr<PythonMemoryTracerBase> make();
+  virtual ~PythonMemoryTracerBase() = default;
+
+  virtual void start() = 0;
+  virtual void stop() = 0;
+  virtual void export_memory_history(const std::string& path) = 0;
+};
+
+using MakeMemoryFn = std::unique_ptr<PythonMemoryTracerBase> (*)();
+TORCH_API void registerMemoryTracer(MakeMemoryFn make_memory_tracer);
+
+} // namespace python_tracer
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/vulkan.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/vulkan.h
new file mode 100644
index 0000000000000000000000000000000000000000..577d1299b83600fecfca5ecf03a399c526ed2fd1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/orchestration/vulkan.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/profiler/stubs/base.h>
+#include <torch/csrc/profiler/util.h>
+#include <cstdint>
+
+namespace torch::profiler::impl::vulkan {
+
+// Using function pointer i.e. [std::tuple<std::string, uint64_t> (*)(int64_t)]
+// doesn't work because we need to capture the QueryPool in the lambda context
+// https://stackoverflow.com/a/28746827
+using GetShaderNameAndDurationNsFn =
+    std::function<std::tuple<std::string, uint64_t>(int64_t)>;
+TORCH_API void registerGetShaderNameAndDurationNs(
+    GetShaderNameAndDurationNsFn get_shader_name_and_duration_ns);
+
+TORCH_API void deregisterGetShaderNameAndDurationNs();
+
+std::tuple<std::string, uint64_t> getShaderNameAndDurationNs(
+    const vulkan_id_t& vulkan_id);
+
+} // namespace torch::profiler::impl::vulkan
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/perf-inl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/perf-inl.h
new file mode 100644
index 0000000000000000000000000000000000000000..448e4747807ee871734722b351a58215ad2a4f60
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/perf-inl.h
@@ -0,0 +1,73 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#if defined(__ANDROID__) || defined(__linux__)
+
+#include <unistd.h>
+
+#include <sys/ioctl.h>
+#include <sys/syscall.h>
+
+#include <linux/perf_event.h>
+
+#endif /* __ANDROID__ || __linux__ */
+
+#include <torch/csrc/profiler/perf.h>
+
+#include <limits>
+
+namespace torch::profiler::impl::linux_perf {
+
+/*
+ * PerfEvent
+ * ---------
+ */
+
+inline void PerfEvent::Disable() const {
+#if defined(__ANDROID__) || defined(__linux__)
+  ioctl(fd_, PERF_EVENT_IOC_DISABLE, 0);
+#endif /* __ANDROID__ || __linux__ */
+}
+
+inline void PerfEvent::Enable() const {
+#if defined(__ANDROID__) || defined(__linux__)
+  ioctl(fd_, PERF_EVENT_IOC_ENABLE, 0);
+#endif /* __ANDROID__ || __linux__ */
+}
+
+inline void PerfEvent::Reset() const {
+#if defined(__ANDROID__) || defined(__linux__)
+  ioctl(fd_, PERF_EVENT_IOC_RESET, 0);
+#endif /* __ANDROID__ || __linux__ */
+}
+
+/*
+ * PerfProfiler
+ * ------------
+ */
+
+inline uint64_t PerfProfiler::CalcDelta(uint64_t start, uint64_t end) const {
+  if (end < start) { // overflow
+    return end + (std::numeric_limits<uint64_t>::max() - start);
+  }
+  // not possible to wrap around start for a 64b cycle counter
+  return end - start;
+}
+
+inline void PerfProfiler::StartCounting() const {
+  for (auto& e : events_) {
+    e.Enable();
+  }
+}
+
+inline void PerfProfiler::StopCounting() const {
+  for (auto& e : events_) {
+    e.Disable();
+  }
+}
+
+} // namespace torch::profiler::impl::linux_perf
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/perf.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/perf.h
new file mode 100644
index 0000000000000000000000000000000000000000..70ea26a46330fdeb556b76a511add5253b0b2314
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/perf.h
@@ -0,0 +1,106 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <array>
+#include <cstdint>
+#include <memory>
+#include <stack>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <torch/csrc/profiler/events.h>
+
+#include <c10/util/Exception.h>
+
+namespace torch::profiler::impl::linux_perf {
+
+/*
+ * Maximum number of events supported
+ * This stems from the hardware limitation on CPU performance counters, and the
+ * fact that we don't support time multiplexing just yet.
+ * Time multiplexing involves scaling the counter values proportional to
+ * the enabled and running time or running the workload multiple times.
+ */
+constexpr uint8_t MAX_EVENTS = 4;
+
+struct PerfCounter {
+  uint64_t value; /* The value of the event */
+  uint64_t time_enabled; /* for TIME_ENABLED */
+  uint64_t time_running; /* for TIME_RUNNING */
+};
+
+/*
+ * Basic perf event handler for Android and Linux
+ */
+class PerfEvent {
+ public:
+  explicit PerfEvent(std::string& name) : name_(name) {}
+
+  PerfEvent(const PerfEvent& other) = delete;
+  PerfEvent& operator=(const PerfEvent&) = delete;
+  PerfEvent& operator=(PerfEvent&& other) noexcept {
+    if (this != &other) {
+      fd_ = other.fd_;
+      other.fd_ = -1;
+      name_ = std::move(other.name_);
+    }
+    return *this;
+  }
+
+  PerfEvent(PerfEvent&& other) noexcept {
+    *this = std::move(other);
+  }
+
+  ~PerfEvent();
+
+  /* Setup perf events with the Linux Kernel, attaches perf to this process
+   * using perf_event_open(2) */
+  void Init();
+
+  /* Stop incrementing hardware counters for this event */
+  void Disable() const;
+
+  /* Start counting hardware event from this point on */
+  void Enable() const;
+
+  /* Zero out the counts for this event */
+  void Reset() const;
+
+  /* Returns PerfCounter values for this event from kernel, on non supported
+   * platforms this always returns zero */
+  uint64_t ReadCounter() const;
+
+ private:
+  /* Name of the event */
+  std::string name_;
+
+  int fd_ = -1;
+};
+
+class PerfProfiler {
+ public:
+  /* Configure all the events and track them as individual PerfEvent */
+  void Configure(std::vector<std::string>& event_names);
+
+  /* Enable events counting from here */
+  void Enable();
+
+  /* Disable counting and fill in the caller supplied container with delta
+   * calculated from the start count values since last Enable() */
+  void Disable(perf_counters_t& /*vals*/);
+
+ private:
+  uint64_t CalcDelta(uint64_t start, uint64_t end) const;
+  void StartCounting() const;
+  void StopCounting() const;
+
+  std::vector<PerfEvent> events_;
+  std::stack<perf_counters_t> start_values_;
+};
+} // namespace torch::profiler::impl::linux_perf
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/combined_traceback.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/combined_traceback.h
new file mode 100644
index 0000000000000000000000000000000000000000..0cdb2e1737d860f40c91814405a39ef310cee5b1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/combined_traceback.h
@@ -0,0 +1,32 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/profiler/combined_traceback.h>
+
+#include <nlohmann/json.hpp>
+#include <pybind11/pybind11.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch {
+
+// symbolize combined traceback objects, converting them into lists of
+// dictionaries that are easily consumed in python.
+
+// returns std::vector because one use is to call it with a batch of
+// tracebacks that come from a larger datastructure (e.g. a memory snapshot)
+// and then have more c++ code to put those objects in the right place.
+TORCH_API std::vector<pybind11::object> py_symbolize(
+    std::vector<CapturedTraceback*>& to_symbolize);
+
+// Return the callback in json format so that it can be used within cpp
+TORCH_API std::vector<nlohmann::json> json_symbolize(
+    std::vector<CapturedTraceback*>& to_symbolize);
+
+// requires GIL to be held, frees any pending free frames
+TORCH_PYTHON_API void freeDeadCapturedTracebackFrames();
+
+TORCH_PYTHON_API void installCapturedTracebackPython();
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..8f5f25ddaf5cb6b647d99e53c45ed81630ce3213
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/init.h
@@ -0,0 +1,36 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <Python.h>
+
+#include <torch/csrc/profiler/collection.h>
+#include <torch/csrc/profiler/python/pybind.h>
+
+namespace pybind11::detail {
+using torch::profiler::impl::TensorID;
+
+#define STRONG_POINTER_TYPE_CASTER(T) \
+  template <>                         \
+  struct type_caster<T> : public strong_pointer_type_caster<T> {};
+
+STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::StorageImplData)
+STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::AllocationID)
+STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::TensorImplAddress)
+STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::PyModuleSelf)
+STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::PyModuleCls)
+STRONG_POINTER_TYPE_CASTER(torch::profiler::impl::PyOptimizerSelf)
+#undef STRONG_POINTER_TYPE_CASTER
+
+template <>
+struct type_caster<TensorID> : public strong_uint_type_caster<TensorID> {};
+} // namespace pybind11::detail
+
+namespace torch::profiler {
+
+void initPythonBindings(PyObject* module);
+
+} // namespace torch::profiler
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/pybind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/pybind.h
new file mode 100644
index 0000000000000000000000000000000000000000..39c63a85349c7493d1d06e922941ff6205cb1770
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/python/pybind.h
@@ -0,0 +1,53 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include <c10/util/strong_type.h>
+#include <torch/csrc/utils/pybind.h>
+#include <torch/csrc/utils/python_numbers.h>
+
+namespace pybind11::detail {
+// Strong typedefs don't make much sense in Python since everything is duck
+// typed. So instead we simply extract the underlying value and let the caller
+// handle correctness.
+template <typename T>
+struct strong_pointer_type_caster {
+  template <typename T_>
+  static handle cast(
+      const T_& src,
+      return_value_policy /*policy*/,
+      handle /*parent*/) {
+    const auto* ptr = reinterpret_cast<const void*>(src.value_of());
+    return ptr ? handle(THPUtils_packUInt64(reinterpret_cast<intptr_t>(ptr)))
+               : none();
+  }
+
+  bool load(handle /*src*/, bool /*convert*/) {
+    return false;
+  }
+
+  PYBIND11_TYPE_CASTER(T, _("strong_pointer"));
+};
+
+template <typename T>
+struct strong_uint_type_caster {
+  template <typename T_>
+  static handle cast(
+      const T_& src,
+      return_value_policy /*policy*/,
+      handle /*parent*/) {
+    return handle(THPUtils_packUInt64(src.value_of()));
+  }
+
+  bool load(handle /*src*/, bool /*convert*/) {
+    return false;
+  }
+
+  PYBIND11_TYPE_CASTER(T, _("strong_uint"));
+};
+} // namespace pybind11::detail
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/execution_trace_observer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/execution_trace_observer.h
new file mode 100644
index 0000000000000000000000000000000000000000..cd5a3addd9cc27717340a29fa268b7f96ae561ba
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/execution_trace_observer.h
@@ -0,0 +1,26 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Export.h>
+#include <string>
+
+namespace torch::profiler::impl {
+
+// Adds the execution trace observer as a global callback function, the data
+// will be written to output file path.
+TORCH_API bool addExecutionTraceObserver(const std::string& output_file_path);
+
+// Remove the execution trace observer from the global callback functions.
+TORCH_API void removeExecutionTraceObserver();
+
+// Enables execution trace observer.
+TORCH_API void enableExecutionTraceObserver();
+
+// Disables execution trace observer.
+TORCH_API void disableExecutionTraceObserver();
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/itt_observer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/itt_observer.h
new file mode 100644
index 0000000000000000000000000000000000000000..e1d340daf8b2917d3b3a4de59f254deb68aefeea
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/itt_observer.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/profiler/api.h>
+
+namespace torch::profiler::impl {
+
+void pushITTCallbacks(
+    const ProfilerConfig& config,
+    const std::unordered_set<at::RecordScope>& scopes);
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/nvtx_observer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/nvtx_observer.h
new file mode 100644
index 0000000000000000000000000000000000000000..75214a0fbd6b9b776bc8788061461057a9843c2a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/nvtx_observer.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/profiler/api.h>
+
+namespace torch::profiler::impl {
+
+void pushNVTXCallbacks(
+    const ProfilerConfig& config,
+    const std::unordered_set<at::RecordScope>& scopes);
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/privateuse1_observer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/privateuse1_observer.h
new file mode 100644
index 0000000000000000000000000000000000000000..09e3f63d535eb7fba56dabd71aaca0ed85e7b968
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/privateuse1_observer.h
@@ -0,0 +1,49 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/profiler/api.h>
+
+namespace torch::profiler::impl {
+
+using CallBackFnPtr = void (*)(
+    const ProfilerConfig& config,
+    const std::unordered_set<at::RecordScope>& scopes);
+
+struct PushPRIVATEUSE1CallbacksStub {
+  PushPRIVATEUSE1CallbacksStub() = default;
+  PushPRIVATEUSE1CallbacksStub(const PushPRIVATEUSE1CallbacksStub&) = delete;
+  PushPRIVATEUSE1CallbacksStub& operator=(const PushPRIVATEUSE1CallbacksStub&) =
+      delete;
+  PushPRIVATEUSE1CallbacksStub(PushPRIVATEUSE1CallbacksStub&&) = default;
+  PushPRIVATEUSE1CallbacksStub& operator=(PushPRIVATEUSE1CallbacksStub&&) =
+      default;
+  ~PushPRIVATEUSE1CallbacksStub() = default;
+
+  template <typename... ArgTypes>
+  void operator()(ArgTypes&&... args) {
+    return (*push_privateuse1_callbacks_fn)(std::forward<ArgTypes>(args)...);
+  }
+
+  void set_privateuse1_dispatch_ptr(CallBackFnPtr fn_ptr) {
+    push_privateuse1_callbacks_fn = fn_ptr;
+  }
+
+ private:
+  CallBackFnPtr push_privateuse1_callbacks_fn = nullptr;
+};
+
+extern TORCH_API struct PushPRIVATEUSE1CallbacksStub
+    pushPRIVATEUSE1CallbacksStub;
+
+struct RegisterPRIVATEUSE1Observer {
+  RegisterPRIVATEUSE1Observer(CallBackFnPtr cb) {
+    pushPRIVATEUSE1CallbacksStub.set_privateuse1_dispatch_ptr(cb);
+  }
+};
+
+#define REGISTER_PRIVATEUSE1_OBSERVER(cb) \
+  static RegisterPRIVATEUSE1Observer privateuse1_callbacks_stub_register(cb);
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/privateuse1_profiler.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/privateuse1_profiler.h
new file mode 100644
index 0000000000000000000000000000000000000000..955bcc6c6623a84c43fe257abe6d55faaa88634a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/standalone/privateuse1_profiler.h
@@ -0,0 +1,118 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#ifdef USE_KINETO
+
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <type_traits>
+
+#include <torch/csrc/Export.h>
+
+#include <IActivityProfiler.h>
+
+namespace torch::profiler::impl {
+
+// Factory function type that creates an IActivityProfiler instance
+using PrivateUse1ProfilerFactory =
+    std::function<std::unique_ptr<libkineto::IActivityProfiler>()>;
+
+// Registry for PrivateUse1 activity profiler factories.
+//
+// This registry allows custom accelerator backends to register their
+// IActivityProfiler implementation with Kineto, enabling full profiling
+// integration without modifying Kineto code.
+//
+// Usage:
+//   1. Backend implements libkineto::IActivityProfiler
+//   2. Backend uses REGISTER_PRIVATEUSE1_PROFILER macro to register
+//   3. PyTorch forwards the factory to Kineto during initialization
+//
+// Example:
+//   class MyAcceleratorProfiler : public libkineto::IActivityProfiler {
+//     const std::string& name() const override { return name_; }
+//     const std::set<libkineto::ActivityType>& availableActivities() const
+//     override; std::unique_ptr<libkineto::IActivityProfilerSession>
+//     configure(...) override;
+//   private:
+//     std::string name_{"my_accelerator"};
+//   };
+//
+//   REGISTER_PRIVATEUSE1_PROFILER(MyAcceleratorProfiler)
+//
+class TORCH_API PrivateUse1ProfilerRegistry {
+ public:
+  static PrivateUse1ProfilerRegistry& instance();
+
+  // Register a factory function for creating the PrivateUse1 profiler.
+  // This should be called during static initialization.
+  void registerFactory(PrivateUse1ProfilerFactory factory);
+
+  // Check if a factory has been registered.
+  bool hasFactory() const;
+
+  // Check if the factory has been registered with Kineto.
+  // Useful for testing to verify the registration logic.
+  bool isRegisteredWithKineto() const;
+
+  // Register the factory with Kineto's activity profiler.
+  // This is called internally when Kineto is ready.
+  // Safe to call multiple times - will only register once.
+  void registerWithKineto();
+
+  // Mark that Kineto has been initialized.
+  // If a factory was registered before Kineto init, it will be forwarded.
+  void onKinetoInit();
+
+ private:
+  PrivateUse1ProfilerRegistry() = default;
+
+  mutable std::mutex mutex_;
+  PrivateUse1ProfilerFactory factory_;
+  bool registered_with_kineto_ = false;
+  bool kineto_initialized_ = false;
+};
+
+// Helper struct for static registration via macro.
+// Enforces at compile-time that ProfilerClass inherits from
+// libkineto::IActivityProfiler.
+template <typename ProfilerClass>
+struct RegisterPrivateUse1Profiler {
+  static_assert(
+      std::is_base_of_v<libkineto::IActivityProfiler, ProfilerClass>,
+      "ProfilerClass must inherit from libkineto::IActivityProfiler. "
+      "Please ensure your profiler class implements the IActivityProfiler interface.");
+
+  RegisterPrivateUse1Profiler() {
+    PrivateUse1ProfilerRegistry::instance().registerFactory(
+        []() -> std::unique_ptr<libkineto::IActivityProfiler> {
+          return std::make_unique<ProfilerClass>();
+        });
+  }
+};
+
+// Macro for registering a PrivateUse1 activity profiler.
+// The profiler class must implement libkineto::IActivityProfiler.
+//
+// Usage:
+//   REGISTER_PRIVATEUSE1_PROFILER(MyAcceleratorProfiler)
+#define REGISTER_PRIVATEUSE1_PROFILER(ProfilerClass)                         \
+  static ::torch::profiler::impl::RegisterPrivateUse1Profiler<ProfilerClass> \
+      privateuse1_profiler_register_##ProfilerClass
+
+} // namespace torch::profiler::impl
+
+#endif // USE_KINETO
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/stubs/base.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/stubs/base.h
new file mode 100644
index 0000000000000000000000000000000000000000..fad34ad6b9daae1c143e9b365ae59a85156a832e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/stubs/base.h
@@ -0,0 +1,58 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <functional>
+#include <memory>
+
+#include <c10/core/Device.h>
+#include <c10/util/strong_type.h>
+#include <torch/csrc/Export.h>
+
+struct CUevent_st;
+
+namespace torch::profiler::impl {
+
+// ----------------------------------------------------------------------------
+// -- Annotation --------------------------------------------------------------
+// ----------------------------------------------------------------------------
+using ProfilerEventStub = std::shared_ptr<CUevent_st>;
+using ProfilerVoidEventStub = std::shared_ptr<void>;
+
+struct TORCH_API ProfilerStubs {
+  virtual void record(
+      c10::DeviceIndex* device,
+      ProfilerVoidEventStub* event,
+      int64_t* cpu_ns) const = 0;
+  virtual float elapsed(
+      const ProfilerVoidEventStub* event,
+      const ProfilerVoidEventStub* event2) const = 0;
+  virtual void mark(const char* name) const = 0;
+  virtual void rangePush(const char* name) const = 0;
+  virtual void rangePop() const = 0;
+  virtual bool enabled() const {
+    return false;
+  }
+  virtual void onEachDevice(std::function<void(int)> op) const = 0;
+  virtual void synchronize() const = 0;
+  virtual ~ProfilerStubs() = default;
+};
+
+TORCH_API void registerCUDAMethods(ProfilerStubs* stubs);
+TORCH_API const ProfilerStubs* cudaStubs();
+TORCH_API void registerITTMethods(ProfilerStubs* stubs);
+TORCH_API const ProfilerStubs* ittStubs();
+TORCH_API void registerPrivateUse1Methods(ProfilerStubs* stubs);
+TORCH_API const ProfilerStubs* privateuse1Stubs();
+
+using vulkan_id_t = strong::type<
+    int64_t,
+    struct _VulkanID,
+    strong::regular,
+    strong::convertible_to<int64_t>,
+    strong::hashable>;
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/action.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/action.h
new file mode 100644
index 0000000000000000000000000000000000000000..b26be822168f491a116db3400e9ad44b78334c5f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/action.h
@@ -0,0 +1,87 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <cstdint>
+#include <ostream>
+
+namespace torch::unwind {
+
+enum {
+  A_UNDEFINED = 0x0,
+  A_REG_PLUS_DATA = 0x1, // exp = REG[reg] + data0
+  A_LOAD_CFA_OFFSET = 0x2, // exp = *(cfa + data0)
+  A_REG_PLUS_DATA_DEREF = 0x3 // exp = *(REG[reg] + data0)
+};
+
+// DWARF register numbers — architecture-specific
+#if defined(__x86_64__)
+enum {
+  D_UNDEFINED = -1,
+  D_RBP = 6,
+  D_RSP = 7,
+  D_RIP = 16,
+  D_REG_SIZE = 17,
+};
+static constexpr int D_FRAME_PTR = D_RBP;
+static constexpr int D_STACK_PTR = D_RSP;
+static constexpr int D_RET_ADDR = D_RIP;
+static constexpr int D_EXPECTED_RA_REG = 16;
+#elif defined(__aarch64__)
+enum {
+  D_UNDEFINED = -1,
+  D_FP = 29,
+  D_LR = 30,
+  D_SP = 31,
+  D_REG_SIZE = 32,
+};
+static constexpr int D_FRAME_PTR = D_FP;
+static constexpr int D_STACK_PTR = D_SP;
+static constexpr int D_RET_ADDR = D_LR;
+static constexpr int D_EXPECTED_RA_REG = 30;
+#else
+enum {
+  D_UNDEFINED = -1,
+  D_REG_SIZE = 1,
+};
+#endif
+
+struct Action {
+  uint8_t kind = A_UNDEFINED;
+  int32_t reg = -1;
+  int64_t data = 0;
+  static Action undefined() {
+    return Action{A_UNDEFINED};
+  }
+  static Action regPlusData(int32_t reg, int64_t offset) {
+    return Action{A_REG_PLUS_DATA, reg, offset};
+  }
+  static Action regPlusDataDeref(int32_t reg, int64_t offset) {
+    return Action{A_REG_PLUS_DATA_DEREF, reg, offset};
+  }
+  static Action loadCfaOffset(int64_t offset) {
+    return Action{A_LOAD_CFA_OFFSET, D_UNDEFINED, offset};
+  }
+
+  friend std::ostream& operator<<(std::ostream& out, const Action& self) {
+    switch (self.kind) {
+      case A_UNDEFINED:
+        out << 'u';
+        break;
+      case A_REG_PLUS_DATA:
+        out << 'r' << (int)self.reg << " + " << self.data;
+        break;
+      case A_REG_PLUS_DATA_DEREF:
+        out << "*(r" << (int)self.reg << " + " << self.data << ')';
+        break;
+      case A_LOAD_CFA_OFFSET:
+        out << "*(cfa + " << self.data << ')';
+        break;
+    }
+    return out;
+  }
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/communicate.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/communicate.h
new file mode 100644
index 0000000000000000000000000000000000000000..a5b2281067ca256e2db7c687c313f29525d556ba
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/communicate.h
@@ -0,0 +1,78 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ext/stdio_filebuf.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <unistd.h>
+#include <array>
+#include <memory>
+
+namespace torch::unwind {
+// helper to open a process with stdin/stdout/stderr streams.
+struct Communicate {
+  Communicate(const char* command, const char** args) {
+    if (pipe(inpipe_.data()) < 0 || pipe(outpipe_.data()) < 0 ||
+        pipe(errpipe_.data()) < 0) {
+      throw UnwindError("pipe() failed");
+    }
+    pid_t pid = fork();
+    if (pid < 0) {
+      throw UnwindError("fork() failed");
+    } else if (pid == 0) { // child process
+      close(inpipe_[1]);
+      close(outpipe_[0]);
+      close(errpipe_[0]);
+
+      dup2(inpipe_[0], STDIN_FILENO);
+      dup2(outpipe_[1], STDOUT_FILENO);
+      dup2(errpipe_[1], STDERR_FILENO);
+      execvp(command, (char* const*)args);
+      throw UnwindError("failed execvp");
+    } else { // parent process
+      close(inpipe_[0]);
+      close(outpipe_[1]);
+      close(errpipe_[1]);
+      outbuf_ = std::make_unique<__gnu_cxx::stdio_filebuf<char>>(
+          inpipe_[1], std::ios::out);
+      inbuf_ = std::make_unique<__gnu_cxx::stdio_filebuf<char>>(
+          outpipe_[0], std::ios::in);
+      errbuf_ = std::make_unique<__gnu_cxx::stdio_filebuf<char>>(
+          errpipe_[0], std::ios::in);
+      in_ = std::make_unique<std::istream>(inbuf_.get());
+      out_ = std::make_unique<std::ostream>(outbuf_.get());
+      err_ = std::make_unique<std::ostream>(errbuf_.get());
+    }
+  }
+  Communicate(const Communicate&) = delete;
+  Communicate(Communicate&&) = delete;
+  Communicate& operator=(const Communicate&) = delete;
+  Communicate& operator=(Communicate&&) = delete;
+  ~Communicate() {
+    close(inpipe_[1]);
+    close(outpipe_[0]);
+    close(errpipe_[0]);
+  }
+  std::ostream& out() {
+    return *out_;
+  }
+  std::ostream& err() {
+    return *err_;
+  }
+  std::istream& in() {
+    return *in_;
+  }
+
+ private:
+  std::array<int, 2> inpipe_{-1, -1};
+  std::array<int, 2> outpipe_{-1, -1};
+  std::array<int, 2> errpipe_{-1, -1};
+  std::unique_ptr<__gnu_cxx::stdio_filebuf<char>> outbuf_, inbuf_, errbuf_;
+  std::unique_ptr<std::istream> in_;
+  std::unique_ptr<std::ostream> out_;
+  std::unique_ptr<std::ostream> err_;
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/debug_info.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/debug_info.h
new file mode 100644
index 0000000000000000000000000000000000000000..9f4d0fe227613489847733783a13872df1015b1f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/debug_info.h
@@ -0,0 +1,285 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/profiler/unwind/dwarf_enums.h>
+#include <torch/csrc/profiler/unwind/dwarf_symbolize_enums.h>
+#include <torch/csrc/profiler/unwind/lexer.h>
+#include <torch/csrc/profiler/unwind/sections.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <cstdint>
+#include <optional>
+
+namespace torch::unwind {
+
+struct DebugInfo {
+  DebugInfo(Sections& s) : s_(s) {}
+
+  void parse(uint64_t offset) {
+    auto L = parseHeader(offset);
+    parseCompileUnit(L);
+  }
+  std::optional<uint64_t> lineNumberProgramOffset() {
+    return line_number_program_offset_;
+  }
+  uint64_t nextOffset() {
+    return end_ - s_.debug_info.data;
+  }
+  std::vector<std::pair<uint64_t, uint64_t>> ranges() {
+    if (range_ptr_) {
+      auto offset = range_ptr_->first;
+      if (range_ptr_->second == DW_FORM_rnglistx) {
+        UNWIND_CHECK(rnglists_base_, "rnglistx but not rnglists_base_ set");
+        LOG_INFO("index for rnglistx {:x} + {:x}\n", *rnglists_base_, offset);
+        CheckedLexer L = s_.debug_rnglists.lexer(
+            *rnglists_base_ + offset * sec_offset_size_);
+        auto read = readSegmentOffset(L);
+        offset = *rnglists_base_ + read;
+      }
+      return version_ == 4 ? readRanges4(offset) : readRanges5(offset);
+    }
+    if (!highpc_) {
+      return {};
+    }
+    return {{lowpc_, lowpc_ + *highpc_}};
+  }
+
+  bool is64bit() {
+    return is_64bit_;
+  }
+
+ private:
+  CheckedLexer parseHeader(uint64_t offset) {
+    offset_ = offset;
+    CheckedLexer L = s_.debug_info.lexer(offset_);
+    std::tie(length_, is_64bit_) = L.readSectionLength();
+    sec_offset_size_ = is_64bit_ ? 8 : 4;
+    end_ = (const char*)L.loc() + length_;
+    version_ = L.read<uint16_t>();
+    UNWIND_CHECK(
+        version_ == 5 || version_ == 4,
+        "unexpected dwarf version {}",
+        version_);
+    uint8_t address_size = 0;
+    if (version_ == 5) {
+      auto unit_type = L.read<uint8_t>();
+      UNWIND_CHECK(unit_type == 0x1, "unexpected unit type {}", unit_type);
+      address_size = L.read<uint8_t>();
+      debug_abbrev_offset_ =
+          is_64bit_ ? L.read<uint64_t>() : L.read<uint32_t>();
+    } else {
+      debug_abbrev_offset_ =
+          is_64bit_ ? L.read<uint64_t>() : L.read<uint32_t>();
+      address_size = L.read<uint8_t>();
+    }
+    LOG_INFO(
+        "compilation unit at offset {:x} with length {:x} and debug_abbrev_offset {:x}\n",
+        offset,
+        length_,
+        debug_abbrev_offset_);
+    UNWIND_CHECK(
+        address_size == 8,
+        "expected 64-bit dwarf but found address size {}",
+        address_size);
+    return L;
+  }
+
+  uint64_t readSegmentOffset(CheckedLexer& L) {
+    return s_.readSegmentOffset(L, is_64bit_);
+  }
+
+  uint64_t readEncoded(CheckedLexer& L, uint64_t encoding) {
+    switch (encoding) {
+      case DW_FORM_data8:
+      case DW_FORM_addr:
+        return L.read<uint64_t>();
+      case DW_FORM_data4:
+        return L.read<uint32_t>();
+      case DW_FORM_addrx: {
+        auto idx = L.readULEB128();
+        return s_.debug_addr.lexer(address_base_ + sizeof(uint64_t) * idx)
+            .read<uint64_t>();
+      }
+      case DW_FORM_sec_offset:
+        return readSegmentOffset(L);
+      case DW_FORM_rnglistx: {
+        return L.readULEB128();
+      }
+      default:
+        UNWIND_CHECK(false, "unexpected encoding");
+    }
+  }
+
+  void parseCompileUnit(CheckedLexer& L) {
+    auto entry = L.readULEB128();
+    auto A = findAbbrev(debug_abbrev_offset_, entry);
+    while (true) {
+      auto attr = A.readULEB128();
+      auto form = A.readULEB128();
+      if (attr == 0 && form == 0) {
+        break;
+      }
+      if (form == DW_FORM_implicit_const) {
+        A.readSLEB128();
+      }
+      if (attr == DW_AT_low_pc) {
+        lowpc_ = readEncoded(L, form);
+        LOG_INFO("  lowpc {:x}\n", lowpc_);
+      } else if (attr == DW_AT_high_pc) {
+        highpc_ = readEncoded(L, form);
+        range_ptr_ = std::nullopt;
+        LOG_INFO("  highpc {:x}\n", *highpc_);
+      } else if (attr == DW_AT_addr_base) {
+        UNWIND_CHECK(form == DW_FORM_sec_offset, "unexpected addr_base form");
+        address_base_ = readSegmentOffset(L);
+        LOG_INFO("  address base {:x}\n", address_base_);
+      } else if (attr == DW_AT_rnglists_base) {
+        UNWIND_CHECK(
+            form == DW_FORM_sec_offset, "unexpected rnglists_base form");
+        rnglists_base_ = readSegmentOffset(L);
+        LOG_INFO("  range base {:x}\n", *rnglists_base_);
+      } else if (form == DW_FORM_string) {
+        L.readCString();
+      } else if (attr == DW_AT_stmt_list) {
+        UNWIND_CHECK(form == DW_FORM_sec_offset, "unexpected stmt_list form");
+        LOG_INFO("  program table offset {:x}\n", *line_number_program_offset_);
+        line_number_program_offset_ = readSegmentOffset(L);
+      } else if (form == DW_FORM_exprloc) {
+        auto sz = L.readULEB128();
+        L.skip(int64_t(sz));
+      } else if (form == DW_FORM_block1) {
+        auto sz = L.read<uint8_t>();
+        L.skip(int64_t(sz));
+      } else if (attr == DW_AT_ranges) {
+        auto range_offset = readEncoded(L, form);
+        LOG_INFO("setting range_ptr to {:x} {:x}\n", range_offset, form);
+        range_ptr_.emplace(range_offset, form);
+      } else if (
+          form == DW_FORM_udata || form == DW_FORM_rnglistx ||
+          form == DW_FORM_strx || form == DW_FORM_loclistx ||
+          form == DW_FORM_addrx) {
+        L.readULEB128();
+      } else if (form == DW_FORM_sdata) {
+        L.readSLEB128();
+      } else {
+        auto sz = formSize(form, sec_offset_size_);
+        UNWIND_CHECK(sz, "unsupported form in compilation unit {:x}", form);
+        L.skip(int64_t(*sz));
+      }
+    }
+  }
+
+  std::vector<std::pair<uint64_t, uint64_t>> readRanges4(uint64_t offset) {
+    CheckedLexer L = s_.debug_ranges.lexer(offset);
+    std::vector<std::pair<uint64_t, uint64_t>> ranges;
+    uint64_t base = lowpc_;
+    while (true) {
+      auto start = L.read<uint64_t>();
+      auto end = L.read<uint64_t>();
+      if (start == 0 && end == 0) {
+        break;
+      }
+      if (start == std::numeric_limits<uint64_t>::max()) {
+        base = end;
+      } else {
+        ranges.emplace_back(base + start, base + end);
+      }
+    }
+    return ranges;
+  }
+
+  std::vector<std::pair<uint64_t, uint64_t>> readRanges5(uint64_t offset) {
+    CheckedLexer L = s_.debug_rnglists.lexer(offset);
+    uint64_t base = 0;
+    LOG_INFO("BEGIN RANGES {:x}\n", offset);
+    std::vector<std::pair<uint64_t, uint64_t>> ranges;
+    while (true) {
+      auto op = L.read<uint8_t>();
+      switch (op) {
+        case DW_RLE_end_of_list:
+          LOG_INFO("END RANGES\n");
+          return ranges;
+        case DW_RLE_base_addressx: {
+          base = readEncoded(L, DW_FORM_addrx);
+          LOG_INFO("BASE ADDRX {:x}\n", base);
+        } break;
+        case DW_RLE_startx_length: {
+          auto s = readEncoded(L, DW_FORM_addrx);
+          auto e = L.readULEB128();
+          LOG_INFO("startx_length {:x} {:x}\n", s, e);
+          ranges.emplace_back(s, s + e);
+        } break;
+        case DW_RLE_base_address:
+          base = L.read<uint64_t>();
+          LOG_INFO("BASE ADDR {:x}\n", base);
+          break;
+        case DW_RLE_offset_pair: {
+          auto s = L.readULEB128();
+          auto e = L.readULEB128();
+          LOG_INFO("offset_pair {:x} {:x}\n", s, e);
+          ranges.emplace_back(base + s, base + e);
+        } break;
+        case DW_RLE_start_length: {
+          auto s = L.read<uint64_t>();
+          auto e = L.readULEB128();
+          LOG_INFO("start_length {:x} {:x}\n", s, e);
+          ranges.emplace_back(s, s + e);
+        } break;
+        default:
+          UNWIND_CHECK(false, "unknown range op: {}", op);
+      }
+    }
+  }
+
+  CheckedLexer findAbbrev(uint64_t offset, uint64_t entry) {
+    CheckedLexer L = s_.debug_abbrev.lexer(offset);
+    while (true) {
+      auto abbrev_code = L.readULEB128();
+      UNWIND_CHECK(
+          abbrev_code != 0,
+          "could not find entry {} at offset {:x}",
+          entry,
+          offset);
+      auto tag = L.readULEB128();
+      L.read<uint8_t>(); // has children
+      if (abbrev_code == entry) {
+        UNWIND_CHECK(
+            tag == DW_TAG_compile_unit,
+            "first entry was not a compile unit but {}",
+            tag);
+        return L;
+      }
+      while (true) {
+        auto attr = L.readULEB128();
+        auto form = L.readULEB128();
+        if (attr == 0 && form == 0) {
+          break;
+        }
+        if (form == DW_FORM_implicit_const) {
+          L.readSLEB128();
+        }
+      }
+    }
+  }
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  Sections& s_;
+  std::optional<uint64_t> line_number_program_offset_;
+  uint64_t offset_ = 0;
+  uint8_t sec_offset_size_ = 0;
+  uint64_t length_ = 0;
+  const char* end_ = nullptr;
+  uint64_t debug_abbrev_offset_ = 0;
+  bool is_64bit_ = false;
+
+  std::optional<std::pair<uint64_t, uint8_t>> range_ptr_;
+  uint64_t lowpc_ = 0;
+  std::optional<uint64_t> highpc_;
+  uint16_t version_ = 0;
+  uint64_t address_base_ = 0;
+  std::optional<uint64_t> rnglists_base_;
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/dwarf_enums.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/dwarf_enums.h
new file mode 100644
index 0000000000000000000000000000000000000000..7254057ca61305f1f12a7b4bf6d94d9eb67ed6a6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/dwarf_enums.h
@@ -0,0 +1,52 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+enum {
+  DW_EH_PE_absptr = 0x00,
+  DW_EH_PE_omit = 0xff,
+  /* FDE data encoding.  */
+  DW_EH_PE_uleb128 = 0x01,
+  DW_EH_PE_udata2 = 0x02,
+  DW_EH_PE_udata4 = 0x03,
+  DW_EH_PE_udata8 = 0x04,
+  DW_EH_PE_sleb128 = 0x09,
+  DW_EH_PE_sdata2 = 0x0a,
+  DW_EH_PE_sdata4 = 0x0b,
+  DW_EH_PE_sdata8 = 0x0c,
+  DW_EH_PE_signed = 0x08,
+  /* FDE flags.  */
+  DW_EH_PE_pcrel = 0x10,
+  DW_EH_PE_textrel = 0x20,
+  DW_EH_PE_datarel = 0x30,
+  DW_EH_PE_funcrel = 0x40,
+  DW_EH_PE_aligned = 0x50,
+  DW_EH_PE_indirect = 0x80,
+};
+
+enum {
+  DW_CFA_nop = 0x0,
+  DW_CFA_advance_loc = 0x01,
+  DW_CFA_offset = 0x02,
+  DW_CFA_restore = 0x03,
+  DW_CFA_advance_loc1 = 0x02,
+  DW_CFA_advance_loc2 = 0x03,
+  DW_CFA_advance_loc4 = 0x04,
+  DW_CFA_offset_extended = 0x05,
+  DW_CFA_restore_extended = 0x06,
+  DW_CFA_undefined = 0x07,
+  DW_CFA_register = 0x09,
+  DW_CFA_remember_state = 0x0a,
+  DW_CFA_restore_state = 0x0b,
+  DW_CFA_def_cfa = 0x0c,
+  DW_CFA_def_cfa_register = 0x0d,
+  DW_CFA_def_cfa_offset = 0x0e,
+  DW_CFA_def_cfa_expression = 0xf,
+  DW_CFA_expression = 0x10,
+  DW_CFA_offset_extended_sf = 0x11,
+  DW_CFA_GNU_args_size = 0x2e,
+  DW_OP_deref = 0x6,
+};
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/dwarf_symbolize_enums.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/dwarf_symbolize_enums.h
new file mode 100644
index 0000000000000000000000000000000000000000..d14a1dfd1d01bcd9ef6e06d63605bb91a8a83ba1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/dwarf_symbolize_enums.h
@@ -0,0 +1,184 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <cstdint>
+#include <optional>
+
+enum {
+  DW_TAG_subprogram = 0x2e,
+  DW_TAG_inlined_subroutine = 0x1d,
+  DW_TAG_compile_unit = 0x11,
+  DW_AT_sibling = 0x1, // reference
+  DW_AT_name = 0x3, // string
+  DW_AT_stmt_list = 0x10, // lineptr
+  DW_AT_addr_base = 0x73, // sec_offset
+  DW_AT_rnglists_base = 0x74, // sec_offset
+  DW_AT_low_pc = 0x11, // address
+  DW_AT_high_pc = 0x12, // address
+  DW_AT_specification = 0x47, // reference
+  DW_AT_abstract_origin = 0x31, // reference
+  DW_AT_linkage_name = 0x6e, // string
+  DW_AT_ranges = 0x55, // rnglist
+  DW_AT_str_offsets_base = 0x72, // sec_offset
+  DW_FORM_addr = 0x01,
+  DW_FORM_block2 = 0x03,
+  DW_FORM_block4 = 0x04,
+  DW_FORM_data2 = 0x05,
+  DW_FORM_data4 = 0x06,
+  DW_FORM_data8 = 0x07,
+  DW_FORM_string = 0x08,
+  DW_FORM_block = 0x09,
+  DW_FORM_block1 = 0x0a,
+  DW_FORM_data1 = 0x0b,
+  DW_FORM_flag = 0x0c,
+  DW_FORM_sdata = 0x0d,
+  DW_FORM_strp = 0x0e,
+  DW_FORM_udata = 0x0f,
+  DW_FORM_ref_addr = 0x10,
+  DW_FORM_ref1 = 0x11,
+  DW_FORM_ref2 = 0x12,
+  DW_FORM_ref4 = 0x13,
+  DW_FORM_ref8 = 0x14,
+  DW_FORM_ref_udata = 0x15,
+  DW_FORM_indirect = 0x16,
+  DW_FORM_sec_offset = 0x17,
+  DW_FORM_exprloc = 0x18,
+  DW_FORM_flag_present = 0x19,
+  DW_FORM_strx = 0x1a,
+  DW_FORM_addrx = 0x1b,
+  DW_FORM_ref_sup4 = 0x1c,
+  DW_FORM_strp_sup = 0x1d,
+  DW_FORM_data16 = 0x1e,
+  DW_FORM_line_strp = 0x1f,
+  DW_FORM_ref_sig8 = 0x20,
+  DW_FORM_implicit_const = 0x21,
+  DW_FORM_loclistx = 0x22,
+  DW_FORM_rnglistx = 0x23,
+  DW_FORM_ref_sup8 = 0x24,
+  DW_FORM_strx1 = 0x25,
+  DW_FORM_strx2 = 0x26,
+  DW_FORM_strx3 = 0x27,
+  DW_FORM_strx4 = 0x28,
+  DW_FORM_addrx1 = 0x29,
+  DW_FORM_addrx2 = 0x2a,
+  DW_FORM_addrx3 = 0x2b,
+  DW_FORM_addrx4 = 0x2c,
+  /* GNU Debug Fission extensions.  */
+  DW_FORM_GNU_addr_index = 0x1f01,
+  DW_FORM_GNU_str_index = 0x1f02,
+  DW_FORM_GNU_ref_alt = 0x1f20, /* offset in alternate .debuginfo.  */
+  DW_FORM_GNU_strp_alt = 0x1f21, /* offset in alternate .debug_str. */
+  DW_LNCT_path = 0x1,
+  DW_LNCT_directory_index = 0x2,
+  DW_LNS_extended_op = 0x00,
+  DW_LNE_end_sequence = 0x01,
+  DW_LNE_set_address = 0x02,
+  DW_LNS_copy = 0x01,
+  DW_LNS_advance_pc = 0x02,
+  DW_LNS_advance_line = 0x03,
+  DW_LNS_set_file = 0x04,
+  DW_LNS_const_add_pc = 0x08,
+  DW_LNS_fixed_advance_pc = 0x09,
+  DW_RLE_end_of_list = 0x0,
+  DW_RLE_base_addressx = 0x1,
+  DW_RLE_startx_endx = 0x2,
+  DW_RLE_startx_length = 0x3,
+  DW_RLE_offset_pair = 0x4,
+  DW_RLE_base_address = 0x5,
+  DW_RLE_start_end = 0x6,
+  DW_RLE_start_length = 0x7
+};
+
+static std::optional<size_t> formSize(uint64_t form, uint8_t sec_offset_size) {
+  switch (form) {
+    case DW_FORM_addr:
+      return sizeof(void*);
+    case DW_FORM_block2:
+    case DW_FORM_block4:
+      return std::nullopt;
+    case DW_FORM_data2:
+      return 2;
+    case DW_FORM_data4:
+      return 4;
+    case DW_FORM_data8:
+      return 8;
+    case DW_FORM_string:
+    case DW_FORM_block:
+    case DW_FORM_block1:
+      return std::nullopt;
+    case DW_FORM_data1:
+    case DW_FORM_flag:
+      return 1;
+    case DW_FORM_sdata:
+      return std::nullopt;
+    case DW_FORM_strp:
+      return sec_offset_size;
+    case DW_FORM_udata:
+      return std::nullopt;
+    case DW_FORM_ref_addr:
+      return sec_offset_size;
+    case DW_FORM_ref1:
+      return 1;
+    case DW_FORM_ref2:
+      return 2;
+    case DW_FORM_ref4:
+      return 4;
+    case DW_FORM_ref8:
+      return 8;
+    case DW_FORM_ref_udata:
+    case DW_FORM_indirect:
+      return std::nullopt;
+    case DW_FORM_sec_offset:
+      return sec_offset_size;
+    case DW_FORM_exprloc:
+      return std::nullopt;
+    case DW_FORM_flag_present:
+      return 0;
+    case DW_FORM_strx:
+    case DW_FORM_addrx:
+      return std::nullopt;
+    case DW_FORM_ref_sup4:
+      return 4;
+    case DW_FORM_strp_sup:
+      return sec_offset_size;
+    case DW_FORM_data16:
+      return 16;
+    case DW_FORM_line_strp:
+      return sec_offset_size;
+    case DW_FORM_ref_sig8:
+      return 8;
+    case DW_FORM_implicit_const:
+      return 0;
+    case DW_FORM_loclistx:
+    case DW_FORM_rnglistx:
+      return std::nullopt;
+    case DW_FORM_ref_sup8:
+      return 8;
+    case DW_FORM_strx1:
+      return 1;
+    case DW_FORM_strx2:
+      return 2;
+    case DW_FORM_strx3:
+      return 3;
+    case DW_FORM_strx4:
+      return 4;
+    case DW_FORM_addrx1:
+      return 1;
+    case DW_FORM_addrx2:
+      return 2;
+    case DW_FORM_addrx3:
+      return 3;
+    case DW_FORM_addrx4:
+      return 4;
+    case DW_FORM_GNU_addr_index:
+    case DW_FORM_GNU_str_index:
+    case DW_FORM_GNU_ref_alt:
+    case DW_FORM_GNU_strp_alt:
+    default:
+      return std::nullopt;
+  }
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/eh_frame_hdr.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/eh_frame_hdr.h
new file mode 100644
index 0000000000000000000000000000000000000000..865ddaecbacc94de1b0995effefbe628f3aefaa8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/eh_frame_hdr.h
@@ -0,0 +1,105 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <cstdint>
+#include <ostream>
+
+#include <torch/csrc/profiler/unwind/lexer.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+
+// Overview of the format described in
+// https://refspecs.linuxfoundation.org/LSB_1.3.0/gLSB/gLSB/ehframehdr.html
+namespace torch::unwind {
+
+struct EHFrameHdr {
+  EHFrameHdr(void* base) : base_(base) {
+    Lexer L(base, base);
+    version_ = L.read<uint8_t>();
+    eh_frame_ptr_enc_ = L.read<uint8_t>();
+    fde_count_enc_ = L.read<uint8_t>();
+    table_enc_ = L.read<uint8_t>();
+    if (table_enc_ == DW_EH_PE_omit) {
+      table_size_ = 0;
+    } else {
+      switch (table_enc_ & 0xF) {
+        case DW_EH_PE_udata2:
+        case DW_EH_PE_sdata2:
+          table_size_ = 2;
+          break;
+        case DW_EH_PE_udata4:
+        case DW_EH_PE_sdata4:
+          table_size_ = 4;
+          break;
+        case DW_EH_PE_udata8:
+        case DW_EH_PE_sdata8:
+          table_size_ = 8;
+          break;
+        case DW_EH_PE_uleb128:
+        case DW_EH_PE_sleb128:
+          throw UnwindError("uleb/sleb table encoding not supported");
+          break;
+        default:
+          throw UnwindError("unknown table encoding");
+      }
+    }
+    // NOLINTNEXTLINE(performance-no-int-to-ptr)
+    eh_frame_ = (void*)L.readEncodedOr(eh_frame_ptr_enc_, 0);
+    fde_count_ = L.readEncodedOr(fde_count_enc_, 0);
+    table_start_ = L.loc();
+  }
+  size_t nentries() const {
+    return fde_count_;
+  }
+
+  uint64_t lowpc(size_t i) const {
+    return Lexer(table_start_, base_)
+        .skip(2 * i * table_size_)
+        .readEncoded(table_enc_);
+  }
+  void* fde(size_t i) const {
+    // NOLINTNEXTLINE(performance-no-int-to-ptr)
+    return (void*)Lexer(table_start_, base_)
+        .skip((2 * i + 1) * table_size_)
+        .readEncoded(table_enc_);
+  }
+
+  void* entryForAddr(uint64_t addr) const {
+    if (!table_size_ || !nentries()) {
+      throw UnwindError("search table not present");
+    }
+    uint64_t low = 0;
+    uint64_t high = nentries();
+    while (low + 1 < high) {
+      auto mid = (low + high) / 2;
+      if (addr < lowpc(mid)) {
+        high = mid;
+      } else {
+        low = mid;
+      }
+    }
+    return fde(low);
+  }
+
+  friend std::ostream& operator<<(std::ostream& out, const EHFrameHdr& self) {
+    out << "EHFrameHeader(version=" << self.version_
+        << ",table_size=" << self.table_size_
+        << ",fde_count=" << self.fde_count_ << ')';
+    return out;
+  }
+
+ private:
+  void* base_;
+  void* table_start_;
+  uint8_t version_;
+  uint8_t eh_frame_ptr_enc_;
+  uint8_t fde_count_enc_;
+  uint8_t table_enc_;
+  void* eh_frame_ = nullptr;
+  int64_t fde_count_;
+  uint32_t table_size_;
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/fast_symbolizer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/fast_symbolizer.h
new file mode 100644
index 0000000000000000000000000000000000000000..2740c50054b41a55ff7190aacc9e7258464f5fa7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/fast_symbolizer.h
@@ -0,0 +1,113 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <fmt/format.h>
+#include <sys/types.h>
+#include <torch/csrc/profiler/unwind/debug_info.h>
+#include <torch/csrc/profiler/unwind/line_number_program.h>
+#include <torch/csrc/profiler/unwind/sections.h>
+#include <torch/csrc/profiler/unwind/unwind.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <memory>
+#include <unordered_map>
+
+namespace torch::unwind {
+
+#define UNWIND_WARN(w, ...)                   \
+  do {                                        \
+    w.emplace_back(fmt::format(__VA_ARGS__)); \
+    LOG_INFO("WARNING: {}\n", w.back());      \
+  } while (0);
+
+struct FastSymbolizer {
+  FastSymbolizer() = default;
+  Frame symbolize(const std::string& library, uint64_t offset) {
+    LOG_INFO("symbolizing {} + 0x{:x}\n", library, offset);
+    Frame frame;
+    frame.funcname = "??";
+    frame.filename = library;
+    frame.lineno = offset;
+    auto s = getOrCreateSections(library);
+    if (auto e = s->findSubprogramName(offset)) {
+      frame.funcname = *e;
+    } else {
+      UNWIND_WARN(
+          warnings_,
+          "failed to find subprogram name for {} 0x{:x}",
+          library,
+          offset);
+    }
+    if (auto e = findLine(s, offset)) {
+      frame.filename = e->first;
+      frame.lineno = e->second;
+    } else {
+      UNWIND_WARN(
+          warnings_, "failed to find file/line for {} 0x{:x}", library, offset);
+    }
+    return frame;
+  }
+  const std::vector<std::string>& warnings() {
+    return warnings_;
+  }
+
+ private:
+  void parseDebugInfo(Sections* s) {
+    uint64_t offset = 0;
+    while (offset < s->debug_info.size) {
+      DebugInfo info(*s);
+      info.parse(offset);
+      if (auto lnp_offset = info.lineNumberProgramOffset()) {
+        for (auto r : info.ranges()) {
+          s->addDebugInfoRange(r.first, r.second, line_number_programs_.size());
+        }
+        line_number_programs_.emplace_back(
+            std::make_unique<LineNumberProgram>(*s, *lnp_offset));
+      }
+      offset = info.nextOffset();
+    }
+  }
+  Sections* getOrCreateSections(const std::string& library) {
+    auto it = libraries_.find(library);
+    if (it == libraries_.end()) {
+      it = libraries_.insert({library, std::make_unique<Sections>()}).first;
+      try {
+        Sections* s = it->second.get();
+        s->parse(library.c_str());
+        parseDebugInfo(s);
+      } catch (UnwindError& err) {
+        UNWIND_WARN(
+            warnings_, "failed to parse library {}: {}", library, err.what());
+      }
+    }
+    return it->second.get();
+  }
+  std::optional<std::pair<std::string, int64_t>> findLine(
+      Sections* s,
+      uint64_t offset) {
+    if (auto idx = s->findDebugInfoOffset(offset)) {
+      auto r = line_number_programs_.at(*idx).get();
+      try {
+        r->parse();
+      } catch (UnwindError& err) {
+        UNWIND_WARN(
+            warnings_,
+            "failed to read line number program [{:x}] {}",
+            r->offset(),
+            err.what());
+      }
+      if (auto e = r->find(offset)) {
+        return std::make_pair(r->filename(e->file), e->line);
+      }
+    }
+    return std::nullopt;
+  }
+  std::unordered_map<std::string, std::unique_ptr<Sections>> libraries_;
+  std::vector<std::unique_ptr<LineNumberProgram>> line_number_programs_;
+  std::vector<std::string> warnings_;
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/fde.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/fde.h
new file mode 100644
index 0000000000000000000000000000000000000000..461363707fc46087a34a8d90014786b95a54a92c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/fde.h
@@ -0,0 +1,423 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/util/Exception.h>
+#include <c10/util/irange.h>
+#include <torch/csrc/profiler/unwind/action.h>
+#include <torch/csrc/profiler/unwind/lexer.h>
+#include <array>
+#include <iostream>
+#include <sstream>
+#include <vector>
+
+namespace torch::unwind {
+
+struct TableState {
+  Action cfa;
+  std::array<Action, D_REG_SIZE> registers;
+  friend std::ostream& operator<<(std::ostream& out, const TableState& self) {
+    out << "cfa = " << self.cfa << "; ";
+    for (auto r : c10::irange(self.registers.size())) {
+      if (self.registers.at(r).kind != A_UNDEFINED) {
+        out << 'r' << r << " = " << self.registers.at(r) << "; ";
+      }
+    }
+    return out;
+  }
+};
+
+// FDE - Frame Description Entry (Concept in ELF spec)
+// This format is explained well by
+// https://www.airs.com/blog/archives/460
+// Details of different dwarf actions are explained
+// in the spec document:
+// https://web.archive.org/web/20221129184704/https://dwarfstd.org/doc/DWARF4.doc
+// An overview of how DWARF unwinding works is given in
+// https://dl.acm.org/doi/pdf/10.1145/3360572
+// A similar implementation written in rust is:
+// https://github.com/mstange/framehop/
+
+template <bool LOG = false>
+struct FDE {
+  FDE(void* data, const char* library_name, uint64_t load_bias)
+      : library_name_(library_name), load_bias_(load_bias) {
+    Lexer L(data);
+    auto length = L.read4or8Length();
+    void* fde_start = L.loc();
+    // NOLINTNEXTLINE(performance-no-int-to-ptr)
+    void* cie_data = (void*)((int64_t)fde_start - L.read<uint32_t>());
+    Lexer LC(cie_data);
+    auto cie_length = LC.read4or8Length();
+    void* cie_start = LC.loc();
+    auto zero = LC.read<uint32_t>();
+    TORCH_INTERNAL_ASSERT(zero == 0, "expected 0 for CIE");
+    auto version = LC.read<uint8_t>();
+    TORCH_INTERNAL_ASSERT(
+        version == 1 || version == 3, "non-1 version for CIE");
+    augmentation_string_ = LC.readCString();
+    if (hasAugmentation("eh")) {
+      throw UnwindError("unsupported 'eh' augmentation string");
+    }
+    code_alignment_factor_ = static_cast<int64_t>(LC.readULEB128());
+    data_alignment_factor_ = LC.readSLEB128();
+    if (version == 1) {
+      ra_register_ = LC.read<uint8_t>();
+    } else {
+      ra_register_ = static_cast<int64_t>(LC.readULEB128());
+    }
+    TORCH_INTERNAL_ASSERT(
+        ra_register_ == D_EXPECTED_RA_REG,
+        "unexpected ra register: ",
+        ra_register_);
+    if (augmentation_string_ && *augmentation_string_ == 'z') {
+      augmentation_length_ = static_cast<int64_t>(LC.readULEB128());
+      Lexer A(LC.loc());
+      for (auto ap = augmentation_string_ + 1; *ap; ap++) {
+        switch (*ap) {
+          case 'L':
+            lsda_enc = A.read<uint8_t>();
+            break;
+          case 'R':
+            fde_enc = A.read<uint8_t>();
+            break;
+          case 'P': {
+            uint8_t personality_enc = A.read<uint8_t>();
+            A.readEncoded(personality_enc);
+          } break;
+          case 'S': {
+            // signal handler
+          } break;
+          default: {
+            throw UnwindError("unknown augmentation string");
+          } break;
+        }
+      }
+    }
+    LC.skip(augmentation_length_);
+    low_pc_ = L.readEncoded(fde_enc);
+    high_pc_ = low_pc_ + L.readEncodedValue(fde_enc);
+
+    if (hasAugmentation("z")) {
+      augmentation_length_fde_ = static_cast<int64_t>(L.readULEB128());
+    }
+    L.readEncodedOr(lsda_enc, 0);
+
+    cie_begin_ = LC.loc();
+    fde_begin_ = L.loc();
+    cie_end_ = (void*)((const char*)cie_start + cie_length);
+    fde_end_ = (void*)((const char*)fde_start + length);
+  }
+
+  // OP Code implementations
+
+  void advance_raw(int64_t amount) {
+    auto previous_pc = current_pc_;
+    current_pc_ += amount;
+    if (LOG) {
+      (*out_) << (void*)(previous_pc - load_bias_) << '-'
+              << (void*)(current_pc_ - load_bias_) << ": " << state() << '\n';
+    }
+  }
+
+  void advance_loc(int64_t amount) {
+    if (LOG) {
+      (*out_) << "advance_loc " << amount << '\n';
+    }
+    advance_raw(amount * code_alignment_factor_);
+  }
+
+  void offset(int64_t reg, int64_t offset) {
+    if (LOG) {
+      (*out_) << "offset " << reg << ' ' << offset << '\n';
+    }
+    if (reg > (int64_t)state().registers.size()) {
+      if (LOG) {
+        (*out_) << "OFFSET OF BIG REGISTER " << reg << "ignored...\n";
+      }
+      return;
+    }
+    state().registers.at(reg) =
+        Action{A_LOAD_CFA_OFFSET, -1, offset * data_alignment_factor_};
+  }
+
+  void restore(int64_t reg) {
+    if (LOG) {
+      (*out_) << "restore " << reg << '\n';
+    }
+    if (reg > (int64_t)state().registers.size()) {
+      if (LOG) {
+        (*out_) << "RESTORE OF BIG REGISTER " << reg << "ignored...\n";
+      }
+      return;
+    }
+    state().registers.at(reg) = initial_state_.registers.at(reg);
+  }
+
+  void def_cfa(int64_t reg, int64_t off) {
+    if (LOG) {
+      (*out_) << "def_cfa " << reg << ' ' << off << '\n';
+    }
+    last_reg_ = reg;
+    last_offset_ = off;
+    state().cfa = Action::regPlusData(static_cast<int32_t>(reg), off);
+  }
+  void def_cfa_register(int64_t reg) {
+    def_cfa(reg, last_offset_);
+  }
+  void def_cfa_offset(int64_t off) {
+    def_cfa(last_reg_, off);
+  }
+
+  void remember_state() {
+    if (LOG) {
+      (*out_) << "remember_state\n";
+    }
+    state_stack_.push_back(state());
+  }
+  void restore_state() {
+    if (LOG) {
+      (*out_) << "restore_state\n";
+    }
+    state_stack_.pop_back();
+  }
+
+  void undefined(int64_t reg) {
+    if (LOG) {
+      (*out_) << "undefined " << reg << '\n';
+    }
+    state().registers.at(reg) = Action::undefined();
+  }
+  void register_(int64_t reg, int64_t rhs_reg) {
+    if (LOG) {
+      (*out_) << "register " << reg << ' ' << rhs_reg << '\n';
+    }
+    state().registers.at(reg) =
+        Action::regPlusData(static_cast<int32_t>(reg), 0);
+  }
+
+  TableState& state() {
+    return state_stack_.back();
+  }
+
+  void dump(std::ostream& out) {
+    out_ = &out;
+    out << "FDE(augmentation_string=" << augmentation_string_
+        << ", low_pc=" << (void*)(low_pc_ - load_bias_)
+        << ",high_pc=" << (void*)(high_pc_ - load_bias_)
+        << ",code_alignment_factor=" << code_alignment_factor_
+        << ", data_alignment_factor=" << data_alignment_factor_
+        << ", ra_register_=" << ra_register_ << ")\n";
+    readUpTo(high_pc_);
+    out_ = &std::cout;
+  }
+
+  TableState readUpTo(uint64_t addr) {
+    if (addr < low_pc_ || addr > high_pc_) {
+      throw UnwindError("Address not in range");
+    }
+    if (LOG) {
+      // NOLINTNEXTLINE(performance-no-int-to-ptr)
+      (*out_) << "readUpTo " << (void*)addr << " for " << library_name_
+              << " at " << (void*)load_bias_ << '\n';
+    }
+    state_stack_.emplace_back();
+    current_pc_ = low_pc_;
+    // parse instructions...
+    Lexer LC(cie_begin_);
+    while (LC.loc() < cie_end_ && current_pc_ <= addr) {
+      readInstruction(LC);
+    }
+    if (current_pc_ > addr) {
+      return state();
+    }
+
+    initial_state_ = state_stack_.back();
+
+    if (LOG) {
+      (*out_) << "--\n";
+    }
+
+    Lexer L(fde_begin_);
+    while (L.loc() < fde_end_ && current_pc_ <= addr) {
+      readInstruction(L);
+    }
+    // so that we print the full range in debugging
+    if (current_pc_ <= addr) {
+      advance_raw(addr - current_pc_);
+    }
+    return state();
+  }
+
+  void dumpAddr2Line() {
+    std::cout << "addr2line -f -e " << library_name_ << ' '
+              << (void*)(low_pc_ - load_bias_) << '\n';
+  }
+
+  void readInstruction(Lexer& L) {
+    uint8_t bc = L.read<uint8_t>();
+    auto op = bc >> 6;
+    auto lowbits = bc & 0x3F;
+    switch (op) {
+      case 0x0: {
+        switch (lowbits) {
+          case DW_CFA_nop: {
+            return; // nop
+          }
+          case DW_CFA_advance_loc1: {
+            auto delta = L.read<uint8_t>();
+            return advance_loc(delta);
+          }
+          case DW_CFA_advance_loc2: {
+            auto delta = L.read<uint16_t>();
+            return advance_loc(delta);
+          }
+          case DW_CFA_advance_loc4: {
+            auto delta = L.read<uint32_t>();
+            return advance_loc(delta);
+          }
+          case DW_CFA_offset_extended: {
+            auto reg = L.readULEB128();
+            auto off = L.readULEB128();
+            return offset(reg, off);
+          }
+          case DW_CFA_restore_extended: {
+            auto reg = L.readULEB128();
+            return restore(reg);
+          }
+          case DW_CFA_undefined: {
+            auto reg = L.readULEB128();
+            return undefined(reg);
+          }
+          case DW_CFA_register: {
+            auto reg = L.readULEB128();
+            auto rhs_reg = L.readULEB128();
+            return register_(reg, rhs_reg);
+          }
+          case DW_CFA_def_cfa: {
+            auto reg = L.readULEB128();
+            auto off = L.readULEB128();
+            return def_cfa(reg, off);
+          }
+          case DW_CFA_def_cfa_register: {
+            auto reg = L.readULEB128();
+            return def_cfa_register(reg);
+          }
+          case DW_CFA_def_cfa_offset: {
+            auto off = L.readULEB128();
+            return def_cfa_offset(off);
+          }
+          case DW_CFA_offset_extended_sf: {
+            auto reg = L.readULEB128();
+            auto off = L.readSLEB128();
+            return offset(reg, off);
+          }
+          case DW_CFA_remember_state: {
+            return remember_state();
+          }
+          case DW_CFA_restore_state: {
+            return restore_state();
+          }
+          case DW_CFA_GNU_args_size: {
+            // GNU_args_size, we do not need to know it..
+            L.readULEB128();
+            return;
+          }
+          case DW_CFA_expression: {
+            auto reg = L.readULEB128();
+            auto len = L.readULEB128();
+            // NOLINTNEXTLINE(performance-no-int-to-ptr)
+            auto end = (void*)((uint64_t)L.loc() + len);
+            auto op = L.read<uint8_t>();
+            if ((op & 0xF0) == 0x70) { // DW_bregX
+              auto rhs_reg = (op & 0xF);
+              auto addend = L.readSLEB128();
+              if (L.loc() == end) {
+                state().registers.at(reg) =
+                    Action::regPlusDataDeref(rhs_reg, addend);
+                return;
+              }
+            }
+            throw UnwindError("Unsupported dwarf expression");
+          }
+          case DW_CFA_def_cfa_expression: {
+            auto len = L.readULEB128();
+            // NOLINTNEXTLINE(performance-no-int-to-ptr)
+            auto end = (void*)((uint64_t)L.loc() + len);
+            auto op = L.read<uint8_t>();
+            if ((op & 0xF0) == 0x70) { // DW_bregX
+              auto rhs_reg = (op & 0xF);
+              auto addend = L.readSLEB128();
+              if (L.loc() != end) {
+                auto op2 = L.read<uint8_t>();
+                if (op2 == DW_OP_deref && L.loc() == end) { // deref
+                  state().cfa = Action::regPlusDataDeref(rhs_reg, addend);
+                  return;
+                }
+              }
+            }
+            throw UnwindError("Unsupported def_cfa dwarf expression");
+          }
+          default: {
+            std::stringstream ss;
+            // NOLINTNEXTLINE(performance-no-int-to-ptr)
+            ss << "unknown op code " << (void*)(uint64_t)lowbits;
+            throw UnwindError(ss.str());
+          }
+        }
+      }
+      case DW_CFA_advance_loc: {
+        return advance_loc(lowbits);
+      }
+      case DW_CFA_offset: {
+        auto off = L.readULEB128();
+        return offset(lowbits, off);
+      }
+      case DW_CFA_restore: {
+        return restore(lowbits);
+      }
+    }
+  }
+  // used for debug printing
+  const char* library_name_;
+  uint64_t load_bias_;
+
+  // parsed from the eh_string data structures:
+  const char* augmentation_string_ = nullptr;
+  int64_t augmentation_length_ = 0;
+  int64_t augmentation_length_fde_ = 0;
+
+  int64_t code_alignment_factor_;
+  int64_t data_alignment_factor_;
+  void* cie_data_{nullptr};
+
+  int64_t ra_register_;
+  uint8_t lsda_enc = DW_EH_PE_omit;
+  uint8_t fde_enc = DW_EH_PE_absptr;
+  uint64_t low_pc_ = UINT64_MAX;
+  uint64_t high_pc_ = UINT64_MAX;
+
+  void* cie_begin_;
+  void* fde_begin_;
+  void* cie_end_;
+  void* fde_end_;
+
+  // state accumulated while parsing instructions
+  int64_t last_reg_ = 0;
+  int64_t last_offset_ = 0;
+  uint64_t current_pc_ = 0;
+
+  TableState
+      initial_state_; // state after the initial instructions, used by restore
+  std::vector<TableState> state_stack_;
+
+  std::ostream* out_ = &std::cout; // for debug dumping
+ private:
+  bool hasAugmentation(const char* s) {
+    return strstr(augmentation_string_, s) != nullptr;
+  }
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/lexer.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/lexer.h
new file mode 100644
index 0000000000000000000000000000000000000000..f9580a88c4574ba9d74464bd61b7d4e0ec527730
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/lexer.h
@@ -0,0 +1,164 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <cstdint>
+#include <cstring>
+#include <utility>
+
+#include <torch/csrc/profiler/unwind/dwarf_enums.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+
+namespace torch::unwind {
+
+template <bool checked>
+struct LexerImpl {
+  LexerImpl(void* data, void* base = nullptr, void* end = nullptr)
+      : next_((const char*)data),
+        base_((int64_t)base),
+        end_((const char*)end) {}
+
+  template <typename T>
+  T read() {
+    T result;
+    auto end = next_ + sizeof(T);
+    UNWIND_CHECK(
+        !checked || end <= end_,
+        "read out of bounds {} >= {}",
+        (void*)end,
+        (void*)end_);
+    memcpy(&result, next_, sizeof(T));
+    next_ = end;
+    return result;
+  }
+
+  // SLEB/ULEB code adapted from LLVM equivalents
+  int64_t readSLEB128() {
+    int64_t Value = 0;
+    unsigned Shift = 0;
+    uint8_t Byte = 0;
+    do {
+      Byte = read<uint8_t>();
+      uint64_t Slice = Byte & 0x7f;
+      if ((Shift >= 64 && Slice != (Value < 0 ? 0x7f : 0x00)) ||
+          (Shift == 63 && Slice != 0 && Slice != 0x7f)) {
+        throw UnwindError("sleb128 too big for int64");
+      }
+      Value |= int64_t(Slice << Shift);
+      Shift += 7;
+    } while (Byte >= 128);
+    // Sign extend negative numbers if needed.
+    if (Shift < 64 && (Byte & 0x40)) {
+      Value |= int64_t((-1ULL) << Shift);
+    }
+    return Value;
+  }
+
+  uint64_t readULEB128() {
+    uint64_t Value = 0;
+    unsigned Shift = 0;
+    uint8_t p = 0;
+    do {
+      p = read<uint8_t>();
+      uint64_t Slice = p & 0x7f;
+      if ((Shift >= 64 && Slice != 0) || Slice << Shift >> Shift != Slice) {
+        throw UnwindError("uleb128 too big for uint64");
+      }
+      Value += Slice << Shift;
+      Shift += 7;
+    } while (p >= 128);
+    return Value;
+  }
+  const char* readCString() {
+    auto result = next_;
+    if (!checked) {
+      next_ += strlen(next_) + 1;
+      return result;
+    }
+    while (next_ < end_) {
+      if (*next_++ == '\0') {
+        return result;
+      }
+    }
+    UNWIND_CHECK(
+        false, "string is out of bounds {} >= {}", (void*)next_, (void*)end_);
+  }
+  int64_t readEncoded(uint8_t enc) {
+    int64_t r = 0;
+    switch (enc & (~DW_EH_PE_indirect & 0xF0)) {
+      case DW_EH_PE_absptr:
+        break;
+      case DW_EH_PE_pcrel:
+        r = (int64_t)next_;
+        break;
+      case DW_EH_PE_datarel:
+        r = base_;
+        break;
+      default:
+        throw UnwindError("unknown encoding");
+    }
+    return r + readEncodedValue(enc);
+  }
+  int64_t readEncodedOr(uint8_t enc, int64_t orelse) {
+    if (enc == DW_EH_PE_omit) {
+      return orelse;
+    }
+    return readEncoded(enc);
+  }
+
+  int64_t read4or8Length() {
+    return readSectionLength().first;
+  }
+
+  std::pair<int64_t, bool> readSectionLength() {
+    int64_t length = read<uint32_t>();
+    if (length == 0xFFFFFFFF) {
+      return std::make_pair(read<int64_t>(), true);
+    }
+    return std::make_pair(length, false);
+  }
+
+  void* loc() const {
+    return (void*)next_;
+  }
+  LexerImpl& skip(size_t bytes) {
+    next_ += bytes;
+    return *this;
+  }
+
+  int64_t readEncodedValue(uint8_t enc) {
+    switch (enc & 0xF) {
+      case DW_EH_PE_udata2:
+        return read<uint16_t>();
+      case DW_EH_PE_sdata2:
+        return read<int16_t>();
+      case DW_EH_PE_udata4:
+        return read<uint32_t>();
+      case DW_EH_PE_sdata4:
+        return read<int32_t>();
+      case DW_EH_PE_udata8:
+        return read<uint64_t>();
+      case DW_EH_PE_sdata8:
+        return read<int64_t>();
+      case DW_EH_PE_uleb128:
+        return readULEB128();
+      case DW_EH_PE_sleb128:
+        return readSLEB128();
+      default:
+        throw UnwindError("not implemented");
+    }
+  }
+
+ private:
+  const char* next_;
+  int64_t base_;
+  const char* end_;
+};
+
+// using Lexer = LexerImpl<false>;
+using CheckedLexer = LexerImpl<true>;
+using Lexer = LexerImpl<false>;
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/line_number_program.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/line_number_program.h
new file mode 100644
index 0000000000000000000000000000000000000000..32fbc59c4655a7d99c2c795c04c88b125b12193c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/line_number_program.h
@@ -0,0 +1,333 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <c10/util/irange.h>
+#include <torch/csrc/profiler/unwind/debug_info.h>
+#include <torch/csrc/profiler/unwind/dwarf_enums.h>
+#include <torch/csrc/profiler/unwind/dwarf_symbolize_enums.h>
+#include <torch/csrc/profiler/unwind/lexer.h>
+#include <torch/csrc/profiler/unwind/sections.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <tuple>
+
+namespace torch::unwind {
+
+struct LineNumberProgram {
+  LineNumberProgram(Sections& s, uint64_t offset) : s_(s), offset_(offset) {}
+
+  uint64_t offset() {
+    return offset_;
+  }
+  void parse() {
+    if (parsed_) {
+      return;
+    }
+    parsed_ = true;
+    CheckedLexer L = s_.debug_line.lexer(offset_);
+    std::tie(length_, is_64bit_) = L.readSectionLength();
+    program_end_ = (char*)L.loc() + length_;
+    auto version = L.read<uint16_t>();
+    UNWIND_CHECK(
+        version == 5 || version == 4,
+        "expected version 4 or 5 but found {}",
+        version);
+    if (version == 5) {
+      auto address_size = L.read<uint8_t>();
+      UNWIND_CHECK(
+          address_size == 8,
+          "expected 64-bit dwarf but found address size {}",
+          address_size);
+      segment_selector_size_ = L.read<uint8_t>();
+    }
+    header_length_ = is_64bit_ ? L.read<uint64_t>() : L.read<uint32_t>();
+    program_ = L;
+    program_.skip(int64_t(header_length_));
+    minimum_instruction_length_ = L.read<uint8_t>();
+    maximum_operations_per_instruction_ = L.read<uint8_t>();
+    default_is_stmt_ = L.read<uint8_t>();
+    line_base_ = L.read<int8_t>();
+    line_range_ = L.read<uint8_t>();
+    opcode_base_ = L.read<uint8_t>();
+    UNWIND_CHECK(line_range_ != 0, "line_range_ must be non-zero");
+    standard_opcode_lengths_.resize(opcode_base_);
+    for (size_t i = 1; i < opcode_base_; i++) {
+      standard_opcode_lengths_[i] = L.read<uint8_t>();
+    }
+    // fmt::print("{:x} {:x} {} {} {} {} {}\n", offset_, header_length_,
+    // minimum_instruction_length_, maximum_operations_per_instruction_,
+    // line_base_, line_range_, opcode_base_);
+    uint8_t directory_entry_format_count = L.read<uint8_t>();
+
+    if (version == 5) {
+      struct Member {
+        uint64_t content_type;
+        uint64_t form;
+      };
+      std::vector<Member> directory_members;
+      directory_members.reserve(directory_entry_format_count);
+      for (size_t i = 0; i < directory_entry_format_count; i++) {
+        directory_members.push_back({L.readULEB128(), L.readULEB128()});
+      }
+      uint64_t directories_count = L.readULEB128();
+      for (size_t i = 0; i < directories_count; i++) {
+        for (auto& member : directory_members) {
+          switch (member.content_type) {
+            case DW_LNCT_path: {
+              include_directories_.emplace_back(
+                  s_.readString(L, member.form, is_64bit_));
+            } break;
+            default: {
+              skipForm(L, member.form);
+            } break;
+          }
+        }
+      }
+
+      for (auto i : c10::irange(directories_count)) {
+        (void)i;
+        LOG_INFO("{} {}\n", i, include_directories_[i]);
+      }
+      auto file_name_entry_format_count = L.read<uint8_t>();
+      std::vector<Member> file_members;
+      file_members.reserve(file_name_entry_format_count);
+      for (size_t i = 0; i < file_name_entry_format_count; i++) {
+        file_members.push_back({L.readULEB128(), L.readULEB128()});
+      }
+      auto files_count = L.readULEB128();
+      for (size_t i = 0; i < files_count; i++) {
+        for (auto& member : file_members) {
+          switch (member.content_type) {
+            case DW_LNCT_path: {
+              file_names_.emplace_back(
+                  s_.readString(L, member.form, is_64bit_));
+            } break;
+            case DW_LNCT_directory_index: {
+              file_directory_index_.emplace_back(readData(L, member.form));
+              UNWIND_CHECK(
+                  file_directory_index_.back() < include_directories_.size(),
+                  "directory index out of range");
+            } break;
+            default: {
+              skipForm(L, member.form);
+            } break;
+          }
+        }
+      }
+      for (auto i : c10::irange(files_count)) {
+        (void)i;
+        LOG_INFO("{} {} {}\n", i, file_names_[i], file_directory_index_[i]);
+      }
+    } else {
+      include_directories_.emplace_back(""); // implicit cwd
+      while (true) {
+        auto str = L.readCString();
+        if (*str == '\0') {
+          break;
+        }
+        include_directories_.emplace_back(str);
+      }
+      file_names_.emplace_back("");
+      file_directory_index_.emplace_back(0);
+      while (true) {
+        auto str = L.readCString();
+        if (*str == '\0') {
+          break;
+        }
+        auto directory_index = L.readULEB128();
+        L.readULEB128(); // mod_time
+        L.readULEB128(); // file_length
+        file_names_.emplace_back(str);
+        file_directory_index_.push_back(directory_index);
+      }
+    }
+    UNWIND_CHECK(
+        maximum_operations_per_instruction_ == 1,
+        "maximum_operations_per_instruction_ must be 1");
+    UNWIND_CHECK(
+        minimum_instruction_length_ == 1,
+        "minimum_instruction_length_ must be 1");
+    readProgram();
+  }
+  struct Entry {
+    uint32_t file = 1;
+    int64_t line = 1;
+  };
+  std::optional<Entry> find(uint64_t address) {
+    auto e = program_index_.find(address);
+    if (!e) {
+      return std::nullopt;
+    }
+    return all_programs_.at(*e).find(address);
+  }
+  std::string filename(uint64_t index) {
+    return fmt::format(
+        "{}/{}",
+        include_directories_.at(file_directory_index_.at(index)),
+        file_names_.at(index));
+  }
+
+ private:
+  void skipForm(CheckedLexer& L, uint64_t form) {
+    auto sz = formSize(form, is_64bit_ ? 8 : 4);
+    UNWIND_CHECK(sz, "unsupported form {}", form);
+    L.skip(int64_t(*sz));
+  }
+
+  uint64_t readData(CheckedLexer& L, uint64_t encoding) {
+    switch (encoding) {
+      case DW_FORM_data1:
+        return L.read<uint8_t>();
+      case DW_FORM_data2:
+        return L.read<uint16_t>();
+      case DW_FORM_data4:
+        return L.read<uint32_t>();
+      case DW_FORM_data8:
+        return L.read<uint64_t>();
+      case DW_FORM_udata:
+        return L.readULEB128();
+      default:
+        UNWIND_CHECK(false, "unsupported data encoding {}", encoding);
+    }
+  }
+
+  void produceEntry() {
+    if (shadow_) {
+      return;
+    }
+    if (ranges_.size() == 1) {
+      start_address_ = address_;
+    }
+    PRINT_LINE_TABLE(
+        "{:x}\t{}\t{}\n", address_, filename(entry_.file), entry_.line);
+    UNWIND_CHECK(
+        entry_.file < file_names_.size(),
+        "file index {} > {} entries",
+        entry_.file,
+        file_names_.size());
+    ranges_.add(address_, entry_, true);
+  }
+  void endSequence() {
+    if (shadow_) {
+      return;
+    }
+    PRINT_LINE_TABLE(
+        "{:x}\tEND\n", address_, filename(entry_.file), entry_.line);
+    program_index_.add(start_address_, all_programs_.size(), false);
+    program_index_.add(address_, std::nullopt, false);
+    all_programs_.emplace_back(std::move(ranges_));
+    ranges_ = RangeTable<Entry>();
+  }
+  void readProgram() {
+    while (program_.loc() < program_end_) {
+      PRINT_INST("{:x}: ", (char*)program_.loc() - (s_.debug_line.data));
+      uint8_t op = program_.read<uint8_t>();
+      if (op >= opcode_base_) {
+        auto op2 = int64_t(op - opcode_base_);
+        address_ += op2 / line_range_;
+        entry_.line += line_base_ + (op2 % line_range_);
+        PRINT_INST(
+            "address += {}, line += {}\n",
+            op2 / line_range_,
+            line_base_ + (op2 % line_range_));
+        produceEntry();
+      } else {
+        switch (op) {
+          case DW_LNS_extended_op: {
+            auto len = program_.readULEB128();
+            auto extended_op = program_.read<uint8_t>();
+            switch (extended_op) {
+              case DW_LNE_end_sequence: {
+                PRINT_INST("end_sequence\n");
+                endSequence();
+                entry_ = Entry{};
+              } break;
+              case DW_LNE_set_address: {
+                address_ = program_.read<uint64_t>();
+                if (!shadow_) {
+                  PRINT_INST(
+                      "set address {:x} {:x} {:x}\n",
+                      address_,
+                      min_address_,
+                      max_address_);
+                }
+                shadow_ = address_ == 0;
+              } break;
+              default: {
+                PRINT_INST("skip extended op {}\n", extended_op);
+                program_.skip(int64_t(len - 1));
+              } break;
+            }
+          } break;
+          case DW_LNS_copy: {
+            PRINT_INST("copy\n");
+            produceEntry();
+          } break;
+          case DW_LNS_advance_pc: {
+            PRINT_INST("advance pc\n");
+            address_ += program_.readULEB128();
+          } break;
+          case DW_LNS_advance_line: {
+            entry_.line += program_.readSLEB128();
+            PRINT_INST("advance line {}\n", entry_.line);
+
+          } break;
+          case DW_LNS_set_file: {
+            PRINT_INST("set file\n");
+            entry_.file = program_.readULEB128();
+          } break;
+          case DW_LNS_const_add_pc: {
+            PRINT_INST("const add pc\n");
+            address_ += (255 - opcode_base_) / line_range_;
+          } break;
+          case DW_LNS_fixed_advance_pc: {
+            PRINT_INST("fixed advance pc\n");
+            address_ += program_.read<uint16_t>();
+          } break;
+          default: {
+            PRINT_INST("other {}\n", op);
+            auto n = standard_opcode_lengths_[op];
+            for (int i = 0; i < n; ++i) {
+              program_.readULEB128();
+            }
+          } break;
+        }
+      }
+    }
+    PRINT_INST(
+        "{:x}: end {:x}\n",
+        ((char*)program_.loc() - s_.debug_line.data),
+        program_end_ - s_.debug_line.data);
+  }
+
+  uint64_t address_ = 0;
+  bool shadow_ = false;
+  bool parsed_ = false;
+  Entry entry_ = {};
+  std::vector<std::string> include_directories_;
+  std::vector<std::string> file_names_;
+  std::vector<uint64_t> file_directory_index_;
+  uint8_t segment_selector_size_ = 0;
+  uint8_t minimum_instruction_length_ = 0;
+  uint8_t maximum_operations_per_instruction_ = 0;
+  int8_t line_base_ = 0;
+  uint8_t line_range_ = 0;
+  uint8_t opcode_base_ = 0;
+  bool default_is_stmt_ = false;
+  CheckedLexer program_ = {nullptr};
+  char* program_end_ = nullptr;
+  uint64_t header_length_ = 0;
+  uint64_t length_ = 0;
+  bool is_64bit_ = false;
+  std::vector<uint8_t> standard_opcode_lengths_;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  Sections& s_;
+  uint64_t offset_;
+  uint64_t start_address_ = 0;
+  RangeTable<uint64_t> program_index_;
+  std::vector<RangeTable<Entry>> all_programs_;
+  RangeTable<Entry> ranges_;
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/mem_file.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/mem_file.h
new file mode 100644
index 0000000000000000000000000000000000000000..15f6821012f0306007a061499867a8d1d97b7ec2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/mem_file.h
@@ -0,0 +1,164 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <c10/util/error.h>
+#include <elf.h>
+#include <fcntl.h>
+#include <fmt/format.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <torch/csrc/profiler/unwind/lexer.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <unistd.h>
+#include <cerrno>
+#include <cstdio>
+#include <cstring>
+
+namespace torch::unwind {
+
+struct Section {
+  char* data = nullptr;
+  size_t size = 0;
+  const char* string(size_t offset) {
+    return lexer(offset).readCString();
+  }
+  CheckedLexer lexer(size_t offset) {
+    return CheckedLexer(data + offset, data, data + size);
+  }
+};
+
+/// Memory maps a file into the address space read-only, and manages the
+/// lifetime of the mapping. Here are a few use cases:
+/// 1. Used in the loader to read in initial image, and to inspect
+// ELF files for dependencies before calling dlopen.
+///
+/// 2. Used in unity to load the elf file.
+struct MemFile {
+  explicit MemFile(const char* filename_)
+      : fd_(open(filename_, O_RDONLY)), name_(filename_) {
+    UNWIND_CHECK(
+        fd_ != -1,
+        "failed to open {}: {}",
+        filename_,
+        c10::utils::str_error(errno));
+    struct stat s{};
+    if (-1 == fstat(fd_, &s)) {
+      close(fd_); // destructors don't run during exceptions
+      UNWIND_CHECK(
+          false,
+          "failed to stat {}: {}",
+          filename_,
+          c10::utils::str_error(errno));
+    }
+    n_bytes_ = s.st_size;
+    UNWIND_CHECK(
+        n_bytes_ > sizeof(Elf64_Ehdr), "empty shared library: {}", filename_);
+    mem_ = (char*)mmap(nullptr, n_bytes_, PROT_READ, MAP_SHARED, fd_, 0);
+    if (MAP_FAILED == mem_) {
+      close(fd_);
+      UNWIND_CHECK(
+          false,
+          "failed to mmap {}: {}",
+          filename_,
+          c10::utils::str_error(errno));
+    }
+    ehdr_ = (Elf64_Ehdr*)mem_;
+#define ELF_CHECK(cond) UNWIND_CHECK(cond, "not an ELF file: {}", filename_)
+    ELF_CHECK(ehdr_->e_ident[EI_MAG0] == ELFMAG0);
+    ELF_CHECK(ehdr_->e_ident[EI_MAG1] == ELFMAG1);
+    ELF_CHECK(ehdr_->e_ident[EI_MAG2] == ELFMAG2);
+    ELF_CHECK(ehdr_->e_ident[EI_MAG3] == ELFMAG3);
+    ELF_CHECK(ehdr_->e_ident[EI_CLASS] == ELFCLASS64);
+    ELF_CHECK(ehdr_->e_ident[EI_VERSION] == EV_CURRENT);
+    ELF_CHECK(ehdr_->e_version == EV_CURRENT);
+    ELF_CHECK(ehdr_->e_machine == EM_X86_64);
+#undef ELF_CHECK
+    UNWIND_CHECK(
+        ehdr_->e_shoff + sizeof(Elf64_Shdr) * ehdr_->e_shnum <= n_bytes_,
+        "invalid section header table {} {} {}",
+        ehdr_->e_shoff + sizeof(Elf64_Shdr) * ehdr_->e_shnum,
+        n_bytes_,
+        ehdr_->e_shnum);
+    shdr_ = (Elf64_Shdr*)(mem_ + ehdr_->e_shoff);
+    UNWIND_CHECK(
+        ehdr_->e_shstrndx < ehdr_->e_shnum, "invalid strtab section offset");
+    auto& strtab_hdr = shdr_[ehdr_->e_shstrndx];
+    strtab_ = getSection(strtab_hdr);
+  }
+
+  MemFile(const MemFile&) = delete;
+  MemFile(MemFile&&) = delete;
+  MemFile& operator=(const MemFile&) = delete;
+  MemFile& operator=(MemFile&&) = delete;
+  [[nodiscard]] const char* data() const {
+    return (const char*)mem_;
+  }
+
+  /// Returns whether or not the file descriptor
+  /// of the underlying file is valid.
+  int valid() {
+    return fcntl(fd_, F_GETFD) != -1 || errno != EBADF;
+  }
+
+  ~MemFile() {
+    if (mem_) {
+      munmap((void*)mem_, n_bytes_);
+    }
+    if (fd_ >= 0) {
+      close(fd_);
+    }
+  }
+
+  /// Returns the size of the underlying file defined by the `MemFile`
+  size_t size() {
+    return n_bytes_;
+  }
+  [[nodiscard]] int fd() const {
+    return fd_;
+  }
+
+  Section getSection(const Elf64_Shdr& shdr) {
+    UNWIND_CHECK(shdr.sh_offset + shdr.sh_size <= n_bytes_, "invalid section");
+    return Section{mem_ + shdr.sh_offset, shdr.sh_size};
+  }
+
+  Section getSection(const char* name, bool optional) {
+    for (int i = 0; i < ehdr_->e_shnum; i++) {
+      if (strcmp(strtab_.string(shdr_[i].sh_name), name) == 0) {
+        return getSection(shdr_[i]);
+      }
+    }
+    UNWIND_CHECK(optional, "{} has no section {}", name_, name);
+    return Section{nullptr, 0};
+  }
+
+  Section strtab() {
+    return strtab_;
+  }
+
+ private:
+  template <typename T>
+  T* load(size_t offset) {
+    UNWIND_CHECK(offset < n_bytes_, "out of range");
+    return (T*)(mem_ + offset);
+  }
+  int fd_;
+  char* mem_{nullptr};
+  size_t n_bytes_{0};
+  std::string name_;
+  Elf64_Ehdr* ehdr_;
+  Elf64_Shdr* shdr_;
+  Section strtab_ = {nullptr, 0};
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/range_table.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/range_table.h
new file mode 100644
index 0000000000000000000000000000000000000000..a08bf00133fe6c4aa5f9fa7ef528b99d1c6b547b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/range_table.h
@@ -0,0 +1,78 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <algorithm>
+#include <memory>
+#include <optional>
+#include <vector>
+
+namespace torch::unwind {
+template <typename T>
+struct RangeTable {
+  RangeTable() {
+    // guarantee that lower_bound[-1] is always valid
+    addresses_.push_back(0);
+    payloads_.emplace_back(std::nullopt);
+  }
+  void add(uint64_t address, std::optional<T> payload, bool sorted) {
+    if (addresses_.back() > address) {
+      UNWIND_CHECK(!sorted, "expected addresses to be sorted");
+      sorted_ = false;
+    }
+    addresses_.push_back(address);
+    payloads_.emplace_back(std::move(payload));
+  }
+  std::optional<T> find(uint64_t address) {
+    maybeSort();
+    auto it = std::upper_bound(addresses_.begin(), addresses_.end(), address);
+    return payloads_.at(it - addresses_.begin() - 1);
+  }
+  void dump() {
+    for (size_t i = 0; i < addresses_.size(); i++) {
+      fmt::print("{} {:x}: {}\n", i, addresses_[i], payloads_[i] ? "" : "END");
+    }
+  }
+  size_t size() const {
+    return addresses_.size();
+  }
+  uint64_t back() {
+    maybeSort();
+    return addresses_.back();
+  }
+
+ private:
+  void maybeSort() {
+    if (sorted_) {
+      return;
+    }
+    std::vector<uint64_t> indices;
+    indices.reserve(addresses_.size());
+    for (size_t i = 0; i < addresses_.size(); i++) {
+      indices.push_back(i);
+    }
+    std::sort(indices.begin(), indices.end(), [&](uint64_t a, uint64_t b) {
+      return addresses_[a] < addresses_[b] ||
+          (addresses_[a] == addresses_[b] &&
+           bool(payloads_[a]) < bool(payloads_[b]));
+    });
+    std::vector<uint64_t> addresses;
+    std::vector<std::optional<T>> payloads;
+    addresses.reserve(addresses_.size());
+    payloads.reserve(addresses_.size());
+    for (auto i : indices) {
+      addresses.push_back(addresses_[i]);
+      payloads.push_back(payloads_[i]);
+    }
+    addresses_ = std::move(addresses);
+    payloads_ = std::move(payloads);
+    sorted_ = true;
+  }
+  bool sorted_ = true;
+  std::vector<uint64_t> addresses_;
+  std::vector<std::optional<T>> payloads_;
+};
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/sections.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/sections.h
new file mode 100644
index 0000000000000000000000000000000000000000..f3f4f63b2d4e92d14ff361c722df2b66f0009d73
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/sections.h
@@ -0,0 +1,125 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <cxxabi.h>
+#include <elf.h>
+#include <torch/csrc/profiler/unwind/dwarf_enums.h>
+#include <torch/csrc/profiler/unwind/dwarf_symbolize_enums.h>
+#include <torch/csrc/profiler/unwind/mem_file.h>
+#include <torch/csrc/profiler/unwind/range_table.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <cstdint>
+
+namespace torch::unwind {
+
+static std::string demangle(const std::string& mangled_name) {
+  int status = 0;
+  char* realname =
+      abi::__cxa_demangle(mangled_name.c_str(), nullptr, nullptr, &status);
+  if (status == 0) {
+    std::string demangled_name(realname);
+    // NOLINTNEXTLINE(cppcoreguidelines-no-malloc)
+    free(realname);
+    return demangled_name;
+  } else {
+    return mangled_name;
+  }
+}
+
+struct Sections {
+  Sections() = default;
+  void parse(const char* name) {
+    library_ = std::make_unique<MemFile>(name);
+    strtab = library_->getSection(".strtab", false);
+
+    symtab = library_->getSection(".symtab", true);
+    debug_info = library_->getSection(".debug_info", true);
+    if (debug_info.size > 0) {
+      debug_abbrev = library_->getSection(".debug_abbrev", false);
+      debug_str = library_->getSection(".debug_str", false);
+      debug_line = library_->getSection(".debug_line", false);
+      // dwarf 5
+      debug_line_str = library_->getSection(".debug_line_str", true);
+      debug_rnglists = library_->getSection(".debug_rnglists", true);
+      debug_addr = library_->getSection(".debug_addr", true);
+      // dwarf 4
+      debug_ranges = library_->getSection(".debug_ranges", true);
+    }
+    parseSymtab();
+  }
+
+  Section debug_info;
+  Section debug_abbrev;
+  Section debug_str;
+  Section debug_line;
+  Section debug_line_str;
+  Section debug_rnglists;
+  Section debug_ranges;
+  Section debug_addr;
+  Section symtab;
+  Section strtab;
+
+  const char* readString(CheckedLexer& data, uint64_t encoding, bool is_64bit) {
+    switch (encoding) {
+      case DW_FORM_string: {
+        return data.readCString();
+      }
+      case DW_FORM_strp: {
+        return debug_str.string(readSegmentOffset(data, is_64bit));
+      }
+      case DW_FORM_line_strp: {
+        return debug_line_str.string(readSegmentOffset(data, is_64bit));
+      }
+      default:
+        UNWIND_CHECK(false, "unsupported string encoding {:x}", encoding);
+    }
+  }
+
+  uint64_t readSegmentOffset(CheckedLexer& data, bool is_64bit) {
+    return is_64bit ? data.read<uint64_t>() : data.read<uint32_t>();
+  }
+
+  std::optional<uint64_t> findDebugInfoOffset(uint64_t address) {
+    return debug_info_offsets_.find(address);
+  }
+  size_t compilationUnitCount() {
+    return debug_info_offsets_.size() / 2;
+  }
+  void addDebugInfoRange(
+      uint64_t start,
+      uint64_t end,
+      uint64_t debug_info_offset) {
+    debug_info_offsets_.add(start, debug_info_offset, false);
+    debug_info_offsets_.add(end, std::nullopt, false);
+  }
+  std::optional<std::string> findSubprogramName(uint64_t address) {
+    if (auto e = symbol_table_.find(address)) {
+      return demangle(strtab.string(*e));
+    }
+    return std::nullopt;
+  }
+
+ private:
+  void parseSymtab() {
+    auto L = symtab.lexer(0);
+    char* end = symtab.data + symtab.size;
+    while (L.loc() < end) {
+      auto symbol = L.read<Elf64_Sym>();
+      if (symbol.st_shndx == SHN_UNDEF ||
+          ELF64_ST_TYPE(symbol.st_info) != STT_FUNC) {
+        continue;
+      }
+      symbol_table_.add(symbol.st_value, symbol.st_name, false);
+      symbol_table_.add(symbol.st_value + symbol.st_size, std::nullopt, false);
+    }
+  }
+
+  std::unique_ptr<MemFile> library_;
+  RangeTable<uint64_t> debug_info_offsets_;
+  RangeTable<uint64_t> symbol_table_;
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwind.h
new file mode 100644
index 0000000000000000000000000000000000000000..5dd90ecbbebb016b28fc3c3581a5e3a4c8f36c6a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwind.h
@@ -0,0 +1,48 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/macros/Export.h>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace torch::unwind {
+// gather current stack, relatively fast.
+// gets faster once the cache of program counter locations is warm.
+TORCH_API std::vector<void*> unwind();
+
+struct Frame {
+  std::string filename;
+  std::string funcname;
+  uint64_t lineno;
+};
+
+enum class Mode { addr2line, fast, dladdr };
+
+// note: symbolize is really slow
+// it will launch an addr2line process that has to parse dwarf
+// information from the libraries that frames point into.
+// Callers should first batch up all the unique void* pointers
+// across a number of unwind states and make a single call to
+// symbolize.
+TORCH_API std::vector<Frame> symbolize(
+    const std::vector<void*>& frames,
+    Mode mode);
+
+// returns path to the library, and the offset of the addr inside the library
+TORCH_API std::optional<std::pair<std::string, uint64_t>> libraryFor(
+    void* addr);
+
+struct Stats {
+  size_t hits = 0;
+  size_t misses = 0;
+  size_t unsupported = 0;
+  size_t resets = 0;
+};
+Stats stats();
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwind_error.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwind_error.h
new file mode 100644
index 0000000000000000000000000000000000000000..9a468e702c2379d2b25d6bd1c251daa517f1cf1f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwind_error.h
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <fmt/format.h>
+#include <optional>
+#include <stdexcept>
+
+namespace torch::unwind {
+
+struct UnwindError : public std::runtime_error {
+  using std::runtime_error::runtime_error;
+};
+
+#define UNWIND_CHECK(cond, fmtstring, ...)                          \
+  do {                                                              \
+    if (!(cond)) {                                                  \
+      throw unwind::UnwindError(fmt::format(                        \
+          "{}:{}: " fmtstring, __FILE__, __LINE__, ##__VA_ARGS__)); \
+    }                                                               \
+  } while (0)
+
+// #define LOG_INFO(...) fmt::print(__VA_ARGS__)
+#define LOG_INFO(...)
+
+// #define PRINT_INST(...) LOG_INFO(__VA_ARGS__)
+#define PRINT_INST(...)
+
+// #define PRINT_LINE_TABLE(...) LOG_INFO(__VA_ARGS__)
+#define PRINT_LINE_TABLE(...)
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwinder.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwinder.h
new file mode 100644
index 0000000000000000000000000000000000000000..2f284e849fe2721fcc0ddb929281d8c8e8844409
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/unwind/unwinder.h
@@ -0,0 +1,88 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/profiler/unwind/action.h>
+#include <torch/csrc/profiler/unwind/unwind_error.h>
+#include <cstdint>
+#include <limits>
+
+namespace torch::unwind {
+
+// Architecture-neutral names: pc (program counter / return address),
+// fp (frame pointer: x86 RBP, aarch64 x29), sp (stack pointer).
+struct UnwindState {
+  int64_t pc, fp, sp;
+};
+
+struct Unwinder {
+  Unwinder(Action cfa, Action ret, Action fp)
+      : kind_(ret.kind == A_UNDEFINED ? END : STANDARD),
+        reg_(cfa.reg),
+        off_(cfa.data),
+        ret_off_(ret.data),
+        fp_off_(
+            fp.kind == A_UNDEFINED ? std::numeric_limits<int64_t>::max()
+                                   : fp.data),
+        deref_(cfa.kind == A_REG_PLUS_DATA_DEREF) {
+    check(cfa.reg == D_STACK_PTR || cfa.reg == D_FRAME_PTR);
+    check(ret.kind == A_UNDEFINED || ret.kind == A_LOAD_CFA_OFFSET);
+    if (cfa.kind == A_REG_PLUS_DATA) {
+      check(fp.kind == A_LOAD_CFA_OFFSET || fp.kind == A_UNDEFINED);
+    } else if (cfa.kind == A_REG_PLUS_DATA_DEREF) {
+      if (fp.kind == A_REG_PLUS_DATA_DEREF) {
+        check(fp.reg == cfa.reg);
+        fp_off_ -= cfa.data;
+      } else {
+        check(fp.kind == A_UNDEFINED);
+      }
+    } else {
+      check(false);
+    }
+  }
+  void check(bool cond) {
+    if (!cond) {
+      throw UnwindError("Unwinding actions do not follow supported patterns");
+    }
+  }
+  bool terminator() const {
+    return kind_ != STANDARD;
+  }
+  bool isUnknown() const {
+    return kind_ == UNKNOWN;
+  }
+  // unwinder representing some pattern unsupported in
+  // current implementation
+  static Unwinder unknown() {
+    return Unwinder();
+  }
+  UnwindState run(const UnwindState& cur) const {
+    UnwindState r = cur;
+    r.sp = (reg_ == D_STACK_PTR ? cur.sp : cur.fp) + off_;
+    r.fp = fp_off_ == std::numeric_limits<int64_t>::max()
+        ? cur.fp
+        // NOLINTNEXTLINE(performance-no-int-to-ptr)
+        : *(int64_t*)(r.sp + fp_off_);
+    if (deref_) {
+      // NOLINTNEXTLINE(performance-no-int-to-ptr)
+      r.sp = *(int64_t*)r.sp;
+    }
+    // NOLINTNEXTLINE(performance-no-int-to-ptr)
+    r.pc = *(int64_t*)(r.sp + ret_off_);
+
+    return r;
+  }
+
+ private:
+  Unwinder() : kind_(UNKNOWN), reg_(0), off_(0), ret_off_(0), fp_off_(0) {}
+  enum Kind { STANDARD, END, UNKNOWN } kind_;
+  uint32_t reg_;
+  int64_t off_;
+  int64_t ret_off_;
+  int64_t fp_off_;
+  bool deref_{false};
+};
+
+} // namespace torch::unwind
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/util.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/util.h
new file mode 100644
index 0000000000000000000000000000000000000000..a8e084ed3fa49bd3606f48a15097ce57b789061b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/profiler/util.h
@@ -0,0 +1,215 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <list>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <ATen/record_function.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/hash.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/jit/frontend/source_range.h>
+#include <optional>
+
+// TODO: replace with pytorch/rfcs#43 when it is ready.
+#define SOFT_ASSERT(cond, ...)                         \
+  [&]() -> bool {                                      \
+    if (C10_UNLIKELY(!(cond))) {                       \
+      torch::profiler::impl::logSoftAssert(            \
+          __func__,                                    \
+          __FILE__,                                    \
+          static_cast<uint32_t>(__LINE__),             \
+          #cond,                                       \
+          ::c10::str(__VA_ARGS__));                    \
+      if (torch::profiler::impl::softAssertRaises()) { \
+        TORCH_INTERNAL_ASSERT(cond, __VA_ARGS__);      \
+      } else {                                         \
+        TORCH_WARN_ONCE(__VA_ARGS__);                  \
+      }                                                \
+      return false;                                    \
+    }                                                  \
+    return true;                                       \
+  }()
+
+namespace torch::profiler::impl {
+TORCH_API bool softAssertRaises();
+TORCH_API void setSoftAssertRaises(std::optional<bool> value);
+TORCH_API void logSoftAssert(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* cond,
+    const char* args);
+inline void logSoftAssert(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* cond,
+    ::c10::detail::CompileTimeEmptyString args) {
+  logSoftAssert(func, file, line, cond, (const char*)args);
+}
+TORCH_API void logSoftAssert(
+    const char* func,
+    const char* file,
+    uint32_t line,
+    const char* cond,
+    const std::string& args);
+
+using shape =
+    std::variant<std::vector<int64_t>, std::vector<std::vector<int64_t>>>;
+constexpr int TENSOR_LIST_DISPLAY_LENGTH_LIMIT = 30;
+
+std::string getNvtxStr(
+    const char* name,
+    int64_t sequence_nr,
+    const std::vector<std::vector<int64_t>>& shapes,
+    at::RecordFunctionHandle op_id = 0,
+    const std::list<std::pair<at::RecordFunctionHandle, int>>& input_op_ids =
+        {});
+
+struct TORCH_API FileLineFunc {
+  std::string filename;
+  size_t line;
+  std::string funcname;
+};
+
+struct TORCH_API SaveNcclMetaConfig {
+  bool truncate;
+  bool introspectMetadata;
+  bool introspectInputs;
+  bool introspectOutputs;
+
+  // Default constructor with default values
+  SaveNcclMetaConfig()
+      : truncate(true),
+        introspectMetadata(true),
+        introspectInputs(false),
+        introspectOutputs(false) {}
+
+  SaveNcclMetaConfig(
+      bool truncate,
+      bool introspectMetadata,
+      bool introspectInputs,
+      bool introspectOutputs)
+      : truncate(truncate),
+        introspectMetadata(introspectMetadata),
+        introspectInputs(introspectInputs),
+        introspectOutputs(introspectOutputs) {}
+};
+
+TORCH_API std::vector<FileLineFunc> prepareCallstack(
+    const std::vector<jit::StackEntry>& cs);
+TORCH_API std::vector<std::string> callstackStr(
+    const std::vector<FileLineFunc>& cs);
+TORCH_API std::string stacksToStr(
+    const std::vector<std::string>& stacks,
+    const char* delim);
+TORCH_API std::vector<std::vector<int64_t>> inputSizes(
+    const at::RecordFunction& fn,
+    const bool flatten_list_enabled = false);
+TORCH_API std::string variantShapesToStr(const std::vector<shape>& shapes);
+TORCH_API std::string shapesToStr(
+    const std::vector<std::vector<int64_t>>& shapes);
+TORCH_API std::string strListToStr(const std::vector<std::string>& types);
+TORCH_API std::string inputOpIdsToStr(
+    const std::list<std::pair<at::RecordFunctionHandle, int>>& input_op_ids);
+TORCH_API std::string ivalueToStr(const c10::IValue& val, bool isString);
+TORCH_API std::string ivalueListToStr(const std::vector<c10::IValue>& list);
+TORCH_API std::vector<std::string> inputTypes(const at::RecordFunction& fn);
+
+std::unordered_map<std::string, c10::IValue> TORCH_API
+saveExtraArgs(const at::RecordFunction& fn);
+std::unordered_map<std::string, std::string> TORCH_API saveNcclMeta(
+    const at::RecordFunction& fn,
+    const SaveNcclMetaConfig& config = SaveNcclMetaConfig());
+int getTensorStartHint(const at::Tensor& t);
+bool checkFunctionOutputsForLogging(const at::RecordFunction& fn);
+bool checkFunctionInputsForLogging(const at::RecordFunction& fn);
+std::pair<bool, std::variant<int, std::vector<int>>> findStartAddrForTensors(
+    const c10::IValue& val);
+uint64_t TORCH_API computeFlops(
+    const std::string& op_name,
+    const std::unordered_map<std::string, c10::IValue>& extra_args);
+
+std::string shapeToStr(const std::vector<int64_t>& shape);
+
+template <typename T>
+class TORCH_API GlobalStateManager {
+ public:
+  static GlobalStateManager& singleton() {
+    /* library-local */ static GlobalStateManager singleton_;
+    return singleton_;
+  }
+
+  static void push(std::shared_ptr<T>&& state) {
+    if (singleton().state_) {
+      LOG(WARNING) << "GlobalStatePtr already exists!";
+    } else {
+      singleton().state_ = std::move(state);
+    }
+  }
+
+  static auto* get() {
+    return singleton().state_.get();
+  }
+
+  static std::shared_ptr<T> pop() {
+    auto out = singleton().state_;
+    singleton().state_.reset();
+    return out;
+  }
+
+ private:
+  GlobalStateManager() = default;
+
+  std::shared_ptr<T> state_;
+};
+
+struct HashCombine {
+  template <typename T0, typename T1>
+  size_t operator()(const std::pair<T0, T1>& i) {
+    return c10::get_hash((*this)(i.first), (*this)(i.second));
+  }
+
+  template <typename... Args>
+  size_t operator()(const std::tuple<Args...>& i) {
+    return c10::get_hash(i);
+  }
+
+  template <typename T>
+  size_t operator()(const T& i) {
+    return c10::get_hash(i);
+  }
+};
+
+#ifdef USE_DISTRIBUTED
+constexpr auto kCommsName = "Collective name";
+constexpr auto kDtype = "dtype";
+constexpr auto kInMsgNelems = "In msg nelems";
+constexpr auto kOutMsgNelems = "Out msg nelems";
+constexpr auto kInSplit = "In split size";
+constexpr auto kOutSplit = "Out split size";
+constexpr auto kGlobalRankStart = "Global rank start";
+constexpr auto kGlobalRankStride = "Global rank stride";
+constexpr auto kGroupSize = "Group size";
+constexpr auto kProcessGroupName = "Process Group Name";
+constexpr auto kProcessGroupDesc = "Process Group Description";
+constexpr auto kGroupRanks = "Process Group Ranks";
+constexpr auto kRank = "Rank";
+constexpr auto kP2pSrc = "Src Rank";
+constexpr auto kP2pDst = "Dst Rank";
+constexpr auto kSeqNum = "Seq";
+constexpr auto kInTensorsStart = "Input Tensors start";
+constexpr auto kOutTensorsStart = "Output Tensors start";
+constexpr auto kIsAsynchronizedOp = "Is asynchronized op";
+#endif // USE_DISTRIBUTED
+
+} // namespace torch::profiler::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/python_dimname.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/python_dimname.h
new file mode 100644
index 0000000000000000000000000000000000000000..34226e53fdfb1dcd5d017c44160b0c1ac8e133db
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/python_dimname.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <ATen/Dimname.h>
+#include <torch/csrc/python_headers.h>
+
+at::Dimname THPDimname_parse(PyObject* obj);
+bool THPUtils_checkDimname(PyObject* obj);
+bool THPUtils_checkDimnameList(PyObject* obj);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/python_headers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/python_headers.h
new file mode 100644
index 0000000000000000000000000000000000000000..1b90571fdc027a49d5c9b0828d2d0ec4a96776cc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/python_headers.h
@@ -0,0 +1,30 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+// workaround for https://github.com/python/cpython/pull/23326
+#include <cmath>
+#include <complex>
+// workaround for Python 2 issue: https://bugs.python.org/issue17120
+// NOTE: It looks like this affects Python 3 as well.
+#pragma push_macro("_XOPEN_SOURCE")
+#pragma push_macro("_POSIX_C_SOURCE")
+#undef _XOPEN_SOURCE
+#undef _POSIX_C_SOURCE
+
+#include <Python.h>
+#include <frameobject.h>
+#include <structseq.h>
+
+#pragma pop_macro("_XOPEN_SOURCE")
+#pragma pop_macro("_POSIX_C_SOURCE")
+
+#ifdef copysign
+#undef copysign
+#endif
+
+#if PY_MAJOR_VERSION < 3
+#error "Python 2 has reached end-of-life and is no longer supported by PyTorch."
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/serialization.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/serialization.h
new file mode 100644
index 0000000000000000000000000000000000000000..7a39f8a320d21183e6a09137e26365ac0d838839
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/serialization.h
@@ -0,0 +1,32 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef THP_SERIALIZATION_INC
+#define THP_SERIALIZATION_INC
+
+#include <c10/core/StorageImpl.h>
+#include <c10/util/intrusive_ptr.h>
+template <class io>
+void doRead(io fildes, void* buf, size_t nbytes);
+
+template <class io>
+void doWrite(io fildes, void* buf, size_t nbytes);
+
+// Note that this takes a mutable storage because it may pass through
+// to at::from_blob.
+template <class io>
+void THPStorage_writeFileRaw(
+    c10::StorageImpl* self,
+    io fd,
+    bool save_size,
+    uint64_t element_size);
+
+template <class io>
+c10::intrusive_ptr<c10::StorageImpl> THPStorage_readFileRaw(
+    io fd,
+    c10::intrusive_ptr<c10::StorageImpl> storage,
+    uint64_t element_size);
+
+#endif
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/shim_conversion_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/shim_conversion_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..1ef8f8ac9329b2640f850026c66b847c755a9f56
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/shim_conversion_utils.h
@@ -0,0 +1,27 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/Exception.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/stable/c/shim.h>
+
+#include <vector>
+
+inline std::vector<StableIValue>* list_handle_to_list_pointer(
+    StableListHandle handle) {
+  return reinterpret_cast<std::vector<StableIValue>*>(handle);
+}
+
+inline StableListHandle list_pointer_to_list_handle(
+    std::vector<StableIValue>* list_ptr) {
+  return reinterpret_cast<StableListHandle>(list_ptr);
+}
+
+inline StableListHandle new_list_handle(std::vector<StableIValue>&& list) {
+  std::vector<StableIValue>* new_list = new std::vector<StableIValue>(list);
+  return list_pointer_to_list_handle(new_list);
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/accelerator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/accelerator.h
new file mode 100644
index 0000000000000000000000000000000000000000..d2a35e58982b84057d85bc047b0d1491ade81071
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/accelerator.h
@@ -0,0 +1,117 @@
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/shim_utils.h>
+
+#include <memory>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable, accelerator)
+
+using DeleterFnPtr = void (*)(void*);
+
+namespace {
+inline void delete_device_guard(void* ptr) {
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_delete_device_guard(reinterpret_cast<DeviceGuardHandle>(ptr)));
+}
+
+} // namespace
+
+// This is bigger than DeviceIndex in c10/core/Device.h but it is the type we
+// can converge on in this world as DeviceIndex in libtorch is not stable.
+/**
+ * @brief Device index type for stable ABI.
+ *
+ * Minimum compatible version: PyTorch 2.9.
+ */
+using DeviceIndex = int32_t;
+
+using StreamId = int64_t; // this is from c10/core/Stream.h
+
+/**
+ * @brief A stable ABI version of c10::DeviceGuard.
+ *
+ * RAII class that sets the current device to the specified device index
+ * on construction and restores the previous device on destruction.
+ *
+ * Minimum compatible version: PyTorch 2.9.
+ */
+class DeviceGuard {
+ public:
+  /// \private
+  explicit DeviceGuard() = delete;
+
+  /**
+   * @brief Constructs a DeviceGuard that sets the current device.
+   *
+   * @param device_index The device index to set as the current device.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  explicit DeviceGuard(DeviceIndex device_index)
+      : guard_(nullptr, delete_device_guard) {
+    DeviceGuardHandle ptr = nullptr;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_create_device_guard(device_index, &ptr));
+    guard_.reset(ptr);
+  }
+
+  /**
+   * @brief Changes the current device to the specified device index.
+   *
+   * @param device_index The new device index to set.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  void set_index(DeviceIndex device_index) {
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_device_guard_set_index(guard_.get(), device_index));
+  }
+
+ private:
+  std::unique_ptr<DeviceGuardOpaque, DeleterFnPtr> guard_;
+};
+
+class Stream {
+ public:
+  explicit Stream() = delete;
+
+  // Construct a stable::Stream from a StreamHandle
+  // Steals ownership from the StreamHandle
+  explicit Stream(StreamHandle stream)
+      : stream_(stream, [](StreamHandle stream) {
+          TORCH_ERROR_CODE_CHECK(aoti_torch_delete_stream(stream));
+        }) {}
+
+  StreamId id() const {
+    StreamId stream_id;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_stream_id(stream_.get(), &stream_id));
+    return stream_id;
+  }
+
+ private:
+  std::shared_ptr<StreamOpaque> stream_;
+};
+
+inline Stream getCurrentStream(DeviceIndex device_index) {
+  StreamHandle stream = nullptr;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_stream(device_index, &stream));
+  return Stream(stream);
+}
+
+/**
+ * @brief Gets the current device index.
+ *
+ * Returns the index of the currently active device for the accelerator.
+ *
+ * @return The current device index.
+ *
+ * Minimum compatible version: PyTorch 2.9.
+ */
+inline DeviceIndex getCurrentDeviceIndex() {
+  DeviceIndex device_index;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_device_index(&device_index));
+  return device_index;
+}
+
+HIDDEN_NAMESPACE_END(torch, stable, accelerator)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/c/shim.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/c/shim.h
new file mode 100644
index 0000000000000000000000000000000000000000..3844a3d5cc21df5163bd1b672019c0e13a070965
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/c/shim.h
@@ -0,0 +1,230 @@
+#ifndef STABLE_TORCH_SHIM
+#define STABLE_TORCH_SHIM
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+
+#include <torch/csrc/stable/version.h>
+
+// This header defines stable C API extensions for backward/forward
+// compatibility when calling ATen operations through the dispatcher.
+//
+// This is separate from the main AOTI shim to provide versioning capabilities
+// for schema changes in native ATen functions.
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+// Has the same semantic as aoti_torch_call_dispatcher, but takes an
+// additional argument for the extension build version. This is
+// needed for backward compatibility when calling native functions via
+// the dispatcher. The caller should pass in the libtorch version the
+// extension is building with (NOT target version).
+AOTI_TORCH_EXPORT AOTITorchError torch_call_dispatcher(
+    const char* opName,
+    const char* overloadName,
+    StableIValue* stack,
+    uint64_t extension_build_version);
+
+// Version-aware variant of aoti_torch_library_impl that takes an
+// extension_build_version parameter for backward compatibility
+AOTI_TORCH_EXPORT AOTITorchError torch_library_impl(
+    TorchLibraryHandle self,
+    const char* name,
+    void (*fn)(StableIValue*, uint64_t, uint64_t),
+    uint64_t extension_build_version);
+
+struct StableListOpaque;
+using StableListHandle = StableListOpaque*;
+
+// returns an owning reference of a StableList. callee is responsible for
+// freeing memory.
+AOTI_TORCH_EXPORT AOTITorchError
+torch_new_list_reserve_size(size_t size, StableListHandle* ret);
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_list_size(StableListHandle list_handle, size_t* size);
+
+AOTI_TORCH_EXPORT AOTITorchError torch_list_get_item(
+    StableListHandle list_handle,
+    size_t index,
+    StableIValue* element);
+
+AOTI_TORCH_EXPORT AOTITorchError torch_list_set_item(
+    StableListHandle list_handle,
+    size_t index,
+    StableIValue element);
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_list_push_back(StableListHandle list_handle, StableIValue element);
+
+// deletes the underlying list referenced by list_handle
+AOTI_TORCH_EXPORT AOTITorchError
+torch_delete_list(StableListHandle list_handle);
+
+// Helper function to parse device string using c10::Device
+// Returns device type and index via output parameters
+AOTI_TORCH_EXPORT AOTITorchError torch_parse_device_string(
+    const char* device_string,
+    uint32_t* out_device_type,
+    int32_t* out_device_index);
+
+// Parallel utility APIs for stable ABI
+// Function pointer type for parallel_for callback
+// The callback receives begin and end indices for a range to process
+typedef void (*ParallelFunc)(int64_t begin, int64_t end, void* ctx);
+
+AOTI_TORCH_EXPORT AOTITorchError torch_parallel_for(
+    int64_t begin,
+    int64_t end,
+    int64_t grain_size,
+    ParallelFunc func,
+    void* ctx);
+
+// Get the current thread index in a parallel region
+// Returns 0 if not in a parallel region
+AOTI_TORCH_EXPORT AOTITorchError torch_get_thread_idx(uint32_t* out_thread_idx);
+
+// Get the number of threads for the parallel backend
+AOTI_TORCH_EXPORT AOTITorchError
+torch_get_num_threads(uint32_t* out_num_threads);
+
+// Get a pointer to the underlying storage data
+AOTI_TORCH_EXPORT AOTITorchError torch_get_mutable_data_ptr(
+    AtenTensorHandle tensor,
+    void** ret_data_ptr // returns borrowed reference
+);
+
+AOTI_TORCH_EXPORT AOTITorchError torch_get_const_data_ptr(
+    AtenTensorHandle tensor,
+    const void** ret_data_ptr // returns borrowed reference
+);
+
+struct StringOpaque;
+using StringHandle = StringOpaque*;
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_new_string_handle(const char* data, size_t length, StringHandle* handle);
+
+AOTI_TORCH_EXPORT AOTITorchError torch_delete_string(StringHandle handle);
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_string_length(StringHandle handle, size_t* length);
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_string_c_str(StringHandle handle, const char** data);
+
+#ifdef USE_CUDA
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_get_current_cuda_blas_handle(void** ret_handle);
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_set_current_cuda_stream(void* stream, int32_t device_index);
+
+AOTI_TORCH_EXPORT AOTITorchError torch_get_cuda_stream_from_pool(
+    bool isHighPriority,
+    int32_t device_index,
+    void** ret_stream);
+
+AOTI_TORCH_EXPORT AOTITorchError
+torch_cuda_stream_synchronize(void* stream, int32_t device_index);
+
+// Wrapper around c10_cuda_check_implementation that captures the error message
+// without propagating the exception. The caller must free error_msg using
+// torch_c10_cuda_free_error_msg if it is non-null.
+AOTI_TORCH_EXPORT AOTITorchError torch_c10_cuda_check_msg(
+    int32_t err,
+    const char* filename,
+    const char* function_name,
+    uint32_t line_number,
+    bool include_device_assertions,
+    char** error_msg);
+
+// Free error message allocated by torch_c10_cuda_check_msg
+AOTI_TORCH_EXPORT void torch_c10_cuda_free_error_msg(char* error_msg);
+
+#endif // USE_CUDA
+
+// Set requires_grad on a tensor
+AOTI_TORCH_EXPORT AOTITorchError
+torch_set_requires_grad(AtenTensorHandle tensor, bool requires_grad);
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+/**
+ * The beginning of all shims added in 2.11.0 onwards.
+ */
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+
+// Shims for the a few dtypes not already in
+// torch/csrc/inductor/aoti_torch/c/shim.h
+AOTI_TORCH_EXPORT int32_t torch_dtype_float8_e8m0fnu();
+AOTI_TORCH_EXPORT int32_t torch_dtype_float4_e2m1fn_x2();
+
+// Creates a tensor from an existing data blob with an optional deleter.
+// The deleter receives both the data pointer and a caller-supplied context
+// pointer, which allows passing capturing lambdas across the C ABI boundary
+// by heap-allocating the callable and passing it as deleter_ctx.
+AOTI_TORCH_EXPORT AOTITorchError torch_from_blob(
+    void* data,
+    int64_t ndim,
+    const int64_t* sizes_ptr,
+    const int64_t* strides_ptr,
+    int64_t storage_offset,
+    int32_t dtype,
+    int32_t device_type,
+    int32_t device_index,
+    AtenTensorHandle* ret, // returns new reference
+    int32_t layout,
+    const uint8_t* opaque_metadata,
+    int64_t opaque_metadata_size,
+    void (*deleter)(void* data, void* ctx),
+    void* deleter_ctx);
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+
+/**
+ * The beginning of all shims added in 2.12.0 onwards.
+ */
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
+// Tag getter functions for ABI-stable tag passing.  By hiding these behind
+// functions, the precise enum ordinal is NOT part of the ABI contract.
+AOTI_TORCH_EXPORT int32_t torch_tag_core();
+AOTI_TORCH_EXPORT int32_t torch_tag_cudagraph_unsafe();
+AOTI_TORCH_EXPORT int32_t torch_tag_data_dependent_output();
+AOTI_TORCH_EXPORT int32_t torch_tag_dynamic_output_shape();
+AOTI_TORCH_EXPORT int32_t torch_tag_flexible_layout();
+AOTI_TORCH_EXPORT int32_t torch_tag_generated();
+AOTI_TORCH_EXPORT int32_t torch_tag_inplace_view();
+AOTI_TORCH_EXPORT int32_t torch_tag_maybe_aliasing_or_mutating();
+AOTI_TORCH_EXPORT int32_t torch_tag_needs_contiguous_strides();
+AOTI_TORCH_EXPORT int32_t torch_tag_needs_exact_strides();
+AOTI_TORCH_EXPORT int32_t torch_tag_needs_fixed_stride_order();
+AOTI_TORCH_EXPORT int32_t torch_tag_nondeterministic_bitwise();
+AOTI_TORCH_EXPORT int32_t torch_tag_nondeterministic_seeded();
+AOTI_TORCH_EXPORT int32_t torch_tag_out_variant();
+AOTI_TORCH_EXPORT int32_t torch_tag_pointwise();
+AOTI_TORCH_EXPORT int32_t torch_tag_pt2_compliant_tag();
+AOTI_TORCH_EXPORT int32_t torch_tag_reduction();
+AOTI_TORCH_EXPORT int32_t torch_tag_view_copy();
+
+// Stable corollary to torch::Library method m.def() with tags.
+// Tags are passed as int32_t values obtained from torch_tag_*() getters,
+// not raw enum ordinals, so the ABI is stable across versions.
+AOTI_TORCH_EXPORT AOTITorchError torch_library_def_with_tags(
+    TorchLibraryHandle self,
+    const char* schema,
+    const int32_t* tags,
+    int32_t num_tags);
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
+#ifdef __cplusplus
+} // extern "C"
+#endif
+
+#endif // STABLE_TORCH_SHIM
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device.h
new file mode 100644
index 0000000000000000000000000000000000000000..223e3320a4fd341ba0b388b1dfe08ff38f2d94df
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device.h
@@ -0,0 +1,4 @@
+#pragma once
+
+#include <torch/csrc/stable/device_inl.h>
+#include <torch/csrc/stable/device_struct.h>
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device_inl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device_inl.h
new file mode 100644
index 0000000000000000000000000000000000000000..8c9685f0d7da7822ddca6fbe80eab065895b8933
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device_inl.h
@@ -0,0 +1,41 @@
+#pragma once
+
+// This file implements device.h. We separated out the Device struct so that
+// other files can depend on the Device struct (like stableivalue_conversions.h)
+// and the implementations of the Device methods can depend on APIs in
+// stableivalue_conversions.h without circular dependencies.
+
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/device_struct.h>
+#include <torch/csrc/stable/stableivalue_conversions.h>
+#include <torch/csrc/stable/version.h>
+#include <torch/headeronly/core/DeviceType.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/shim_utils.h>
+
+#include <string>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable)
+
+using DeviceType = torch::headeronly::DeviceType;
+using DeviceIndex = torch::stable::accelerator::DeviceIndex;
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+inline Device::Device(const std::string& device_string) {
+  uint32_t device_type;
+  int32_t device_index;
+
+  TORCH_ERROR_CODE_CHECK(torch_parse_device_string(
+      device_string.c_str(), &device_type, &device_index));
+
+  DeviceType dt = torch::stable::detail::to<DeviceType>(
+      torch::stable::detail::from(device_type));
+  DeviceIndex di = static_cast<DeviceIndex>(device_index);
+
+  *this = Device(dt, di);
+}
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+HIDDEN_NAMESPACE_END(torch, stable)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device_struct.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device_struct.h
new file mode 100644
index 0000000000000000000000000000000000000000..e7c478578cc09e198c2f6bc8d7d4dbfed8b4957b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/device_struct.h
@@ -0,0 +1,189 @@
+#pragma once
+
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/version.h>
+#include <torch/headeronly/core/DeviceType.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Exception.h>
+#include <torch/headeronly/util/shim_utils.h>
+
+#include <string>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable)
+
+using DeviceType = torch::headeronly::DeviceType;
+using DeviceIndex = torch::stable::accelerator::DeviceIndex;
+
+// The torch::stable::Device class is an approximate copy of c10::Device.
+// It has some slight modifications:
+// 1. TORCH_INTERNAL_ASSERT_DEBUG_ONLY -> STD_TORCH_CHECK
+// 2. Has a string constructor that uses a shim function
+// 3. does not include some is_{device} variants that we can add later
+//
+// We chose to copy it rather than moving it to headeronly as
+// 1. Device is < 8 bytes so the *Handle approach used for tensor doesn't make
+// sense
+// 2. c10::Device is not header-only due to its string constructor.
+//
+// StableIValue conversions handle conversion between c10::Device (in libtorch)
+// and torch::stable::Device (in stable user extensions)
+
+/**
+ * @brief A stable version of c10::Device.
+ *
+ * Minimum compatible version: PyTorch 2.9.
+ */
+class Device {
+ private:
+  DeviceType type_;
+  DeviceIndex index_ = -1;
+
+  void validate() {
+    STD_TORCH_CHECK(
+        index_ >= -1,
+        "Device index must be -1 or non-negative, got ",
+        static_cast<int>(index_));
+    STD_TORCH_CHECK(
+        type_ != DeviceType::CPU || index_ <= 0,
+        "CPU device index must be -1 or zero, got ",
+        static_cast<int>(index_));
+  }
+
+ public:
+  /**
+   * @brief Constructs a Device from a DeviceType and optional device index.
+   *
+   * @param type The type of device (e.g., DeviceType::CPU, DeviceType::CUDA).
+   * @param index The device index. Default is -1 (current device).
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  /* implicit */ Device(DeviceType type, DeviceIndex index = -1)
+      : type_(type), index_(index) {
+    validate();
+  }
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  // Defined in device_inl.h to avoid circular dependencies.
+  /**
+   * @brief Constructs a stable::Device from a string description.
+   *
+   * The string must follow the schema: (cpu|cuda|...)[:<device-index>]
+   *
+   * @param device_string A string describing the device (e.g., "cuda:0",
+   * "cpu").
+   *
+   * Minimum compatible version: PyTorch 2.10.
+   */
+  /* implicit */ Device(const std::string& device_string);
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+  // Copy and move constructors can be default
+  /// \private
+  Device(const Device& other) = default;
+  /// \private
+  Device(Device&& other) noexcept = default;
+
+  // Copy and move assignment operators can be default
+  /// \private
+  Device& operator=(const Device& other) = default;
+  /// \private
+  Device& operator=(Device&& other) noexcept = default;
+
+  // Destructor can be default
+  /// \private
+  ~Device() = default;
+
+  /**
+   * @brief Checks if two devices are equal.
+   *
+   * @param other The device to compare with.
+   * @return true if both type and index match, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool operator==(const Device& other) const noexcept {
+    return type() == other.type() && index() == other.index();
+  }
+
+  /**
+   * @brief Checks if two devices are not equal.
+   *
+   * @param other The device to compare with.
+   * @return true if type or index differ, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool operator!=(const Device& other) const noexcept {
+    return !(*this == other);
+  }
+
+  /**
+   * @brief Sets the device index.
+   *
+   * @param index The new device index.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  void set_index(DeviceIndex index) {
+    index_ = index;
+  }
+
+  /**
+   * @brief Returns the device type.
+   *
+   * @return The DeviceType of this device.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  DeviceType type() const noexcept {
+    return type_;
+  }
+
+  /**
+   * @brief Returns the device index.
+   *
+   * @return The device index, or -1 if no specific index is set.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  DeviceIndex index() const noexcept {
+    return index_;
+  }
+
+  /**
+   * @brief Checks if this device has a specific index.
+   *
+   * @return true if index is not -1, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool has_index() const noexcept {
+    return index_ != -1;
+  }
+
+  /**
+   * @brief Checks if this is a CUDA device.
+   *
+   * @return true if the device type is CUDA, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool is_cuda() const noexcept {
+    return type_ == DeviceType::CUDA;
+  }
+
+  /**
+   * @brief Checks if this is a CPU device.
+   *
+   * @return true if the device type is CPU, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool is_cpu() const noexcept {
+    return type_ == DeviceType::CPU;
+  }
+};
+
+HIDDEN_NAMESPACE_END(torch, stable)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/library.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/library.h
new file mode 100644
index 0000000000000000000000000000000000000000..f12f0d768acfc361825875fb50565a0c5f8343ad
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/library.h
@@ -0,0 +1,372 @@
+#pragma once
+// this file can only have stable stuff! Akin to shim.h
+// but unlike shim.h, this file can contain header-only C++
+// code for better UX.
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/headeronly/core/enum_tag.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Metaprogramming.h>
+
+// Technically, this file doesn't use anything from stableivalue_conversions.h,
+// but we need to include it here as the contents of stableivalue_conversions.h
+// used to live here and so we need to expose them for backwards compatibility.
+#include <torch/csrc/stable/stableivalue_conversions.h>
+#include <torch/csrc/stable/version.h>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable, detail)
+
+class StableLibrary final {
+ private:
+  TorchLibraryHandle lib_;
+
+ public:
+  enum class Kind {
+    DEF,
+    IMPL,
+    FRAGMENT,
+  };
+
+  // constructor
+  /// \private
+  ///
+  /// Use STABLE_TORCH_LIBRARY or STABLE_TORCH_LIBRARY_IMPL() instead of using
+  /// these constructors directly
+  StableLibrary(
+      Kind kind,
+      const char* ns,
+      const char* k,
+      const char* file,
+      uint32_t line) {
+    if (kind == Kind::IMPL) {
+      aoti_torch_library_init_impl(ns, k, file, line, &lib_);
+    } else if (kind == Kind::DEF) {
+      aoti_torch_library_init_def(ns, file, line, &lib_);
+    } else { // kind == FRAGMENT
+      aoti_torch_library_init_fragment(ns, file, line, &lib_);
+    }
+  }
+
+  // do not permit copy
+  StableLibrary(const StableLibrary&) = delete;
+  StableLibrary& operator=(const StableLibrary&) = delete;
+
+  // do not permit move
+  StableLibrary(StableLibrary&& other) = delete;
+  StableLibrary& operator=(StableLibrary&& other) = delete;
+
+  ~StableLibrary() {
+    aoti_torch_delete_library_object(lib_);
+  }
+
+  // corresponds to a limited, stable version of torch::library::impl()
+  // Inputs:
+  //   name: the name of the function to implement
+  //   fn: a boxed function with schema
+  //       (StableIValue* stack, uint64_t num_inputs, uint64_t num_outputs) ->
+  //       void
+  // fn should follow the calling convention of our boxed kernels that convert
+  // to IValues. fn will be called with a StableIValue* array of length
+  // max(num_inputs, num_outputs), where the first num_inputs entries are
+  // populated with inputs. fn is responsible for stealing the memory of the
+  // inputs, in effect "popping" them off the stack, and then populating the
+  // stack with StableIValue outputs. Concretely, fn should:
+  //    1. read StableIValue inputs from the given stack
+  //    2. convert the inputs to the proper types
+  //    3. call the function corresponding to name with the inputs
+  //    4. convert the outputs to StableIValues
+  //    5. populate the now empty stack with StableIValue outputs
+  // If the operation corresponding to name takes in 4 inputs and returns 2
+  // outputs, fn should expect stack to contain 4 StableIValues:
+  //    [stable_arg1, stable_arg2, stable_arg3, stable_arg4]
+  // to end, fn should fill the stack with 2 StableIValues representing outputs:
+  //    [stable_ret1, stable_ret2, -, -]
+  StableLibrary& impl(
+      const char* name,
+      void (*fn)(StableIValue*, uint64_t, uint64_t)) {
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+    torch_library_impl(lib_, name, fn, TORCH_ABI_VERSION);
+#else
+    aoti_torch_library_impl(lib_, name, fn);
+#endif
+    return *this;
+  }
+
+  // corresponds to a limited, stable version of torch::library::def()
+  StableLibrary& def(const char* schema) {
+    aoti_torch_library_def(lib_, schema);
+    return *this;
+  }
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+  // corresponds to a limited, stable version of torch::library::def() with tags
+  StableLibrary& def(const char* schema, const std::vector<at::Tag>& tags) {
+    std::vector<int32_t> tag_ints;
+    tag_ints.reserve(tags.size());
+    for (auto t : tags) {
+      tag_ints.push_back(
+          torch::stable::detail::to<int32_t>(torch::stable::detail::from(t)));
+    }
+    torch_library_def_with_tags(
+        lib_, schema, tag_ints.data(), static_cast<int32_t>(tag_ints.size()));
+    return *this;
+  }
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+};
+
+class StableTorchLibraryInit final {
+ private:
+  using InitFn = void(StableLibrary&);
+  StableLibrary lib_;
+
+ public:
+  StableTorchLibraryInit(
+      StableLibrary::Kind kind,
+      InitFn* fn,
+      const char* ns,
+      const char* k,
+      const char* file,
+      uint32_t line)
+      : lib_(kind, ns, k, file, line) {
+    fn(lib_);
+  }
+};
+
+// type mapper: since to<HeaderOnlyArrayRef<T>> cannot exist,
+// we map that to to<std::vector<T>> to preserve ownership semantics.
+// note that unbox_type_t is used to convert ParamTypes, so that
+// the tuple holding the arguments will have proper ownership too.
+template <typename T>
+struct UnboxType {
+  using type = T;
+};
+
+template <typename T>
+struct UnboxType<torch::headeronly::HeaderOnlyArrayRef<T>> {
+  using type = std::vector<T>;
+};
+
+template <typename T>
+struct UnboxType<std::optional<torch::headeronly::HeaderOnlyArrayRef<T>>> {
+  using type = std::optional<std::vector<T>>;
+};
+
+template <>
+struct UnboxType<std::string_view> {
+  using type = std::string;
+};
+
+// const and reference are stripped before UnboxType is applied
+// in order to avoid ambiguous template matches
+template <typename T>
+using unbox_type_t =
+    typename UnboxType<std::remove_cv_t<std::remove_reference_t<T>>>::type;
+
+template <class... T, std::size_t... I>
+std::tuple<T...> unbox_to_tuple_impl(
+    StableIValue* stack,
+    std::index_sequence<I...> /*unused*/) {
+  return std::make_tuple(to<T>(stack[I])...);
+}
+
+template <class... T>
+std::tuple<T...> unbox_to_tuple(StableIValue* stack) {
+  return unbox_to_tuple_impl<T...>(
+      stack, std::make_index_sequence<sizeof...(T)>());
+}
+
+template <class... T, std::size_t... I>
+void box_from_tuple_impl(
+    StableIValue* stack,
+    std::tuple<T...> vals,
+    std::index_sequence<I...> /*unused*/) {
+  ((stack[I] = from<T>(std::get<I>(vals))), ...);
+}
+
+template <class... T>
+void box_from_tuple(StableIValue* stack, std::tuple<T...> vals) {
+  box_from_tuple_impl<T...>(
+      stack, vals, std::make_index_sequence<sizeof...(T)>());
+}
+
+template <
+    typename ReturnType,
+    typename ParameterTypeList,
+    typename FuncT,
+    FuncT* func>
+struct boxer_impl {
+  static_assert(
+      torch::headeronly::guts::false_t<ReturnType>::value,
+      "Unsupported function schema for TORCH_BOX.");
+};
+
+// Multiple returns
+template <
+    typename... ReturnTypes,
+    typename... ParameterTypes,
+    typename FuncT,
+    FuncT* func>
+struct boxer_impl<
+    std::tuple<ReturnTypes...>,
+    torch::headeronly::guts::typelist::typelist<ParameterTypes...>,
+    FuncT,
+    func> {
+  static void boxed_fn(
+      StableIValue* stack,
+      uint64_t num_args,
+      uint64_t num_outputs) {
+    STD_TORCH_CHECK(
+        num_args == sizeof...(ParameterTypes),
+        "Registered schema has ",
+        num_args,
+        " args, but the kernel to box has ",
+        sizeof...(ParameterTypes));
+    STD_TORCH_CHECK(
+        num_outputs == sizeof...(ReturnTypes),
+        "Registered schema has ",
+        num_outputs,
+        " outputs, but the kernel to box has ",
+        sizeof...(ReturnTypes));
+    std::tuple<unbox_type_t<ParameterTypes>...> args =
+        unbox_to_tuple<unbox_type_t<ParameterTypes>...>(stack);
+    auto res = std::apply(func, args);
+    box_from_tuple<ReturnTypes...>(stack, res);
+  }
+};
+
+// Single return
+template <
+    typename ReturnType,
+    typename... ParameterTypes,
+    typename FuncT,
+    FuncT* func>
+struct boxer_impl<
+    ReturnType,
+    torch::headeronly::guts::typelist::typelist<ParameterTypes...>,
+    FuncT,
+    func> {
+  static void boxed_fn(
+      StableIValue* stack,
+      uint64_t num_args,
+      uint64_t num_outputs) {
+    STD_TORCH_CHECK(
+        num_args == sizeof...(ParameterTypes),
+        "Registered schema has ",
+        num_args,
+        " args, but the kernel to box has ",
+        sizeof...(ParameterTypes));
+    STD_TORCH_CHECK(
+        num_outputs == 1,
+        "Registered schema has ",
+        num_outputs,
+        " outputs, but the kernel to box has ",
+        1);
+    std::tuple<unbox_type_t<ParameterTypes>...> args =
+        unbox_to_tuple<unbox_type_t<ParameterTypes>...>(stack);
+    auto res = std::apply(func, args);
+    stack[0] = from<ReturnType>(res);
+  }
+};
+
+// No/void return
+template <typename... ParameterTypes, typename FuncT, FuncT* func>
+struct boxer_impl<
+    void,
+    torch::headeronly::guts::typelist::typelist<ParameterTypes...>,
+    FuncT,
+    func> {
+  static void boxed_fn(
+      StableIValue* stack,
+      uint64_t num_args,
+      uint64_t num_outputs) {
+    STD_TORCH_CHECK(
+        num_args == sizeof...(ParameterTypes),
+        "Registered schema has ",
+        num_args,
+        " args, but the kernel to box has ",
+        sizeof...(ParameterTypes));
+    STD_TORCH_CHECK(
+        num_outputs == 0,
+        "Registered schema has ",
+        num_outputs,
+        " outputs, but the kernel to box has ",
+        0);
+    std::tuple<unbox_type_t<ParameterTypes>...> args =
+        unbox_to_tuple<unbox_type_t<ParameterTypes>...>(stack);
+    std::apply(func, args);
+  }
+};
+
+template <typename FuncT, FuncT* func>
+struct boxer {
+  using FunctionTraits =
+      torch::headeronly::guts::infer_function_traits_t<FuncT>;
+
+  static void boxed_fn(
+      StableIValue* stack,
+      uint64_t num_args,
+      uint64_t num_outputs) {
+    boxer_impl<
+        typename FunctionTraits::return_type,
+        typename FunctionTraits::parameter_types,
+        FuncT,
+        func>::boxed_fn(stack, num_args, num_outputs);
+  }
+};
+
+HIDDEN_NAMESPACE_END(torch, stable, detail)
+
+#define TORCH_BOX(func)                                               \
+  torch::stable::detail::boxer<                                       \
+      std::remove_pointer_t<std::remove_reference_t<decltype(func)>>, \
+      (func)>::boxed_fn
+
+#define STABLE_TORCH_LIBRARY_IMPL(ns, k, m) \
+  _STABLE_TORCH_LIBRARY_IMPL(ns, k, m, C10_UID)
+
+#define _STABLE_TORCH_LIBRARY_IMPL(ns, k, m, uid)                             \
+  static void C10_CONCATENATE(                                                \
+      STABLE_TORCH_LIBRARY_IMPL_init_##ns##_##k##_,                           \
+      uid)(torch::stable::detail::StableLibrary&);                            \
+  static const torch::stable::detail::StableTorchLibraryInit C10_CONCATENATE( \
+      STABLE_TORCH_LIBRARY_IMPL_static_init_##ns##_##k##_, uid)(              \
+      torch::stable::detail::StableLibrary::Kind::IMPL,                       \
+      &C10_CONCATENATE(STABLE_TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid),    \
+      C10_STRINGIZE(ns),                                                      \
+      C10_STRINGIZE(k),                                                       \
+      __FILE__,                                                               \
+      __LINE__);                                                              \
+  void C10_CONCATENATE(STABLE_TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid)(    \
+      torch::stable::detail::StableLibrary & m)
+
+#define STABLE_TORCH_LIBRARY(ns, m)                          \
+  static void STABLE_TORCH_LIBRARY_init_##ns(                \
+      torch::stable::detail::StableLibrary&);                \
+  static const torch::stable::detail::StableTorchLibraryInit \
+      STABLE_TORCH_LIBRARY_static_init_##ns(                 \
+          torch::stable::detail::StableLibrary::Kind::DEF,   \
+          &STABLE_TORCH_LIBRARY_init_##ns,                   \
+          C10_STRINGIZE(ns),                                 \
+          nullptr,                                           \
+          __FILE__,                                          \
+          __LINE__);                                         \
+  void STABLE_TORCH_LIBRARY_init_##ns(torch::stable::detail::StableLibrary& m)
+
+#define STABLE_TORCH_LIBRARY_FRAGMENT(ns, m) \
+  _STABLE_TORCH_LIBRARY_FRAGMENT(ns, m, C10_UID)
+
+#define _STABLE_TORCH_LIBRARY_FRAGMENT(ns, m, uid)                            \
+  static void C10_CONCATENATE(                                                \
+      STABLE_TORCH_LIBRARY_FRAGMENT_init_##ns##_,                             \
+      uid)(torch::stable::detail::StableLibrary&);                            \
+  static const torch::stable::detail::StableTorchLibraryInit C10_CONCATENATE( \
+      STABLE_TORCH_LIBRARY_FRAGMENT_static_init_##ns##_, uid)(                \
+      torch::stable::detail::StableLibrary::Kind::FRAGMENT,                   \
+      &C10_CONCATENATE(STABLE_TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid),      \
+      C10_STRINGIZE(ns),                                                      \
+      nullptr,                                                                \
+      __FILE__,                                                               \
+      __LINE__);                                                              \
+  void C10_CONCATENATE(STABLE_TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid)(      \
+      torch::stable::detail::StableLibrary & m)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/macros.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/macros.h
new file mode 100644
index 0000000000000000000000000000000000000000..52f108d500683299ca44ee0dddc772a68edafa21
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/macros.h
@@ -0,0 +1,30 @@
+#include <torch/csrc/stable/c/shim.h>
+
+#include <stdexcept>
+#include <string>
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+// Users of this macro are expected to include cuda_runtime.h
+#define STD_CUDA_CHECK(EXPR)                      \
+  do {                                            \
+    const cudaError_t __err = EXPR;               \
+    char* __error_msg = nullptr;                  \
+    torch_c10_cuda_check_msg(                     \
+        static_cast<int32_t>(__err),              \
+        __FILE__,                                 \
+        __func__,                                 \
+        static_cast<uint32_t>(__LINE__),          \
+        true,                                     \
+        &__error_msg);                            \
+    if (__error_msg != nullptr) {                 \
+      std::string __msg(__error_msg);             \
+      torch_c10_cuda_free_error_msg(__error_msg); \
+      throw std::runtime_error(__msg);            \
+    }                                             \
+  } while (0)
+
+// Users of this macro are expected to include cuda_runtime.h
+#define STD_CUDA_KERNEL_LAUNCH_CHECK() STD_CUDA_CHECK(cudaGetLastError())
+
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/ops.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/ops.h
new file mode 100644
index 0000000000000000000000000000000000000000..dbede4faba49eca9984488c0d7ccad1879c6f82f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/ops.h
@@ -0,0 +1,1134 @@
+#pragma once
+
+#include <torch/csrc/stable/stableivalue_conversions.h>
+#include <array>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <type_traits>
+
+#include <torch/csrc/inductor/aoti_torch/generated/c_shim_aten.h>
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/version.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/HeaderOnlyArrayRef.h>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable)
+
+/// A function pointer type for data deleters used with from_blob.
+/// The deleter is called with the data pointer when the tensor's storage
+/// is deallocated.
+using DeleterFnPtr = void (*)(void*);
+
+/// Stable version of the empty_like op.
+///
+/// Creates a new uninitialized tensor with the same size, dtype, layout, and
+/// device as the input tensor. This version does not support kwargs (device,
+/// dtype, layout, memory_format) - kwargs support may be added in the future.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The input tensor whose properties will be used for the new
+/// tensor.
+/// @return A new uninitialized tensor with the same properties as self.
+inline torch::stable::Tensor empty_like(const torch::stable::Tensor& self) {
+  const auto num_args = 6;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(std::nullopt),
+      torch::stable::detail::from(std::nullopt),
+      torch::stable::detail::from(std::nullopt),
+      torch::stable::detail::from(std::nullopt),
+      torch::stable::detail::from(std::nullopt)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::empty_like", "", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::empty_like", "", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the fill_.Scalar op.
+///
+/// Fills the input tensor with the specified scalar value in-place and returns
+/// it. This has identical semantics to the existing fill_.Scalar op.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @note The value parameter is typed as double
+///       This is because Scalar.h is currently not header-only.
+///
+/// @param self The tensor to fill.
+/// @param value The scalar value to fill the tensor with.
+/// @return The input tensor, now filled with the specified value.
+inline torch::stable::Tensor fill_(
+    const torch::stable::Tensor& self,
+    double value) {
+  TORCH_ERROR_CODE_CHECK(aoti_torch_aten_fill__Scalar(self.get(), value));
+  return self;
+}
+
+/// Stable version of the narrow.default op.
+///
+/// Returns a new tensor that is a narrowed version of the input tensor. The
+/// dimension dim is narrowed from start to start + length.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @note The start and length parameters
+///       is not yet header-only.
+///
+/// @param self The input tensor to narrow.
+/// @param dim The dimension along which to narrow.
+/// @param start The starting index for the narrowed dimension.
+/// @param length The length of the narrowed dimension.
+/// @return A new tensor that is a narrowed view of the input.
+inline torch::stable::Tensor narrow(
+    torch::stable::Tensor& self,
+    int64_t dim,
+    int64_t start,
+    int64_t length) {
+  AtenTensorHandle ret0 = nullptr;
+
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_aten_narrow(self.get(), dim, start, length, &ret0));
+  return torch::stable::Tensor(ret0);
+}
+
+#if TORCH_FEATURE_VERSION < TORCH_VERSION_2_10_0
+/// Stable version of the new_empty op (2.9 version).
+///
+/// Creates a new uninitialized tensor with the specified size, inheriting
+/// device and layout from the input tensor. This version only supports the
+/// dtype kwarg. For the full kwargs version, use PyTorch 2.10+.
+///
+/// Minimum compatible version: PyTorch 2.9. For full kwargs support, use
+/// PyTorch 2.10+.
+///
+/// @param self The input tensor whose device
+/// @param size The desired size of the output tensor.
+/// @param dtype Optional scalar type for the tensor elements. If not provided,
+///              inherits from self.
+/// @return A new uninitialized tensor with the specified properties.
+inline torch::stable::Tensor new_empty(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef size,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt) {
+  int32_t device_type;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_device_type(self.get(), &device_type));
+
+  int32_t device_index;
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_get_device_index(self.get(), &device_index));
+
+  int32_t target_dtype;
+  if (dtype.has_value()) {
+    target_dtype = torch::stable::detail::to<int32_t>(
+        torch::stable::detail::from(dtype.value()));
+  } else {
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(self.get(), &target_dtype));
+  }
+
+  int32_t layout;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_layout(self.get(), &layout));
+
+  AtenTensorHandle ret0;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_aten_new_empty(
+      self.get(),
+      size.data(),
+      static_cast<int64_t>(size.size()),
+      &target_dtype,
+      &layout,
+      &device_type,
+      device_index,
+      nullptr, // pin_memory (nullptr for default)
+      &ret0));
+
+  return torch::stable::Tensor(ret0);
+}
+
+/// Stable version of the new_zeros op (2.9 version).
+///
+/// Creates a new tensor filled with zeros with the specified size, inheriting
+/// device and layout from the input tensor. This version only supports the
+/// dtype kwarg. For the full kwargs version, use PyTorch 2.10+.
+///
+/// Minimum compatible version: PyTorch 2.9. For full kwargs support, use
+/// PyTorch 2.10+.
+///
+/// @param self The input tensor whose device and layout will be inherited.
+/// @param size The desired size of the output tensor.
+/// @param dtype Optional scalar type for the tensor elements. If not provided,
+///              inherits from self.
+/// @return A new zero-filled tensor with the specified properties.
+inline torch::stable::Tensor new_zeros(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef size,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt) {
+  int32_t device_type;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_device_type(self.get(), &device_type));
+
+  int32_t device_index;
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_get_device_index(self.get(), &device_index));
+
+  int32_t target_dtype;
+  if (dtype.has_value()) {
+    target_dtype = torch::stable::detail::to<int32_t>(
+        torch::stable::detail::from(dtype.value()));
+  } else {
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(self.get(), &target_dtype));
+  }
+
+  int32_t layout;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_layout(self.get(), &layout));
+
+  AtenTensorHandle ath;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_aten_new_zeros(
+      self.get(),
+      size.data(),
+      static_cast<int64_t>(size.size()),
+      &target_dtype,
+      &layout,
+      &device_type,
+      device_index,
+      nullptr, // pin_memory (nullptr for default)
+      &ath));
+
+  return torch::stable::Tensor(ath);
+}
+#endif // TORCH_FEATURE_VERSION < TORCH_VERSION_2_10_0
+
+/// Stable version of the pad.default op.
+///
+/// Pads the input tensor according to the specified padding sizes. The padding
+/// is applied symmetrically to each dimension, with the padding sizes specified
+/// in reverse order (last dimension first).
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @note The pad parameter is typed
+///       not yet header-only.
+///
+/// @param self The input tensor to pad.
+/// @param pad The padding sizes for each dimension (in pairs, starting from
+///            the last dimension).
+/// @param mode The padding mode: "constant", "reflect", "replicate", or
+///             "circular". Defaults to "constant".
+/// @param value The fill value for constant padding. Defaults to 0.0.
+/// @return A new padded tensor.
+inline torch::stable::Tensor pad(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef pad,
+    const std::string& mode = "constant",
+    double value = 0.0) {
+  AtenTensorHandle ret0 = nullptr;
+
+  TORCH_ERROR_CODE_CHECK(aoti_torch_aten_pad(
+      self.get(), pad.data(), pad.size(), mode.c_str(), &value, &ret0));
+  return torch::stable::Tensor(ret0);
+}
+
+/// Stable version of the amax.default op (single dimension).
+///
+/// Computes the maximum value along the specified dimension. If keepdim is
+/// true, the output tensor has the same number of dimensions as the input,
+/// with the reduced dimension having size 1. Otherwise, the reduced dimension
+/// is removed.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The input tensor.
+/// @param dim The dimension along which to compute the maximum.
+/// @param keepdim Whether to retain
+/// @return A tensor containing the maximum values along the specified
+/// dimension.
+inline torch::stable::Tensor amax(
+    const torch::stable::Tensor& self,
+    int64_t dim,
+    bool keepdim = false) {
+  AtenTensorHandle ret = nullptr;
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_aten_amax(self.get(), &dim, 1, keepdim, &ret));
+  return torch::stable::Tensor(ret);
+}
+
+/// Stable version of the amax.default op (multiple dimensions).
+///
+/// Computes the maximum value reducing over all the specified dimensions. If
+/// keepdim is true, the output tensor has the same number of dimensions as the
+/// input, with the reduced dimensions having size 1. Otherwise, the reduced
+/// dimensions are removed.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @note The dims parameter is typed
+///       is not yet header-only.
+///
+/// @param self The input tensor.
+/// @param dims The dimensions along which to compute the maximum.
+/// @param keepdim Whether to retain the reduced dimensions. Defaults to false.
+/// @return A tensor containing the maximum values.
+inline torch::stable::Tensor amax(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef dims,
+    bool keepdim = false) {
+  AtenTensorHandle ret = nullptr;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_aten_amax(
+      self.get(),
+      dims.data(),
+      static_cast<int64_t>(dims.size()),
+      keepdim,
+      &ret));
+  return torch::stable::Tensor(ret);
+}
+
+/// Stable version of the transpose.int op.
+///
+/// Returns a tensor that is a transposed version of the input, with dimensions
+/// dim0 and dim1 swapped. The returned tensor shares storage with the input.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The input tensor.
+/// @param dim0 The first dimension to transpose.
+/// @param dim1 The second dimension to transpose.
+/// @return A transposed view of the input tensor.
+inline torch::stable::Tensor transpose(
+    const torch::stable::Tensor& self,
+    int64_t dim0,
+    int64_t dim1) {
+  const auto num_args = 3;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(dim0),
+      torch::stable::detail::from(dim1)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::transpose", "int", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::transpose", "int", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the zero_ op.
+///
+/// Fills the input tensor with zeros in-place and returns it. Unlike the
+/// tensor method version (t.zero_()), this is called as a function: zero_(t).
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The tensor to fill with zeros.
+/// @return The input tensor, now filled with zeros.
+inline torch::stable::Tensor zero_(torch::stable::Tensor& self) {
+  const auto num_args = 1;
+  std::array<StableIValue, num_args> stack{torch::stable::detail::from(self)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::zero_", "", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::zero_", "", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the copy_ op.
+///
+/// Copies the elements from the source tensor into the destination tensor
+/// in-place and returns the destination tensor. The tensors must be
+/// broadcastable.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The destination tensor (modified in-place).
+/// @param src The source tensor to copy from.
+/// @param non_blocking If true, the copy may occur asynchronously with respect
+///                     to the host. Defaults to false.
+/// @return The destination tensor with copied values.
+inline torch::stable::Tensor copy_(
+    torch::stable::Tensor& self,
+    const torch::stable::Tensor& src,
+    std::optional<bool> non_blocking = std::nullopt) {
+  const auto num_args = 3;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(src),
+      torch::stable::detail::from(non_blocking.value_or(false))};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::copy_", "", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::copy_", "", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the clone op.
+///
+/// Returns a copy of the input tensor. The returned tensor has the same data
+/// and type as the input, but is stored in a new memory location.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @note Optional memory_format kwarg support
+///
+/// @param self The input tensor to clone.
+/// @return A new tensor with copied data.
+inline torch::stable::Tensor clone(const torch::stable::Tensor& self) {
+  const auto num_args = 2;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(std::nullopt)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::clone", "", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::clone", "", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the flatten.using_ints op.
+///
+/// Flattens the input tensor by reshaping it into a one-dimensional tensor.
+/// If start_dim or end_dim are specified, only dimensions starting from
+/// start_dim to end_dim are flattened.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The input tensor to flatten.
+/// @param start_dim The first dimension to flatten. Defaults to 0.
+/// @param end_dim The last dimension to flatten. Defaults to -1 (last dim).
+/// @return A flattened tensor.
+inline torch::stable::Tensor flatten(
+    const torch::stable::Tensor& self,
+    int64_t start_dim = 0,
+    int64_t end_dim = -1) {
+  const auto num_args = 3;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(start_dim),
+      torch::stable::detail::from(end_dim)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::flatten", "using_ints", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::flatten", "using_ints", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the unsqueeze op.
+///
+/// Returns a new tensor with a dimension of size one inserted at the specified
+/// position. The returned tensor shares the same underlying data with the input
+/// tensor.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The input tensor.
+/// @param dim The index at which to insert
+///            values are supported.
+/// @return A tensor with an additional dimension.
+inline torch::stable::Tensor unsqueeze(
+    const torch::stable::Tensor& self,
+    int64_t dim) {
+  const auto num_args = 2;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self), torch::stable::detail::from(dim)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::unsqueeze", "", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::unsqueeze", "", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the squeeze.dim op.
+///
+/// Returns a tensor with the dimension of size one at the specified position
+/// removed. The returned tensor shares the same underlying data with the input
+/// tensor.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The input tensor.
+/// @param dim The dimension to squeeze.
+///            the tensor is returned unchanged.
+/// @return A tensor with the specified dimension removed (if size was 1).
+inline torch::stable::Tensor squeeze(
+    const torch::stable::Tensor& self,
+    int64_t dim) {
+  const auto num_args = 2;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self), torch::stable::detail::from(dim)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::squeeze", "dim", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::squeeze", "dim", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the select.int op.
+///
+/// Slices the input tensor along the specified dimension at the given index.
+/// This function returns a view of the original tensor with the given dimension
+/// removed.
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @note The index parameter is typed
+///       header-only.
+///
+/// @param self The input tensor.
+/// @param dim The dimension to slice.
+/// @param index The index to select along the dimension.
+/// @return A tensor with one fewer dimension.
+inline torch::stable::Tensor select(
+    const torch::stable::Tensor& self,
+    int64_t dim,
+    int64_t index) {
+  const auto num_args = 3;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(dim),
+      torch::stable::detail::from(index)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::select", "int", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::select", "int", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the matmul op.
+///
+/// Performs matrix multiplication between two tensors. The behavior depends on
+/// the dimensionality of the tensors (see PyTorch documentation for details on
+/// broadcasting rules for matmul).
+///
+/// Minimum compatible version: PyTorch 2.9.
+///
+/// @param self The first input tensor.
+/// @param other The second input tensor.
+/// @return The result of matrix multiplication.
+inline torch::stable::Tensor matmul(
+    const torch::stable::Tensor& self,
+    const torch::stable::Tensor& other) {
+  const auto num_args = 2;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self), torch::stable::detail::from(other)};
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::matmul", "", stack.data(), TORCH_ABI_VERSION));
+#else
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_call_dispatcher("aten::matmul", "", stack.data()));
+#endif
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+/// Stable parallel_for utility.
+///
+/// Provides a stable interface to at::parallel_for for parallel execution.
+/// The function f will be called with (begin, end) ranges to process in
+/// parallel. grain_size controls the minimum work size per thread for efficient
+/// parallelization.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @tparam F The callable type
+/// @param begin The start of the iteration range.
+/// @param end The end of the iteration range (exclusive).
+/// @param grain_size The minimum number of iterations per thread.
+/// @param f The function to execute in parallel.
+template <class F>
+inline void parallel_for(
+    const int64_t begin,
+    const int64_t end,
+    const int64_t grain_size,
+    const F& f) {
+  auto callback = [](int64_t cb_begin, int64_t cb_end, void* ctx) {
+    const F* func = static_cast<const F*>(ctx);
+    (*func)(cb_begin, cb_end);
+  };
+  TORCH_ERROR_CODE_CHECK(torch_parallel_for(
+      begin,
+      end,
+      grain_size,
+      callback,
+      const_cast<void*>(static_cast<const void*>(&f))));
+}
+
+/// Gets the number of threads for the parallel backend.
+///
+/// Provides a stable interface to at::get_num_threads.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @return The number of threads
+inline uint32_t get_num_threads() {
+  uint32_t num_threads;
+  TORCH_ERROR_CODE_CHECK(torch_get_num_threads(&num_threads));
+  return num_threads;
+}
+
+/// Stable version of the empty.memory_format op.
+///
+/// Creates a new uninitialized tensor with the specified size and options.
+/// This function supports full tensor creation options including device,
+/// dtype, layout, and memory format.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param size The desired size of the output tensor.
+/// @param dtype Optional scalar type for the tensor elements.
+/// @param layout Optional memory layout (e.g., strided, sparse).
+/// @param device Optional device to place the tensor on.
+/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors).
+/// @param memory_format Optional memory format for the tensor.
+/// @return A new uninitialized tensor with the specified properties.
+inline torch::stable::Tensor empty(
+    torch::headeronly::IntHeaderOnlyArrayRef size,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt,
+    std::optional<torch::headeronly::Layout> layout = std::nullopt,
+    std::optional<torch::stable::Device> device = std::nullopt,
+    std::optional<bool> pin_memory = std::nullopt,
+    std::optional<torch::headeronly::MemoryFormat> memory_format =
+        std::nullopt) {
+  const auto num_args = 6;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(size),
+      torch::stable::detail::from(dtype),
+      torch::stable::detail::from(layout),
+      torch::stable::detail::from(device),
+      torch::stable::detail::from(pin_memory),
+      torch::stable::detail::from(memory_format)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::empty", "memory_format", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the reshape op.
+///
+/// Returns a tensor with the same data and number of elements as the input,
+/// but with the specified shape. When possible, the returned tensor will be
+/// a view of the input.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor.
+/// @param shape The desired output shape.
+/// @return A tensor with the specified shape.
+inline torch::stable::Tensor reshape(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef shape) {
+  const auto num_args = 2;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self), torch::stable::detail::from(shape)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::reshape", "", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the view op.
+///
+/// Returns a new tensor with the same data as the input tensor but with a
+/// different shape. The returned tensor shares the same data and must have
+/// the same number of elements.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor.
+/// @param size The desired output shape.
+/// @return A view tensor with the specified shape.
+inline torch::stable::Tensor view(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef size) {
+  const auto num_args = 2;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self), torch::stable::detail::from(size)};
+  TORCH_ERROR_CODE_CHECK(
+      torch_call_dispatcher("aten::view", "", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Creates a tensor from an existing data blob.
+///
+/// Creates a tensor that uses the provided data pointer as its storage.
+/// The tensor does not own the data, so the caller must ensure the data
+/// remains valid for the lifetime of the tensor.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param data Pointer to the data buffer.
+/// @param sizes The size of each dimension of the tensor.
+/// @param strides The stride for each dimension.
+/// @param device The device where the data resides.
+/// @param dtype The scalar type of the data.
+/// @param storage_offset The offset into the data buffer. Defaults to 0.
+/// @param layout The memory layout. Defaults to Strided.
+/// @return A tensor backed by the provided data.
+inline torch::stable::Tensor from_blob(
+    void* data,
+    torch::headeronly::IntHeaderOnlyArrayRef sizes,
+    torch::headeronly::IntHeaderOnlyArrayRef strides,
+    torch::stable::Device device,
+    torch::headeronly::ScalarType dtype,
+    int64_t storage_offset = 0,
+    torch::headeronly::Layout layout = torch::headeronly::Layout::Strided) {
+  auto shim_dtype =
+      torch::stable::detail::to<int32_t>(torch::stable::detail::from(dtype));
+  auto shim_device_type = torch::stable::detail::to<int32_t>(
+      torch::stable::detail::from(device.type()));
+  auto shim_layout =
+      torch::stable::detail::to<int32_t>(torch::stable::detail::from(layout));
+  AtenTensorHandle ath;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_v2(
+      data,
+      sizes.size(),
+      sizes.data(),
+      strides.data(),
+      storage_offset,
+      shim_dtype,
+      shim_device_type,
+      device.index(),
+      &ath,
+      shim_layout,
+      nullptr,
+      0));
+  return torch::stable::Tensor(ath);
+}
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+/// Creates a tensor from an existing data blob with a custom deleter.
+///
+/// This is the same as the from_blob function above, but allows specifying a
+/// custom deleter function that will be called when the tensor's storage is
+/// deallocated. Accepts both plain function pointers and capturing lambdas.
+///
+/// Minimum compatible version: PyTorch 2.11.
+///
+/// @tparam F The callable type. Must be invocable with (void*).
+/// @param data Pointer to the data buffer.
+/// @param sizes The size of each dimension of the tensor.
+/// @param strides The stride for each dimension.
+/// @param device The device where the data resides.
+/// @param dtype The scalar type of the data.
+/// @param deleter Callable to invoke when the tensor is deallocated.
+/// @param storage_offset The offset into the data buffer. Defaults to 0.
+/// @param layout The memory layout. Defaults to Strided.
+/// @return A tensor backed by the provided data.
+template <class F, std::enable_if_t<std::is_invocable_v<F, void*>, int> = 0>
+inline torch::stable::Tensor from_blob(
+    void* data,
+    torch::headeronly::IntHeaderOnlyArrayRef sizes,
+    torch::headeronly::IntHeaderOnlyArrayRef strides,
+    torch::stable::Device device,
+    torch::headeronly::ScalarType dtype,
+    F deleter,
+    int64_t storage_offset = 0,
+    torch::headeronly::Layout layout = torch::headeronly::Layout::Strided) {
+  auto shim_dtype =
+      torch::stable::detail::to<int32_t>(torch::stable::detail::from(dtype));
+  auto shim_device_type = torch::stable::detail::to<int32_t>(
+      torch::stable::detail::from(device.type()));
+  auto shim_layout =
+      torch::stable::detail::to<int32_t>(torch::stable::detail::from(layout));
+
+  AtenTensorHandle ath;
+  if constexpr (std::is_convertible_v<F, DeleterFnPtr>) {
+    // Simple function pointer: pass it as ctx, no heap allocation.
+    auto deleter_callback = [](void* data, void* ctx) {
+      auto fn = reinterpret_cast<DeleterFnPtr>(ctx);
+      fn(data);
+    };
+    TORCH_ERROR_CODE_CHECK(torch_from_blob(
+        data,
+        sizes.size(),
+        sizes.data(),
+        strides.data(),
+        storage_offset,
+        shim_dtype,
+        shim_device_type,
+        device.index(),
+        &ath,
+        shim_layout,
+        nullptr,
+        0,
+        deleter_callback,
+        reinterpret_cast<void*>(static_cast<DeleterFnPtr>(deleter))));
+  } else {
+    // Capturing lambda: heap-allocate and type-erase.
+    F* heap_allocated_deleter = new F(std::move(deleter));
+    auto deleter_callback = [](void* data, void* ctx) {
+      F* func = static_cast<F*>(ctx);
+      (*func)(data);
+      delete func;
+    };
+    TORCH_ERROR_CODE_CHECK(torch_from_blob(
+        data,
+        sizes.size(),
+        sizes.data(),
+        strides.data(),
+        storage_offset,
+        shim_dtype,
+        shim_device_type,
+        device.index(),
+        &ath,
+        shim_layout,
+        nullptr,
+        0,
+        deleter_callback,
+        static_cast<void*>(heap_allocated_deleter)));
+  }
+  return torch::stable::Tensor(ath);
+}
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+
+/// Stable version of the to.dtype_layout op.
+///
+/// Converts a tensor to the specified dtype, layout, device, and/or memory
+/// format. Returns a new tensor with the specified properties.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor.
+/// @param dtype Optional target scalar type.
+/// @param layout Optional target memory layout.
+/// @param device Optional target device.
+/// @param pin_memory Optional flag to use pinned memory.
+/// @param non_blocking If true, the operation may be asynchronous. Defaults to
+///                     false.
+/// @param copy If true, always create a copy. Defaults to false.
+/// @param memory_format Optional target memory format.
+/// @return A tensor with the specified properties.
+inline torch::stable::Tensor to(
+    const torch::stable::Tensor& self,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt,
+    std::optional<torch::headeronly::Layout> layout = std::nullopt,
+    std::optional<torch::stable::Device> device = std::nullopt,
+    std::optional<bool> pin_memory = std::nullopt,
+    bool non_blocking = false,
+    bool copy = false,
+    std::optional<torch::headeronly::MemoryFormat> memory_format =
+        std::nullopt) {
+  const auto num_args = 8;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(dtype),
+      torch::stable::detail::from(layout),
+      torch::stable::detail::from(device),
+      torch::stable::detail::from(pin_memory),
+      torch::stable::detail::from(non_blocking),
+      torch::stable::detail::from(copy),
+      torch::stable::detail::from(memory_format)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::to", "dtype_layout", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Convenience overload for moving a tensor to a device.
+///
+/// Moves the tensor to the specified device. This is a convenience wrapper
+/// around the full to() function.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor.
+/// @param device The target device.
+/// @param non_blocking If true, the operation may be asynchronous. Defaults to
+///                     false.
+/// @param copy If true, always create a copy. Defaults to false.
+/// @return A tensor on the specified device.
+inline torch::stable::Tensor to(
+    const torch::stable::Tensor& self,
+    torch::stable::Device device,
+    bool non_blocking = false,
+    bool copy = false) {
+  return to(
+      self,
+      std::nullopt,
+      std::nullopt,
+      device,
+      std::nullopt,
+      non_blocking,
+      copy,
+      std::nullopt);
+}
+
+/// Stable version of the contiguous op.
+///
+/// Returns a contiguous in memory tensor containing the same data as the input
+/// tensor. If the input tensor is already contiguous in the specified memory
+/// format, the input tensor is returned.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor.
+/// @param memory_format The desired memory format.
+/// @return A contiguous tensor.
+inline torch::stable::Tensor contiguous(
+    const torch::stable::Tensor& self,
+    torch::headeronly::MemoryFormat memory_format =
+        torch::headeronly::MemoryFormat::Contiguous) {
+  const auto num_args = 2;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(memory_format)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::contiguous", "", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the new_empty op (2.10 version with full kwargs).
+///
+/// Creates a new uninitialized tensor with the specified size and options.
+/// This version supports all tensor creation kwargs. For versions < 2.10,
+/// a simpler overload that only takes dtype is available.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor whose properties may be inherited if kwargs are
+///             not provided.
+/// @param size The desired size of the output tensor.
+/// @param dtype Optional scalar type for the tensor elements.
+/// @param layout Optional memory layout (e.g., strided, sparse).
+/// @param device Optional device to place the tensor on.
+/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors).
+/// @return A new uninitialized tensor with the specified properties.
+inline torch::stable::Tensor new_empty(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef size,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt,
+    std::optional<torch::headeronly::Layout> layout = std::nullopt,
+    std::optional<torch::stable::Device> device = std::nullopt,
+    std::optional<bool> pin_memory = std::nullopt) {
+  const auto num_args = 6;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(size),
+      torch::stable::detail::from(dtype),
+      torch::stable::detail::from(layout),
+      torch::stable::detail::from(device),
+      torch::stable::detail::from(pin_memory)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::new_empty", "", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the new_zeros op (2.10 version with full kwargs).
+///
+/// Creates a new zero-filled tensor with the specified size and options.
+/// This version supports all tensor creation kwargs. For versions < 2.10,
+/// a simpler overload that only takes dtype is available.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor whose properties may be inherited if kwargs are
+///             not provided.
+/// @param size The desired size of the output tensor.
+/// @param dtype Optional scalar type for the tensor elements.
+/// @param layout Optional memory layout (e.g., strided, sparse).
+/// @param device Optional device to place the tensor on.
+/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors).
+/// @return A new zero-filled tensor with the specified properties.
+inline torch::stable::Tensor new_zeros(
+    const torch::stable::Tensor& self,
+    torch::headeronly::IntHeaderOnlyArrayRef size,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt,
+    std::optional<torch::headeronly::Layout> layout = std::nullopt,
+    std::optional<torch::stable::Device> device = std::nullopt,
+    std::optional<bool> pin_memory = std::nullopt) {
+  const auto num_args = 6;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(size),
+      torch::stable::detail::from(dtype),
+      torch::stable::detail::from(layout),
+      torch::stable::detail::from(device),
+      torch::stable::detail::from(pin_memory)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::new_zeros", "", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the sum.dim_IntList op.
+///
+/// Computes the sum of the input tensor along the specified dimensions.
+/// If dim is not provided, sums over all dimensions.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param self The input tensor.
+/// @param dim Optional dimensions to reduce. If not provided, reduces all
+///            dimensions.
+/// @param keepdim Whether to retain the reduced dimensions. Defaults to false.
+/// @param dtype Optional output dtype. If not provided, uses the input dtype.
+/// @return A tensor containing the sum.
+inline torch::stable::Tensor sum(
+    const torch::stable::Tensor& self,
+    std::optional<torch::headeronly::IntHeaderOnlyArrayRef> dim = std::nullopt,
+    bool keepdim = false,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt) {
+  const auto num_args = 4;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(dim),
+      torch::stable::detail::from(keepdim),
+      torch::stable::detail::from(dtype)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::sum", "dim_IntList", stack.data(), TORCH_ABI_VERSION));
+  return torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+}
+
+/// Stable version of the sum.IntList_out op.
+///
+/// Computes the sum of the input tensor along the specified dimensions,
+/// storing the result in the provided output tensor. Following C++ convention,
+/// the out parameter comes first.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @param out The output tensor (modified in-place).
+/// @param self The input tensor.
+/// @param dim Optional dimensions to reduce.
+/// @param keepdim Whether to retain the reduced dimensions. Defaults to false.
+/// @param dtype Optional output dtype.
+/// @return Reference to the output tensor.
+inline torch::stable::Tensor& sum_out(
+    torch::stable::Tensor& out,
+    const torch::stable::Tensor& self,
+    std::optional<torch::headeronly::IntHeaderOnlyArrayRef> dim = std::nullopt,
+    bool keepdim = false,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt) {
+  const auto num_args = 5;
+  std::array<StableIValue, num_args> stack{
+      torch::stable::detail::from(self),
+      torch::stable::detail::from(dim),
+      torch::stable::detail::from(keepdim),
+      torch::stable::detail::from(dtype),
+      torch::stable::detail::from(out)};
+  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
+      "aten::sum", "IntList_out", stack.data(), TORCH_ABI_VERSION));
+  // Clean up the handle in stack[0], discard the temporary
+  (void)torch::stable::detail::to<torch::stable::Tensor>(stack[0]);
+  return out;
+}
+
+/// Stable version of the subtract.Tensor op.
+///
+/// Subtracts the other tensor from self, with an optional scaling factor alpha.
+/// Computes: self - alpha * other.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @note The alpha parameter is typed as double
+///       API uses double for the Scalar parameter.
+///
+/// @param self The input tensor.
+/// @param other The tensor to subtract.
+/// @param alpha The scaling factor for other. Defaults to 1.0.
+/// @return The result of self - alpha * other.
+inline torch::stable::Tensor subtract(
+    const torch::stable::Tensor& self,
+    const torch::stable::Tensor& other,
+    double alpha = 1.0) {
+  AtenTensorHandle ret0;
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_aten_subtract_Tensor(self.get(), other.get(), alpha, &ret0));
+  return torch::stable::Tensor(ret0);
+}
+
+/// Stable version of the full.default op.
+///
+/// Creates a tensor of the specified size filled with the given value.
+///
+/// Minimum compatible version: PyTorch 2.10.
+///
+/// @note The fill_value parameter is typed
+///       C shim API uses double for the Scalar parameter.
+///
+/// @param size The desired size of the output tensor.
+/// @param fill_value The value to fill the tensor with.
+/// @param dtype Optional scalar type for the tensor elements.
+/// @param layout Optional memory layout.
+/// @param device Optional device to place the tensor on.
+/// @param pin_memory Optional flag to use pinned memory.
+/// @return A new tensor filled with the specified value.
+inline torch::stable::Tensor full(
+    torch::headeronly::IntHeaderOnlyArrayRef size,
+    double fill_value,
+    std::optional<torch::headeronly::ScalarType> dtype = std::nullopt,
+    std::optional<torch::headeronly::Layout> layout = std::nullopt,
+    std::optional<torch::stable::Device> device = std::nullopt,
+    std::optional<bool> pin_memory = std::nullopt) {
+  int32_t* dtype_ptr = nullptr;
+  int32_t dtype_val;
+  if (dtype.has_value()) {
+    dtype_val = torch::stable::detail::to<int32_t>(
+        torch::stable::detail::from(dtype.value()));
+    dtype_ptr = &dtype_val;
+  }
+
+  int32_t* layout_ptr = nullptr;
+  int32_t layout_val;
+  if (layout.has_value()) {
+    layout_val = torch::stable::detail::to<int32_t>(
+        torch::stable::detail::from(layout.value()));
+    layout_ptr = &layout_val;
+  }
+
+  int32_t* device_type_ptr = nullptr;
+  int32_t device_type_val;
+  int32_t device_index = 0;
+  if (device.has_value()) {
+    device_type_val = torch::stable::detail::to<int32_t>(
+        torch::stable::detail::from(device.value().type()));
+    device_type_ptr = &device_type_val;
+    device_index = device.value().index();
+  }
+
+  int32_t* pin_memory_ptr = nullptr;
+  int32_t pin_memory_val;
+  if (pin_memory.has_value()) {
+    pin_memory_val = pin_memory.value() ? 1 : 0;
+    pin_memory_ptr = &pin_memory_val;
+  }
+
+  AtenTensorHandle ret0;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_aten_full(
+      size.data(),
+      static_cast<int64_t>(size.size()),
+      fill_value,
+      dtype_ptr,
+      layout_ptr,
+      device_type_ptr,
+      device_index,
+      pin_memory_ptr,
+      &ret0));
+
+  return torch::stable::Tensor(ret0);
+}
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+HIDDEN_NAMESPACE_END(torch, stable)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/stableivalue_conversions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/stableivalue_conversions.h
new file mode 100644
index 0000000000000000000000000000000000000000..eece32c7eb33defaed09117fd53e89e475a3d4e7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/stableivalue_conversions.h
@@ -0,0 +1,996 @@
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/device_struct.h>
+#include <torch/csrc/stable/tensor_struct.h>
+#include <torch/headeronly/core/DeviceType.h>
+#include <torch/headeronly/core/Layout.h>
+#include <torch/headeronly/core/MemoryFormat.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/core/enum_tag.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Deprecated.h>
+#include <torch/headeronly/util/Exception.h>
+#include <torch/headeronly/util/shim_utils.h>
+
+#include <optional>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable, detail)
+
+// Helper variable templates to detect 2.10+ types for better compile-time error
+// messages
+template <typename T>
+inline constexpr bool is_header_only_array_ref_v = false;
+
+template <typename T>
+inline constexpr bool
+    is_header_only_array_ref_v<torch::headeronly::HeaderOnlyArrayRef<T>> = true;
+
+template <typename T>
+inline constexpr bool is_std_vector_v = false;
+
+template <typename T>
+inline constexpr bool is_std_vector_v<std::vector<T>> = true;
+
+// forward declare so that the from/to() implementations in the detail
+// namespace of library.h where the real work is done can compile.
+template <typename T>
+StableIValue from(T val);
+template <typename T>
+T to(StableIValue val);
+
+// =============================================================================
+//  Below are the helpers for converting between StableIValue and T
+// =============================================================================
+// =============================================================================
+// FROM CONVERSIONS (T -> StableIValue)
+// ======================================================================
+
+// Specialization for general copyable types (catch-all) => StableIValue
+template <typename T>
+struct FromImpl {
+  static StableIValue call(
+      T val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    // Ensure 2.10+ types don't accidentally use the base case - provide clear
+    // compile-time errors.
+    static_assert(
+        !std::is_same_v<T, torch::stable::Device>,
+        "torch::stable::Device requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(
+        !is_header_only_array_ref_v<T>,
+        "HeaderOnlyArrayRef<T> requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(
+        !is_std_vector_v<T>,
+        "std::vector<T> requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(
+        !std::is_same_v<T, std::string>,
+        "std::string requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(
+        sizeof(T) <= sizeof(StableIValue),
+        "StableLibrary stack does not support parameter types larger than 64 bits.");
+    static_assert(std::is_trivially_copyable_v<T>);
+    // Initialization should be cheap enough; let's give people well-specified
+    // reproducible behavior.
+    StableIValue result = 0;
+    // NOTE [ -Wclass-memaccess ]: reinterpret_cast to suppress
+    // overzealous -Wclass-memaccess. (see
+    // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=107361) We have a
+    // static_assert above that T is trivially copyable, which should be
+    // enough.
+#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
+    std::memcpy(&result, reinterpret_cast<const void*>(&val), sizeof(val));
+#elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+    // if value has size less than sizeof(StableIValue), then only lowest bytes
+    // have to be updated
+    std::memcpy(
+        reinterpret_cast<unsigned char*>(&result) + sizeof(StableIValue) -
+            sizeof(val),
+        reinterpret_cast<const void*>(&val),
+        sizeof(val));
+#else
+#error "Unexpected or undefined __BYTE_ORDER__"
+#endif
+    return result;
+  }
+};
+
+// Specialization for torch::headeronly::ScalarType => StableIValue
+// Note that we call into the shim to translate between the user's
+// ScalarType and libtorch's ScalarType, which can be different!
+// Also note that the list below is not comprehensive, as it does not
+// include types that are no longer really used and should probably be
+// deprecated (like qint8).
+using torch::headeronly::ScalarType;
+template <>
+struct FromImpl<ScalarType> {
+  static StableIValue call(
+      ScalarType val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    switch (val) {
+      case ScalarType::Byte:
+        return torch::stable::detail::from(aoti_torch_dtype_uint8());
+      case ScalarType::Char:
+        return torch::stable::detail::from(aoti_torch_dtype_int8());
+      case ScalarType::Short:
+        return torch::stable::detail::from(aoti_torch_dtype_int16());
+      case ScalarType::Int:
+        return torch::stable::detail::from(aoti_torch_dtype_int32());
+      case ScalarType::Long:
+        return torch::stable::detail::from(aoti_torch_dtype_int64());
+      case ScalarType::Half:
+        return torch::stable::detail::from(aoti_torch_dtype_float16());
+      case ScalarType::Float:
+        return torch::stable::detail::from(aoti_torch_dtype_float32());
+      case ScalarType::Double:
+        return torch::stable::detail::from(aoti_torch_dtype_float64());
+      case ScalarType::ComplexHalf:
+        return torch::stable::detail::from(aoti_torch_dtype_complex32());
+      case ScalarType::ComplexFloat:
+        return torch::stable::detail::from(aoti_torch_dtype_complex64());
+      case ScalarType::ComplexDouble:
+        return torch::stable::detail::from(aoti_torch_dtype_complex128());
+      case ScalarType::Bool:
+        return torch::stable::detail::from(aoti_torch_dtype_bool());
+      case ScalarType::BFloat16:
+        return torch::stable::detail::from(aoti_torch_dtype_bfloat16());
+      case ScalarType::Float8_e5m2:
+        return torch::stable::detail::from(aoti_torch_dtype_float8_e5m2());
+      case ScalarType::Float8_e4m3fn:
+        return torch::stable::detail::from(aoti_torch_dtype_float8_e4m3fn());
+      case ScalarType::Float8_e5m2fnuz:
+        return torch::stable::detail::from(aoti_torch_dtype_float8_e5m2fnuz());
+      case ScalarType::Float8_e4m3fnuz:
+        return torch::stable::detail::from(aoti_torch_dtype_float8_e4m3fnuz());
+      case ScalarType::UInt16:
+        return torch::stable::detail::from(aoti_torch_dtype_uint16());
+      case ScalarType::UInt32:
+        return torch::stable::detail::from(aoti_torch_dtype_uint32());
+      case ScalarType::UInt64:
+        return torch::stable::detail::from(aoti_torch_dtype_uint64());
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+      case ScalarType::Float8_e8m0fnu:
+        return torch::stable::detail::from(torch_dtype_float8_e8m0fnu());
+      case ScalarType::Float4_e2m1fn_x2:
+        return torch::stable::detail::from(torch_dtype_float4_e2m1fn_x2());
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+      default:
+        STD_TORCH_CHECK(
+            false,
+            "Not yet supported ScalarType ",
+            toString(val),
+            ", please file an issue describing your use case.");
+    }
+  }
+};
+
+// [Note DeviceType version guard]
+// This conversion was introduced in 2.10. However, we do not gate it
+// with TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 because this
+// conversion is not actually used to pass DeviceType between user
+// extensions and libtorch (i.e. there is no c10::TypeKind::DeviceType).
+// The purpose of gating other conversions is to ensure that user
+// extensions do not try to pass a StableIValue that libtorch is
+// unable to interpret.
+// This conversion is only used
+// (1) In the conversion for torch::stable::Device (already gated)
+// (2) Within the user extension to translate between libtorch/extension's
+//     DeviceType (no gating needed)
+// Specialization for torch::headeronly::DeviceType => StableIValue
+// Note that we call into the shim to translate between the user's
+// DeviceType and libtorch's DeviceType, which can be different!
+using torch::headeronly::DeviceType;
+template <>
+struct FromImpl<DeviceType> {
+  static StableIValue call(
+      DeviceType val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    switch (val) {
+      case DeviceType::CPU:
+        return torch::stable::detail::from(aoti_torch_device_type_cpu());
+      case DeviceType::CUDA:
+        return torch::stable::detail::from(aoti_torch_device_type_cuda());
+      case DeviceType::Meta:
+        return torch::stable::detail::from(aoti_torch_device_type_meta());
+      case DeviceType::XPU:
+        return torch::stable::detail::from(aoti_torch_device_type_xpu());
+      case DeviceType::MPS:
+        return torch::stable::detail::from(aoti_torch_device_type_mps());
+      case DeviceType::PrivateUse1:
+        return torch::stable::detail::from(
+            aoti_torch_device_type_privateuse1());
+      default:
+        STD_TORCH_CHECK(
+            false,
+            "Not yet supported DeviceType, please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for std::nullopt_t => StableIValue
+template <>
+struct FromImpl<std::nullopt_t> {
+  static StableIValue call(
+      std::nullopt_t val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    return torch::stable::detail::from(nullptr);
+  }
+};
+
+// Specialization for std::optional => StableIValue
+// [Handling std::optional]
+// When the schema is represented by an optional type, say int?, then we
+// expect the custom extension representation to be a std::optional<int>
+// (critically NOT int!). In order for all parameters to be stably parsed and
+// handled by our dispatcher, we liaison custom extension parameters through
+// boxed kernels, meaning that every value will make its way to be an IValue:
+//
+// custom extension value --(from)-> StableIValue --(to_ivalue)-> IValue
+//
+// When the custom extension value is a literal that can be trivially
+// casted to StableIValue, e.g., an int, a float, a pointer, this route is
+// ...trivial. The below specialization is for a case when the custom
+// extension value would NOT fit within a StableIValue: a std::optional.
+//
+// If the std::optional has no value, it is treated as std::nullopt,
+// whose StableIValue representation is from(nullptr). Otherwise, we:
+// 1. unwrap the std::optional<T>
+// 2. recursively convert its value of type T to a StableIValue
+// 3. allocate heap space for said StableIValue
+// 4. convert the resulting StableIValue* into a StableIValue
+//
+// note that this allocates heap memory! which we expect to be cleaned
+// up in the to_ivalue() function defined in shim_common.cpp. We
+// purposefully hide this implementation detail from the user so that
+// all the user needs to know is:
+//
+// The schema requests an optional (T?) so I must call `from` on a
+// std::optional<T> or a std::nullopt.
+template <typename T>
+struct FromImpl<std::optional<T>> {
+  static StableIValue call(
+      const std::optional<T>& val,
+      uint64_t extension_build_version,
+      bool is_internal) {
+    if (!val.has_value()) {
+      return torch::stable::detail::from(std::nullopt);
+    }
+    return torch::stable::detail::from(
+        new StableIValue(detail::FromImpl<T>::call(
+            val.value(), extension_build_version, is_internal)));
+  }
+};
+
+// Specialization for torch::stable::Tensor => StableIValue
+// Returns a new owning reference of the underlying Tensor.
+template <>
+struct FromImpl<torch::stable::Tensor> {
+  static StableIValue call(
+      const torch::stable::Tensor& val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    AtenTensorHandle new_ath;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_new_tensor_handle(val.get(), &new_ath));
+    return torch::stable::detail::from(new_ath);
+  }
+};
+
+// =============================================================================
+// FROM CONVERSIONS requiring TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+// =============================================================================
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+// Specialization for torch::headeronly::Layout => StableIValue
+// Note that we call into the shim to translate between the user's
+// Layout and libtorch's Layout, which can be different!
+using torch::headeronly::Layout;
+template <>
+struct FromImpl<Layout> {
+  static StableIValue call(
+      Layout val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    switch (val) {
+      case Layout::Strided:
+        return torch::stable::detail::from(aoti_torch_layout_strided());
+      case Layout::Sparse:
+        return torch::stable::detail::from(aoti_torch_layout_sparse_coo());
+      case Layout::SparseCsr:
+        return torch::stable::detail::from(aoti_torch_layout_sparse_csr());
+      case Layout::SparseCsc:
+        return torch::stable::detail::from(aoti_torch_layout_sparse_csc());
+      case Layout::SparseBsr:
+        return torch::stable::detail::from(aoti_torch_layout_sparse_bsr());
+      case Layout::SparseBsc:
+        return torch::stable::detail::from(aoti_torch_layout_sparse_bsc());
+      case Layout::Mkldnn:
+        return torch::stable::detail::from(aoti_torch_layout__mkldnn());
+      case Layout::Jagged:
+        return torch::stable::detail::from(aoti_torch_layout_jagged());
+      default:
+        STD_TORCH_CHECK(
+            false,
+            "Not yet supported Layout, please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for torch::headeronly::MemoryFormat => StableIValue
+// Note that we call into the shim to translate between the user's
+// MemoryFormat and libtorch's MemoryFormat, which can be different!
+using torch::headeronly::MemoryFormat;
+template <>
+struct FromImpl<MemoryFormat> {
+  static StableIValue call(
+      MemoryFormat val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    switch (val) {
+      case MemoryFormat::Contiguous:
+        return torch::stable::detail::from(
+            aoti_torch_memory_format_contiguous_format());
+      case MemoryFormat::Preserve:
+        return torch::stable::detail::from(
+            aoti_torch_memory_format_preserve_format());
+      case MemoryFormat::ChannelsLast:
+        return torch::stable::detail::from(
+            aoti_torch_memory_format_channels_last());
+      case MemoryFormat::ChannelsLast3d:
+        return torch::stable::detail::from(
+            aoti_torch_memory_format_channels_last_3d());
+      default:
+        STD_TORCH_CHECK(
+            false,
+            "Not yet supported MemoryFormat, please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for torch::headeronly::HeaderOnlyArrayRef<T> => StableIValue
+// Returns a new owning reference of the underlying list.
+template <typename T>
+struct FromImpl<torch::headeronly::HeaderOnlyArrayRef<T>> {
+  static StableIValue call(
+      const torch::headeronly::HeaderOnlyArrayRef<T>& val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    StableListHandle new_list_handle;
+    try {
+      TORCH_ERROR_CODE_CHECK(
+          torch_new_list_reserve_size(val.size(), &new_list_handle));
+      for (const auto& elem : val) {
+        TORCH_ERROR_CODE_CHECK(torch_list_push_back(
+            new_list_handle, torch::stable::detail::from(elem)));
+      }
+      return torch::stable::detail::from(new_list_handle);
+    } catch (const std::runtime_error&) {
+      if (new_list_handle != nullptr) {
+        // clean up memory if an error was thrown
+        TORCH_ERROR_CODE_CHECK(torch_delete_list(new_list_handle));
+      }
+      throw;
+    }
+  }
+};
+
+// Specialization for std::vector<T> => StableIValue, which is implemented the
+// same way as HeaderOnlyArrayRef<T> => StableIValue
+// Returns a new owning reference of the underlying list.
+template <typename T>
+struct FromImpl<std::vector<T>> {
+  static StableIValue call(
+      const std::vector<T>& val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    return torch::stable::detail::from<
+        torch::headeronly::HeaderOnlyArrayRef<T>>(val);
+  }
+};
+
+// Specialization for torch::stable::Device => StableIValue
+// Pack the device type and index into a StableIValue in a platform-independent
+// format. We use the shim representation for DeviceType (int32_t) for ABI
+// stability. StableIValue layout: DeviceIndex in lower 32 bits,
+// DeviceType (shim int32_t) in upper 32 bits
+template <>
+struct FromImpl<torch::stable::Device> {
+  static StableIValue call(
+      const torch::stable::Device& val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    // Convert DeviceType to shim representation (int32_t)
+    StableIValue device_type_shim = torch::stable::detail::from(val.type());
+    // Pack: lower 32 bits = device index, upper 32 bits = device type (shim)
+    uint64_t device_index_bits =
+        static_cast<uint64_t>(static_cast<uint32_t>(val.index()));
+    uint64_t device_type_bits =
+        static_cast<uint64_t>(static_cast<uint32_t>(device_type_shim)) << 32;
+    return device_index_bits | device_type_bits;
+  }
+};
+
+// Specialization for std::string, which should return a new owning reference of
+// the string
+template <>
+struct FromImpl<std::string> {
+  static StableIValue call(
+      const std::string& val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    StringHandle handle;
+    TORCH_ERROR_CODE_CHECK(
+        torch_new_string_handle(val.c_str(), val.length(), &handle))
+    return torch::stable::detail::from(handle);
+  }
+};
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+// =============================================================================
+// TO CONVERSIONS (StableIValue -> T)
+// =============================================================================
+
+// Specialization for StableIValue => general copyable types (catch-all)
+template <typename T>
+struct ToImpl {
+  static T call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    // Ensure 2.10+ types don't accidentally use the base case - provide clear
+    // compile-time errors.
+    static_assert(
+        !std::is_same_v<T, torch::stable::Device>,
+        "torch::stable::Device requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(
+        !is_header_only_array_ref_v<T>,
+        "HeaderOnlyArrayRef<T> requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(
+        !is_std_vector_v<T>,
+        "std::vector<T> requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(
+        !std::is_same_v<T, std::string>,
+        "std::string requires TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0");
+    static_assert(std::is_trivially_copyable_v<T>);
+    // T may not have a default constructor. (For example, it might be
+    // c10::Device.) However, std::memcpy implicitly creates a T at the
+    // destination. So, we can use a union to work around this lack of
+    // default constructor.
+    union Result {
+      Result() {}
+      T t;
+    };
+    Result result;
+    // See NOTE[ -Wclass-memaccess ] above.
+#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
+    std::memcpy(reinterpret_cast<void*>(&result.t), &val, sizeof(result));
+#elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+    static_assert(
+        sizeof(T) <= sizeof(StableIValue),
+        "StableLibrary stack does not support parameter types larger than 64 bits.");
+    // if value has size less than sizeof(StableIValue), then only lowest bytes
+    // have to be updated
+    std::memcpy(
+        reinterpret_cast<void*>(&result.t),
+        reinterpret_cast<unsigned char*>(&val) + sizeof(StableIValue) -
+            sizeof(result),
+        sizeof(result));
+#else
+#error "Unexpected or undefined __BYTE_ORDER__"
+#endif
+    return result.t;
+  }
+};
+
+// Specialization for StableIValue => torch::headeronly::ScalarType
+template <>
+struct ToImpl<ScalarType> {
+  static ScalarType call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    int32_t shim_scalartype = torch::stable::detail::to<int32_t>(val);
+    if (shim_scalartype == aoti_torch_dtype_uint8()) {
+      return ScalarType::Byte;
+    } else if (shim_scalartype == aoti_torch_dtype_int8()) {
+      return ScalarType::Char;
+    } else if (shim_scalartype == aoti_torch_dtype_int16()) {
+      return ScalarType::Short;
+    } else if (shim_scalartype == aoti_torch_dtype_int32()) {
+      return ScalarType::Int;
+    } else if (shim_scalartype == aoti_torch_dtype_int64()) {
+      return ScalarType::Long;
+    } else if (shim_scalartype == aoti_torch_dtype_float16()) {
+      return ScalarType::Half;
+    } else if (shim_scalartype == aoti_torch_dtype_float32()) {
+      return ScalarType::Float;
+    } else if (shim_scalartype == aoti_torch_dtype_float64()) {
+      return ScalarType::Double;
+    } else if (shim_scalartype == aoti_torch_dtype_complex32()) {
+      return ScalarType::ComplexHalf;
+    } else if (shim_scalartype == aoti_torch_dtype_complex64()) {
+      return ScalarType::ComplexFloat;
+    } else if (shim_scalartype == aoti_torch_dtype_complex128()) {
+      return ScalarType::ComplexDouble;
+    } else if (shim_scalartype == aoti_torch_dtype_bool()) {
+      return ScalarType::Bool;
+    } else if (shim_scalartype == aoti_torch_dtype_bfloat16()) {
+      return ScalarType::BFloat16;
+    } else if (shim_scalartype == aoti_torch_dtype_float8_e5m2()) {
+      return ScalarType::Float8_e5m2;
+    } else if (shim_scalartype == aoti_torch_dtype_float8_e4m3fn()) {
+      return ScalarType::Float8_e4m3fn;
+    } else if (shim_scalartype == aoti_torch_dtype_float8_e5m2fnuz()) {
+      return ScalarType::Float8_e5m2fnuz;
+    } else if (shim_scalartype == aoti_torch_dtype_float8_e4m3fnuz()) {
+      return ScalarType::Float8_e4m3fnuz;
+    } else if (shim_scalartype == aoti_torch_dtype_uint16()) {
+      return ScalarType::UInt16;
+    } else if (shim_scalartype == aoti_torch_dtype_uint32()) {
+      return ScalarType::UInt32;
+    } else if (shim_scalartype == aoti_torch_dtype_uint64()) {
+      return ScalarType::UInt64;
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+    } else if (shim_scalartype == torch_dtype_float8_e8m0fnu()) {
+      return ScalarType::Float8_e8m0fnu;
+    } else if (shim_scalartype == torch_dtype_float4_e2m1fn_x2()) {
+      return ScalarType::Float4_e2m1fn_x2;
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0
+    } else {
+      STD_TORCH_CHECK(
+          false,
+          "Not yet supported ScalarType ",
+          std::to_string(shim_scalartype),
+          ", please file an issue describing your use case.");
+    }
+  }
+};
+
+// See [Note DeviceType version guard]
+// Specialization for StableIValue => torch::headeronly::DeviceType
+template <>
+struct ToImpl<DeviceType> {
+  static DeviceType call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    int32_t shim_devicetype = torch::stable::detail::to<int32_t>(val);
+    if (shim_devicetype == aoti_torch_device_type_cpu()) {
+      return DeviceType::CPU;
+    } else if (shim_devicetype == aoti_torch_device_type_cuda()) {
+      return DeviceType::CUDA;
+    } else if (shim_devicetype == aoti_torch_device_type_meta()) {
+      return DeviceType::Meta;
+    } else if (shim_devicetype == aoti_torch_device_type_xpu()) {
+      return DeviceType::XPU;
+    } else if (shim_devicetype == aoti_torch_device_type_mps()) {
+      return DeviceType::MPS;
+    } else if (shim_devicetype == aoti_torch_device_type_privateuse1()) {
+      return DeviceType::PrivateUse1;
+    } else {
+      STD_TORCH_CHECK(
+          false,
+          "Not yet supported DeviceType ",
+          std::to_string(shim_devicetype),
+          ", please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for StableIValue => std::nullopt_t
+template <>
+struct ToImpl<std::nullopt_t> {
+  static std::nullopt_t call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    // val should be equivalent to from(nullptr)
+    return std::nullopt;
+  }
+};
+
+// Specialization for StableIValue => std::optional, see [Handling
+// std::optional] as the semantic is the same but in reverse direction as we go
+// from IValue --(from_ivalue)-> StableIValue --(to<T>)-> T in custom extension
+template <typename T>
+struct ToImpl<std::optional<T>> {
+  static std::optional<T> call(
+      StableIValue val,
+      uint64_t extension_build_version,
+      bool is_internal) {
+    auto sivp = torch::stable::detail::to<StableIValue*>(val);
+
+    // sivp is either nullptr or a pointer to a StableIValue
+    if (sivp == nullptr) {
+      return {};
+    }
+    auto inner_val =
+        detail::ToImpl<T>::call(*sivp, extension_build_version, is_internal);
+
+    // free the memory associated with StableIValue* sivp
+    delete sivp;
+
+    return std::make_optional(inner_val);
+  }
+};
+
+// Specialization for StableIValue => torch::stable::Tensor
+// The resulting stable::Tensor steals ownership of the input's
+// underlying AtenTensorHandle.
+template <>
+struct ToImpl<torch::stable::Tensor> {
+  static torch::stable::Tensor call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    return torch::stable::Tensor(
+        torch::stable::detail::to<AtenTensorHandle>(val));
+  }
+};
+
+// =============================================================================
+// TO CONVERSIONS requiring TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+// =============================================================================
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+// Specialization for StableIValue => torch::headeronly::Layout
+template <>
+struct ToImpl<Layout> {
+  static Layout call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    int32_t shim_layout = torch::stable::detail::to<int32_t>(val);
+    if (shim_layout == aoti_torch_layout_strided()) {
+      return Layout::Strided;
+    } else if (shim_layout == aoti_torch_layout_sparse_coo()) {
+      return Layout::Sparse;
+    } else if (shim_layout == aoti_torch_layout_sparse_csr()) {
+      return Layout::SparseCsr;
+    } else if (shim_layout == aoti_torch_layout_sparse_csc()) {
+      return Layout::SparseCsc;
+    } else if (shim_layout == aoti_torch_layout_sparse_bsr()) {
+      return Layout::SparseBsr;
+    } else if (shim_layout == aoti_torch_layout_sparse_bsc()) {
+      return Layout::SparseBsc;
+    } else if (shim_layout == aoti_torch_layout__mkldnn()) {
+      return Layout::Mkldnn;
+    } else if (shim_layout == aoti_torch_layout_jagged()) {
+      return Layout::Jagged;
+    } else {
+      STD_TORCH_CHECK(
+          false,
+          "Not yet supported Layout ",
+          std::to_string(shim_layout),
+          ", please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for StableIValue => torch::headeronly::MemoryFormat
+template <>
+struct ToImpl<MemoryFormat> {
+  static MemoryFormat call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    int32_t shim_memory_format = torch::stable::detail::to<int32_t>(val);
+    if (shim_memory_format == aoti_torch_memory_format_contiguous_format()) {
+      return MemoryFormat::Contiguous;
+    } else if (
+        shim_memory_format == aoti_torch_memory_format_preserve_format()) {
+      return MemoryFormat::Preserve;
+    } else if (shim_memory_format == aoti_torch_memory_format_channels_last()) {
+      return MemoryFormat::ChannelsLast;
+    } else if (
+        shim_memory_format == aoti_torch_memory_format_channels_last_3d()) {
+      return MemoryFormat::ChannelsLast3d;
+    } else {
+      STD_TORCH_CHECK(
+          false,
+          "Not yet supported MemoryFormat ",
+          std::to_string(shim_memory_format),
+          ", please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for StableIValue => std::vector<T>
+// std::vector<T> should be represented as a StableListHandle
+// filled with StableIValues
+// The new std::vector steals ownership of the underlying elements
+// and we free the underlying list referred by the input StableListHandle.
+template <typename T>
+struct ToImpl<std::vector<T>> {
+  static std::vector<T> call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    auto list_handle = torch::stable::detail::to<StableListHandle>(val);
+    size_t size;
+    try {
+      TORCH_ERROR_CODE_CHECK(torch_list_size(list_handle, &size));
+      std::vector<T> result;
+      result.reserve(size);
+      for (size_t i = 0; i < size; i++) {
+        StableIValue element;
+        TORCH_ERROR_CODE_CHECK(torch_list_get_item(list_handle, i, &element));
+        result.push_back(torch::stable::detail::to<T>(element));
+      }
+      TORCH_ERROR_CODE_CHECK(torch_delete_list(list_handle));
+      return result;
+    } catch (const std::runtime_error&) {
+      // clean up memory if an exception is thrown, and rethrow
+      TORCH_ERROR_CODE_CHECK(torch_delete_list(list_handle));
+      throw;
+    }
+  }
+};
+
+// Specialization for StableIValue => torch::stable::Device
+// Unpack device type and index from StableIValue in platform-independent
+// format. StableIValue layout: DeviceIndex in lower 32 bits,
+// DeviceType (shim int32_t) in upper 32 bits
+template <>
+struct ToImpl<torch::stable::Device> {
+  static torch::stable::Device call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    // Unpack: lower 32 bits = device index, upper 32 bits = device type (shim)
+    int32_t device_index = static_cast<int32_t>(val & 0xFFFFFFFF);
+    StableIValue device_type_shim = (val >> 32) & 0xFFFFFFFF;
+    DeviceType device_type =
+        torch::stable::detail::to<DeviceType>(device_type_shim);
+    return torch::stable::Device(device_type, device_index);
+  }
+};
+
+// Specialization for std::string
+// Returns a new std::string; the string in val is deleted.
+template <>
+struct ToImpl<std::string> {
+  static std::string call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    StringHandle handle = torch::stable::detail::to<StringHandle>(val);
+    size_t length;
+    TORCH_ERROR_CODE_CHECK(torch_string_length(handle, &length));
+    const char* data;
+    TORCH_ERROR_CODE_CHECK(torch_string_c_str(handle, &data));
+    auto strptr = new std::string(data, length);
+
+    // delete the old string before returning new string
+    TORCH_ERROR_CODE_CHECK(torch_delete_string(handle));
+    return *strptr;
+  }
+};
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+// =============================================================================
+// FROM/TO CONVERSIONS requiring TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+// =============================================================================
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
+// Specialization for torch::headeronly::Tag => StableIValue
+// Uses shim getter functions so the integer representation is resolved at
+// runtime from libtorch, not baked in at extension compile time.
+using torch::headeronly::Tag;
+template <>
+struct FromImpl<Tag> {
+  static StableIValue call(
+      Tag val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    switch (val) {
+      case Tag::core:
+        return torch::stable::detail::from(torch_tag_core());
+      case Tag::cudagraph_unsafe:
+        return torch::stable::detail::from(torch_tag_cudagraph_unsafe());
+      case Tag::data_dependent_output:
+        return torch::stable::detail::from(torch_tag_data_dependent_output());
+      case Tag::dynamic_output_shape:
+        return torch::stable::detail::from(torch_tag_dynamic_output_shape());
+      case Tag::flexible_layout:
+        return torch::stable::detail::from(torch_tag_flexible_layout());
+      case Tag::generated:
+        return torch::stable::detail::from(torch_tag_generated());
+      case Tag::inplace_view:
+        return torch::stable::detail::from(torch_tag_inplace_view());
+      case Tag::maybe_aliasing_or_mutating:
+        return torch::stable::detail::from(
+            torch_tag_maybe_aliasing_or_mutating());
+      case Tag::needs_contiguous_strides:
+        return torch::stable::detail::from(
+            torch_tag_needs_contiguous_strides());
+      case Tag::needs_exact_strides:
+        return torch::stable::detail::from(torch_tag_needs_exact_strides());
+      case Tag::needs_fixed_stride_order:
+        return torch::stable::detail::from(
+            torch_tag_needs_fixed_stride_order());
+      case Tag::nondeterministic_bitwise:
+        return torch::stable::detail::from(
+            torch_tag_nondeterministic_bitwise());
+      case Tag::nondeterministic_seeded:
+        return torch::stable::detail::from(torch_tag_nondeterministic_seeded());
+      case Tag::out_variant:
+        return torch::stable::detail::from(torch_tag_out_variant());
+      case Tag::pointwise:
+        return torch::stable::detail::from(torch_tag_pointwise());
+      case Tag::pt2_compliant_tag:
+        return torch::stable::detail::from(torch_tag_pt2_compliant_tag());
+      case Tag::reduction:
+        return torch::stable::detail::from(torch_tag_reduction());
+      case Tag::view_copy:
+        return torch::stable::detail::from(torch_tag_view_copy());
+      default:
+        STD_TORCH_CHECK(
+            false,
+            "Not yet supported Tag, please file an issue describing your use case.");
+    }
+  }
+};
+
+// Specialization for StableIValue => torch::headeronly::Tag
+template <>
+struct ToImpl<Tag> {
+  static Tag call(
+      StableIValue val,
+      [[maybe_unused]] uint64_t extension_build_version,
+      [[maybe_unused]] bool is_internal) {
+    int32_t shim_tag = torch::stable::detail::to<int32_t>(val);
+    if (shim_tag == torch_tag_core()) {
+      return Tag::core;
+    } else if (shim_tag == torch_tag_cudagraph_unsafe()) {
+      return Tag::cudagraph_unsafe;
+    } else if (shim_tag == torch_tag_data_dependent_output()) {
+      return Tag::data_dependent_output;
+    } else if (shim_tag == torch_tag_dynamic_output_shape()) {
+      return Tag::dynamic_output_shape;
+    } else if (shim_tag == torch_tag_flexible_layout()) {
+      return Tag::flexible_layout;
+    } else if (shim_tag == torch_tag_generated()) {
+      return Tag::generated;
+    } else if (shim_tag == torch_tag_inplace_view()) {
+      return Tag::inplace_view;
+    } else if (shim_tag == torch_tag_maybe_aliasing_or_mutating()) {
+      return Tag::maybe_aliasing_or_mutating;
+    } else if (shim_tag == torch_tag_needs_contiguous_strides()) {
+      return Tag::needs_contiguous_strides;
+    } else if (shim_tag == torch_tag_needs_exact_strides()) {
+      return Tag::needs_exact_strides;
+    } else if (shim_tag == torch_tag_needs_fixed_stride_order()) {
+      return Tag::needs_fixed_stride_order;
+    } else if (shim_tag == torch_tag_nondeterministic_bitwise()) {
+      return Tag::nondeterministic_bitwise;
+    } else if (shim_tag == torch_tag_nondeterministic_seeded()) {
+      return Tag::nondeterministic_seeded;
+    } else if (shim_tag == torch_tag_out_variant()) {
+      return Tag::out_variant;
+    } else if (shim_tag == torch_tag_pointwise()) {
+      return Tag::pointwise;
+    } else if (shim_tag == torch_tag_pt2_compliant_tag()) {
+      return Tag::pt2_compliant_tag;
+    } else if (shim_tag == torch_tag_reduction()) {
+      return Tag::reduction;
+    } else if (shim_tag == torch_tag_view_copy()) {
+      return Tag::view_copy;
+    } else {
+      STD_TORCH_CHECK(
+          false,
+          "Not yet supported Tag ",
+          std::to_string(shim_tag),
+          ", please file an issue describing your use case.");
+    }
+  }
+};
+
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_12_0
+
+// =============================================================================
+//  end to helpers for converting between StableIValue and T
+// =============================================================================
+
+// Expose the partially templated class functions through single functions
+// The non-private versions will be used by the extension or headers that
+// the extension includes.
+template <typename T>
+inline StableIValue from(T val) {
+  return detail::FromImpl<T>::call(
+      val, aoti_torch_abi_version(), /*is_internal=*/false);
+}
+
+template <typename T>
+inline StableIValue from(const std::optional<T>& val) {
+  return detail::FromImpl<std::optional<T>>::call(
+      val, aoti_torch_abi_version(), /*is_internal=*/false);
+}
+
+// The below overload is used! See https://godbolt.org/z/859cshxrW
+// We are suppressing the warning for versions clang12- and gcc11-
+[[maybe_unused]] inline StableIValue from(const torch::stable::Tensor& val) {
+  return detail::FromImpl<torch::stable::Tensor>::call(
+      val, aoti_torch_abi_version(), /*is_internal=*/false);
+}
+
+template <typename T>
+inline T to(StableIValue val) {
+  return detail::ToImpl<T>::call(
+      val, aoti_torch_abi_version(), /*is_internal=*/false);
+}
+
+// Internal conversion functions used by from_ivalue and to_ivalue.
+// These are used in libtorch
+template <typename T>
+inline StableIValue _from(T val, uint64_t extension_build_version) {
+  return detail::FromImpl<T>::call(
+      val, extension_build_version, /*is_internal=*/true);
+}
+
+template <typename T>
+inline StableIValue _from(
+    const std::optional<T>& val,
+    uint64_t extension_build_version) {
+  return detail::FromImpl<std::optional<T>>::call(
+      val, extension_build_version, /*is_internal=*/true);
+}
+
+[[maybe_unused]] inline StableIValue _from(
+    const torch::stable::Tensor& val,
+    uint64_t extension_build_version) {
+  return detail::FromImpl<torch::stable::Tensor>::call(
+      val, extension_build_version, /*is_internal=*/true);
+}
+
+template <typename T>
+inline T _to(StableIValue val, uint64_t extension_build_version) {
+  return detail::ToImpl<T>::call(
+      val, extension_build_version, /*is_internal=*/true);
+}
+
+HIDDEN_NAMESPACE_END(torch, stable, detail)
+
+// [global from/to deprecation note]
+// WARNING! the following APIs will be removed!! We deprecated global from/to
+// (in 2.10) in favor of torch::stable::detail from/to to not pollute the global
+// namespace. We are only including the following wrappers for backwards
+// compatibility.
+
+// WARNING! Will be removed. Only exists for BC. See [global from/to deprecation
+// note]
+template <typename T>
+[[deprecated("Use torch::stable::detail::from instead.")]]
+inline StableIValue from(T val) {
+  return torch::stable::detail::from(val);
+}
+
+// WARNING! Will be removed. Only exists for BC. See [global from/to deprecation
+// note]
+template <typename T>
+[[deprecated("Use torch::stable::detail::from instead.")]]
+inline StableIValue from(const std::optional<T>& val) {
+  return torch::stable::detail::from(val);
+}
+
+// WARNING! Will be removed. Only exists for BC. See [global from/to deprecation
+// note]
+[[deprecated(
+    "Use torch::stable::detail::from instead.")]] [[maybe_unused]] inline StableIValue
+from(const torch::stable::Tensor& val) {
+  return torch::stable::detail::from(val);
+}
+
+// WARNING! Will be removed. Only exists for BC. See [global from/to deprecation
+// note]
+template <typename T>
+[[deprecated("Use torch::stable::detail::to instead.")]]
+inline T to(StableIValue val) {
+  return torch::stable::detail::to<T>(val);
+}
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..8762372a415cf30f429808980e0e2f2af7942b1d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor.h
@@ -0,0 +1,4 @@
+#pragma once
+
+#include <torch/csrc/stable/tensor_inl.h>
+#include <torch/csrc/stable/tensor_struct.h>
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor_inl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor_inl.h
new file mode 100644
index 0000000000000000000000000000000000000000..f6e7ecaf7456fdb9e9ae589f57b1c115b1857a58
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor_inl.h
@@ -0,0 +1,77 @@
+#pragma once
+
+// This file implements tensor.h. We separated out the Tensor struct so that
+// other files can depend on the Tensor struct (like library.h) and the
+// implementations of the Tensor methods can depend on APIs in library.h
+// without circular dependencies.
+
+#include <torch/csrc/stable/stableivalue_conversions.h>
+#include <torch/csrc/stable/tensor_struct.h>
+#include <torch/headeronly/core/Layout.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/shim_utils.h>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable)
+
+using torch::headeronly::Layout;
+using torch::headeronly::ScalarType;
+
+inline ScalarType Tensor::scalar_type() const {
+  int32_t dtype;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(ath_.get(), &dtype));
+  return torch::stable::detail::to<ScalarType>(
+      torch::stable::detail::from(dtype));
+}
+
+inline Device Tensor::device() const {
+  int32_t device_type;
+  int32_t device_index;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_device_type(ath_.get(), &device_type));
+  TORCH_ERROR_CODE_CHECK(
+      aoti_torch_get_device_index(ath_.get(), &device_index));
+  DeviceType extension_device_type = torch::stable::detail::to<DeviceType>(
+      torch::stable::detail::from(device_type));
+  return Device(extension_device_type, static_cast<DeviceIndex>(device_index));
+}
+
+inline Layout Tensor::layout() const {
+  int32_t layout;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_layout(ath_.get(), &layout));
+  return torch::stable::detail::to<Layout>(torch::stable::detail::from(layout));
+}
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+// The following data ptr cast methods mirror the methods defined in
+// aten/src/ATen/templates/TensorMethods.cpp
+#define DEFINE_DATA_PTR_CAST(T, name, PRED)               \
+  template <>                                             \
+  inline T* Tensor::mutable_data_ptr() const {            \
+    auto stype = scalar_type();                           \
+    STD_TORCH_CHECK(                                      \
+        PRED(stype, torch::headeronly::ScalarType::name), \
+        "expected scalar type " #name " but found ",      \
+        torch::headeronly::toString(stype));              \
+    return static_cast<T*>(mutable_data_ptr());           \
+  }                                                       \
+  template <>                                             \
+  inline const T* Tensor::const_data_ptr() const {        \
+    auto stype = scalar_type();                           \
+    STD_TORCH_CHECK(                                      \
+        PRED(stype, torch::headeronly::ScalarType::name), \
+        "expected scalar type " #name " but found ",      \
+        torch::headeronly::toString(stype));              \
+    return static_cast<const T*>(const_data_ptr());       \
+  }
+
+#define _PRED(S1, S2) S1 == S2
+#define DEFINE_CAST(T, name) DEFINE_DATA_PTR_CAST(T, name, _PRED)
+AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CAST)
+DEFINE_CAST(uint16_t, UInt16)
+DEFINE_CAST(uint32_t, UInt32)
+DEFINE_CAST(uint64_t, UInt64)
+#undef DEFINE_CAST
+#undef _PRED
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+HIDDEN_NAMESPACE_END(torch, stable)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor_struct.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor_struct.h
new file mode 100644
index 0000000000000000000000000000000000000000..45461211955990e8c022bd6026ed6d1ab259858e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/tensor_struct.h
@@ -0,0 +1,447 @@
+#pragma once
+
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/headeronly/core/Layout.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Exception.h>
+#include <torch/headeronly/util/HeaderOnlyArrayRef.h>
+#include <torch/headeronly/util/shim_utils.h>
+#include <climits>
+#include <memory>
+
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/device_struct.h>
+
+HIDDEN_NAMESPACE_BEGIN(torch, stable)
+
+using accelerator::DeviceIndex;
+using torch::headeronly::IntHeaderOnlyArrayRef;
+using torch::headeronly::Layout;
+using torch::headeronly::ScalarType;
+
+// The torch::stable::Tensor class is a highlevel C++ wrapper around
+// the C shim Tensor APIs. We've modeled this class after TensorBase, as custom
+// op kernels only really need to interact with Tensor metadata (think sizes,
+// strides, device, dtype). Other functions on Tensor (like empty_like) should
+// live like the ATen op that they are and exist outside of this struct.
+//
+// There are several goals of this class over AtenTensorHandle and
+// RAIIAtenTensorHandle:
+// 1. torch::stable::Tensor is a nicer UX much closer to torch::Tensor than the
+//    C APIs with AtenTensorHandle. Under the hood we still call to these C shim
+//    APIs to preserve stability.
+// 2. RAIIAtenTensorHandle boils down to a uniq_ptr that forces the user to pass
+//    around ownership. This makes it difficult to pass one input into 2
+//    different functions, e.g., doing something like c = a(t) + b(t) for
+//    stable::Tensor t. Thus, we use a shared_ptr here.
+
+/**
+ * @brief An ABI stable wrapper around PyTorch tensors.
+ *
+ * This class is modeled after TensorBase, as custom
+ * op kernels primarily need to interact with Tensor metadata (sizes,
+ * strides, device, dtype). Other tensor operations (like ``empty_like``) exist
+ * as standalone functions outside of this struct.
+ *
+ * Minimum compatible version: PyTorch 2.9.
+ */
+class Tensor {
+ private:
+  std::shared_ptr<AtenTensorOpaque> ath_;
+
+ public:
+  /**
+   * @brief Constructs a Tensor with an uninitialized AtenTensorHandle.
+   *
+   * Creates a new stable::Tensor by allocating an uninitialized tensor handle.
+   * The ownership of the handle is managed internally via shared_ptr.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  Tensor() {
+    AtenTensorHandle ret;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_new_uninitialized_tensor(&ret));
+    ath_ = std::shared_ptr<AtenTensorOpaque>(ret, [](AtenTensorHandle ath) {
+      TORCH_ERROR_CODE_CHECK(aoti_torch_delete_tensor_object(ath));
+    });
+  }
+
+  /**
+   * @brief Constructs a Tensor from an existing AtenTensorHandle.
+   *
+   * Steals ownership of the provided AtenTensorHandle.
+   *
+   * @param ath The AtenTensorHandle to wrap. Ownership is transferred to this
+   *            Tensor.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  explicit Tensor(AtenTensorHandle ath)
+      : ath_(ath, [](AtenTensorHandle ath) {
+          TORCH_ERROR_CODE_CHECK(aoti_torch_delete_tensor_object(ath));
+        }) {}
+
+  // Copy and move constructors can be default cuz the underlying handle is a
+  // shared_ptr
+  /// \private
+  Tensor(const Tensor& other) = default;
+  /// \private
+  Tensor(Tensor&& other) noexcept = default;
+
+  // Copy and move assignment operators can be default cuz the underlying handle
+  // is a shared_ptr
+  /// \private
+  Tensor& operator=(const Tensor& other) = default;
+  /// \private
+  Tensor& operator=(Tensor&& other) noexcept = default;
+
+  // Destructor can be default: shared ptr has custom deletion logic
+  /// \private
+  ~Tensor() = default;
+
+  /**
+   * @brief Returns a borrowed reference to the underlying AtenTensorHandle.
+   *
+   * @return The underlying AtenTensorHandle.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  AtenTensorHandle get() const {
+    return ath_.get();
+  }
+
+  // =============================================================================
+  // C-shimified TensorBase APIs: the below APIs have the same signatures and
+  // semantics as their counterparts in TensorBase.h.
+  // =============================================================================
+
+  /**
+   * @brief Returns a pointer to the tensor's data.
+   *
+   * @return A void pointer to the tensor's data storage.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  void* data_ptr() const {
+    void* data_ptr;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr(ath_.get(), &data_ptr));
+    return data_ptr;
+  }
+
+#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+  /**
+   * @brief Returns a mutable pointer to the tensor's data.
+   *
+   * @return A mutable void pointer to the tensor's data storage.
+   *
+   * Minimum compatible version: PyTorch 2.10.
+   */
+  void* mutable_data_ptr() const {
+    void* data_ptr{};
+    TORCH_ERROR_CODE_CHECK(torch_get_mutable_data_ptr(ath_.get(), &data_ptr));
+    return data_ptr;
+  }
+
+  /**
+   * @brief Returns a const pointer to the tensor's data.
+   *
+   * @return A const void pointer to the tensor's data storage.
+   *
+   * Minimum compatible version: PyTorch 2.10.
+   */
+  const void* const_data_ptr() const {
+    const void* data_ptr{};
+    TORCH_ERROR_CODE_CHECK(torch_get_const_data_ptr(ath_.get(), &data_ptr));
+    return data_ptr;
+  }
+
+  /**
+   * @brief Returns a typed mutable pointer to the tensor's data.
+   *
+   * @tparam T The type to cast the data pointer to.
+   * @return A mutable pointer to the tensor's data cast to type T*.
+   *
+   * Minimum compatible version: PyTorch 2.10.
+   */
+  template <typename T>
+  T* mutable_data_ptr() const;
+
+  /**
+   * @brief Returns a typed const pointer to the tensor's data.
+   *
+   * @tparam T The type to cast the data pointer to. Must not be
+   * const-qualified.
+   * @return A const pointer to the tensor's data cast to type const T*.
+   *
+   * Minimum compatible version: PyTorch 2.10.
+   */
+  template <typename T, std::enable_if_t<!std::is_const_v<T>, int> = 0>
+  const T* const_data_ptr() const;
+
+  /**
+   * @brief Sets whether this tensor requires gradient computation.
+   *
+   * @param requires_grad If true, gradients will be computed for this tensor
+   *                      during backpropagation.
+   * @return A reference to this Tensor.
+   *
+   * Minimum compatible version: PyTorch 2.10.
+   */
+  const Tensor& set_requires_grad(bool requires_grad) const {
+    TORCH_ERROR_CODE_CHECK(torch_set_requires_grad(ath_.get(), requires_grad));
+    return *this;
+  }
+#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0
+
+  /**
+   * @brief Returns the number of dimensions of the tensor.
+   *
+   * @return The number of dimensions (rank) of the tensor.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  int64_t dim() const {
+    int64_t dim;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_dim(ath_.get(), &dim));
+    return dim;
+  }
+
+  /**
+   * @brief Returns the total number of elements in the tensor.
+   *
+   * @return The total number of elements across all dimensions.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  int64_t numel() const {
+    int64_t numel;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_numel(ath_.get(), &numel));
+    return numel;
+  }
+
+  // note: sizes and strides, for all intents and purposes, the same as in
+  // TensorBase.h: it returns a borrowed reference of the dimension sizes of
+  // a Tensor.
+  //
+  // The only difference is that it returns a header-only IntHeaderOnlyArrayRef,
+  // which has slightly less functionality than a regular IntArrayRef. See
+  // [HeaderOnlyArrayRef vs ArrayRef note] for more details.
+  /**
+   * @brief Returns the sizes (shape) of the tensor.
+   *
+   * Returns a borrowed reference of the dimension sizes of the tensor.
+   *
+   * @return An IntHeaderOnlyArrayRef containing the size of each dimension.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  IntHeaderOnlyArrayRef sizes() const {
+    int64_t* sizes;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes(ath_.get(), &sizes));
+    return IntHeaderOnlyArrayRef(sizes, dim());
+  }
+
+  /**
+   * @brief Returns the strides of the tensor.
+   *
+   * Returns a borrowed reference of the strides of the tensor.
+   *
+   *
+   * @return An IntHeaderOnlyArrayRef containing the stride of each dimension.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  IntHeaderOnlyArrayRef strides() const {
+    int64_t* strides;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides(ath_.get(), &strides));
+    return IntHeaderOnlyArrayRef(strides, dim());
+  }
+
+  /**
+   * @brief Checks if the tensor is contiguous in memory.
+   *
+   * @note This is a subset of the original TensorBase API. It takes no
+   *       arguments whereas the original API takes a memory format argument.
+   *       Here, we assume the default contiguous memory format.
+   *
+   * @return true if the tensor is contiguous, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool is_contiguous() const {
+    bool is_contiguous;
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_is_contiguous(ath_.get(), &is_contiguous));
+    return is_contiguous;
+  }
+
+  /**
+   * @brief Returns the stride of a specific dimension.
+   *
+   * @param dim The dimension index to query.
+   * @return The stride of the specified dimension.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  int64_t stride(int64_t dim) const {
+    int64_t stride;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_stride(ath_.get(), dim, &stride));
+    return stride;
+  }
+
+  // This is almost the same API as the one in TensorBase.h, except
+  // we add a check that the returned device_index is within the
+  // range of int8_t.
+  /// \private
+  int8_t get_device() const {
+    int32_t device_index;
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_device_index(ath_.get(), &device_index));
+    STD_TORCH_CHECK(
+        device_index >= std::numeric_limits<int8_t>::min() &&
+            device_index <= std::numeric_limits<int8_t>::max(),
+        "Device index is out of range of return type int8_t, please use get_device_index() instead.");
+    return static_cast<int8_t>(device_index);
+  }
+
+  // The same as get_device but with two differences:
+  // 1. it has a more suiting name
+  // 2. it returns a DeviceIndex, which is int32_t in this world
+  //    that should be more stable than the likely shifting
+  //    DeviceIndex in libtorch (it is int8_t that might become int16_t)
+  /**
+   * @brief Returns the device index of the tensor.
+   *
+   * @return The device index as DeviceIndex (int32_t).
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  DeviceIndex get_device_index() const {
+    int32_t device_index;
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_device_index(ath_.get(), &device_index));
+    return device_index;
+  }
+
+  /**
+   * @brief Checks if the tensor is on a CUDA device.
+   *
+   * @return true if the tensor is on a CUDA device, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool is_cuda() const {
+    int32_t device_type;
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_device_type(ath_.get(), &device_type));
+    return device_type == aoti_torch_device_type_cuda();
+  }
+
+  /**
+   * @brief Checks if the tensor is on the CPU.
+   *
+   * @return true if the tensor is on the CPU, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool is_cpu() const {
+    int32_t device_type;
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_device_type(ath_.get(), &device_type));
+    return device_type == aoti_torch_device_type_cpu();
+  }
+
+  /**
+   * @brief Returns the size of a specific dimension.
+   *
+   * @param dim The dimension index to query.
+   * @return The size of the specified dimension.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  int64_t size(int64_t dim) const {
+    int64_t size;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_size(ath_.get(), dim, &size));
+    return size;
+  }
+
+  /**
+   * @brief Checks if the tensor is defined (not null).
+   *
+   * @return true if the tensor is defined, false otherwise.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  bool defined() const {
+    bool defined;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_is_defined(ath_.get(), &defined));
+    return defined;
+  }
+
+  /**
+   * @brief Returns the storage offset of the tensor.
+   *
+   * The storage offset is the number of elements from the beginning of the
+   * underlying storage to the first element of the tensor.
+   *
+   * @return The storage offset in number of elements.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  int64_t storage_offset() const {
+    int64_t storage_offset;
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_storage_offset(ath_.get(), &storage_offset));
+    return storage_offset;
+  }
+
+  /**
+   * @brief Returns the size in bytes of each element in the tensor.
+   *
+   * @return The element size in bytes.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  size_t element_size() const {
+    int32_t dtype;
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_dtype(ath_.get(), &dtype));
+    return aoti_torch_dtype_element_size(dtype);
+  }
+
+  // defined in tensor-inl.h to avoid circular dependencies
+  /**
+   * @brief Returns the scalar type (dtype) of the tensor.
+   *
+   * @return The ScalarType of the tensor.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  ScalarType scalar_type() const;
+
+  // defined in tensor-inl.h to avoid circular dependencies
+  /**
+   * @brief Returns the device of the tensor.
+   *
+   * @return The Device on which the tensor resides.
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  Device device() const;
+
+  // defined in tensor_inl.h to avoid circular dependencies
+  /**
+   * @brief Returns the layout of the tensor.
+   *
+   * @return The Layout of the tensor (e.g., Strided, Sparse).
+   *
+   * Minimum compatible version: PyTorch 2.9.
+   */
+  Layout layout() const;
+
+  // =============================================================================
+  // END of C-shimified TensorBase APIs
+  // =============================================================================
+};
+
+HIDDEN_NAMESPACE_END(torch, stable)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/version.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/version.h
new file mode 100644
index 0000000000000000000000000000000000000000..eeb4f9152546aab161d6ac468eb8f1d2221982e8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/stable/version.h
@@ -0,0 +1,31 @@
+#pragma once
+
+#include <torch/headeronly/version.h>
+
+// Stable ABI Version Targeting
+//
+// This header provides version targeting capabilities for the PyTorch Stable
+// ABI. Users can define TORCH_TARGET_VERSION to target a specific stable ABI
+// version instead of using the current TORCH_ABI_VERSION of libtorch at
+// compile time.
+//
+// Usage:
+//   Default behavior (uses current ABI version):
+//     #include <torch/csrc/stable/library.h>
+//
+//   Target a specific stable version (major.minor) (e.g. PyTorch 2.9):
+//   (1) Pass a compiler flag -DTORCH_TARGET_VERSION=0x0209000000000000
+//   (2) Alternatively, define TORCH_TARGET_VERSION in the source code before
+//   including any header files:
+//     #define TORCH_TARGET_VERSION (((0ULL + 2) << 56) | ((0ULL + 9) << 48))
+//     #include <torch/csrc/stable/library.h>
+
+#ifdef TORCH_TARGET_VERSION
+#define TORCH_FEATURE_VERSION TORCH_TARGET_VERSION
+#else
+#define TORCH_FEATURE_VERSION TORCH_ABI_VERSION
+#endif
+
+#define TORCH_VERSION_2_10_0 (((0ULL + 2) << 56) | ((0ULL + 10) << 48))
+#define TORCH_VERSION_2_11_0 (((0ULL + 2) << 56) | ((0ULL + 11) << 48))
+#define TORCH_VERSION_2_12_0 (((0ULL + 2) << 56) | ((0ULL + 12) << 48))
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/tensor/python_tensor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/tensor/python_tensor.h
new file mode 100644
index 0000000000000000000000000000000000000000..87183f3f4eed50b44407c6df72e9e39ead754db8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/tensor/python_tensor.h
@@ -0,0 +1,40 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/Device.h>
+#include <c10/core/DispatchKey.h>
+#include <c10/core/ScalarType.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/python_headers.h>
+
+namespace at {
+class Tensor;
+} // namespace at
+
+namespace torch::tensors {
+
+// Initializes the Python tensor type objects: torch.FloatTensor,
+// torch.DoubleTensor, etc. and binds them in their containing modules.
+TORCH_PYTHON_API void initialize_python_bindings();
+
+// Same as set_default_tensor_type() but takes a PyObject*
+TORCH_PYTHON_API void py_set_default_tensor_type(PyObject* type_obj);
+
+// Same as py_set_default_tensor_type, but only changes the dtype (ScalarType).
+TORCH_PYTHON_API void py_set_default_dtype(PyObject* dtype_obj);
+
+// Gets the DispatchKey for the default tensor type.
+//
+// TODO: This is nuts!  There is no reason to let the default tensor type id
+// change.  Probably only store ScalarType, as that's the only flex point
+// we support.
+TORCH_PYTHON_API c10::DispatchKey get_default_dispatch_key();
+TORCH_PYTHON_API at::Device get_default_device();
+
+// Gets the ScalarType for the default tensor type.
+TORCH_PYTHON_API at::ScalarType get_default_scalar_type();
+} // namespace torch::tensors
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..36869bdc99dd4fcf67f5c5b7a10bc495d1d81f0c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils.h
@@ -0,0 +1,130 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/util/Exception.h>
+#include <torch/csrc/Storage.h>
+#include <torch/csrc/THConcat.h>
+#include <torch/csrc/utils/object_ptr.h>
+#include <torch/csrc/utils/python_compat.h>
+#include <torch/csrc/utils/python_numbers.h>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#define THPUtils_(NAME) TH_CONCAT_4(THP, Real, Utils_, NAME)
+
+#define THPUtils_typename(obj) (Py_TYPE(obj)->tp_name)
+
+#if defined(__GNUC__) || defined(__ICL) || defined(__clang__)
+#define THP_EXPECT(x, y) (__builtin_expect((x), (y)))
+#else
+#define THP_EXPECT(x, y) (x)
+#endif
+
+#define THPUtils_unpackReal_FLOAT(object)           \
+  (PyFloat_Check(object) ? PyFloat_AsDouble(object) \
+       : PyLong_Check(object)                       \
+       ? PyLong_AsLongLong(object)                  \
+       : (throw std::runtime_error("Could not parse real"), 0))
+
+#define THPUtils_checkReal_INT(object) PyLong_Check(object)
+
+#define THPUtils_unpackReal_INT(object) \
+  (PyLong_Check(object)                 \
+       ? PyLong_AsLongLong(object)      \
+       : (throw std::runtime_error("Could not parse real"), 0))
+
+#define THPUtils_unpackReal_BOOL(object) \
+  (PyBool_Check(object)                  \
+       ? object                          \
+       : (throw std::runtime_error("Could not parse real"), Py_False))
+
+#define THPUtils_unpackReal_COMPLEX(object)                                   \
+  (PyComplex_Check(object)                                                    \
+       ? (c10::complex<double>(                                               \
+             PyComplex_RealAsDouble(object), PyComplex_ImagAsDouble(object))) \
+       : PyFloat_Check(object)                                                \
+       ? (c10::complex<double>(PyFloat_AsDouble(object), 0))                  \
+       : PyLong_Check(object)                                                 \
+       ? (c10::complex<double>(PyLong_AsLongLong(object), 0))                 \
+       : (throw std::runtime_error("Could not parse real"),                   \
+          c10::complex<double>(0, 0)))
+
+#define THPBoolUtils_unpackReal(object) THPUtils_unpackReal_BOOL(object)
+#define THPBoolUtils_checkAccreal(object) THPUtils_checkReal_BOOL(object)
+#define THPByteUtils_checkReal(object) THPUtils_checkReal_INT(object)
+#define THPByteUtils_unpackReal(object) \
+  (unsigned char)THPUtils_unpackReal_INT(object)
+
+/*
+   From https://github.com/python/cpython/blob/v3.7.0/Modules/xxsubtype.c
+   If compiled as a shared library, some compilers don't allow addresses of
+   Python objects defined in other libraries to be used in static PyTypeObject
+   initializers. The DEFERRED_ADDRESS macro is used to tag the slots where such
+   addresses appear; the module init function that adds the PyTypeObject to the
+   module must fill in the tagged slots at runtime. The argument is for
+   documentation -- the macro ignores it.
+*/
+#define DEFERRED_ADDRESS(ADDR) nullptr
+
+TORCH_PYTHON_API void THPUtils_setError(const char* format, ...);
+TORCH_PYTHON_API void THPUtils_invalidArguments(
+    PyObject* given_args,
+    PyObject* given_kwargs,
+    const char* function_name,
+    size_t num_options,
+    ...);
+
+bool THPUtils_checkIntTuple(PyObject* arg);
+std::vector<int> THPUtils_unpackIntTuple(PyObject* arg);
+
+TORCH_PYTHON_API void THPUtils_addPyMethodDefs(
+    std::vector<PyMethodDef>& vector,
+    const PyMethodDef* methods);
+
+int THPUtils_getCallable(PyObject* arg, PyObject** result);
+
+typedef THPPointer<THPGenerator> THPGeneratorPtr;
+typedef class THPPointer<THPStorage> THPStoragePtr;
+
+TORCH_PYTHON_API std::vector<int64_t> THPUtils_unpackLongs(PyObject* arg);
+PyObject* THPUtils_dispatchStateless(
+    PyObject* tensor,
+    const char* name,
+    PyObject* args,
+    PyObject* kwargs);
+
+template <typename _real, typename = void>
+struct mod_traits {};
+
+template <typename _real>
+struct mod_traits<_real, std::enable_if_t<std::is_floating_point_v<_real>>> {
+  static _real mod(_real a, _real b) {
+    return fmod(a, b);
+  }
+};
+
+template <typename _real>
+struct mod_traits<_real, std::enable_if_t<std::is_integral_v<_real>>> {
+  static _real mod(_real a, _real b) {
+    return a % b;
+  }
+};
+
+void setBackCompatBroadcastWarn(bool warn);
+bool getBackCompatBroadcastWarn();
+
+void setBackCompatKeepdimWarn(bool warn);
+bool getBackCompatKeepdimWarn();
+bool maybeThrowBackCompatKeepdimWarn(char* func);
+
+void storage_fill(const at::Storage& self, uint8_t value);
+void storage_set(const at::Storage& self, ptrdiff_t idx, uint8_t value);
+uint8_t storage_get(const at::Storage& self, ptrdiff_t idx);
+
+std::string uuid_to_string(const char* uuid_bytes);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/byte_order.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/byte_order.h
new file mode 100644
index 0000000000000000000000000000000000000000..74ce8c0dadbe4e7baa6f9237520edce5921cb19c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/byte_order.h
@@ -0,0 +1,86 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Float8_e4m3fn.h>
+#include <c10/util/Float8_e4m3fnuz.h>
+#include <c10/util/Float8_e5m2.h>
+#include <c10/util/Float8_e5m2fnuz.h>
+#include <c10/util/Half.h>
+#include <torch/csrc/Export.h>
+#include <cstddef>
+#include <cstdint>
+
+#ifdef __FreeBSD__
+#include <sys/endian.h>
+#include <sys/types.h>
+#define thp_bswap16(x) bswap16(x)
+#define thp_bswap32(x) bswap32(x)
+#define thp_bswap64(x) bswap64(x)
+#elif defined(__APPLE__)
+#include <libkern/OSByteOrder.h>
+#define thp_bswap16(x) OSSwapInt16(x)
+#define thp_bswap32(x) OSSwapInt32(x)
+#define thp_bswap64(x) OSSwapInt64(x)
+#elif defined(__GNUC__) && !defined(__MINGW32__)
+#include <byteswap.h>
+#define thp_bswap16(x) bswap_16(x)
+#define thp_bswap32(x) bswap_32(x)
+#define thp_bswap64(x) bswap_64(x)
+#elif defined _WIN32 || defined _WIN64
+#define thp_bswap16(x) _byteswap_ushort(x)
+#define thp_bswap32(x) _byteswap_ulong(x)
+#define thp_bswap64(x) _byteswap_uint64(x)
+#endif
+
+#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
+#define to_be16(x) thp_bswap16(x)
+#define from_be16(x) thp_bswap16(x)
+#define to_be32(x) thp_bswap32(x)
+#define from_be32(x) thp_bswap32(x)
+#define to_be64(x) thp_bswap64(x)
+#define from_be64(x) thp_bswap64(x)
+#define to_le16(x) (x)
+#define from_le16(x) (x)
+#define to_le32(x) (x)
+#define from_le32(x) (x)
+#define to_le64(x) (x)
+#define from_le64(x) (x)
+#elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+#define to_be16(x) (x)
+#define from_be16(x) (x)
+#define to_be32(x) (x)
+#define from_be32(x) (x)
+#define to_be64(x) (x)
+#define from_be64(x) (x)
+#define to_le16(x) thp_bswap16(x)
+#define from_le16(x) thp_bswap16(x)
+#define to_le32(x) thp_bswap32(x)
+#define from_le32(x) thp_bswap32(x)
+#define to_le64(x) thp_bswap64(x)
+#define from_le64(x) thp_bswap64(x)
+#else
+#error Unexpected or undefined __BYTE_ORDER__
+#endif
+
+namespace torch::utils {
+
+enum THPByteOrder { THP_LITTLE_ENDIAN = 0, THP_BIG_ENDIAN = 1 };
+
+TORCH_API THPByteOrder THP_nativeByteOrder();
+
+template <typename T, typename U>
+TORCH_API void THP_decodeBuffer(T* dst, const uint8_t* src, U type, size_t len);
+
+template <typename T>
+TORCH_API void THP_encodeBuffer(
+    uint8_t* dst,
+    const T* src,
+    THPByteOrder order,
+    size_t len);
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/cpp_stacktraces.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/cpp_stacktraces.h
new file mode 100644
index 0000000000000000000000000000000000000000..2e24da7a55d82a86777253263813ff58b1264ed6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/cpp_stacktraces.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/profiler/unwind/unwind.h>
+
+namespace torch {
+TORCH_API bool get_cpp_stacktraces_enabled();
+TORCH_API torch::unwind::Mode get_symbolize_mode();
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/cuda_enabled.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/cuda_enabled.h
new file mode 100644
index 0000000000000000000000000000000000000000..7704171f7c8e889a2680ef2815f1b2b0c04c7d92
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/cuda_enabled.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::utils {
+
+inline constexpr bool cuda_enabled() {
+#ifdef USE_CUDA
+  return true;
+#else
+  return false;
+#endif
+}
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/device_lazy_init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/device_lazy_init.h
new file mode 100644
index 0000000000000000000000000000000000000000..0f48437820daeb566418292c1de24c21b3f136b1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/device_lazy_init.h
@@ -0,0 +1,92 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/TensorOptions.h>
+#include <torch/csrc/Export.h>
+
+// device_lazy_init() is always compiled, even for CPU-only builds.
+
+namespace torch::utils {
+
+/**
+ * This mechanism of lazy initialization is designed for each device backend.
+ * Currently, CUDA and XPU follow this design. This function `device_lazy_init`
+ * MUST be called before you attempt to access any Type(CUDA or XPU) object
+ * from ATen, in any way. It guarantees that the device runtime status is lazily
+ * initialized when the first runtime API is requested.
+ *
+ * Here are some common ways that a device object may be retrieved:
+ *   - You call getNonVariableType or getNonVariableTypeOpt
+ *   - You call toBackend() on a Type
+ *
+ * It's important to do this correctly, because if you forget to add it you'll
+ * get an oblique error message seems like "Cannot initialize CUDA without
+ * ATen_cuda library" or "Cannot initialize XPU without ATen_xpu library" if you
+ * try to use CUDA or XPU functionality from a CPU-only build, which is not good
+ * UX.
+ */
+TORCH_PYTHON_API void device_lazy_init(at::DeviceType device_type);
+TORCH_PYTHON_API void set_requires_device_init(
+    at::DeviceType device_type,
+    bool value);
+
+inline bool is_device_lazy_init_supported(at::DeviceType device_type) {
+  // Add more devices here to enable lazy initialization.
+  return (
+      device_type == at::DeviceType::CUDA ||
+      device_type == at::DeviceType::XPU ||
+      device_type == at::DeviceType::HPU ||
+      device_type == at::DeviceType::MTIA ||
+      device_type == at::DeviceType::PrivateUse1);
+}
+
+inline void maybe_initialize_device(at::Device& device) {
+  if (is_device_lazy_init_supported(device.type())) {
+    device_lazy_init(device.type());
+  }
+}
+
+inline void maybe_initialize_device(std::optional<at::Device>& device) {
+  if (!device.has_value()) {
+    return;
+  }
+  maybe_initialize_device(device.value());
+}
+
+inline void maybe_initialize_device(const at::TensorOptions& options) {
+  auto device = options.device();
+  maybe_initialize_device(device);
+}
+
+inline void maybe_initialize_device(
+    std::optional<at::DeviceType>& device_type) {
+  if (!device_type.has_value()) {
+    return;
+  }
+  maybe_initialize_device(device_type.value());
+}
+
+bool is_device_initialized(at::DeviceType device_type);
+
+TORCH_PYTHON_API bool is_device_in_bad_fork(at::DeviceType device_type);
+
+TORCH_PYTHON_API void set_device_in_bad_fork(
+    at::DeviceType device_type,
+    bool value);
+
+TORCH_PYTHON_API void register_fork_handler_for_device_init(
+    at::DeviceType device_type);
+
+inline void maybe_register_fork_handler_for_device_init(
+    std::optional<at::DeviceType>& device_type) {
+  if (!device_type.has_value()) {
+    return;
+  }
+  register_fork_handler_for_device_init(device_type.value());
+}
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/disable_torch_function.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/disable_torch_function.h
new file mode 100644
index 0000000000000000000000000000000000000000..694a47e8fd6cf4af36aa1e749a26979eee1799c8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/disable_torch_function.h
@@ -0,0 +1,52 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <c10/core/DispatchKey.h>
+#include <c10/core/impl/LocalDispatchKeySet.h>
+#include <torch/csrc/python_headers.h>
+
+namespace torch {
+// Sometimes we don't want infinite recursion for subclasses,
+// Or a way to achieve the old behaviour.
+
+// This is an internal utility, not exposed to users.
+bool torch_function_enabled();
+PyObject* disabled_torch_function_impl();
+PyObject* disabled_torch_dispatch_impl();
+void set_disabled_torch_function_impl(PyObject* value);
+void set_disabled_torch_dispatch_impl(PyObject* value);
+// Set ignore_mode to true if you're trying to collect overloaded arguments;
+// using mode here will improperly cause you to add ALL objects to the
+// overloaded list even if they don't actually have __torch_function__
+bool check_has_torch_function(PyObject* obj, bool ignore_mode = false);
+
+struct DisableTorchDispatch {
+  DisableTorchDispatch()
+      : guard_(c10::DispatchKeySet(
+            {c10::DispatchKey::Python, c10::DispatchKey::PreDispatch})),
+        guard_tls_snapshot_(c10::DispatchKey::PythonTLSSnapshot) {}
+  c10::impl::ExcludeDispatchKeyGuard guard_;
+  c10::impl::ExcludeDispatchKeyGuard guard_tls_snapshot_;
+};
+
+} // namespace torch
+
+PyObject* THPModule_isEnabledTorchFunction(PyObject* self, PyObject* unused);
+PyObject* THPModule_isAllDisabledTorchFunction(
+    PyObject* self,
+    PyObject* unused);
+PyObject* THPModule_DisableTorchFunctionType();
+PyObject* THPModule_DisableTorchFunctionSubclassType();
+PyObject* THPModule_disable_torch_function(PyObject* self, PyObject* args);
+PyObject* THPModule_disable_torch_dispatch(PyObject* self, PyObject* args);
+PyObject* THPModule_has_torch_function(PyObject* /*unused*/, PyObject* arg);
+PyObject* THPModule_has_torch_function_unary(
+    PyObject* /*unused*/,
+    PyObject* obj);
+PyObject* THPModule_has_torch_function_variadic(
+    PyObject* /*unused*/,
+    PyObject* const* args,
+    Py_ssize_t nargs);
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/generated_serialization_types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/generated_serialization_types.h
new file mode 100644
index 0000000000000000000000000000000000000000..55d7be4cba1b71babfa500055577852a1a64a32a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/generated_serialization_types.h
@@ -0,0 +1,3937 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// @generated by update_schema.py
+// checksum<<eb9645ae5f539caace53ada70fb262097ce780e6edf685149dafc5a5abe8d31b>>
+// clang-format off
+
+#pragma once
+
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+#include <nlohmann/json.hpp>
+
+#ifndef NLOHMANN_JSON_NAMESPACE_BEGIN
+#define NLOHMANN_JSON_NAMESPACE_BEGIN namespace nlohmann {
+#endif
+
+#ifndef NLOHMANN_JSON_NAMESPACE_END
+#define NLOHMANN_JSON_NAMESPACE_END }
+#endif
+
+// https://github.com/nlohmann/json/pull/2117
+NLOHMANN_JSON_NAMESPACE_BEGIN
+template <typename T>
+struct adl_serializer<std::optional<T>> {
+  static void to_json(json& j, const std::optional<T>& opt) {
+    if (opt == std::nullopt) {
+      j = nullptr;
+    } else {
+      j = *opt; // this will call adl_serializer<T>::to_json which will
+                // find the free function to_json in T's namespace!
+    }
+  }
+
+  static void from_json(const json& j, std::optional<T>& opt) {
+    if (j.is_null()) {
+      opt = std::nullopt;
+    } else {
+      opt = j.template get<T>(); // same as above, but with
+                                 // adl_serializer<T>::from_json
+    }
+  }
+};
+NLOHMANN_JSON_NAMESPACE_END
+
+namespace torch {
+namespace _export {
+
+template <typename T>
+class ForwardRef {
+  static_assert(!std::is_reference_v<T>, "ForwardRef cannot be a reference type");
+
+ public:
+  ForwardRef(): ptr_(std::make_unique<T>()) {}
+  ForwardRef(ForwardRef<T>&&);
+  ForwardRef(const ForwardRef<T>& other): ptr_(std::make_unique<T>(*other.ptr_)) {}
+  ForwardRef<T>& operator=(ForwardRef<T>&&);
+  ForwardRef<T>& operator=(const ForwardRef<T>& other) {
+    ptr_ = std::make_unique<T>(*other.ptr_);
+    return *this;
+  }
+  ~ForwardRef();
+  const T& operator*() const {
+    return *ptr_;
+  }
+
+  const T* operator->() const {
+    return ptr_.get();
+  }
+
+  void emplace(T&& t) {
+    ptr_ = std::make_unique<T>(std::move(t));
+  }
+
+ private:
+  std::unique_ptr<T> ptr_;
+};
+
+template <typename T>
+void to_json(nlohmann::json& j, const ForwardRef<T>& p) {
+  j = *p;
+}
+
+template <typename T>
+void from_json(const nlohmann::json& j, ForwardRef<T>& p) {
+  p.emplace(j.template get<T>());
+}
+
+class F64 {
+ public:
+  double get() const {
+    return value_;
+  }
+
+  void set(double value) {
+    value_ = value;
+  }
+
+ private:
+  double value_;
+};
+
+inline void to_json(nlohmann::json& j, const F64& f) {
+  if (std::isinf(f.get())) {
+    j = "Infinity";
+  } else if (std::isinf(-f.get())) {
+    j = "-Infinity";
+  } else if (std::isnan(f.get())) {
+    j = "NaN";
+  } else {
+    j = f.get();
+  }
+}
+
+inline void from_json(const nlohmann::json& j, F64& f) {
+  if (j == "Infinity") {
+    f.set(std::numeric_limits<double>::infinity());
+  } else if (j == "-Infinity") {
+    f.set(-std::numeric_limits<double>::infinity());
+  } else if (j == "NaN") {
+    f.set(std::numeric_limits<double>::quiet_NaN());
+  } else {
+    f.set(j.get<double>());
+  }
+}
+
+class AOTInductorModelPickleData;
+class Argument;
+class BufferMutationSpec;
+class ComplexValue;
+class ConstantValue;
+class CustomObjArgument;
+class Device;
+class ExportedProgram;
+class ExternKernelNode;
+class ExternKernelNodes;
+class GradientToParameterSpec;
+class GradientToUserInputSpec;
+class Graph;
+class GraphArgument;
+class GraphModule;
+class GraphSignature;
+class InputSpec;
+class InputToBufferSpec;
+class InputToConstantInputSpec;
+class InputToCustomObjSpec;
+class InputToParameterSpec;
+class InputToTensorConstantSpec;
+class InputTokenSpec;
+class LossOutputSpec;
+class ModuleCallEntry;
+class ModuleCallSignature;
+class NamedArgument;
+class NamedTupleDef;
+class Node;
+class OptionalTensorArgument;
+class OutputSpec;
+class OutputTokenSpec;
+class ParameterMutationSpec;
+class PayloadConfig;
+class PayloadMeta;
+class RangeConstraint;
+class SchemaVersion;
+class SymBool;
+class SymBoolArgument;
+class SymExpr;
+class SymExprHint;
+class SymFloat;
+class SymFloatArgument;
+class SymInt;
+class SymIntArgument;
+class TensorArgument;
+class TensorMeta;
+class TokenArgument;
+class UserInputMutationSpec;
+class UserInputSpec;
+class UserOutputSpec;
+
+enum class ArgumentKind {
+  UNKNOWN = 0,
+  POSITIONAL = 1,
+  KEYWORD = 2,
+};
+
+inline std::string_view printEnum(const ArgumentKind& e) {
+  switch (e) {
+    case ArgumentKind::UNKNOWN: return "UNKNOWN";
+    case ArgumentKind::POSITIONAL: return "POSITIONAL";
+    case ArgumentKind::KEYWORD: return "KEYWORD";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, ArgumentKind& t) {
+  if (s == "UNKNOWN") { t = ArgumentKind::UNKNOWN; return; }
+  if (s == "POSITIONAL") { t = ArgumentKind::POSITIONAL; return; }
+  if (s == "KEYWORD") { t = ArgumentKind::KEYWORD; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+enum class Layout {
+  Unknown = 0,
+  SparseCoo = 1,
+  SparseCsr = 2,
+  SparseCsc = 3,
+  SparseBsr = 4,
+  SparseBsc = 5,
+  _mkldnn = 6,
+  Strided = 7,
+};
+
+inline std::string_view printEnum(const Layout& e) {
+  switch (e) {
+    case Layout::Unknown: return "Unknown";
+    case Layout::SparseCoo: return "SparseCoo";
+    case Layout::SparseCsr: return "SparseCsr";
+    case Layout::SparseCsc: return "SparseCsc";
+    case Layout::SparseBsr: return "SparseBsr";
+    case Layout::SparseBsc: return "SparseBsc";
+    case Layout::_mkldnn: return "_mkldnn";
+    case Layout::Strided: return "Strided";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, Layout& t) {
+  if (s == "Unknown") { t = Layout::Unknown; return; }
+  if (s == "SparseCoo") { t = Layout::SparseCoo; return; }
+  if (s == "SparseCsr") { t = Layout::SparseCsr; return; }
+  if (s == "SparseCsc") { t = Layout::SparseCsc; return; }
+  if (s == "SparseBsr") { t = Layout::SparseBsr; return; }
+  if (s == "SparseBsc") { t = Layout::SparseBsc; return; }
+  if (s == "_mkldnn") { t = Layout::_mkldnn; return; }
+  if (s == "Strided") { t = Layout::Strided; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+enum class MemoryFormat {
+  Unknown = 0,
+  ContiguousFormat = 1,
+  ChannelsLast = 2,
+  ChannelsLast3d = 3,
+  PreserveFormat = 4,
+};
+
+inline std::string_view printEnum(const MemoryFormat& e) {
+  switch (e) {
+    case MemoryFormat::Unknown: return "Unknown";
+    case MemoryFormat::ContiguousFormat: return "ContiguousFormat";
+    case MemoryFormat::ChannelsLast: return "ChannelsLast";
+    case MemoryFormat::ChannelsLast3d: return "ChannelsLast3d";
+    case MemoryFormat::PreserveFormat: return "PreserveFormat";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, MemoryFormat& t) {
+  if (s == "Unknown") { t = MemoryFormat::Unknown; return; }
+  if (s == "ContiguousFormat") { t = MemoryFormat::ContiguousFormat; return; }
+  if (s == "ChannelsLast") { t = MemoryFormat::ChannelsLast; return; }
+  if (s == "ChannelsLast3d") { t = MemoryFormat::ChannelsLast3d; return; }
+  if (s == "PreserveFormat") { t = MemoryFormat::PreserveFormat; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+enum class ScalarType {
+  UNKNOWN = 0,
+  BYTE = 1,
+  CHAR = 2,
+  SHORT = 3,
+  INT = 4,
+  LONG = 5,
+  HALF = 6,
+  FLOAT = 7,
+  DOUBLE = 8,
+  COMPLEXHALF = 9,
+  COMPLEXFLOAT = 10,
+  COMPLEXDOUBLE = 11,
+  BOOL = 12,
+  BFLOAT16 = 13,
+  UINT16 = 28,
+  FLOAT8E4M3FN = 29,
+  FLOAT8E5M2 = 30,
+  FLOAT8E4M3FNUZ = 31,
+  FLOAT8E5M2FNUZ = 32,
+  FLOAT8E8M0FNU = 33,
+  UINT32 = 34,
+  UINT64 = 35,
+};
+
+inline std::string_view printEnum(const ScalarType& e) {
+  switch (e) {
+    case ScalarType::UNKNOWN: return "UNKNOWN";
+    case ScalarType::BYTE: return "BYTE";
+    case ScalarType::CHAR: return "CHAR";
+    case ScalarType::SHORT: return "SHORT";
+    case ScalarType::INT: return "INT";
+    case ScalarType::LONG: return "LONG";
+    case ScalarType::HALF: return "HALF";
+    case ScalarType::FLOAT: return "FLOAT";
+    case ScalarType::DOUBLE: return "DOUBLE";
+    case ScalarType::COMPLEXHALF: return "COMPLEXHALF";
+    case ScalarType::COMPLEXFLOAT: return "COMPLEXFLOAT";
+    case ScalarType::COMPLEXDOUBLE: return "COMPLEXDOUBLE";
+    case ScalarType::BOOL: return "BOOL";
+    case ScalarType::BFLOAT16: return "BFLOAT16";
+    case ScalarType::UINT16: return "UINT16";
+    case ScalarType::FLOAT8E4M3FN: return "FLOAT8E4M3FN";
+    case ScalarType::FLOAT8E5M2: return "FLOAT8E5M2";
+    case ScalarType::FLOAT8E4M3FNUZ: return "FLOAT8E4M3FNUZ";
+    case ScalarType::FLOAT8E5M2FNUZ: return "FLOAT8E5M2FNUZ";
+    case ScalarType::FLOAT8E8M0FNU: return "FLOAT8E8M0FNU";
+    case ScalarType::UINT32: return "UINT32";
+    case ScalarType::UINT64: return "UINT64";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, ScalarType& t) {
+  if (s == "UNKNOWN") { t = ScalarType::UNKNOWN; return; }
+  if (s == "BYTE") { t = ScalarType::BYTE; return; }
+  if (s == "CHAR") { t = ScalarType::CHAR; return; }
+  if (s == "SHORT") { t = ScalarType::SHORT; return; }
+  if (s == "INT") { t = ScalarType::INT; return; }
+  if (s == "LONG") { t = ScalarType::LONG; return; }
+  if (s == "HALF") { t = ScalarType::HALF; return; }
+  if (s == "FLOAT") { t = ScalarType::FLOAT; return; }
+  if (s == "DOUBLE") { t = ScalarType::DOUBLE; return; }
+  if (s == "COMPLEXHALF") { t = ScalarType::COMPLEXHALF; return; }
+  if (s == "COMPLEXFLOAT") { t = ScalarType::COMPLEXFLOAT; return; }
+  if (s == "COMPLEXDOUBLE") { t = ScalarType::COMPLEXDOUBLE; return; }
+  if (s == "BOOL") { t = ScalarType::BOOL; return; }
+  if (s == "BFLOAT16") { t = ScalarType::BFLOAT16; return; }
+  if (s == "UINT16") { t = ScalarType::UINT16; return; }
+  if (s == "FLOAT8E4M3FN") { t = ScalarType::FLOAT8E4M3FN; return; }
+  if (s == "FLOAT8E5M2") { t = ScalarType::FLOAT8E5M2; return; }
+  if (s == "FLOAT8E4M3FNUZ") { t = ScalarType::FLOAT8E4M3FNUZ; return; }
+  if (s == "FLOAT8E5M2FNUZ") { t = ScalarType::FLOAT8E5M2FNUZ; return; }
+  if (s == "FLOAT8E8M0FNU") { t = ScalarType::FLOAT8E8M0FNU; return; }
+  if (s == "UINT32") { t = ScalarType::UINT32; return; }
+  if (s == "UINT64") { t = ScalarType::UINT64; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class Device {
+ private:
+  std::string type;
+  std::optional<int64_t> index = std::nullopt;
+
+ public:
+
+  const std::string& get_type() const {
+    return type;
+  }
+
+  void set_type(std::string def) {
+    type = std::move(def);
+  }
+
+  const std::optional<int64_t>& get_index() const {
+    return index;
+  }
+
+  void set_index(std::optional<int64_t> def) {
+    index = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const Device& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, Device& nlohmann_json_t);
+};
+
+class SymExprHint {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_INT, AS_BOOL, AS_FLOAT
+  };
+
+ private:
+  std::variant<Void, int64_t, bool, F64> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const int64_t& get_as_int() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_int(int64_t def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_INT;
+  }
+
+  const bool& get_as_bool() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_bool(bool def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_BOOL;
+  }
+
+  const F64& get_as_float() const {
+    return std::get<3>(variant_);
+  }
+
+  void set_as_float(F64 def) {
+    variant_.emplace<3>(std::move(def));
+    tag_ = Tag::AS_FLOAT;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymExprHint& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_INT) {
+      nlohmann_json_j["as_int"] = nlohmann_json_t.get_as_int();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_BOOL) {
+      nlohmann_json_j["as_bool"] = nlohmann_json_t.get_as_bool();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOAT) {
+      nlohmann_json_j["as_float"] = nlohmann_json_t.get_as_float();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymExprHint& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_int")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_int").template get<int64_t>());
+      nlohmann_json_t.tag_ = Tag::AS_INT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_bool")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_bool").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_BOOL;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_float")) {
+      nlohmann_json_t.variant_.emplace<3>(nlohmann_json_j.at("as_float").template get<F64>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOAT;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const SymExprHint::Tag& e) {
+  switch (e) {
+    case SymExprHint::Tag::AS_INT: return "AS_INT";
+    case SymExprHint::Tag::AS_BOOL: return "AS_BOOL";
+    case SymExprHint::Tag::AS_FLOAT: return "AS_FLOAT";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, SymExprHint::Tag& t) {
+  if (s == "AS_INT") { t = SymExprHint::Tag::AS_INT; return; }
+  if (s == "AS_BOOL") { t = SymExprHint::Tag::AS_BOOL; return; }
+  if (s == "AS_FLOAT") { t = SymExprHint::Tag::AS_FLOAT; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class SymExpr {
+ private:
+  std::string expr_str;
+  std::optional<SymExprHint> hint = std::nullopt;
+
+ public:
+
+  const std::string& get_expr_str() const {
+    return expr_str;
+  }
+
+  void set_expr_str(std::string def) {
+    expr_str = std::move(def);
+  }
+
+  const std::optional<SymExprHint>& get_hint() const {
+    return hint;
+  }
+
+  void set_hint(std::optional<SymExprHint> def) {
+    hint = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymExpr& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymExpr& nlohmann_json_t);
+};
+
+class SymInt {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_EXPR, AS_INT
+  };
+
+ private:
+  std::variant<Void, SymExpr, int64_t> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const SymExpr& get_as_expr() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_expr(SymExpr def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_EXPR;
+  }
+
+  const int64_t& get_as_int() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_int(int64_t def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_INT;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymInt& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_EXPR) {
+      nlohmann_json_j["as_expr"] = nlohmann_json_t.get_as_expr();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_INT) {
+      nlohmann_json_j["as_int"] = nlohmann_json_t.get_as_int();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymInt& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_expr")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_expr").template get<SymExpr>());
+      nlohmann_json_t.tag_ = Tag::AS_EXPR;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_int")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_int").template get<int64_t>());
+      nlohmann_json_t.tag_ = Tag::AS_INT;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const SymInt::Tag& e) {
+  switch (e) {
+    case SymInt::Tag::AS_EXPR: return "AS_EXPR";
+    case SymInt::Tag::AS_INT: return "AS_INT";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, SymInt::Tag& t) {
+  if (s == "AS_EXPR") { t = SymInt::Tag::AS_EXPR; return; }
+  if (s == "AS_INT") { t = SymInt::Tag::AS_INT; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class SymFloat {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_EXPR, AS_FLOAT
+  };
+
+ private:
+  std::variant<Void, SymExpr, F64> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const SymExpr& get_as_expr() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_expr(SymExpr def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_EXPR;
+  }
+
+  const F64& get_as_float() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_float(F64 def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_FLOAT;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymFloat& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_EXPR) {
+      nlohmann_json_j["as_expr"] = nlohmann_json_t.get_as_expr();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOAT) {
+      nlohmann_json_j["as_float"] = nlohmann_json_t.get_as_float();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymFloat& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_expr")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_expr").template get<SymExpr>());
+      nlohmann_json_t.tag_ = Tag::AS_EXPR;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_float")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_float").template get<F64>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOAT;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const SymFloat::Tag& e) {
+  switch (e) {
+    case SymFloat::Tag::AS_EXPR: return "AS_EXPR";
+    case SymFloat::Tag::AS_FLOAT: return "AS_FLOAT";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, SymFloat::Tag& t) {
+  if (s == "AS_EXPR") { t = SymFloat::Tag::AS_EXPR; return; }
+  if (s == "AS_FLOAT") { t = SymFloat::Tag::AS_FLOAT; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class SymBool {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_EXPR, AS_BOOL
+  };
+
+ private:
+  std::variant<Void, SymExpr, bool> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const SymExpr& get_as_expr() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_expr(SymExpr def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_EXPR;
+  }
+
+  const bool& get_as_bool() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_bool(bool def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_BOOL;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymBool& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_EXPR) {
+      nlohmann_json_j["as_expr"] = nlohmann_json_t.get_as_expr();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_BOOL) {
+      nlohmann_json_j["as_bool"] = nlohmann_json_t.get_as_bool();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymBool& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_expr")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_expr").template get<SymExpr>());
+      nlohmann_json_t.tag_ = Tag::AS_EXPR;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_bool")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_bool").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_BOOL;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const SymBool::Tag& e) {
+  switch (e) {
+    case SymBool::Tag::AS_EXPR: return "AS_EXPR";
+    case SymBool::Tag::AS_BOOL: return "AS_BOOL";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, SymBool::Tag& t) {
+  if (s == "AS_EXPR") { t = SymBool::Tag::AS_EXPR; return; }
+  if (s == "AS_BOOL") { t = SymBool::Tag::AS_BOOL; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class TensorMeta {
+ private:
+  int64_t dtype;
+  std::vector<SymInt> sizes;
+  bool requires_grad;
+  Device device;
+  std::vector<SymInt> strides;
+  SymInt storage_offset;
+  int64_t layout;
+
+ public:
+
+  ScalarType get_dtype() const {
+    return static_cast<ScalarType>(dtype);
+  }
+
+  void set_dtype(ScalarType def) {
+    dtype = static_cast<int64_t>(def);
+  }
+
+  const std::vector<SymInt>& get_sizes() const {
+    return sizes;
+  }
+
+  void set_sizes(std::vector<SymInt> def) {
+    sizes = std::move(def);
+  }
+
+  const bool& get_requires_grad() const {
+    return requires_grad;
+  }
+
+  void set_requires_grad(bool def) {
+    requires_grad = std::move(def);
+  }
+
+  const Device& get_device() const {
+    return device;
+  }
+
+  void set_device(Device def) {
+    device = std::move(def);
+  }
+
+  const std::vector<SymInt>& get_strides() const {
+    return strides;
+  }
+
+  void set_strides(std::vector<SymInt> def) {
+    strides = std::move(def);
+  }
+
+  const SymInt& get_storage_offset() const {
+    return storage_offset;
+  }
+
+  void set_storage_offset(SymInt def) {
+    storage_offset = std::move(def);
+  }
+
+  Layout get_layout() const {
+    return static_cast<Layout>(layout);
+  }
+
+  void set_layout(Layout def) {
+    layout = static_cast<int64_t>(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const TensorMeta& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, TensorMeta& nlohmann_json_t);
+};
+
+class SymIntArgument {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_NAME, AS_INT
+  };
+
+ private:
+  std::variant<Void, std::string, int64_t> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const std::string& get_as_name() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_name(std::string def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_NAME;
+  }
+
+  const int64_t& get_as_int() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_int(int64_t def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_INT;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymIntArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_NAME) {
+      nlohmann_json_j["as_name"] = nlohmann_json_t.get_as_name();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_INT) {
+      nlohmann_json_j["as_int"] = nlohmann_json_t.get_as_int();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymIntArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_name")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_name").template get<std::string>());
+      nlohmann_json_t.tag_ = Tag::AS_NAME;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_int")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_int").template get<int64_t>());
+      nlohmann_json_t.tag_ = Tag::AS_INT;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const SymIntArgument::Tag& e) {
+  switch (e) {
+    case SymIntArgument::Tag::AS_NAME: return "AS_NAME";
+    case SymIntArgument::Tag::AS_INT: return "AS_INT";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, SymIntArgument::Tag& t) {
+  if (s == "AS_NAME") { t = SymIntArgument::Tag::AS_NAME; return; }
+  if (s == "AS_INT") { t = SymIntArgument::Tag::AS_INT; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class SymFloatArgument {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_NAME, AS_FLOAT
+  };
+
+ private:
+  std::variant<Void, std::string, F64> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const std::string& get_as_name() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_name(std::string def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_NAME;
+  }
+
+  const F64& get_as_float() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_float(F64 def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_FLOAT;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymFloatArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_NAME) {
+      nlohmann_json_j["as_name"] = nlohmann_json_t.get_as_name();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOAT) {
+      nlohmann_json_j["as_float"] = nlohmann_json_t.get_as_float();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymFloatArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_name")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_name").template get<std::string>());
+      nlohmann_json_t.tag_ = Tag::AS_NAME;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_float")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_float").template get<F64>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOAT;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const SymFloatArgument::Tag& e) {
+  switch (e) {
+    case SymFloatArgument::Tag::AS_NAME: return "AS_NAME";
+    case SymFloatArgument::Tag::AS_FLOAT: return "AS_FLOAT";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, SymFloatArgument::Tag& t) {
+  if (s == "AS_NAME") { t = SymFloatArgument::Tag::AS_NAME; return; }
+  if (s == "AS_FLOAT") { t = SymFloatArgument::Tag::AS_FLOAT; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class SymBoolArgument {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_NAME, AS_BOOL
+  };
+
+ private:
+  std::variant<Void, std::string, bool> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const std::string& get_as_name() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_name(std::string def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_NAME;
+  }
+
+  const bool& get_as_bool() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_bool(bool def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_BOOL;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SymBoolArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_NAME) {
+      nlohmann_json_j["as_name"] = nlohmann_json_t.get_as_name();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_BOOL) {
+      nlohmann_json_j["as_bool"] = nlohmann_json_t.get_as_bool();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SymBoolArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_name")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_name").template get<std::string>());
+      nlohmann_json_t.tag_ = Tag::AS_NAME;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_bool")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_bool").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_BOOL;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const SymBoolArgument::Tag& e) {
+  switch (e) {
+    case SymBoolArgument::Tag::AS_NAME: return "AS_NAME";
+    case SymBoolArgument::Tag::AS_BOOL: return "AS_BOOL";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, SymBoolArgument::Tag& t) {
+  if (s == "AS_NAME") { t = SymBoolArgument::Tag::AS_NAME; return; }
+  if (s == "AS_BOOL") { t = SymBoolArgument::Tag::AS_BOOL; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class TensorArgument {
+ private:
+  std::string name;
+
+ public:
+
+  const std::string& get_name() const {
+    return name;
+  }
+
+  void set_name(std::string def) {
+    name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const TensorArgument& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, TensorArgument& nlohmann_json_t);
+};
+
+class TokenArgument {
+ private:
+  std::string name;
+
+ public:
+
+  const std::string& get_name() const {
+    return name;
+  }
+
+  void set_name(std::string def) {
+    name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const TokenArgument& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, TokenArgument& nlohmann_json_t);
+};
+
+class OptionalTensorArgument {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_TENSOR, AS_NONE
+  };
+
+ private:
+  std::variant<Void, TensorArgument, bool> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const TensorArgument& get_as_tensor() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_tensor(TensorArgument def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_TENSOR;
+  }
+
+  const bool& get_as_none() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_none(bool def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_NONE;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const OptionalTensorArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_TENSOR) {
+      nlohmann_json_j["as_tensor"] = nlohmann_json_t.get_as_tensor();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_NONE) {
+      nlohmann_json_j["as_none"] = nlohmann_json_t.get_as_none();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, OptionalTensorArgument& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_tensor")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_tensor").template get<TensorArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_TENSOR;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_none")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_none").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_NONE;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const OptionalTensorArgument::Tag& e) {
+  switch (e) {
+    case OptionalTensorArgument::Tag::AS_TENSOR: return "AS_TENSOR";
+    case OptionalTensorArgument::Tag::AS_NONE: return "AS_NONE";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, OptionalTensorArgument::Tag& t) {
+  if (s == "AS_TENSOR") { t = OptionalTensorArgument::Tag::AS_TENSOR; return; }
+  if (s == "AS_NONE") { t = OptionalTensorArgument::Tag::AS_NONE; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class GraphArgument {
+ private:
+  std::string name;
+  ForwardRef<Graph> graph;
+
+ public:
+
+  const std::string& get_name() const {
+    return name;
+  }
+
+  void set_name(std::string def) {
+    name = std::move(def);
+  }
+
+  const ForwardRef<Graph>& get_graph() const {
+    return graph;
+  }
+
+  void set_graph(ForwardRef<Graph> def) {
+    graph = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const GraphArgument& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, GraphArgument& nlohmann_json_t);
+};
+
+class CustomObjArgument {
+ private:
+  std::string name;
+  std::string class_fqn;
+
+ public:
+
+  const std::string& get_name() const {
+    return name;
+  }
+
+  void set_name(std::string def) {
+    name = std::move(def);
+  }
+
+  const std::string& get_class_fqn() const {
+    return class_fqn;
+  }
+
+  void set_class_fqn(std::string def) {
+    class_fqn = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const CustomObjArgument& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, CustomObjArgument& nlohmann_json_t);
+};
+
+class ComplexValue {
+ private:
+  F64 real;
+  F64 imag;
+
+ public:
+
+  const F64& get_real() const {
+    return real;
+  }
+
+  void set_real(F64 def) {
+    real = std::move(def);
+  }
+
+  const F64& get_imag() const {
+    return imag;
+  }
+
+  void set_imag(F64 def) {
+    imag = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ComplexValue& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ComplexValue& nlohmann_json_t);
+};
+
+class Argument {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_NONE, AS_TENSOR, AS_TENSORS, AS_INT, AS_INTS, AS_FLOAT, AS_FLOATS, AS_STRING, AS_STRINGS, AS_SYM_INT, AS_SYM_INTS, AS_SCALAR_TYPE, AS_MEMORY_FORMAT, AS_LAYOUT, AS_DEVICE, AS_BOOL, AS_BOOLS, AS_SYM_BOOL, AS_SYM_BOOLS, AS_GRAPH, AS_OPTIONAL_TENSORS, AS_CUSTOM_OBJ, AS_OPERATOR, AS_SYM_FLOAT, AS_SYM_FLOATS, AS_OPTIONAL_TENSOR, AS_COMPLEX, AS_NESTED_TENSORS, AS_INT_LISTS, AS_STRING_TO_ARGUMENT, AS_FLOAT_LISTS
+  };
+
+ private:
+  std::variant<Void, bool, TensorArgument, std::vector<TensorArgument>, int64_t, std::vector<int64_t>, F64, std::vector<F64>, std::string, std::vector<std::string>, SymIntArgument, std::vector<SymIntArgument>, ScalarType, MemoryFormat, Layout, Device, bool, std::vector<bool>, SymBoolArgument, std::vector<SymBoolArgument>, GraphArgument, std::vector<OptionalTensorArgument>, CustomObjArgument, std::string, SymFloatArgument, std::vector<SymFloatArgument>, OptionalTensorArgument, ComplexValue, std::vector<std::vector<TensorArgument>>, std::vector<std::vector<int64_t>>, std::unordered_map<std::string, ForwardRef<Argument>>, std::vector<std::vector<F64>>> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const bool& get_as_none() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_none(bool def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_NONE;
+  }
+
+  const TensorArgument& get_as_tensor() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_tensor(TensorArgument def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_TENSOR;
+  }
+
+  const std::vector<TensorArgument>& get_as_tensors() const {
+    return std::get<3>(variant_);
+  }
+
+  void set_as_tensors(std::vector<TensorArgument> def) {
+    variant_.emplace<3>(std::move(def));
+    tag_ = Tag::AS_TENSORS;
+  }
+
+  const int64_t& get_as_int() const {
+    return std::get<4>(variant_);
+  }
+
+  void set_as_int(int64_t def) {
+    variant_.emplace<4>(std::move(def));
+    tag_ = Tag::AS_INT;
+  }
+
+  const std::vector<int64_t>& get_as_ints() const {
+    return std::get<5>(variant_);
+  }
+
+  void set_as_ints(std::vector<int64_t> def) {
+    variant_.emplace<5>(std::move(def));
+    tag_ = Tag::AS_INTS;
+  }
+
+  const F64& get_as_float() const {
+    return std::get<6>(variant_);
+  }
+
+  void set_as_float(F64 def) {
+    variant_.emplace<6>(std::move(def));
+    tag_ = Tag::AS_FLOAT;
+  }
+
+  const std::vector<F64>& get_as_floats() const {
+    return std::get<7>(variant_);
+  }
+
+  void set_as_floats(std::vector<F64> def) {
+    variant_.emplace<7>(std::move(def));
+    tag_ = Tag::AS_FLOATS;
+  }
+
+  const std::string& get_as_string() const {
+    return std::get<8>(variant_);
+  }
+
+  void set_as_string(std::string def) {
+    variant_.emplace<8>(std::move(def));
+    tag_ = Tag::AS_STRING;
+  }
+
+  const std::vector<std::string>& get_as_strings() const {
+    return std::get<9>(variant_);
+  }
+
+  void set_as_strings(std::vector<std::string> def) {
+    variant_.emplace<9>(std::move(def));
+    tag_ = Tag::AS_STRINGS;
+  }
+
+  const SymIntArgument& get_as_sym_int() const {
+    return std::get<10>(variant_);
+  }
+
+  void set_as_sym_int(SymIntArgument def) {
+    variant_.emplace<10>(std::move(def));
+    tag_ = Tag::AS_SYM_INT;
+  }
+
+  const std::vector<SymIntArgument>& get_as_sym_ints() const {
+    return std::get<11>(variant_);
+  }
+
+  void set_as_sym_ints(std::vector<SymIntArgument> def) {
+    variant_.emplace<11>(std::move(def));
+    tag_ = Tag::AS_SYM_INTS;
+  }
+
+  const ScalarType& get_as_scalar_type() const {
+    return std::get<12>(variant_);
+  }
+
+  void set_as_scalar_type(ScalarType def) {
+    variant_.emplace<12>(std::move(def));
+    tag_ = Tag::AS_SCALAR_TYPE;
+  }
+
+  const MemoryFormat& get_as_memory_format() const {
+    return std::get<13>(variant_);
+  }
+
+  void set_as_memory_format(MemoryFormat def) {
+    variant_.emplace<13>(std::move(def));
+    tag_ = Tag::AS_MEMORY_FORMAT;
+  }
+
+  const Layout& get_as_layout() const {
+    return std::get<14>(variant_);
+  }
+
+  void set_as_layout(Layout def) {
+    variant_.emplace<14>(std::move(def));
+    tag_ = Tag::AS_LAYOUT;
+  }
+
+  const Device& get_as_device() const {
+    return std::get<15>(variant_);
+  }
+
+  void set_as_device(Device def) {
+    variant_.emplace<15>(std::move(def));
+    tag_ = Tag::AS_DEVICE;
+  }
+
+  const bool& get_as_bool() const {
+    return std::get<16>(variant_);
+  }
+
+  void set_as_bool(bool def) {
+    variant_.emplace<16>(std::move(def));
+    tag_ = Tag::AS_BOOL;
+  }
+
+  const std::vector<bool>& get_as_bools() const {
+    return std::get<17>(variant_);
+  }
+
+  void set_as_bools(std::vector<bool> def) {
+    variant_.emplace<17>(std::move(def));
+    tag_ = Tag::AS_BOOLS;
+  }
+
+  const SymBoolArgument& get_as_sym_bool() const {
+    return std::get<18>(variant_);
+  }
+
+  void set_as_sym_bool(SymBoolArgument def) {
+    variant_.emplace<18>(std::move(def));
+    tag_ = Tag::AS_SYM_BOOL;
+  }
+
+  const std::vector<SymBoolArgument>& get_as_sym_bools() const {
+    return std::get<19>(variant_);
+  }
+
+  void set_as_sym_bools(std::vector<SymBoolArgument> def) {
+    variant_.emplace<19>(std::move(def));
+    tag_ = Tag::AS_SYM_BOOLS;
+  }
+
+  const GraphArgument& get_as_graph() const {
+    return std::get<20>(variant_);
+  }
+
+  void set_as_graph(GraphArgument def) {
+    variant_.emplace<20>(std::move(def));
+    tag_ = Tag::AS_GRAPH;
+  }
+
+  const std::vector<OptionalTensorArgument>& get_as_optional_tensors() const {
+    return std::get<21>(variant_);
+  }
+
+  void set_as_optional_tensors(std::vector<OptionalTensorArgument> def) {
+    variant_.emplace<21>(std::move(def));
+    tag_ = Tag::AS_OPTIONAL_TENSORS;
+  }
+
+  const CustomObjArgument& get_as_custom_obj() const {
+    return std::get<22>(variant_);
+  }
+
+  void set_as_custom_obj(CustomObjArgument def) {
+    variant_.emplace<22>(std::move(def));
+    tag_ = Tag::AS_CUSTOM_OBJ;
+  }
+
+  const std::string& get_as_operator() const {
+    return std::get<23>(variant_);
+  }
+
+  void set_as_operator(std::string def) {
+    variant_.emplace<23>(std::move(def));
+    tag_ = Tag::AS_OPERATOR;
+  }
+
+  const SymFloatArgument& get_as_sym_float() const {
+    return std::get<24>(variant_);
+  }
+
+  void set_as_sym_float(SymFloatArgument def) {
+    variant_.emplace<24>(std::move(def));
+    tag_ = Tag::AS_SYM_FLOAT;
+  }
+
+  const std::vector<SymFloatArgument>& get_as_sym_floats() const {
+    return std::get<25>(variant_);
+  }
+
+  void set_as_sym_floats(std::vector<SymFloatArgument> def) {
+    variant_.emplace<25>(std::move(def));
+    tag_ = Tag::AS_SYM_FLOATS;
+  }
+
+  const OptionalTensorArgument& get_as_optional_tensor() const {
+    return std::get<26>(variant_);
+  }
+
+  void set_as_optional_tensor(OptionalTensorArgument def) {
+    variant_.emplace<26>(std::move(def));
+    tag_ = Tag::AS_OPTIONAL_TENSOR;
+  }
+
+  const ComplexValue& get_as_complex() const {
+    return std::get<27>(variant_);
+  }
+
+  void set_as_complex(ComplexValue def) {
+    variant_.emplace<27>(std::move(def));
+    tag_ = Tag::AS_COMPLEX;
+  }
+
+  const std::vector<std::vector<TensorArgument>>& get_as_nested_tensors() const {
+    return std::get<28>(variant_);
+  }
+
+  void set_as_nested_tensors(std::vector<std::vector<TensorArgument>> def) {
+    variant_.emplace<28>(std::move(def));
+    tag_ = Tag::AS_NESTED_TENSORS;
+  }
+
+  const std::vector<std::vector<int64_t>>& get_as_int_lists() const {
+    return std::get<29>(variant_);
+  }
+
+  void set_as_int_lists(std::vector<std::vector<int64_t>> def) {
+    variant_.emplace<29>(std::move(def));
+    tag_ = Tag::AS_INT_LISTS;
+  }
+
+  const std::unordered_map<std::string, ForwardRef<Argument>>& get_as_string_to_argument() const {
+    return std::get<30>(variant_);
+  }
+
+  void set_as_string_to_argument(std::unordered_map<std::string, ForwardRef<Argument>> def) {
+    variant_.emplace<30>(std::move(def));
+    tag_ = Tag::AS_STRING_TO_ARGUMENT;
+  }
+
+  const std::vector<std::vector<F64>>& get_as_float_lists() const {
+    return std::get<31>(variant_);
+  }
+
+  void set_as_float_lists(std::vector<std::vector<F64>> def) {
+    variant_.emplace<31>(std::move(def));
+    tag_ = Tag::AS_FLOAT_LISTS;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const Argument& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_NONE) {
+      nlohmann_json_j["as_none"] = nlohmann_json_t.get_as_none();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_TENSOR) {
+      nlohmann_json_j["as_tensor"] = nlohmann_json_t.get_as_tensor();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_TENSORS) {
+      nlohmann_json_j["as_tensors"] = nlohmann_json_t.get_as_tensors();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_INT) {
+      nlohmann_json_j["as_int"] = nlohmann_json_t.get_as_int();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_INTS) {
+      nlohmann_json_j["as_ints"] = nlohmann_json_t.get_as_ints();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOAT) {
+      nlohmann_json_j["as_float"] = nlohmann_json_t.get_as_float();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOATS) {
+      nlohmann_json_j["as_floats"] = nlohmann_json_t.get_as_floats();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_STRING) {
+      nlohmann_json_j["as_string"] = nlohmann_json_t.get_as_string();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_STRINGS) {
+      nlohmann_json_j["as_strings"] = nlohmann_json_t.get_as_strings();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_SYM_INT) {
+      nlohmann_json_j["as_sym_int"] = nlohmann_json_t.get_as_sym_int();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_SYM_INTS) {
+      nlohmann_json_j["as_sym_ints"] = nlohmann_json_t.get_as_sym_ints();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_SCALAR_TYPE) {
+      nlohmann_json_j["as_scalar_type"] = nlohmann_json_t.get_as_scalar_type();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_MEMORY_FORMAT) {
+      nlohmann_json_j["as_memory_format"] = nlohmann_json_t.get_as_memory_format();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_LAYOUT) {
+      nlohmann_json_j["as_layout"] = nlohmann_json_t.get_as_layout();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_DEVICE) {
+      nlohmann_json_j["as_device"] = nlohmann_json_t.get_as_device();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_BOOL) {
+      nlohmann_json_j["as_bool"] = nlohmann_json_t.get_as_bool();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_BOOLS) {
+      nlohmann_json_j["as_bools"] = nlohmann_json_t.get_as_bools();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_SYM_BOOL) {
+      nlohmann_json_j["as_sym_bool"] = nlohmann_json_t.get_as_sym_bool();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_SYM_BOOLS) {
+      nlohmann_json_j["as_sym_bools"] = nlohmann_json_t.get_as_sym_bools();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_GRAPH) {
+      nlohmann_json_j["as_graph"] = nlohmann_json_t.get_as_graph();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_OPTIONAL_TENSORS) {
+      nlohmann_json_j["as_optional_tensors"] = nlohmann_json_t.get_as_optional_tensors();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_CUSTOM_OBJ) {
+      nlohmann_json_j["as_custom_obj"] = nlohmann_json_t.get_as_custom_obj();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_OPERATOR) {
+      nlohmann_json_j["as_operator"] = nlohmann_json_t.get_as_operator();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_SYM_FLOAT) {
+      nlohmann_json_j["as_sym_float"] = nlohmann_json_t.get_as_sym_float();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_SYM_FLOATS) {
+      nlohmann_json_j["as_sym_floats"] = nlohmann_json_t.get_as_sym_floats();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_OPTIONAL_TENSOR) {
+      nlohmann_json_j["as_optional_tensor"] = nlohmann_json_t.get_as_optional_tensor();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_COMPLEX) {
+      nlohmann_json_j["as_complex"] = nlohmann_json_t.get_as_complex();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_NESTED_TENSORS) {
+      nlohmann_json_j["as_nested_tensors"] = nlohmann_json_t.get_as_nested_tensors();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_INT_LISTS) {
+      nlohmann_json_j["as_int_lists"] = nlohmann_json_t.get_as_int_lists();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_STRING_TO_ARGUMENT) {
+      nlohmann_json_j["as_string_to_argument"] = nlohmann_json_t.get_as_string_to_argument();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOAT_LISTS) {
+      nlohmann_json_j["as_float_lists"] = nlohmann_json_t.get_as_float_lists();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, Argument& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_none")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_none").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_NONE;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_tensor")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_tensor").template get<TensorArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_TENSOR;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_tensors")) {
+      nlohmann_json_t.variant_.emplace<3>(nlohmann_json_j.at("as_tensors").template get<std::vector<TensorArgument>>());
+      nlohmann_json_t.tag_ = Tag::AS_TENSORS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_int")) {
+      nlohmann_json_t.variant_.emplace<4>(nlohmann_json_j.at("as_int").template get<int64_t>());
+      nlohmann_json_t.tag_ = Tag::AS_INT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_ints")) {
+      nlohmann_json_t.variant_.emplace<5>(nlohmann_json_j.at("as_ints").template get<std::vector<int64_t>>());
+      nlohmann_json_t.tag_ = Tag::AS_INTS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_float")) {
+      nlohmann_json_t.variant_.emplace<6>(nlohmann_json_j.at("as_float").template get<F64>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOAT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_floats")) {
+      nlohmann_json_t.variant_.emplace<7>(nlohmann_json_j.at("as_floats").template get<std::vector<F64>>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOATS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_string")) {
+      nlohmann_json_t.variant_.emplace<8>(nlohmann_json_j.at("as_string").template get<std::string>());
+      nlohmann_json_t.tag_ = Tag::AS_STRING;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_strings")) {
+      nlohmann_json_t.variant_.emplace<9>(nlohmann_json_j.at("as_strings").template get<std::vector<std::string>>());
+      nlohmann_json_t.tag_ = Tag::AS_STRINGS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_sym_int")) {
+      nlohmann_json_t.variant_.emplace<10>(nlohmann_json_j.at("as_sym_int").template get<SymIntArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_SYM_INT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_sym_ints")) {
+      nlohmann_json_t.variant_.emplace<11>(nlohmann_json_j.at("as_sym_ints").template get<std::vector<SymIntArgument>>());
+      nlohmann_json_t.tag_ = Tag::AS_SYM_INTS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_scalar_type")) {
+      nlohmann_json_t.variant_.emplace<12>(nlohmann_json_j.at("as_scalar_type").template get<ScalarType>());
+      nlohmann_json_t.tag_ = Tag::AS_SCALAR_TYPE;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_memory_format")) {
+      nlohmann_json_t.variant_.emplace<13>(nlohmann_json_j.at("as_memory_format").template get<MemoryFormat>());
+      nlohmann_json_t.tag_ = Tag::AS_MEMORY_FORMAT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_layout")) {
+      nlohmann_json_t.variant_.emplace<14>(nlohmann_json_j.at("as_layout").template get<Layout>());
+      nlohmann_json_t.tag_ = Tag::AS_LAYOUT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_device")) {
+      nlohmann_json_t.variant_.emplace<15>(nlohmann_json_j.at("as_device").template get<Device>());
+      nlohmann_json_t.tag_ = Tag::AS_DEVICE;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_bool")) {
+      nlohmann_json_t.variant_.emplace<16>(nlohmann_json_j.at("as_bool").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_BOOL;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_bools")) {
+      nlohmann_json_t.variant_.emplace<17>(nlohmann_json_j.at("as_bools").template get<std::vector<bool>>());
+      nlohmann_json_t.tag_ = Tag::AS_BOOLS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_sym_bool")) {
+      nlohmann_json_t.variant_.emplace<18>(nlohmann_json_j.at("as_sym_bool").template get<SymBoolArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_SYM_BOOL;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_sym_bools")) {
+      nlohmann_json_t.variant_.emplace<19>(nlohmann_json_j.at("as_sym_bools").template get<std::vector<SymBoolArgument>>());
+      nlohmann_json_t.tag_ = Tag::AS_SYM_BOOLS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_graph")) {
+      nlohmann_json_t.variant_.emplace<20>(nlohmann_json_j.at("as_graph").template get<GraphArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_GRAPH;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_optional_tensors")) {
+      nlohmann_json_t.variant_.emplace<21>(nlohmann_json_j.at("as_optional_tensors").template get<std::vector<OptionalTensorArgument>>());
+      nlohmann_json_t.tag_ = Tag::AS_OPTIONAL_TENSORS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_custom_obj")) {
+      nlohmann_json_t.variant_.emplace<22>(nlohmann_json_j.at("as_custom_obj").template get<CustomObjArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_CUSTOM_OBJ;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_operator")) {
+      nlohmann_json_t.variant_.emplace<23>(nlohmann_json_j.at("as_operator").template get<std::string>());
+      nlohmann_json_t.tag_ = Tag::AS_OPERATOR;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_sym_float")) {
+      nlohmann_json_t.variant_.emplace<24>(nlohmann_json_j.at("as_sym_float").template get<SymFloatArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_SYM_FLOAT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_sym_floats")) {
+      nlohmann_json_t.variant_.emplace<25>(nlohmann_json_j.at("as_sym_floats").template get<std::vector<SymFloatArgument>>());
+      nlohmann_json_t.tag_ = Tag::AS_SYM_FLOATS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_optional_tensor")) {
+      nlohmann_json_t.variant_.emplace<26>(nlohmann_json_j.at("as_optional_tensor").template get<OptionalTensorArgument>());
+      nlohmann_json_t.tag_ = Tag::AS_OPTIONAL_TENSOR;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_complex")) {
+      nlohmann_json_t.variant_.emplace<27>(nlohmann_json_j.at("as_complex").template get<ComplexValue>());
+      nlohmann_json_t.tag_ = Tag::AS_COMPLEX;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_nested_tensors")) {
+      nlohmann_json_t.variant_.emplace<28>(nlohmann_json_j.at("as_nested_tensors").template get<std::vector<std::vector<TensorArgument>>>());
+      nlohmann_json_t.tag_ = Tag::AS_NESTED_TENSORS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_int_lists")) {
+      nlohmann_json_t.variant_.emplace<29>(nlohmann_json_j.at("as_int_lists").template get<std::vector<std::vector<int64_t>>>());
+      nlohmann_json_t.tag_ = Tag::AS_INT_LISTS;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_string_to_argument")) {
+      nlohmann_json_t.variant_.emplace<30>(nlohmann_json_j.at("as_string_to_argument").template get<std::unordered_map<std::string, ForwardRef<Argument>>>());
+      nlohmann_json_t.tag_ = Tag::AS_STRING_TO_ARGUMENT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_float_lists")) {
+      nlohmann_json_t.variant_.emplace<31>(nlohmann_json_j.at("as_float_lists").template get<std::vector<std::vector<F64>>>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOAT_LISTS;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const Argument::Tag& e) {
+  switch (e) {
+    case Argument::Tag::AS_NONE: return "AS_NONE";
+    case Argument::Tag::AS_TENSOR: return "AS_TENSOR";
+    case Argument::Tag::AS_TENSORS: return "AS_TENSORS";
+    case Argument::Tag::AS_INT: return "AS_INT";
+    case Argument::Tag::AS_INTS: return "AS_INTS";
+    case Argument::Tag::AS_FLOAT: return "AS_FLOAT";
+    case Argument::Tag::AS_FLOATS: return "AS_FLOATS";
+    case Argument::Tag::AS_STRING: return "AS_STRING";
+    case Argument::Tag::AS_STRINGS: return "AS_STRINGS";
+    case Argument::Tag::AS_SYM_INT: return "AS_SYM_INT";
+    case Argument::Tag::AS_SYM_INTS: return "AS_SYM_INTS";
+    case Argument::Tag::AS_SCALAR_TYPE: return "AS_SCALAR_TYPE";
+    case Argument::Tag::AS_MEMORY_FORMAT: return "AS_MEMORY_FORMAT";
+    case Argument::Tag::AS_LAYOUT: return "AS_LAYOUT";
+    case Argument::Tag::AS_DEVICE: return "AS_DEVICE";
+    case Argument::Tag::AS_BOOL: return "AS_BOOL";
+    case Argument::Tag::AS_BOOLS: return "AS_BOOLS";
+    case Argument::Tag::AS_SYM_BOOL: return "AS_SYM_BOOL";
+    case Argument::Tag::AS_SYM_BOOLS: return "AS_SYM_BOOLS";
+    case Argument::Tag::AS_GRAPH: return "AS_GRAPH";
+    case Argument::Tag::AS_OPTIONAL_TENSORS: return "AS_OPTIONAL_TENSORS";
+    case Argument::Tag::AS_CUSTOM_OBJ: return "AS_CUSTOM_OBJ";
+    case Argument::Tag::AS_OPERATOR: return "AS_OPERATOR";
+    case Argument::Tag::AS_SYM_FLOAT: return "AS_SYM_FLOAT";
+    case Argument::Tag::AS_SYM_FLOATS: return "AS_SYM_FLOATS";
+    case Argument::Tag::AS_OPTIONAL_TENSOR: return "AS_OPTIONAL_TENSOR";
+    case Argument::Tag::AS_COMPLEX: return "AS_COMPLEX";
+    case Argument::Tag::AS_NESTED_TENSORS: return "AS_NESTED_TENSORS";
+    case Argument::Tag::AS_INT_LISTS: return "AS_INT_LISTS";
+    case Argument::Tag::AS_STRING_TO_ARGUMENT: return "AS_STRING_TO_ARGUMENT";
+    case Argument::Tag::AS_FLOAT_LISTS: return "AS_FLOAT_LISTS";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, Argument::Tag& t) {
+  if (s == "AS_NONE") { t = Argument::Tag::AS_NONE; return; }
+  if (s == "AS_TENSOR") { t = Argument::Tag::AS_TENSOR; return; }
+  if (s == "AS_TENSORS") { t = Argument::Tag::AS_TENSORS; return; }
+  if (s == "AS_INT") { t = Argument::Tag::AS_INT; return; }
+  if (s == "AS_INTS") { t = Argument::Tag::AS_INTS; return; }
+  if (s == "AS_FLOAT") { t = Argument::Tag::AS_FLOAT; return; }
+  if (s == "AS_FLOATS") { t = Argument::Tag::AS_FLOATS; return; }
+  if (s == "AS_STRING") { t = Argument::Tag::AS_STRING; return; }
+  if (s == "AS_STRINGS") { t = Argument::Tag::AS_STRINGS; return; }
+  if (s == "AS_SYM_INT") { t = Argument::Tag::AS_SYM_INT; return; }
+  if (s == "AS_SYM_INTS") { t = Argument::Tag::AS_SYM_INTS; return; }
+  if (s == "AS_SCALAR_TYPE") { t = Argument::Tag::AS_SCALAR_TYPE; return; }
+  if (s == "AS_MEMORY_FORMAT") { t = Argument::Tag::AS_MEMORY_FORMAT; return; }
+  if (s == "AS_LAYOUT") { t = Argument::Tag::AS_LAYOUT; return; }
+  if (s == "AS_DEVICE") { t = Argument::Tag::AS_DEVICE; return; }
+  if (s == "AS_BOOL") { t = Argument::Tag::AS_BOOL; return; }
+  if (s == "AS_BOOLS") { t = Argument::Tag::AS_BOOLS; return; }
+  if (s == "AS_SYM_BOOL") { t = Argument::Tag::AS_SYM_BOOL; return; }
+  if (s == "AS_SYM_BOOLS") { t = Argument::Tag::AS_SYM_BOOLS; return; }
+  if (s == "AS_GRAPH") { t = Argument::Tag::AS_GRAPH; return; }
+  if (s == "AS_OPTIONAL_TENSORS") { t = Argument::Tag::AS_OPTIONAL_TENSORS; return; }
+  if (s == "AS_CUSTOM_OBJ") { t = Argument::Tag::AS_CUSTOM_OBJ; return; }
+  if (s == "AS_OPERATOR") { t = Argument::Tag::AS_OPERATOR; return; }
+  if (s == "AS_SYM_FLOAT") { t = Argument::Tag::AS_SYM_FLOAT; return; }
+  if (s == "AS_SYM_FLOATS") { t = Argument::Tag::AS_SYM_FLOATS; return; }
+  if (s == "AS_OPTIONAL_TENSOR") { t = Argument::Tag::AS_OPTIONAL_TENSOR; return; }
+  if (s == "AS_COMPLEX") { t = Argument::Tag::AS_COMPLEX; return; }
+  if (s == "AS_NESTED_TENSORS") { t = Argument::Tag::AS_NESTED_TENSORS; return; }
+  if (s == "AS_INT_LISTS") { t = Argument::Tag::AS_INT_LISTS; return; }
+  if (s == "AS_STRING_TO_ARGUMENT") { t = Argument::Tag::AS_STRING_TO_ARGUMENT; return; }
+  if (s == "AS_FLOAT_LISTS") { t = Argument::Tag::AS_FLOAT_LISTS; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class NamedArgument {
+ private:
+  std::string name;
+  Argument arg;
+  std::optional<int64_t> kind = std::nullopt;
+
+ public:
+
+  const std::string& get_name() const {
+    return name;
+  }
+
+  void set_name(std::string def) {
+    name = std::move(def);
+  }
+
+  const Argument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(Argument def) {
+    arg = std::move(def);
+  }
+
+  const std::optional<int64_t>& get_kind() const {
+    return kind;
+  }
+
+  void set_kind(std::optional<int64_t> def) {
+    kind = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const NamedArgument& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, NamedArgument& nlohmann_json_t);
+};
+
+class Node {
+ private:
+  std::string target;
+  std::vector<NamedArgument> inputs;
+  std::vector<Argument> outputs;
+  std::unordered_map<std::string, std::string> metadata;
+  std::optional<bool> is_hop_single_tensor_return = std::nullopt;
+  std::optional<std::string> name = std::nullopt;
+
+ public:
+
+  const std::string& get_target() const {
+    return target;
+  }
+
+  void set_target(std::string def) {
+    target = std::move(def);
+  }
+
+  const std::vector<NamedArgument>& get_inputs() const {
+    return inputs;
+  }
+
+  void set_inputs(std::vector<NamedArgument> def) {
+    inputs = std::move(def);
+  }
+
+  const std::vector<Argument>& get_outputs() const {
+    return outputs;
+  }
+
+  void set_outputs(std::vector<Argument> def) {
+    outputs = std::move(def);
+  }
+
+  const std::unordered_map<std::string, std::string>& get_metadata() const {
+    return metadata;
+  }
+
+  void set_metadata(std::unordered_map<std::string, std::string> def) {
+    metadata = std::move(def);
+  }
+
+  const std::optional<bool>& get_is_hop_single_tensor_return() const {
+    return is_hop_single_tensor_return;
+  }
+
+  void set_is_hop_single_tensor_return(std::optional<bool> def) {
+    is_hop_single_tensor_return = std::move(def);
+  }
+
+  const std::optional<std::string>& get_name() const {
+    return name;
+  }
+
+  void set_name(std::optional<std::string> def) {
+    name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const Node& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, Node& nlohmann_json_t);
+};
+
+class Graph {
+ private:
+  std::vector<Argument> inputs;
+  std::vector<Argument> outputs;
+  std::vector<Node> nodes;
+  std::unordered_map<std::string, TensorMeta> tensor_values;
+  std::unordered_map<std::string, SymInt> sym_int_values;
+  std::unordered_map<std::string, SymBool> sym_bool_values;
+  bool is_single_tensor_return = false;
+  std::unordered_map<std::string, CustomObjArgument> custom_obj_values = {};
+  std::unordered_map<std::string, SymFloat> sym_float_values = {};
+
+ public:
+
+  const std::vector<Argument>& get_inputs() const {
+    return inputs;
+  }
+
+  void set_inputs(std::vector<Argument> def) {
+    inputs = std::move(def);
+  }
+
+  const std::vector<Argument>& get_outputs() const {
+    return outputs;
+  }
+
+  void set_outputs(std::vector<Argument> def) {
+    outputs = std::move(def);
+  }
+
+  const std::vector<Node>& get_nodes() const {
+    return nodes;
+  }
+
+  void set_nodes(std::vector<Node> def) {
+    nodes = std::move(def);
+  }
+
+  const std::unordered_map<std::string, TensorMeta>& get_tensor_values() const {
+    return tensor_values;
+  }
+
+  void set_tensor_values(std::unordered_map<std::string, TensorMeta> def) {
+    tensor_values = std::move(def);
+  }
+
+  const std::unordered_map<std::string, SymInt>& get_sym_int_values() const {
+    return sym_int_values;
+  }
+
+  void set_sym_int_values(std::unordered_map<std::string, SymInt> def) {
+    sym_int_values = std::move(def);
+  }
+
+  const std::unordered_map<std::string, SymBool>& get_sym_bool_values() const {
+    return sym_bool_values;
+  }
+
+  void set_sym_bool_values(std::unordered_map<std::string, SymBool> def) {
+    sym_bool_values = std::move(def);
+  }
+
+  const bool& get_is_single_tensor_return() const {
+    return is_single_tensor_return;
+  }
+
+  void set_is_single_tensor_return(bool def) {
+    is_single_tensor_return = std::move(def);
+  }
+
+  const std::unordered_map<std::string, CustomObjArgument>& get_custom_obj_values() const {
+    return custom_obj_values;
+  }
+
+  void set_custom_obj_values(std::unordered_map<std::string, CustomObjArgument> def) {
+    custom_obj_values = std::move(def);
+  }
+
+  const std::unordered_map<std::string, SymFloat>& get_sym_float_values() const {
+    return sym_float_values;
+  }
+
+  void set_sym_float_values(std::unordered_map<std::string, SymFloat> def) {
+    sym_float_values = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const Graph& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, Graph& nlohmann_json_t);
+};
+
+class UserInputSpec {
+ private:
+  Argument arg;
+
+ public:
+
+  const Argument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(Argument def) {
+    arg = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const UserInputSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, UserInputSpec& nlohmann_json_t);
+};
+
+class ConstantValue {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    AS_NONE, AS_INT, AS_FLOAT, AS_STRING, AS_BOOL
+  };
+
+ private:
+  std::variant<Void, bool, int64_t, F64, std::string, bool> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const bool& get_as_none() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_as_none(bool def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::AS_NONE;
+  }
+
+  const int64_t& get_as_int() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_as_int(int64_t def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::AS_INT;
+  }
+
+  const F64& get_as_float() const {
+    return std::get<3>(variant_);
+  }
+
+  void set_as_float(F64 def) {
+    variant_.emplace<3>(std::move(def));
+    tag_ = Tag::AS_FLOAT;
+  }
+
+  const std::string& get_as_string() const {
+    return std::get<4>(variant_);
+  }
+
+  void set_as_string(std::string def) {
+    variant_.emplace<4>(std::move(def));
+    tag_ = Tag::AS_STRING;
+  }
+
+  const bool& get_as_bool() const {
+    return std::get<5>(variant_);
+  }
+
+  void set_as_bool(bool def) {
+    variant_.emplace<5>(std::move(def));
+    tag_ = Tag::AS_BOOL;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ConstantValue& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::AS_NONE) {
+      nlohmann_json_j["as_none"] = nlohmann_json_t.get_as_none();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_INT) {
+      nlohmann_json_j["as_int"] = nlohmann_json_t.get_as_int();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_FLOAT) {
+      nlohmann_json_j["as_float"] = nlohmann_json_t.get_as_float();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_STRING) {
+      nlohmann_json_j["as_string"] = nlohmann_json_t.get_as_string();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::AS_BOOL) {
+      nlohmann_json_j["as_bool"] = nlohmann_json_t.get_as_bool();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ConstantValue& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("as_none")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("as_none").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_NONE;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_int")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("as_int").template get<int64_t>());
+      nlohmann_json_t.tag_ = Tag::AS_INT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_float")) {
+      nlohmann_json_t.variant_.emplace<3>(nlohmann_json_j.at("as_float").template get<F64>());
+      nlohmann_json_t.tag_ = Tag::AS_FLOAT;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_string")) {
+      nlohmann_json_t.variant_.emplace<4>(nlohmann_json_j.at("as_string").template get<std::string>());
+      nlohmann_json_t.tag_ = Tag::AS_STRING;
+      return;
+    }
+    if (nlohmann_json_j.contains("as_bool")) {
+      nlohmann_json_t.variant_.emplace<5>(nlohmann_json_j.at("as_bool").template get<bool>());
+      nlohmann_json_t.tag_ = Tag::AS_BOOL;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const ConstantValue::Tag& e) {
+  switch (e) {
+    case ConstantValue::Tag::AS_NONE: return "AS_NONE";
+    case ConstantValue::Tag::AS_INT: return "AS_INT";
+    case ConstantValue::Tag::AS_FLOAT: return "AS_FLOAT";
+    case ConstantValue::Tag::AS_STRING: return "AS_STRING";
+    case ConstantValue::Tag::AS_BOOL: return "AS_BOOL";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, ConstantValue::Tag& t) {
+  if (s == "AS_NONE") { t = ConstantValue::Tag::AS_NONE; return; }
+  if (s == "AS_INT") { t = ConstantValue::Tag::AS_INT; return; }
+  if (s == "AS_FLOAT") { t = ConstantValue::Tag::AS_FLOAT; return; }
+  if (s == "AS_STRING") { t = ConstantValue::Tag::AS_STRING; return; }
+  if (s == "AS_BOOL") { t = ConstantValue::Tag::AS_BOOL; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class InputToConstantInputSpec {
+ private:
+  std::string name;
+  ConstantValue value;
+
+ public:
+
+  const std::string& get_name() const {
+    return name;
+  }
+
+  void set_name(std::string def) {
+    name = std::move(def);
+  }
+
+  const ConstantValue& get_value() const {
+    return value;
+  }
+
+  void set_value(ConstantValue def) {
+    value = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const InputToConstantInputSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, InputToConstantInputSpec& nlohmann_json_t);
+};
+
+class InputToParameterSpec {
+ private:
+  TensorArgument arg;
+  std::string parameter_name;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_parameter_name() const {
+    return parameter_name;
+  }
+
+  void set_parameter_name(std::string def) {
+    parameter_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const InputToParameterSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, InputToParameterSpec& nlohmann_json_t);
+};
+
+class InputToBufferSpec {
+ private:
+  TensorArgument arg;
+  std::string buffer_name;
+  bool persistent;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_buffer_name() const {
+    return buffer_name;
+  }
+
+  void set_buffer_name(std::string def) {
+    buffer_name = std::move(def);
+  }
+
+  const bool& get_persistent() const {
+    return persistent;
+  }
+
+  void set_persistent(bool def) {
+    persistent = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const InputToBufferSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, InputToBufferSpec& nlohmann_json_t);
+};
+
+class InputToTensorConstantSpec {
+ private:
+  TensorArgument arg;
+  std::string tensor_constant_name;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_tensor_constant_name() const {
+    return tensor_constant_name;
+  }
+
+  void set_tensor_constant_name(std::string def) {
+    tensor_constant_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const InputToTensorConstantSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, InputToTensorConstantSpec& nlohmann_json_t);
+};
+
+class InputToCustomObjSpec {
+ private:
+  CustomObjArgument arg;
+  std::string custom_obj_name;
+
+ public:
+
+  const CustomObjArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(CustomObjArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_custom_obj_name() const {
+    return custom_obj_name;
+  }
+
+  void set_custom_obj_name(std::string def) {
+    custom_obj_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const InputToCustomObjSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, InputToCustomObjSpec& nlohmann_json_t);
+};
+
+class InputTokenSpec {
+ private:
+  TokenArgument arg;
+
+ public:
+
+  const TokenArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TokenArgument def) {
+    arg = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const InputTokenSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, InputTokenSpec& nlohmann_json_t);
+};
+
+class InputSpec {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    USER_INPUT, PARAMETER, BUFFER, TENSOR_CONSTANT, CUSTOM_OBJ, TOKEN, CONSTANT_INPUT
+  };
+
+ private:
+  std::variant<Void, UserInputSpec, InputToParameterSpec, InputToBufferSpec, InputToTensorConstantSpec, InputToCustomObjSpec, InputTokenSpec, InputToConstantInputSpec> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const UserInputSpec& get_user_input() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_user_input(UserInputSpec def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::USER_INPUT;
+  }
+
+  const InputToParameterSpec& get_parameter() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_parameter(InputToParameterSpec def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::PARAMETER;
+  }
+
+  const InputToBufferSpec& get_buffer() const {
+    return std::get<3>(variant_);
+  }
+
+  void set_buffer(InputToBufferSpec def) {
+    variant_.emplace<3>(std::move(def));
+    tag_ = Tag::BUFFER;
+  }
+
+  const InputToTensorConstantSpec& get_tensor_constant() const {
+    return std::get<4>(variant_);
+  }
+
+  void set_tensor_constant(InputToTensorConstantSpec def) {
+    variant_.emplace<4>(std::move(def));
+    tag_ = Tag::TENSOR_CONSTANT;
+  }
+
+  const InputToCustomObjSpec& get_custom_obj() const {
+    return std::get<5>(variant_);
+  }
+
+  void set_custom_obj(InputToCustomObjSpec def) {
+    variant_.emplace<5>(std::move(def));
+    tag_ = Tag::CUSTOM_OBJ;
+  }
+
+  const InputTokenSpec& get_token() const {
+    return std::get<6>(variant_);
+  }
+
+  void set_token(InputTokenSpec def) {
+    variant_.emplace<6>(std::move(def));
+    tag_ = Tag::TOKEN;
+  }
+
+  const InputToConstantInputSpec& get_constant_input() const {
+    return std::get<7>(variant_);
+  }
+
+  void set_constant_input(InputToConstantInputSpec def) {
+    variant_.emplace<7>(std::move(def));
+    tag_ = Tag::CONSTANT_INPUT;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const InputSpec& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::USER_INPUT) {
+      nlohmann_json_j["user_input"] = nlohmann_json_t.get_user_input();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::PARAMETER) {
+      nlohmann_json_j["parameter"] = nlohmann_json_t.get_parameter();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::BUFFER) {
+      nlohmann_json_j["buffer"] = nlohmann_json_t.get_buffer();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::TENSOR_CONSTANT) {
+      nlohmann_json_j["tensor_constant"] = nlohmann_json_t.get_tensor_constant();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::CUSTOM_OBJ) {
+      nlohmann_json_j["custom_obj"] = nlohmann_json_t.get_custom_obj();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::TOKEN) {
+      nlohmann_json_j["token"] = nlohmann_json_t.get_token();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::CONSTANT_INPUT) {
+      nlohmann_json_j["constant_input"] = nlohmann_json_t.get_constant_input();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, InputSpec& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("user_input")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("user_input").template get<UserInputSpec>());
+      nlohmann_json_t.tag_ = Tag::USER_INPUT;
+      return;
+    }
+    if (nlohmann_json_j.contains("parameter")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("parameter").template get<InputToParameterSpec>());
+      nlohmann_json_t.tag_ = Tag::PARAMETER;
+      return;
+    }
+    if (nlohmann_json_j.contains("buffer")) {
+      nlohmann_json_t.variant_.emplace<3>(nlohmann_json_j.at("buffer").template get<InputToBufferSpec>());
+      nlohmann_json_t.tag_ = Tag::BUFFER;
+      return;
+    }
+    if (nlohmann_json_j.contains("tensor_constant")) {
+      nlohmann_json_t.variant_.emplace<4>(nlohmann_json_j.at("tensor_constant").template get<InputToTensorConstantSpec>());
+      nlohmann_json_t.tag_ = Tag::TENSOR_CONSTANT;
+      return;
+    }
+    if (nlohmann_json_j.contains("custom_obj")) {
+      nlohmann_json_t.variant_.emplace<5>(nlohmann_json_j.at("custom_obj").template get<InputToCustomObjSpec>());
+      nlohmann_json_t.tag_ = Tag::CUSTOM_OBJ;
+      return;
+    }
+    if (nlohmann_json_j.contains("token")) {
+      nlohmann_json_t.variant_.emplace<6>(nlohmann_json_j.at("token").template get<InputTokenSpec>());
+      nlohmann_json_t.tag_ = Tag::TOKEN;
+      return;
+    }
+    if (nlohmann_json_j.contains("constant_input")) {
+      nlohmann_json_t.variant_.emplace<7>(nlohmann_json_j.at("constant_input").template get<InputToConstantInputSpec>());
+      nlohmann_json_t.tag_ = Tag::CONSTANT_INPUT;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const InputSpec::Tag& e) {
+  switch (e) {
+    case InputSpec::Tag::USER_INPUT: return "USER_INPUT";
+    case InputSpec::Tag::PARAMETER: return "PARAMETER";
+    case InputSpec::Tag::BUFFER: return "BUFFER";
+    case InputSpec::Tag::TENSOR_CONSTANT: return "TENSOR_CONSTANT";
+    case InputSpec::Tag::CUSTOM_OBJ: return "CUSTOM_OBJ";
+    case InputSpec::Tag::TOKEN: return "TOKEN";
+    case InputSpec::Tag::CONSTANT_INPUT: return "CONSTANT_INPUT";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, InputSpec::Tag& t) {
+  if (s == "USER_INPUT") { t = InputSpec::Tag::USER_INPUT; return; }
+  if (s == "PARAMETER") { t = InputSpec::Tag::PARAMETER; return; }
+  if (s == "BUFFER") { t = InputSpec::Tag::BUFFER; return; }
+  if (s == "TENSOR_CONSTANT") { t = InputSpec::Tag::TENSOR_CONSTANT; return; }
+  if (s == "CUSTOM_OBJ") { t = InputSpec::Tag::CUSTOM_OBJ; return; }
+  if (s == "TOKEN") { t = InputSpec::Tag::TOKEN; return; }
+  if (s == "CONSTANT_INPUT") { t = InputSpec::Tag::CONSTANT_INPUT; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class UserOutputSpec {
+ private:
+  Argument arg;
+
+ public:
+
+  const Argument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(Argument def) {
+    arg = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const UserOutputSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, UserOutputSpec& nlohmann_json_t);
+};
+
+class LossOutputSpec {
+ private:
+  TensorArgument arg;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const LossOutputSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, LossOutputSpec& nlohmann_json_t);
+};
+
+class BufferMutationSpec {
+ private:
+  TensorArgument arg;
+  std::string buffer_name;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_buffer_name() const {
+    return buffer_name;
+  }
+
+  void set_buffer_name(std::string def) {
+    buffer_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const BufferMutationSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, BufferMutationSpec& nlohmann_json_t);
+};
+
+class ParameterMutationSpec {
+ private:
+  TensorArgument arg;
+  std::string parameter_name;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_parameter_name() const {
+    return parameter_name;
+  }
+
+  void set_parameter_name(std::string def) {
+    parameter_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ParameterMutationSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ParameterMutationSpec& nlohmann_json_t);
+};
+
+class GradientToParameterSpec {
+ private:
+  TensorArgument arg;
+  std::string parameter_name;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_parameter_name() const {
+    return parameter_name;
+  }
+
+  void set_parameter_name(std::string def) {
+    parameter_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const GradientToParameterSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, GradientToParameterSpec& nlohmann_json_t);
+};
+
+class GradientToUserInputSpec {
+ private:
+  TensorArgument arg;
+  std::string user_input_name;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_user_input_name() const {
+    return user_input_name;
+  }
+
+  void set_user_input_name(std::string def) {
+    user_input_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const GradientToUserInputSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, GradientToUserInputSpec& nlohmann_json_t);
+};
+
+class UserInputMutationSpec {
+ private:
+  TensorArgument arg;
+  std::string user_input_name;
+
+ public:
+
+  const TensorArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TensorArgument def) {
+    arg = std::move(def);
+  }
+
+  const std::string& get_user_input_name() const {
+    return user_input_name;
+  }
+
+  void set_user_input_name(std::string def) {
+    user_input_name = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const UserInputMutationSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, UserInputMutationSpec& nlohmann_json_t);
+};
+
+class OutputTokenSpec {
+ private:
+  TokenArgument arg;
+
+ public:
+
+  const TokenArgument& get_arg() const {
+    return arg;
+  }
+
+  void set_arg(TokenArgument def) {
+    arg = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const OutputTokenSpec& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, OutputTokenSpec& nlohmann_json_t);
+};
+
+class OutputSpec {
+  struct Void {};
+
+ public:
+  enum class Tag {
+    USER_OUTPUT, LOSS_OUTPUT, BUFFER_MUTATION, GRADIENT_TO_PARAMETER, GRADIENT_TO_USER_INPUT, USER_INPUT_MUTATION, TOKEN, PARAMETER_MUTATION
+  };
+
+ private:
+  std::variant<Void, UserOutputSpec, LossOutputSpec, BufferMutationSpec, GradientToParameterSpec, GradientToUserInputSpec, UserInputMutationSpec, OutputTokenSpec, ParameterMutationSpec> variant_;
+  Tag tag_;
+
+ public:
+  Tag tag() const {
+    return tag_;
+  }
+
+  const UserOutputSpec& get_user_output() const {
+    return std::get<1>(variant_);
+  }
+
+  void set_user_output(UserOutputSpec def) {
+    variant_.emplace<1>(std::move(def));
+    tag_ = Tag::USER_OUTPUT;
+  }
+
+  const LossOutputSpec& get_loss_output() const {
+    return std::get<2>(variant_);
+  }
+
+  void set_loss_output(LossOutputSpec def) {
+    variant_.emplace<2>(std::move(def));
+    tag_ = Tag::LOSS_OUTPUT;
+  }
+
+  const BufferMutationSpec& get_buffer_mutation() const {
+    return std::get<3>(variant_);
+  }
+
+  void set_buffer_mutation(BufferMutationSpec def) {
+    variant_.emplace<3>(std::move(def));
+    tag_ = Tag::BUFFER_MUTATION;
+  }
+
+  const GradientToParameterSpec& get_gradient_to_parameter() const {
+    return std::get<4>(variant_);
+  }
+
+  void set_gradient_to_parameter(GradientToParameterSpec def) {
+    variant_.emplace<4>(std::move(def));
+    tag_ = Tag::GRADIENT_TO_PARAMETER;
+  }
+
+  const GradientToUserInputSpec& get_gradient_to_user_input() const {
+    return std::get<5>(variant_);
+  }
+
+  void set_gradient_to_user_input(GradientToUserInputSpec def) {
+    variant_.emplace<5>(std::move(def));
+    tag_ = Tag::GRADIENT_TO_USER_INPUT;
+  }
+
+  const UserInputMutationSpec& get_user_input_mutation() const {
+    return std::get<6>(variant_);
+  }
+
+  void set_user_input_mutation(UserInputMutationSpec def) {
+    variant_.emplace<6>(std::move(def));
+    tag_ = Tag::USER_INPUT_MUTATION;
+  }
+
+  const OutputTokenSpec& get_token() const {
+    return std::get<7>(variant_);
+  }
+
+  void set_token(OutputTokenSpec def) {
+    variant_.emplace<7>(std::move(def));
+    tag_ = Tag::TOKEN;
+  }
+
+  const ParameterMutationSpec& get_parameter_mutation() const {
+    return std::get<8>(variant_);
+  }
+
+  void set_parameter_mutation(ParameterMutationSpec def) {
+    variant_.emplace<8>(std::move(def));
+    tag_ = Tag::PARAMETER_MUTATION;
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const OutputSpec& nlohmann_json_t) {
+
+    if (nlohmann_json_t.tag_ == Tag::USER_OUTPUT) {
+      nlohmann_json_j["user_output"] = nlohmann_json_t.get_user_output();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::LOSS_OUTPUT) {
+      nlohmann_json_j["loss_output"] = nlohmann_json_t.get_loss_output();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::BUFFER_MUTATION) {
+      nlohmann_json_j["buffer_mutation"] = nlohmann_json_t.get_buffer_mutation();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::GRADIENT_TO_PARAMETER) {
+      nlohmann_json_j["gradient_to_parameter"] = nlohmann_json_t.get_gradient_to_parameter();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::GRADIENT_TO_USER_INPUT) {
+      nlohmann_json_j["gradient_to_user_input"] = nlohmann_json_t.get_gradient_to_user_input();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::USER_INPUT_MUTATION) {
+      nlohmann_json_j["user_input_mutation"] = nlohmann_json_t.get_user_input_mutation();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::TOKEN) {
+      nlohmann_json_j["token"] = nlohmann_json_t.get_token();
+      return;
+    }
+    if (nlohmann_json_t.tag_ == Tag::PARAMETER_MUTATION) {
+      nlohmann_json_j["parameter_mutation"] = nlohmann_json_t.get_parameter_mutation();
+      return;
+    }
+  }
+
+  friend void from_json(const nlohmann::json& nlohmann_json_j, OutputSpec& nlohmann_json_t) {
+
+    if (nlohmann_json_j.contains("user_output")) {
+      nlohmann_json_t.variant_.emplace<1>(nlohmann_json_j.at("user_output").template get<UserOutputSpec>());
+      nlohmann_json_t.tag_ = Tag::USER_OUTPUT;
+      return;
+    }
+    if (nlohmann_json_j.contains("loss_output")) {
+      nlohmann_json_t.variant_.emplace<2>(nlohmann_json_j.at("loss_output").template get<LossOutputSpec>());
+      nlohmann_json_t.tag_ = Tag::LOSS_OUTPUT;
+      return;
+    }
+    if (nlohmann_json_j.contains("buffer_mutation")) {
+      nlohmann_json_t.variant_.emplace<3>(nlohmann_json_j.at("buffer_mutation").template get<BufferMutationSpec>());
+      nlohmann_json_t.tag_ = Tag::BUFFER_MUTATION;
+      return;
+    }
+    if (nlohmann_json_j.contains("gradient_to_parameter")) {
+      nlohmann_json_t.variant_.emplace<4>(nlohmann_json_j.at("gradient_to_parameter").template get<GradientToParameterSpec>());
+      nlohmann_json_t.tag_ = Tag::GRADIENT_TO_PARAMETER;
+      return;
+    }
+    if (nlohmann_json_j.contains("gradient_to_user_input")) {
+      nlohmann_json_t.variant_.emplace<5>(nlohmann_json_j.at("gradient_to_user_input").template get<GradientToUserInputSpec>());
+      nlohmann_json_t.tag_ = Tag::GRADIENT_TO_USER_INPUT;
+      return;
+    }
+    if (nlohmann_json_j.contains("user_input_mutation")) {
+      nlohmann_json_t.variant_.emplace<6>(nlohmann_json_j.at("user_input_mutation").template get<UserInputMutationSpec>());
+      nlohmann_json_t.tag_ = Tag::USER_INPUT_MUTATION;
+      return;
+    }
+    if (nlohmann_json_j.contains("token")) {
+      nlohmann_json_t.variant_.emplace<7>(nlohmann_json_j.at("token").template get<OutputTokenSpec>());
+      nlohmann_json_t.tag_ = Tag::TOKEN;
+      return;
+    }
+    if (nlohmann_json_j.contains("parameter_mutation")) {
+      nlohmann_json_t.variant_.emplace<8>(nlohmann_json_j.at("parameter_mutation").template get<ParameterMutationSpec>());
+      nlohmann_json_t.tag_ = Tag::PARAMETER_MUTATION;
+      return;
+    }
+  }
+};
+
+inline std::string_view printEnum(const OutputSpec::Tag& e) {
+  switch (e) {
+    case OutputSpec::Tag::USER_OUTPUT: return "USER_OUTPUT";
+    case OutputSpec::Tag::LOSS_OUTPUT: return "LOSS_OUTPUT";
+    case OutputSpec::Tag::BUFFER_MUTATION: return "BUFFER_MUTATION";
+    case OutputSpec::Tag::GRADIENT_TO_PARAMETER: return "GRADIENT_TO_PARAMETER";
+    case OutputSpec::Tag::GRADIENT_TO_USER_INPUT: return "GRADIENT_TO_USER_INPUT";
+    case OutputSpec::Tag::USER_INPUT_MUTATION: return "USER_INPUT_MUTATION";
+    case OutputSpec::Tag::TOKEN: return "TOKEN";
+    case OutputSpec::Tag::PARAMETER_MUTATION: return "PARAMETER_MUTATION";
+    default:
+      throw std::runtime_error("Unknown enum value");
+  }
+}
+
+inline void parseEnum(std::string_view s, OutputSpec::Tag& t) {
+  if (s == "USER_OUTPUT") { t = OutputSpec::Tag::USER_OUTPUT; return; }
+  if (s == "LOSS_OUTPUT") { t = OutputSpec::Tag::LOSS_OUTPUT; return; }
+  if (s == "BUFFER_MUTATION") { t = OutputSpec::Tag::BUFFER_MUTATION; return; }
+  if (s == "GRADIENT_TO_PARAMETER") { t = OutputSpec::Tag::GRADIENT_TO_PARAMETER; return; }
+  if (s == "GRADIENT_TO_USER_INPUT") { t = OutputSpec::Tag::GRADIENT_TO_USER_INPUT; return; }
+  if (s == "USER_INPUT_MUTATION") { t = OutputSpec::Tag::USER_INPUT_MUTATION; return; }
+  if (s == "TOKEN") { t = OutputSpec::Tag::TOKEN; return; }
+  if (s == "PARAMETER_MUTATION") { t = OutputSpec::Tag::PARAMETER_MUTATION; return; }
+  throw std::runtime_error("Unknown enum value: " + std::string{s});
+}
+
+
+class GraphSignature {
+ private:
+  std::vector<InputSpec> input_specs;
+  std::vector<OutputSpec> output_specs;
+
+ public:
+
+  const std::vector<InputSpec>& get_input_specs() const {
+    return input_specs;
+  }
+
+  void set_input_specs(std::vector<InputSpec> def) {
+    input_specs = std::move(def);
+  }
+
+  const std::vector<OutputSpec>& get_output_specs() const {
+    return output_specs;
+  }
+
+  void set_output_specs(std::vector<OutputSpec> def) {
+    output_specs = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const GraphSignature& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, GraphSignature& nlohmann_json_t);
+};
+
+class RangeConstraint {
+ private:
+  std::optional<int64_t> min_val;
+  std::optional<int64_t> max_val;
+
+ public:
+
+  const std::optional<int64_t>& get_min_val() const {
+    return min_val;
+  }
+
+  void set_min_val(std::optional<int64_t> def) {
+    min_val = std::move(def);
+  }
+
+  const std::optional<int64_t>& get_max_val() const {
+    return max_val;
+  }
+
+  void set_max_val(std::optional<int64_t> def) {
+    max_val = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const RangeConstraint& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, RangeConstraint& nlohmann_json_t);
+};
+
+class ModuleCallSignature {
+ private:
+  std::vector<Argument> inputs;
+  std::vector<Argument> outputs;
+  std::string in_spec;
+  std::string out_spec;
+  std::optional<std::vector<std::string>> forward_arg_names = std::nullopt;
+
+ public:
+
+  const std::vector<Argument>& get_inputs() const {
+    return inputs;
+  }
+
+  void set_inputs(std::vector<Argument> def) {
+    inputs = std::move(def);
+  }
+
+  const std::vector<Argument>& get_outputs() const {
+    return outputs;
+  }
+
+  void set_outputs(std::vector<Argument> def) {
+    outputs = std::move(def);
+  }
+
+  const std::string& get_in_spec() const {
+    return in_spec;
+  }
+
+  void set_in_spec(std::string def) {
+    in_spec = std::move(def);
+  }
+
+  const std::string& get_out_spec() const {
+    return out_spec;
+  }
+
+  void set_out_spec(std::string def) {
+    out_spec = std::move(def);
+  }
+
+  const std::optional<std::vector<std::string>>& get_forward_arg_names() const {
+    return forward_arg_names;
+  }
+
+  void set_forward_arg_names(std::optional<std::vector<std::string>> def) {
+    forward_arg_names = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ModuleCallSignature& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ModuleCallSignature& nlohmann_json_t);
+};
+
+class ModuleCallEntry {
+ private:
+  std::string fqn;
+  std::optional<ModuleCallSignature> signature = std::nullopt;
+
+ public:
+
+  const std::string& get_fqn() const {
+    return fqn;
+  }
+
+  void set_fqn(std::string def) {
+    fqn = std::move(def);
+  }
+
+  const std::optional<ModuleCallSignature>& get_signature() const {
+    return signature;
+  }
+
+  void set_signature(std::optional<ModuleCallSignature> def) {
+    signature = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ModuleCallEntry& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ModuleCallEntry& nlohmann_json_t);
+};
+
+class NamedTupleDef {
+ private:
+  std::vector<std::string> field_names;
+
+ public:
+
+  const std::vector<std::string>& get_field_names() const {
+    return field_names;
+  }
+
+  void set_field_names(std::vector<std::string> def) {
+    field_names = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const NamedTupleDef& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, NamedTupleDef& nlohmann_json_t);
+};
+
+class GraphModule {
+ private:
+  Graph graph;
+  GraphSignature signature;
+  std::vector<ModuleCallEntry> module_call_graph;
+  std::unordered_map<std::string, std::string> metadata = {};
+  std::unordered_map<std::string, NamedTupleDef> treespec_namedtuple_fields = {};
+
+ public:
+
+  const Graph& get_graph() const {
+    return graph;
+  }
+
+  void set_graph(Graph def) {
+    graph = std::move(def);
+  }
+
+  const GraphSignature& get_signature() const {
+    return signature;
+  }
+
+  void set_signature(GraphSignature def) {
+    signature = std::move(def);
+  }
+
+  const std::vector<ModuleCallEntry>& get_module_call_graph() const {
+    return module_call_graph;
+  }
+
+  void set_module_call_graph(std::vector<ModuleCallEntry> def) {
+    module_call_graph = std::move(def);
+  }
+
+  const std::unordered_map<std::string, std::string>& get_metadata() const {
+    return metadata;
+  }
+
+  void set_metadata(std::unordered_map<std::string, std::string> def) {
+    metadata = std::move(def);
+  }
+
+  const std::unordered_map<std::string, NamedTupleDef>& get_treespec_namedtuple_fields() const {
+    return treespec_namedtuple_fields;
+  }
+
+  void set_treespec_namedtuple_fields(std::unordered_map<std::string, NamedTupleDef> def) {
+    treespec_namedtuple_fields = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const GraphModule& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, GraphModule& nlohmann_json_t);
+};
+
+class SchemaVersion {
+ private:
+  int64_t major;
+  int64_t minor;
+
+ public:
+
+  const int64_t& get_major() const {
+    return major;
+  }
+
+  void set_major(int64_t def) {
+    major = std::move(def);
+  }
+
+  const int64_t& get_minor() const {
+    return minor;
+  }
+
+  void set_minor(int64_t def) {
+    minor = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const SchemaVersion& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, SchemaVersion& nlohmann_json_t);
+};
+
+class ExportedProgram {
+ private:
+  GraphModule graph_module;
+  std::unordered_map<std::string, int64_t> opset_version;
+  std::unordered_map<std::string, RangeConstraint> range_constraints;
+  SchemaVersion schema_version;
+  std::vector<std::string> verifiers = {};
+  std::string torch_version = "<=2.4";
+  std::vector<std::string> guards_code = {};
+
+ public:
+
+  const GraphModule& get_graph_module() const {
+    return graph_module;
+  }
+
+  void set_graph_module(GraphModule def) {
+    graph_module = std::move(def);
+  }
+
+  const std::unordered_map<std::string, int64_t>& get_opset_version() const {
+    return opset_version;
+  }
+
+  void set_opset_version(std::unordered_map<std::string, int64_t> def) {
+    opset_version = std::move(def);
+  }
+
+  const std::unordered_map<std::string, RangeConstraint>& get_range_constraints() const {
+    return range_constraints;
+  }
+
+  void set_range_constraints(std::unordered_map<std::string, RangeConstraint> def) {
+    range_constraints = std::move(def);
+  }
+
+  const SchemaVersion& get_schema_version() const {
+    return schema_version;
+  }
+
+  void set_schema_version(SchemaVersion def) {
+    schema_version = std::move(def);
+  }
+
+  const std::vector<std::string>& get_verifiers() const {
+    return verifiers;
+  }
+
+  void set_verifiers(std::vector<std::string> def) {
+    verifiers = std::move(def);
+  }
+
+  const std::string& get_torch_version() const {
+    return torch_version;
+  }
+
+  void set_torch_version(std::string def) {
+    torch_version = std::move(def);
+  }
+
+  const std::vector<std::string>& get_guards_code() const {
+    return guards_code;
+  }
+
+  void set_guards_code(std::vector<std::string> def) {
+    guards_code = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ExportedProgram& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ExportedProgram& nlohmann_json_t);
+};
+
+class PayloadMeta {
+ private:
+  std::string path_name;
+  bool is_param;
+  bool use_pickle;
+  std::optional<TensorMeta> tensor_meta;
+
+ public:
+
+  const std::string& get_path_name() const {
+    return path_name;
+  }
+
+  void set_path_name(std::string def) {
+    path_name = std::move(def);
+  }
+
+  const bool& get_is_param() const {
+    return is_param;
+  }
+
+  void set_is_param(bool def) {
+    is_param = std::move(def);
+  }
+
+  const bool& get_use_pickle() const {
+    return use_pickle;
+  }
+
+  void set_use_pickle(bool def) {
+    use_pickle = std::move(def);
+  }
+
+  const std::optional<TensorMeta>& get_tensor_meta() const {
+    return tensor_meta;
+  }
+
+  void set_tensor_meta(std::optional<TensorMeta> def) {
+    tensor_meta = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const PayloadMeta& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, PayloadMeta& nlohmann_json_t);
+};
+
+class PayloadConfig {
+ private:
+  std::unordered_map<std::string, PayloadMeta> config;
+
+ public:
+
+  const std::unordered_map<std::string, PayloadMeta>& get_config() const {
+    return config;
+  }
+
+  void set_config(std::unordered_map<std::string, PayloadMeta> def) {
+    config = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const PayloadConfig& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, PayloadConfig& nlohmann_json_t);
+};
+
+class AOTInductorModelPickleData {
+ private:
+  std::string library_basename;
+  std::vector<std::string> input_names;
+  std::vector<std::string> output_names;
+  std::optional<int64_t> floating_point_input_dtype = std::nullopt;
+  std::optional<int64_t> floating_point_output_dtype = std::nullopt;
+  std::optional<bool> aot_inductor_model_is_cpu = std::nullopt;
+
+ public:
+
+  const std::string& get_library_basename() const {
+    return library_basename;
+  }
+
+  void set_library_basename(std::string def) {
+    library_basename = std::move(def);
+  }
+
+  const std::vector<std::string>& get_input_names() const {
+    return input_names;
+  }
+
+  void set_input_names(std::vector<std::string> def) {
+    input_names = std::move(def);
+  }
+
+  const std::vector<std::string>& get_output_names() const {
+    return output_names;
+  }
+
+  void set_output_names(std::vector<std::string> def) {
+    output_names = std::move(def);
+  }
+
+  const std::optional<int64_t>& get_floating_point_input_dtype() const {
+    return floating_point_input_dtype;
+  }
+
+  void set_floating_point_input_dtype(std::optional<int64_t> def) {
+    floating_point_input_dtype = std::move(def);
+  }
+
+  const std::optional<int64_t>& get_floating_point_output_dtype() const {
+    return floating_point_output_dtype;
+  }
+
+  void set_floating_point_output_dtype(std::optional<int64_t> def) {
+    floating_point_output_dtype = std::move(def);
+  }
+
+  const std::optional<bool>& get_aot_inductor_model_is_cpu() const {
+    return aot_inductor_model_is_cpu;
+  }
+
+  void set_aot_inductor_model_is_cpu(std::optional<bool> def) {
+    aot_inductor_model_is_cpu = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const AOTInductorModelPickleData& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, AOTInductorModelPickleData& nlohmann_json_t);
+};
+
+class ExternKernelNode {
+ private:
+  std::string name;
+  Node node;
+
+ public:
+
+  const std::string& get_name() const {
+    return name;
+  }
+
+  void set_name(std::string def) {
+    name = std::move(def);
+  }
+
+  const Node& get_node() const {
+    return node;
+  }
+
+  void set_node(Node def) {
+    node = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ExternKernelNode& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ExternKernelNode& nlohmann_json_t);
+};
+
+class ExternKernelNodes {
+ private:
+  std::vector<ExternKernelNode> nodes;
+
+ public:
+
+  const std::vector<ExternKernelNode>& get_nodes() const {
+    return nodes;
+  }
+
+  void set_nodes(std::vector<ExternKernelNode> def) {
+    nodes = std::move(def);
+  }
+
+  friend void to_json(nlohmann::json& nlohmann_json_j, const ExternKernelNodes& nlohmann_json_t);
+  friend void from_json(const nlohmann::json& nlohmann_json_j, ExternKernelNodes& nlohmann_json_t);
+};
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const AOTInductorModelPickleData& nlohmann_json_t) {
+  nlohmann_json_j["library_basename"] = nlohmann_json_t.library_basename;
+  nlohmann_json_j["input_names"] = nlohmann_json_t.input_names;
+  nlohmann_json_j["output_names"] = nlohmann_json_t.output_names;
+  nlohmann_json_j["floating_point_input_dtype"] = nlohmann_json_t.floating_point_input_dtype;
+  nlohmann_json_j["floating_point_output_dtype"] = nlohmann_json_t.floating_point_output_dtype;
+  nlohmann_json_j["aot_inductor_model_is_cpu"] = nlohmann_json_t.aot_inductor_model_is_cpu;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, AOTInductorModelPickleData& nlohmann_json_t) {
+  AOTInductorModelPickleData nlohmann_json_default_obj;
+  nlohmann_json_t.library_basename = nlohmann_json_j.value("library_basename", nlohmann_json_default_obj.library_basename);
+  nlohmann_json_t.input_names = nlohmann_json_j.value("input_names", nlohmann_json_default_obj.input_names);
+  nlohmann_json_t.output_names = nlohmann_json_j.value("output_names", nlohmann_json_default_obj.output_names);
+  nlohmann_json_t.floating_point_input_dtype = nlohmann_json_j.value("floating_point_input_dtype", nlohmann_json_default_obj.floating_point_input_dtype);
+  nlohmann_json_t.floating_point_output_dtype = nlohmann_json_j.value("floating_point_output_dtype", nlohmann_json_default_obj.floating_point_output_dtype);
+  nlohmann_json_t.aot_inductor_model_is_cpu = nlohmann_json_j.value("aot_inductor_model_is_cpu", nlohmann_json_default_obj.aot_inductor_model_is_cpu);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const BufferMutationSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["buffer_name"] = nlohmann_json_t.buffer_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, BufferMutationSpec& nlohmann_json_t) {
+  BufferMutationSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.buffer_name = nlohmann_json_j.value("buffer_name", nlohmann_json_default_obj.buffer_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const ComplexValue& nlohmann_json_t) {
+  nlohmann_json_j["real"] = nlohmann_json_t.real;
+  nlohmann_json_j["imag"] = nlohmann_json_t.imag;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, ComplexValue& nlohmann_json_t) {
+  ComplexValue nlohmann_json_default_obj;
+  nlohmann_json_t.real = nlohmann_json_j.value("real", nlohmann_json_default_obj.real);
+  nlohmann_json_t.imag = nlohmann_json_j.value("imag", nlohmann_json_default_obj.imag);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const CustomObjArgument& nlohmann_json_t) {
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+  nlohmann_json_j["class_fqn"] = nlohmann_json_t.class_fqn;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, CustomObjArgument& nlohmann_json_t) {
+  CustomObjArgument nlohmann_json_default_obj;
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+  nlohmann_json_t.class_fqn = nlohmann_json_j.value("class_fqn", nlohmann_json_default_obj.class_fqn);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const Device& nlohmann_json_t) {
+  nlohmann_json_j["type"] = nlohmann_json_t.type;
+  nlohmann_json_j["index"] = nlohmann_json_t.index;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, Device& nlohmann_json_t) {
+  Device nlohmann_json_default_obj;
+  nlohmann_json_t.type = nlohmann_json_j.value("type", nlohmann_json_default_obj.type);
+  nlohmann_json_t.index = nlohmann_json_j.value("index", nlohmann_json_default_obj.index);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const ExportedProgram& nlohmann_json_t) {
+  nlohmann_json_j["graph_module"] = nlohmann_json_t.graph_module;
+  nlohmann_json_j["opset_version"] = nlohmann_json_t.opset_version;
+  nlohmann_json_j["range_constraints"] = nlohmann_json_t.range_constraints;
+  nlohmann_json_j["schema_version"] = nlohmann_json_t.schema_version;
+  nlohmann_json_j["verifiers"] = nlohmann_json_t.verifiers;
+  nlohmann_json_j["torch_version"] = nlohmann_json_t.torch_version;
+  nlohmann_json_j["guards_code"] = nlohmann_json_t.guards_code;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, ExportedProgram& nlohmann_json_t) {
+  ExportedProgram nlohmann_json_default_obj;
+  nlohmann_json_t.graph_module = nlohmann_json_j.value("graph_module", nlohmann_json_default_obj.graph_module);
+  nlohmann_json_t.opset_version = nlohmann_json_j.value("opset_version", nlohmann_json_default_obj.opset_version);
+  nlohmann_json_t.range_constraints = nlohmann_json_j.value("range_constraints", nlohmann_json_default_obj.range_constraints);
+  nlohmann_json_t.schema_version = nlohmann_json_j.value("schema_version", nlohmann_json_default_obj.schema_version);
+  nlohmann_json_t.verifiers = nlohmann_json_j.value("verifiers", nlohmann_json_default_obj.verifiers);
+  nlohmann_json_t.torch_version = nlohmann_json_j.value("torch_version", nlohmann_json_default_obj.torch_version);
+  nlohmann_json_t.guards_code = nlohmann_json_j.value("guards_code", nlohmann_json_default_obj.guards_code);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const ExternKernelNode& nlohmann_json_t) {
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+  nlohmann_json_j["node"] = nlohmann_json_t.node;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, ExternKernelNode& nlohmann_json_t) {
+  ExternKernelNode nlohmann_json_default_obj;
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+  nlohmann_json_t.node = nlohmann_json_j.value("node", nlohmann_json_default_obj.node);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const ExternKernelNodes& nlohmann_json_t) {
+  nlohmann_json_j["nodes"] = nlohmann_json_t.nodes;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, ExternKernelNodes& nlohmann_json_t) {
+  ExternKernelNodes nlohmann_json_default_obj;
+  nlohmann_json_t.nodes = nlohmann_json_j.value("nodes", nlohmann_json_default_obj.nodes);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const GradientToParameterSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["parameter_name"] = nlohmann_json_t.parameter_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, GradientToParameterSpec& nlohmann_json_t) {
+  GradientToParameterSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.parameter_name = nlohmann_json_j.value("parameter_name", nlohmann_json_default_obj.parameter_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const GradientToUserInputSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["user_input_name"] = nlohmann_json_t.user_input_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, GradientToUserInputSpec& nlohmann_json_t) {
+  GradientToUserInputSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.user_input_name = nlohmann_json_j.value("user_input_name", nlohmann_json_default_obj.user_input_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const Graph& nlohmann_json_t) {
+  nlohmann_json_j["inputs"] = nlohmann_json_t.inputs;
+  nlohmann_json_j["outputs"] = nlohmann_json_t.outputs;
+  nlohmann_json_j["nodes"] = nlohmann_json_t.nodes;
+  nlohmann_json_j["tensor_values"] = nlohmann_json_t.tensor_values;
+  nlohmann_json_j["sym_int_values"] = nlohmann_json_t.sym_int_values;
+  nlohmann_json_j["sym_bool_values"] = nlohmann_json_t.sym_bool_values;
+  nlohmann_json_j["is_single_tensor_return"] = nlohmann_json_t.is_single_tensor_return;
+  nlohmann_json_j["custom_obj_values"] = nlohmann_json_t.custom_obj_values;
+  nlohmann_json_j["sym_float_values"] = nlohmann_json_t.sym_float_values;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, Graph& nlohmann_json_t) {
+  Graph nlohmann_json_default_obj;
+  nlohmann_json_t.inputs = nlohmann_json_j.value("inputs", nlohmann_json_default_obj.inputs);
+  nlohmann_json_t.outputs = nlohmann_json_j.value("outputs", nlohmann_json_default_obj.outputs);
+  nlohmann_json_t.nodes = nlohmann_json_j.value("nodes", nlohmann_json_default_obj.nodes);
+  nlohmann_json_t.tensor_values = nlohmann_json_j.value("tensor_values", nlohmann_json_default_obj.tensor_values);
+  nlohmann_json_t.sym_int_values = nlohmann_json_j.value("sym_int_values", nlohmann_json_default_obj.sym_int_values);
+  nlohmann_json_t.sym_bool_values = nlohmann_json_j.value("sym_bool_values", nlohmann_json_default_obj.sym_bool_values);
+  nlohmann_json_t.is_single_tensor_return = nlohmann_json_j.value("is_single_tensor_return", nlohmann_json_default_obj.is_single_tensor_return);
+  nlohmann_json_t.custom_obj_values = nlohmann_json_j.value("custom_obj_values", nlohmann_json_default_obj.custom_obj_values);
+  nlohmann_json_t.sym_float_values = nlohmann_json_j.value("sym_float_values", nlohmann_json_default_obj.sym_float_values);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const GraphArgument& nlohmann_json_t) {
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+  nlohmann_json_j["graph"] = nlohmann_json_t.graph;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, GraphArgument& nlohmann_json_t) {
+  GraphArgument nlohmann_json_default_obj;
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+  nlohmann_json_t.graph = nlohmann_json_j.value("graph", nlohmann_json_default_obj.graph);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const GraphModule& nlohmann_json_t) {
+  nlohmann_json_j["graph"] = nlohmann_json_t.graph;
+  nlohmann_json_j["signature"] = nlohmann_json_t.signature;
+  nlohmann_json_j["module_call_graph"] = nlohmann_json_t.module_call_graph;
+  nlohmann_json_j["metadata"] = nlohmann_json_t.metadata;
+  nlohmann_json_j["treespec_namedtuple_fields"] = nlohmann_json_t.treespec_namedtuple_fields;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, GraphModule& nlohmann_json_t) {
+  GraphModule nlohmann_json_default_obj;
+  nlohmann_json_t.graph = nlohmann_json_j.value("graph", nlohmann_json_default_obj.graph);
+  nlohmann_json_t.signature = nlohmann_json_j.value("signature", nlohmann_json_default_obj.signature);
+  nlohmann_json_t.module_call_graph = nlohmann_json_j.value("module_call_graph", nlohmann_json_default_obj.module_call_graph);
+  nlohmann_json_t.metadata = nlohmann_json_j.value("metadata", nlohmann_json_default_obj.metadata);
+  nlohmann_json_t.treespec_namedtuple_fields = nlohmann_json_j.value("treespec_namedtuple_fields", nlohmann_json_default_obj.treespec_namedtuple_fields);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const GraphSignature& nlohmann_json_t) {
+  nlohmann_json_j["input_specs"] = nlohmann_json_t.input_specs;
+  nlohmann_json_j["output_specs"] = nlohmann_json_t.output_specs;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, GraphSignature& nlohmann_json_t) {
+  GraphSignature nlohmann_json_default_obj;
+  nlohmann_json_t.input_specs = nlohmann_json_j.value("input_specs", nlohmann_json_default_obj.input_specs);
+  nlohmann_json_t.output_specs = nlohmann_json_j.value("output_specs", nlohmann_json_default_obj.output_specs);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const InputToBufferSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["buffer_name"] = nlohmann_json_t.buffer_name;
+  nlohmann_json_j["persistent"] = nlohmann_json_t.persistent;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, InputToBufferSpec& nlohmann_json_t) {
+  InputToBufferSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.buffer_name = nlohmann_json_j.value("buffer_name", nlohmann_json_default_obj.buffer_name);
+  nlohmann_json_t.persistent = nlohmann_json_j.value("persistent", nlohmann_json_default_obj.persistent);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const InputToConstantInputSpec& nlohmann_json_t) {
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+  nlohmann_json_j["value"] = nlohmann_json_t.value;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, InputToConstantInputSpec& nlohmann_json_t) {
+  InputToConstantInputSpec nlohmann_json_default_obj;
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+  nlohmann_json_t.value = nlohmann_json_j.value("value", nlohmann_json_default_obj.value);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const InputToCustomObjSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["custom_obj_name"] = nlohmann_json_t.custom_obj_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, InputToCustomObjSpec& nlohmann_json_t) {
+  InputToCustomObjSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.custom_obj_name = nlohmann_json_j.value("custom_obj_name", nlohmann_json_default_obj.custom_obj_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const InputToParameterSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["parameter_name"] = nlohmann_json_t.parameter_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, InputToParameterSpec& nlohmann_json_t) {
+  InputToParameterSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.parameter_name = nlohmann_json_j.value("parameter_name", nlohmann_json_default_obj.parameter_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const InputToTensorConstantSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["tensor_constant_name"] = nlohmann_json_t.tensor_constant_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, InputToTensorConstantSpec& nlohmann_json_t) {
+  InputToTensorConstantSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.tensor_constant_name = nlohmann_json_j.value("tensor_constant_name", nlohmann_json_default_obj.tensor_constant_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const InputTokenSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, InputTokenSpec& nlohmann_json_t) {
+  InputTokenSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const LossOutputSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, LossOutputSpec& nlohmann_json_t) {
+  LossOutputSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const ModuleCallEntry& nlohmann_json_t) {
+  nlohmann_json_j["fqn"] = nlohmann_json_t.fqn;
+  nlohmann_json_j["signature"] = nlohmann_json_t.signature;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, ModuleCallEntry& nlohmann_json_t) {
+  ModuleCallEntry nlohmann_json_default_obj;
+  nlohmann_json_t.fqn = nlohmann_json_j.value("fqn", nlohmann_json_default_obj.fqn);
+  nlohmann_json_t.signature = nlohmann_json_j.value("signature", nlohmann_json_default_obj.signature);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const ModuleCallSignature& nlohmann_json_t) {
+  nlohmann_json_j["inputs"] = nlohmann_json_t.inputs;
+  nlohmann_json_j["outputs"] = nlohmann_json_t.outputs;
+  nlohmann_json_j["in_spec"] = nlohmann_json_t.in_spec;
+  nlohmann_json_j["out_spec"] = nlohmann_json_t.out_spec;
+  nlohmann_json_j["forward_arg_names"] = nlohmann_json_t.forward_arg_names;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, ModuleCallSignature& nlohmann_json_t) {
+  ModuleCallSignature nlohmann_json_default_obj;
+  nlohmann_json_t.inputs = nlohmann_json_j.value("inputs", nlohmann_json_default_obj.inputs);
+  nlohmann_json_t.outputs = nlohmann_json_j.value("outputs", nlohmann_json_default_obj.outputs);
+  nlohmann_json_t.in_spec = nlohmann_json_j.value("in_spec", nlohmann_json_default_obj.in_spec);
+  nlohmann_json_t.out_spec = nlohmann_json_j.value("out_spec", nlohmann_json_default_obj.out_spec);
+  nlohmann_json_t.forward_arg_names = nlohmann_json_j.value("forward_arg_names", nlohmann_json_default_obj.forward_arg_names);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const NamedArgument& nlohmann_json_t) {
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["kind"] = nlohmann_json_t.kind;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, NamedArgument& nlohmann_json_t) {
+  NamedArgument nlohmann_json_default_obj;
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.kind = nlohmann_json_j.value("kind", nlohmann_json_default_obj.kind);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const NamedTupleDef& nlohmann_json_t) {
+  nlohmann_json_j["field_names"] = nlohmann_json_t.field_names;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, NamedTupleDef& nlohmann_json_t) {
+  NamedTupleDef nlohmann_json_default_obj;
+  nlohmann_json_t.field_names = nlohmann_json_j.value("field_names", nlohmann_json_default_obj.field_names);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const Node& nlohmann_json_t) {
+  nlohmann_json_j["target"] = nlohmann_json_t.target;
+  nlohmann_json_j["inputs"] = nlohmann_json_t.inputs;
+  nlohmann_json_j["outputs"] = nlohmann_json_t.outputs;
+  nlohmann_json_j["metadata"] = nlohmann_json_t.metadata;
+  nlohmann_json_j["is_hop_single_tensor_return"] = nlohmann_json_t.is_hop_single_tensor_return;
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, Node& nlohmann_json_t) {
+  Node nlohmann_json_default_obj;
+  nlohmann_json_t.target = nlohmann_json_j.value("target", nlohmann_json_default_obj.target);
+  nlohmann_json_t.inputs = nlohmann_json_j.value("inputs", nlohmann_json_default_obj.inputs);
+  nlohmann_json_t.outputs = nlohmann_json_j.value("outputs", nlohmann_json_default_obj.outputs);
+  nlohmann_json_t.metadata = nlohmann_json_j.value("metadata", nlohmann_json_default_obj.metadata);
+  nlohmann_json_t.is_hop_single_tensor_return = nlohmann_json_j.value("is_hop_single_tensor_return", nlohmann_json_default_obj.is_hop_single_tensor_return);
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const OutputTokenSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, OutputTokenSpec& nlohmann_json_t) {
+  OutputTokenSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const ParameterMutationSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["parameter_name"] = nlohmann_json_t.parameter_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, ParameterMutationSpec& nlohmann_json_t) {
+  ParameterMutationSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.parameter_name = nlohmann_json_j.value("parameter_name", nlohmann_json_default_obj.parameter_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const PayloadConfig& nlohmann_json_t) {
+  nlohmann_json_j["config"] = nlohmann_json_t.config;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, PayloadConfig& nlohmann_json_t) {
+  PayloadConfig nlohmann_json_default_obj;
+  nlohmann_json_t.config = nlohmann_json_j.value("config", nlohmann_json_default_obj.config);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const PayloadMeta& nlohmann_json_t) {
+  nlohmann_json_j["path_name"] = nlohmann_json_t.path_name;
+  nlohmann_json_j["is_param"] = nlohmann_json_t.is_param;
+  nlohmann_json_j["use_pickle"] = nlohmann_json_t.use_pickle;
+  nlohmann_json_j["tensor_meta"] = nlohmann_json_t.tensor_meta;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, PayloadMeta& nlohmann_json_t) {
+  PayloadMeta nlohmann_json_default_obj;
+  nlohmann_json_t.path_name = nlohmann_json_j.value("path_name", nlohmann_json_default_obj.path_name);
+  nlohmann_json_t.is_param = nlohmann_json_j.value("is_param", nlohmann_json_default_obj.is_param);
+  nlohmann_json_t.use_pickle = nlohmann_json_j.value("use_pickle", nlohmann_json_default_obj.use_pickle);
+  nlohmann_json_t.tensor_meta = nlohmann_json_j.value("tensor_meta", nlohmann_json_default_obj.tensor_meta);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const RangeConstraint& nlohmann_json_t) {
+  nlohmann_json_j["min_val"] = nlohmann_json_t.min_val;
+  nlohmann_json_j["max_val"] = nlohmann_json_t.max_val;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, RangeConstraint& nlohmann_json_t) {
+  RangeConstraint nlohmann_json_default_obj;
+  nlohmann_json_t.min_val = nlohmann_json_j.value("min_val", nlohmann_json_default_obj.min_val);
+  nlohmann_json_t.max_val = nlohmann_json_j.value("max_val", nlohmann_json_default_obj.max_val);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const SchemaVersion& nlohmann_json_t) {
+  nlohmann_json_j["major"] = nlohmann_json_t.major;
+  nlohmann_json_j["minor"] = nlohmann_json_t.minor;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, SchemaVersion& nlohmann_json_t) {
+  SchemaVersion nlohmann_json_default_obj;
+  nlohmann_json_t.major = nlohmann_json_j.value("major", nlohmann_json_default_obj.major);
+  nlohmann_json_t.minor = nlohmann_json_j.value("minor", nlohmann_json_default_obj.minor);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const SymExpr& nlohmann_json_t) {
+  nlohmann_json_j["expr_str"] = nlohmann_json_t.expr_str;
+  nlohmann_json_j["hint"] = nlohmann_json_t.hint;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, SymExpr& nlohmann_json_t) {
+  SymExpr nlohmann_json_default_obj;
+  nlohmann_json_t.expr_str = nlohmann_json_j.value("expr_str", nlohmann_json_default_obj.expr_str);
+  nlohmann_json_t.hint = nlohmann_json_j.value("hint", nlohmann_json_default_obj.hint);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const TensorArgument& nlohmann_json_t) {
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, TensorArgument& nlohmann_json_t) {
+  TensorArgument nlohmann_json_default_obj;
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const TensorMeta& nlohmann_json_t) {
+  nlohmann_json_j["dtype"] = nlohmann_json_t.dtype;
+  nlohmann_json_j["sizes"] = nlohmann_json_t.sizes;
+  nlohmann_json_j["requires_grad"] = nlohmann_json_t.requires_grad;
+  nlohmann_json_j["device"] = nlohmann_json_t.device;
+  nlohmann_json_j["strides"] = nlohmann_json_t.strides;
+  nlohmann_json_j["storage_offset"] = nlohmann_json_t.storage_offset;
+  nlohmann_json_j["layout"] = nlohmann_json_t.layout;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, TensorMeta& nlohmann_json_t) {
+  TensorMeta nlohmann_json_default_obj;
+  nlohmann_json_t.dtype = nlohmann_json_j.value("dtype", nlohmann_json_default_obj.dtype);
+  nlohmann_json_t.sizes = nlohmann_json_j.value("sizes", nlohmann_json_default_obj.sizes);
+  nlohmann_json_t.requires_grad = nlohmann_json_j.value("requires_grad", nlohmann_json_default_obj.requires_grad);
+  nlohmann_json_t.device = nlohmann_json_j.value("device", nlohmann_json_default_obj.device);
+  nlohmann_json_t.strides = nlohmann_json_j.value("strides", nlohmann_json_default_obj.strides);
+  nlohmann_json_t.storage_offset = nlohmann_json_j.value("storage_offset", nlohmann_json_default_obj.storage_offset);
+  nlohmann_json_t.layout = nlohmann_json_j.value("layout", nlohmann_json_default_obj.layout);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const TokenArgument& nlohmann_json_t) {
+  nlohmann_json_j["name"] = nlohmann_json_t.name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, TokenArgument& nlohmann_json_t) {
+  TokenArgument nlohmann_json_default_obj;
+  nlohmann_json_t.name = nlohmann_json_j.value("name", nlohmann_json_default_obj.name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const UserInputMutationSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+  nlohmann_json_j["user_input_name"] = nlohmann_json_t.user_input_name;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, UserInputMutationSpec& nlohmann_json_t) {
+  UserInputMutationSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+  nlohmann_json_t.user_input_name = nlohmann_json_j.value("user_input_name", nlohmann_json_default_obj.user_input_name);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const UserInputSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, UserInputSpec& nlohmann_json_t) {
+  UserInputSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+}
+
+inline void to_json(nlohmann::json& nlohmann_json_j, const UserOutputSpec& nlohmann_json_t) {
+  nlohmann_json_j["arg"] = nlohmann_json_t.arg;
+}
+
+inline void from_json(const nlohmann::json& nlohmann_json_j, UserOutputSpec& nlohmann_json_t) {
+  UserOutputSpec nlohmann_json_default_obj;
+  nlohmann_json_t.arg = nlohmann_json_j.value("arg", nlohmann_json_default_obj.arg);
+}
+
+
+template <typename T> ForwardRef<T>::ForwardRef(ForwardRef<T>&&) = default;
+template <typename T> ForwardRef<T>& ForwardRef<T>::operator=(ForwardRef<T>&&) = default;
+template <typename T> ForwardRef<T>::~ForwardRef() = default;
+} // namespace _export
+} // namespace torch
+
+// clang-format on
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/init.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/init.h
new file mode 100644
index 0000000000000000000000000000000000000000..0496865beeb547590c24aa992f4484333fe02230
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/init.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::throughput_benchmark {
+
+void initThroughputBenchmarkBindings(PyObject* module);
+
+} // namespace torch::throughput_benchmark
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/invalid_arguments.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/invalid_arguments.h
new file mode 100644
index 0000000000000000000000000000000000000000..d2e818adf738250bc1a707af802d2915e9711067
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/invalid_arguments.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+#include <string>
+#include <vector>
+
+namespace torch {
+
+std::string format_invalid_args(
+    PyObject* given_args,
+    PyObject* given_kwargs,
+    const std::string& function_name,
+    const std::vector<std::string>& options);
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/nested.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/nested.h
new file mode 100644
index 0000000000000000000000000000000000000000..0d06a65acbe2b076b72dd01d6ceed3f10cc51dc6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/nested.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/python_arg_parser.h>
+
+#include <ATen/core/Tensor.h>
+
+namespace torch::utils {
+
+at::Tensor nested_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/numpy_stub.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/numpy_stub.h
new file mode 100644
index 0000000000000000000000000000000000000000..11846d0879f0b2b95c46190d997ac9a0e011e840
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/numpy_stub.h
@@ -0,0 +1,26 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+#ifdef USE_NUMPY
+
+#if !defined(NO_IMPORT_ARRAY) && !defined(WITH_NUMPY_IMPORT_ARRAY)
+#define NO_IMPORT_ARRAY
+#endif
+
+#ifndef PY_ARRAY_UNIQUE_SYMBOL
+#define PY_ARRAY_UNIQUE_SYMBOL __numpy_array_api
+#endif
+
+#ifndef NPY_NO_DEPRECATED_API
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#endif
+
+#include <numpy/arrayobject.h>
+
+#endif // USE_NUMPY
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/object_ptr.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/object_ptr.h
new file mode 100644
index 0000000000000000000000000000000000000000..aabd25e545d5df22e2480dce0cd01d9f69c68d3b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/object_ptr.h
@@ -0,0 +1,86 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <torch/csrc/python_headers.h>
+#include <utility>
+
+template <class T>
+class TORCH_PYTHON_API THPPointer {
+ public:
+  THPPointer() : ptr(nullptr) {}
+  explicit THPPointer(T* ptr) noexcept : ptr(ptr) {}
+  THPPointer(THPPointer&& p) noexcept : ptr(std::exchange(p.ptr, nullptr)) {}
+  THPPointer(const THPPointer& p) = delete;
+  THPPointer& operator=(const THPPointer&) = delete;
+
+  ~THPPointer() {
+    free();
+  }
+  T* get() {
+    return ptr;
+  }
+  const T* get() const {
+    return ptr;
+  }
+  THPPointer dup() const {
+    return dup(ptr);
+  }
+  static THPPointer dup(const T* ptr) {
+    Py_XINCREF(ptr);
+    return THPPointer(
+        const_cast<T*>(ptr)); // NOLINT(cppcoreguidelines-pro-type-const-cast)
+  }
+  static THPPointer none() {
+    Py_INCREF(Py_None);
+    return THPPointer(reinterpret_cast<T*>(Py_None));
+  }
+  T* release() {
+    T* tmp = ptr;
+    ptr = nullptr;
+    return tmp;
+  }
+  operator T*() {
+    return ptr;
+  }
+  THPPointer& operator=(T* new_ptr) noexcept {
+    free();
+    ptr = new_ptr;
+    return *this;
+  }
+  THPPointer& operator=(THPPointer&& p) noexcept {
+    free();
+    ptr = p.ptr;
+    p.ptr = nullptr;
+    return *this;
+  }
+  T* operator->() {
+    return ptr;
+  }
+  explicit operator bool() const {
+    return ptr != nullptr;
+  }
+
+ private:
+  void free();
+  T* ptr = nullptr;
+};
+
+/**
+ * An RAII-style, owning pointer to a PyObject.  You must protect
+ * destruction of this object with the GIL.
+ *
+ * WARNING: Think twice before putting this as a field in a C++
+ * struct.  This class does NOT take out the GIL on destruction,
+ * so if you will need to ensure that the destructor of your struct
+ * is either (a) always invoked when the GIL is taken or (b) takes
+ * out the GIL itself.  Easiest way to avoid this problem is to
+ * not use THPPointer in this situation.
+ */
+using THPObjectPtr = THPPointer<PyObject>;
+using THPCodeObjectPtr = THPPointer<PyCodeObject>;
+using THPFrameObjectPtr = THPPointer<PyFrameObject>;
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/out_types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/out_types.h
new file mode 100644
index 0000000000000000000000000000000000000000..baa59bad3c1fce24c4f8dcaca5a3a9bf2ac99d89
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/out_types.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/Tensor.h>
+
+namespace torch::utils {
+
+TORCH_API void check_out_type_matches(
+    const at::Tensor& result,
+    std::optional<at::ScalarType> scalarType,
+    bool scalarType_is_none,
+    std::optional<at::Layout> layout,
+    std::optional<at::Device> device,
+    bool device_is_none);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pybind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pybind.h
new file mode 100644
index 0000000000000000000000000000000000000000..4c399988a8d8a05859c9b71917e9e94c8b4147f6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pybind.h
@@ -0,0 +1,425 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/pythoncapi_compat.h>
+
+#include <ATen/core/Tensor.h>
+#include <ATen/core/jit_type_base.h>
+#include <c10/util/irange.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <torch/csrc/Device.h>
+#include <torch/csrc/Dtype.h>
+#include <torch/csrc/DynamicTypes.h>
+#include <torch/csrc/Generator.h>
+#include <torch/csrc/MemoryFormat.h>
+#include <torch/csrc/Stream.h>
+#include <torch/csrc/utils/tensor_memoryformats.h>
+
+namespace py = pybind11;
+
+#define IS_PYBIND_2_13_PLUS PYBIND11_VERSION_HEX >= 0x020D0000
+
+// This makes intrusive_ptr to be available as a custom pybind11 holder type,
+// see
+// https://pybind11.readthedocs.io/en/stable/advanced/smart_ptrs.html#custom-smart-pointers
+PYBIND11_DECLARE_HOLDER_TYPE(T, c10::intrusive_ptr<T>, true)
+
+PYBIND11_DECLARE_HOLDER_TYPE(T, c10::SingletonOrSharedTypePtr<T>)
+PYBIND11_DECLARE_HOLDER_TYPE(T, c10::SingletonTypePtr<T>, true)
+
+namespace pybind11::detail {
+
+// torch.Tensor <-> at::Tensor conversions (without unwrapping)
+template <>
+struct TORCH_PYTHON_API type_caster<at::Tensor> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::Tensor, _("torch.Tensor"));
+
+  bool load(handle src, bool /*unused*/);
+
+  static handle cast(
+      const at::Tensor& src,
+      return_value_policy /* policy */,
+      handle /* parent */);
+};
+
+// torch._StorageBase <-> at::Storage
+template <>
+struct type_caster<at::Storage> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::Storage, _("torch.StorageBase"));
+
+  bool load(handle src, bool /*unused*/) {
+    PyObject* obj = src.ptr();
+    if (torch::isStorage(obj)) {
+      value = torch::createStorage(obj);
+      return true;
+    }
+    return false;
+  }
+
+  static handle cast(
+      const at::Storage& src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return handle(torch::createPyObject(src));
+  }
+};
+
+template <>
+struct type_caster<at::Generator> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::Generator, _("torch.Generator"));
+
+  bool load(handle src, bool /*unused*/) {
+    PyObject* obj = src.ptr();
+    if (THPGenerator_Check(obj)) {
+      value = reinterpret_cast<THPGenerator*>(obj)->cdata;
+      return true;
+    }
+    return false;
+  }
+
+  static handle cast(
+      const at::Generator& src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return handle(THPGenerator_Wrap(src));
+  }
+};
+
+template <>
+struct TORCH_PYTHON_API type_caster<at::IntArrayRef> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::IntArrayRef, _("Tuple[int, ...]"));
+
+  bool load(handle src, bool /*unused*/);
+  static handle cast(
+      at::IntArrayRef src,
+      return_value_policy /* policy */,
+      handle /* parent */);
+
+ private:
+  std::vector<int64_t> v_value;
+};
+
+template <>
+struct TORCH_PYTHON_API type_caster<at::SymIntArrayRef> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::SymIntArrayRef, _("List[int]"));
+
+  bool load(handle src, bool /*unused*/);
+  static handle cast(
+      at::SymIntArrayRef src,
+      return_value_policy /* policy */,
+      handle /* parent */);
+
+ private:
+  std::vector<c10::SymInt> v_value;
+};
+
+template <>
+struct TORCH_PYTHON_API type_caster<at::ArrayRef<c10::SymNode>> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::ArrayRef<c10::SymNode>, _("List[SymNode]"));
+
+  bool load(handle src, bool /*unused*/);
+  static handle cast(
+      at::ArrayRef<c10::SymNode> src,
+      return_value_policy /* policy */,
+      handle /* parent */);
+
+ private:
+  std::vector<c10::SymNode> v_value;
+};
+
+template <>
+struct type_caster<at::MemoryFormat> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::MemoryFormat, _("torch.memory_format"));
+
+  bool load(handle src, bool /*unused*/) {
+    PyObject* obj = src.ptr();
+    if (THPMemoryFormat_Check(obj)) {
+      value = reinterpret_cast<THPMemoryFormat*>(obj)->memory_format;
+      return true;
+    }
+    return false;
+  }
+  static handle cast(
+      at::MemoryFormat src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return handle(Py_NewRef(torch::utils::getTHPMemoryFormat(src)));
+  }
+};
+
+template <>
+struct type_caster<at::Device> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::Device, _("torch.device"));
+
+  // PYBIND11_TYPE_CASTER defines a member field called value. Since at::Device
+  // cannot be default-initialized, we provide this constructor to explicitly
+  // initialize that field. The value doesn't matter as it will be overwritten
+  // after a successful call to load.
+  type_caster() : value(c10::kCPU) {}
+
+  bool load(handle src, bool /*unused*/) {
+    PyObject* obj = src.ptr();
+    if (THPDevice_Check(obj)) {
+      value = reinterpret_cast<THPDevice*>(obj)->device;
+      return true;
+    }
+    return false;
+  }
+
+  static handle cast(
+      const at::Device& src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return handle(THPDevice_New(src));
+  }
+};
+
+template <>
+struct type_caster<at::ScalarType> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(at::ScalarType, _("torch.dtype"));
+
+  // PYBIND11_TYPE_CASTER defines a member field called value. at::ScalarType
+  // cannot be default-initialized, we provide this constructor to explicitly
+  // initialize that field. The value doesn't matter as it will be overwritten
+  // after a successful call to load.
+  type_caster() : value(at::kFloat) {}
+
+  bool load(handle src, bool /*unused*/) {
+    PyObject* obj = src.ptr();
+    if (THPDtype_Check(obj)) {
+      value = reinterpret_cast<THPDtype*>(obj)->scalar_type;
+      return true;
+    }
+    return false;
+  }
+
+  static handle cast(
+      const at::ScalarType& src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return Py_NewRef(torch::getTHPDtype(src));
+  }
+};
+
+template <>
+struct type_caster<c10::Stream> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(c10::Stream, _("torch.Stream"));
+
+  // PYBIND11_TYPE_CASTER defines a member field called value. Since c10::Stream
+  // cannot be default-initialized, we provide this constructor to explicitly
+  // initialize that field. The value doesn't matter as it will be overwritten
+  // after a successful call to load.
+  type_caster() : value(c10::Stream::DEFAULT, c10::Device(c10::kCPU, 0)) {}
+
+  bool load(handle src, bool /*unused*/) {
+    PyObject* obj = src.ptr();
+    if (THPStream_Check(obj)) {
+      value = c10::Stream::unpack3(
+          ((THPStream*)obj)->stream_id,
+          static_cast<c10::DeviceIndex>(((THPStream*)obj)->device_index),
+          static_cast<c10::DeviceType>(((THPStream*)obj)->device_type));
+      return true;
+    }
+    return false;
+  }
+
+  static handle cast(
+      const c10::Stream& src,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    return handle(THPStream_Wrap(src));
+  }
+};
+
+template <>
+struct type_caster<c10::DispatchKey>
+    : public type_caster_base<c10::DispatchKey> {
+  using base = type_caster_base<c10::DispatchKey>;
+  c10::DispatchKey tmp{};
+
+ public:
+  bool load(handle src, bool convert) {
+    if (base::load(src, convert)) {
+      return true;
+    } else if (py::isinstance(
+                   src, py::module_::import("builtins").attr("str"))) {
+      tmp = c10::parseDispatchKey(py::cast<std::string>(src));
+      value = &tmp;
+      return true;
+    }
+    return false;
+  }
+
+  static handle cast(
+      c10::DispatchKey src,
+      return_value_policy policy,
+      handle parent) {
+    return base::cast(src, policy, parent);
+  }
+};
+
+template <>
+struct TORCH_PYTHON_API type_caster<c10::Scalar> {
+ public:
+  PYBIND11_TYPE_CASTER(
+      c10::Scalar,
+      _("Union[Number, torch.SymInt, torch.SymFloat, torch.SymBool]"));
+  bool load(py::handle src, bool /*unused*/);
+
+  static py::handle cast(
+      const c10::Scalar& si,
+      return_value_policy /* policy */,
+      handle /* parent */);
+};
+
+template <>
+struct TORCH_PYTHON_API type_caster<c10::SymInt> {
+ public:
+  PYBIND11_TYPE_CASTER(c10::SymInt, _("Union[int, torch.SymInt]"));
+  bool load(py::handle src, bool /*unused*/);
+
+  static py::handle cast(
+      const c10::SymInt& si,
+      return_value_policy /* policy */,
+      handle /* parent */);
+};
+
+template <>
+struct TORCH_PYTHON_API type_caster<c10::SymFloat> {
+ public:
+  PYBIND11_TYPE_CASTER(c10::SymFloat, _("float"));
+  bool load(py::handle src, bool /*unused*/);
+
+  static py::handle cast(
+      const c10::SymFloat& si,
+      return_value_policy /* policy */,
+      handle /* parent */);
+};
+
+template <>
+struct TORCH_PYTHON_API type_caster<c10::SymBool> {
+ public:
+  PYBIND11_TYPE_CASTER(c10::SymBool, _("Union[bool, torch.SymBool]"));
+  bool load(py::handle src, bool /*unused*/);
+
+  static py::handle cast(
+      const c10::SymBool& si,
+      return_value_policy /* policy */,
+      handle /* parent */);
+};
+
+template <typename T>
+struct type_caster<c10::complex<T>> {
+ public:
+  // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
+  PYBIND11_TYPE_CASTER(c10::complex<T>, _("complex"));
+
+  bool load(handle src, bool /*unused*/) {
+    PyObject* obj = src.ptr();
+
+    // Referred from `THPUtils_unpackComplexDouble`
+    Py_complex py_complex = PyComplex_AsCComplex(obj);
+    if (py_complex.real == -1.0 && PyErr_Occurred()) {
+      return false;
+    }
+
+    // Python's Complex is always double precision.
+    value = c10::complex<double>(py_complex.real, py_complex.imag);
+    return true;
+  }
+
+  static handle cast(
+      const c10::complex<T>& complex,
+      return_value_policy /* policy */,
+      handle /* parent */) {
+    // Python only knows double precision complex.
+    return handle(PyComplex_FromDoubles(complex.real(), complex.imag()));
+  }
+};
+
+} // namespace pybind11::detail
+
+namespace torch::impl {
+
+// Use this function if you have a C++ object that is used from both C++
+// and Python contexts, and you need its GIL to be released when you
+// destruct it in the Python context.
+//
+// This function is a valid shared_ptr destructor and can be used to
+// conveniently allocate a shared_ptr to an object whose destructor will be run
+// without the GIL.  Pass it as the second argument to shared_ptr, e.g.,
+//
+//    shared_ptr<T>(new T(), destroy_without_gil<T>)
+//
+// Attaching the GIL release logic to the holder pointer rather than the
+// actual destructor of T is helpful when T is Python-agnostic and
+// shouldn't refer to the PYthon API.
+//
+// Note there are limitations to the correctness of code that makes use of this.
+// In particular, if a shared_ptr is constructed from C++ code without this
+// destructor and then passed to pybind11, pybind11 will happily take ownership
+// of the shared_ptr (and be willing to destruct it from a context where it is
+// holding the GIL).  unique_ptr with a type branded deleter is less prone to
+// this problem, because a stock deleter unique_ptr is not convertible with it.
+// I plan to mitigate this problem by adding DEBUG-only asserts to the true C++
+// destructors that the GIL is not held (using a virtual call to get to the
+// Python interpreter); alternately, we could use a virtual call to simply
+// ensure we release the GIL in the C++ destructor, however, this is a layering
+// violation (why does code that is ostensibly Python agnostic calling into the
+// GIL).
+//
+// Adapted from
+// https://github.com/pybind/pybind11/issues/1446#issuecomment-406341510
+template <typename T>
+inline void destroy_without_gil(T* ptr) {
+  // Because the ownership of a shared_ptr is diffuse, it's not possible to
+  // necessarily predict whether or not the last reference to an object will
+  // be destructed from Python or C++.  This means that in the destructor here,
+  // we don't necessarily know if we actually have the GIL or not; in fact,
+  // we don't even know if the Python interpreter still exists!  Thus, we have
+  // to test for it before releasing the GIL.
+  //
+  // PyGILState_Check is hopefully self explanatory.  But Py_IsInitialized or
+  // _PyIsFinalizing?  Both get set at the same time during the Python
+  // destruction process:
+  // https://github.com/python/cpython/blob/d92513390a1a0da781bb08c284136f4d7abea36d/Python/pylifecycle.c#L1716-L1717
+  // so the operant question is whether or not you want to release the GIL after
+  // finalization has completed (and there is just no Python interpreter).
+  // Clearly there is no need to release GIL in that state, so we want
+  // Py_IsInitialized.
+  if (Py_IsInitialized() && PyGILState_Check()) {
+    pybind11::gil_scoped_release nogil;
+    delete ptr;
+  } else {
+    delete ptr;
+  }
+}
+
+} // namespace torch::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pycfunction_helpers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pycfunction_helpers.h
new file mode 100644
index 0000000000000000000000000000000000000000..a4c40edc4d326b36ef1ab8f162e08a4fa70f01be
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pycfunction_helpers.h
@@ -0,0 +1,31 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Macros.h>
+#include <torch/csrc/utils/python_compat.h>
+
+#include <Python.h>
+
+inline PyCFunction castPyCFunctionWithKeywords(PyCFunctionWithKeywords func) {
+  C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wcast-function-type")
+  C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wcast-function-type-strict")
+  return reinterpret_cast<PyCFunction>(func);
+  C10_DIAGNOSTIC_POP()
+  C10_DIAGNOSTIC_POP()
+}
+
+#if !IS_PYTHON_3_13_PLUS
+using PyCFunctionFast = _PyCFunctionFast;
+#endif
+
+inline PyCFunction castPyCFunctionFast(PyCFunctionFast func) {
+  C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wcast-function-type")
+  C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wcast-function-type-strict")
+  return reinterpret_cast<PyCFunction>(func);
+  C10_DIAGNOSTIC_POP()
+  C10_DIAGNOSTIC_POP()
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pyobject_preservation.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pyobject_preservation.h
new file mode 100644
index 0000000000000000000000000000000000000000..5c8183ae78cac120a7423a8917c9761cb9b79cf4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pyobject_preservation.h
@@ -0,0 +1,36 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+// This file contains utilities used for handling PyObject preservation
+
+namespace c10 {
+class intrusive_ptr_target;
+namespace impl {
+struct PyObjectSlot;
+} // namespace impl
+} // namespace c10
+
+namespace torch::utils {
+
+class PyObjectPreservation {
+ public:
+  // Store a PyObject wrapper on a fresh c10 wrapper. The caller must hold
+  // a unique reference to `target`.
+  static void init_fresh_nonatomic(
+      c10::intrusive_ptr_target* target,
+      c10::impl::PyObjectSlot* slot,
+      PyObject* pyobj);
+
+  static PyObject* init_once(
+      c10::intrusive_ptr_target* target,
+      c10::impl::PyObjectSlot* slot,
+      PyObject* pyobj);
+};
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_arg_parser.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_arg_parser.h
new file mode 100644
index 0000000000000000000000000000000000000000..4e170dcdbdf839890d9686d01c6d6ca7f841fb87
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_arg_parser.h
@@ -0,0 +1,1377 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+// Parse arguments to Python functions implemented in C++
+// This is similar to PyArg_ParseTupleAndKeywords(), but specifically handles
+// the types relevant to PyTorch and distinguishes between overloaded function
+// signatures.
+//
+// Example:
+//
+//   static PythonArgParser parser({
+//     "norm(Scalar p, int64_t dim, bool keepdim=False)",
+//     "norm(Scalar p=2)",
+//   });
+//   ParsedArgs<3> parsed_args;
+//   auto r = parser.parse(args, kwargs, parsed_args);
+//   if (r.idx == 0) {
+//     norm(r.scalar(0), r.int64(1), r.bool(0));
+//   } else {
+//     norm(r.scalar(0));
+//   }
+//
+// We auto-generate most uses of PythonArgParser; the generated files
+// are torch/csrc/autograd/generated/python_*.cpp
+//
+// Some gotchas that you should watch out for:
+//
+//    - Note [Order of overloads matters]
+//      Order of overloads matters.  A set of input arguments may
+//      bind to multiple argument specs; we will always pick the
+//      first one in PythonArgParser.  However, when you are writing
+//      overloads in, e.g., native_functions.yaml, you don't have to
+//      worry about what order you write them, because the code
+//      generation logic always gives the overloads a canonical
+//      order, where Tensor overloads come first, before Scalar overloads.
+//      This logic is in sort_declarations in
+//      tools/autograd/gen_python_functions.py
+//
+//    - Zero-dim tensors (e.g., torch.tensor(2)) bind to both
+//      Scalar and Tensor, UNLESS they require grad (in which case
+//      they only bind to Tensor).
+
+#include <fmt/format.h>
+#include <pybind11/pytypes.h>
+#include <torch/csrc/python_headers.h>
+
+#include <torch/csrc/Device.h>
+#include <torch/csrc/Dtype.h>
+#include <torch/csrc/DynamicTypes.h>
+#include <torch/csrc/Exceptions.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/Generator.h>
+#include <torch/csrc/Layout.h>
+#include <torch/csrc/MemoryFormat.h>
+#include <torch/csrc/QScheme.h>
+#include <torch/csrc/Stream.h>
+#include <torch/csrc/autograd/python_variable.h>
+#include <torch/csrc/autograd/variable.h>
+#include <torch/csrc/dynamo/eval_frame.h>
+#include <torch/csrc/jit/frontend/tracer.h>
+#include <torch/csrc/python_dimname.h>
+#include <torch/csrc/tensor/python_tensor.h>
+#include <torch/csrc/utils/disable_torch_function.h>
+#include <torch/csrc/utils/object_ptr.h>
+#include <torch/csrc/utils/pybind.h>
+#include <torch/csrc/utils/python_numbers.h>
+#include <torch/csrc/utils/python_strings.h>
+#include <torch/csrc/utils/python_symnode.h>
+#include <torch/csrc/utils/structseq.h>
+
+#include <ATen/DeviceAccelerator.h>
+#include <ATen/PythonTorchFunctionTLS.h>
+#include <ATen/core/Tensor.h>
+#include <c10/util/Exception.h>
+#include <c10/util/irange.h>
+
+#include <c10/core/SymFloat.h>
+#include <c10/core/SymNodeImpl.h>
+
+#include <c10/core/DispatchKeySet.h>
+#include <array>
+#include <cstddef>
+#include <string>
+#include <vector>
+
+inline bool THPUtils_checkScalar(PyObject* obj) {
+#ifdef USE_NUMPY
+  if (torch::utils::is_numpy_scalar(obj)) {
+    return true;
+  }
+#endif
+  return PyFloat_Check(obj) || PyLong_Check(obj) || PyComplex_Check(obj) ||
+      torch::is_symint(py::handle(obj)) || torch::is_dynint(py::handle(obj)) ||
+      torch::is_symfloat(py::handle(obj)) || torch::is_symbool(py::handle(obj));
+}
+
+namespace torch {
+
+TORCH_PYTHON_API bool should_allow_numbers_as_tensors(const std::string& name);
+
+enum class ParameterType {
+  TENSOR,
+  SCALAR,
+  INT64,
+  SYM_INT,
+  DOUBLE,
+  COMPLEX,
+  TENSOR_LIST,
+  INT_LIST,
+  GENERATOR,
+  BOOL,
+  STORAGE,
+  PYOBJECT,
+  SCALARTYPE,
+  LAYOUT,
+  MEMORY_FORMAT,
+  DEVICE,
+  STREAM,
+  STRING,
+  DIMNAME,
+  DIMNAME_LIST,
+  QSCHEME,
+  FLOAT_LIST,
+  SCALAR_LIST,
+  SYM_INT_LIST,
+  DISPATCH_KEY_SET
+};
+
+struct PythonArgs;
+
+// Contains bound Python arguments in declaration order
+template <int N>
+struct ParsedArgs {
+  ParsedArgs() : args() {}
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+  PyObject* args[N];
+};
+
+// FunctionParameter is a single formal parameter of a Python function.
+// It is immutable once constructed.
+struct FunctionParameter {
+  FunctionParameter(const std::string& fmt, bool keyword_only);
+
+  bool check(
+      PyObject* obj,
+      std::vector<PyObject*>& overloaded_args,
+      int argnum,
+      int64_t* failed_idx = nullptr);
+
+  bool _check(
+      PyObject* obj,
+      std::vector<PyObject*>& overloaded_args,
+      int argnum,
+      int64_t* failed_idx = nullptr);
+
+  void set_default_str(const std::string& str);
+  TORCH_PYTHON_API std::string type_name() const;
+
+  ParameterType type_;
+  bool optional{false};
+  bool allow_none{false};
+  bool keyword_only;
+  bool allow_numbers_as_tensors = false;
+  int size{0};
+  std::string name;
+  // having this as a raw PyObject * will presumably leak it, but these are only
+  // held by static objects anyway, and Py_Finalize can already be called when
+  // this is destructed.
+  PyObject* python_name;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers)
+  at::SmallVector<PyObject*, 5> numpy_python_names;
+  at::Scalar default_scalar;
+  std::vector<int64_t> default_intlist;
+  std::string default_string;
+  union {
+    bool default_bool;
+    int64_t default_int;
+    double default_double;
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+    double default_complex[2]; // see Scalar
+    at::ScalarType default_scalartype;
+    at::Layout default_layout;
+  };
+  std::string default_value;
+};
+
+// FunctionSignature represents a single valid signature for a Python function.
+// It is immutable once constructed. The contained data can be concurrently
+// accessed by multiple calls.
+struct FunctionSignature {
+  explicit FunctionSignature(const std::string& fmt, int index);
+
+  bool parse(
+      PyObject* self,
+      PyObject* args,
+      PyObject* kwargs,
+      // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+      PyObject* dst[],
+      std::vector<PyObject*>& overloaded_args,
+      bool raise_exception);
+
+  std::string toString() const;
+
+  std::string name;
+  std::vector<FunctionParameter> params;
+  size_t min_args{0};
+  size_t max_args{0};
+  size_t max_pos_args{0};
+  int index;
+  bool hidden{false};
+  bool deprecated{false};
+};
+
+// A PythonArgParser contains a list of valid signatures. Instances are
+// typically global variables and should be immutable.
+struct PYBIND11_EXPORT PythonArgParser {
+  explicit PythonArgParser(
+      const std::vector<std::string>& fmts,
+      bool traceable = false);
+
+  // meant only for `torch` functions.
+  template <int N>
+  inline PythonArgs parse(
+      PyObject* self,
+      PyObject* args,
+      PyObject* kwargs,
+      ParsedArgs<N>& dst);
+
+  template <int N>
+  inline PythonArgs parse(PyObject* args, PyObject* kwargs, ParsedArgs<N>& dst);
+
+  inline PythonArgs parse(PyObject* self, ParsedArgs<0>& dst);
+
+  // Formatted strings of non-hidden signatures
+  std::vector<std::string> get_signatures() const;
+
+ private:
+  [[noreturn]] void print_error(
+      PyObject* self,
+      PyObject* args,
+      PyObject* kwargs,
+      // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+      PyObject* parsed_args[]);
+  void check_deprecated(const FunctionSignature& signature);
+  PythonArgs raw_parse(
+      PyObject* self,
+      PyObject* args,
+      PyObject* kwargs,
+      // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+      PyObject* parsed_args[]);
+
+  std::vector<FunctionSignature> signatures_;
+  std::string function_name;
+  size_t max_args{0};
+  bool traceable;
+};
+
+// PythonArgs contains bound Python arguments for an actual invocation
+// along with references to the matched signature.
+struct TORCH_PYTHON_API PythonArgs {
+  PythonArgs(
+      bool traceable,
+      const FunctionSignature& signature,
+      PyObject** args,
+      std::vector<PyObject*> overloaded_args)
+      : idx(signature.index),
+        traceable(traceable),
+        signature(signature),
+        args(args),
+        overloaded_args(std::move(overloaded_args)) {}
+
+  int idx;
+  bool traceable;
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const FunctionSignature& signature;
+  PyObject** args;
+  std::vector<PyObject*> overloaded_args; // NOTE: borrowed references
+
+  inline bool has_torch_function();
+  inline std::string get_func_name();
+  inline at::Tensor tensor(int i);
+  inline std::optional<at::Tensor> optionalTensor(int i);
+  inline at::Scalar scalar(int i);
+  inline at::Scalar scalarWithDefault(int i, const at::Scalar& default_scalar);
+  inline std::vector<at::Scalar> scalarlist(int i);
+  inline std::vector<at::Tensor> tensorlist(int i);
+  inline torch::List<std::optional<at::Tensor>> list_of_optional_tensors(int i);
+  template <int N>
+  inline std::array<at::Tensor, N> tensorlist_n(int i);
+  inline std::vector<int64_t> intlist(int i);
+  inline std::vector<c10::SymInt> symintlist(int i);
+  inline c10::OptionalArray<int64_t> intlistOptional(int i);
+  inline c10::OptionalArray<c10::SymInt> symintlistOptional(int i);
+  inline std::vector<int64_t> intlistWithDefault(
+      int i,
+      std::vector<int64_t> default_intlist);
+  inline std::optional<at::Generator> generator(int i);
+  inline at::Storage storage(int i);
+  inline at::Storage storage(
+      int i,
+      at::ScalarType& storage_scalar_type,
+      bool& is_typed_storage);
+  inline c10::Stream stream(int i);
+  inline at::ScalarType scalartype(int i);
+  inline at::ScalarType scalartypeWithDefault(
+      int i,
+      at::ScalarType default_scalartype);
+  inline std::optional<at::ScalarType> scalartypeOptional(int i);
+  inline std::optional<at::Scalar> scalarOptional(int i);
+  inline std::optional<int64_t> toInt64Optional(int i);
+  inline std::optional<c10::SymInt> toSymIntOptional(int i);
+  inline std::optional<bool> toBoolOptional(int i);
+  inline std::optional<double> toDoubleOptional(int i);
+  inline c10::OptionalArray<double> doublelistOptional(int i);
+  inline std::vector<double> doublelist(int i);
+  inline std::vector<double> getDoublelist(int i);
+  inline at::Layout layout(int i);
+  inline at::Layout layoutWithDefault(int i, at::Layout default_layout);
+  inline std::optional<at::Layout> layoutOptional(int i);
+  inline at::Device device(int i);
+  inline at::Device deviceWithDefault(int i, const at::Device& default_device);
+  inline std::optional<at::Device> deviceOptional(int i);
+  inline at::Dimname dimname(int i);
+  inline std::vector<at::Dimname> dimnamelist(int i);
+  inline std::optional<std::vector<at::Dimname>> toDimnameListOptional(int i);
+  inline at::MemoryFormat memoryformat(int i);
+  inline std::optional<at::MemoryFormat> memoryformatOptional(int i);
+  inline at::QScheme toQScheme(int i);
+  inline std::string string(int i);
+  inline std::string stringWithDefault(int i, const std::string& default_str);
+  inline std::optional<std::string> stringOptional(int i);
+  inline std::string_view stringView(int i);
+  inline std::string_view stringViewWithDefault(
+      int i,
+      const std::string_view default_str);
+  inline std::optional<std::string_view> stringViewOptional(int i);
+  inline PyObject* pyobject(int i);
+  inline int64_t toInt64(int i);
+  inline c10::SymInt toSymInt(int i);
+  inline c10::SymBool toSymBool(int i);
+  inline int64_t toInt64WithDefault(int i, int64_t default_int);
+  inline double toDouble(int i);
+  inline double toDoubleWithDefault(int i, double default_double);
+  inline c10::complex<double> toComplex(int i);
+  inline c10::complex<double> toComplexWithDefault(
+      int i,
+      c10::complex<double> default_complex);
+  inline bool toBool(int i);
+  inline bool toBoolWithDefault(int i, bool default_bool);
+  inline bool isNone(int i);
+  inline std::optional<c10::DispatchKeySet> toDispatchKeySetOptional(int i);
+
+ private:
+  // Non-inline functions' symbols are exposed to torch_python DLL
+  // via TORCH_PYTHON_API tag at struct level.
+  at::Tensor tensor_slow(int i);
+  at::Scalar scalar_slow(int i);
+  at::Scalar scalar_slow(PyObject* arg);
+};
+
+template <int N>
+inline PythonArgs PythonArgParser::parse(
+    PyObject* self,
+    PyObject* args,
+    PyObject* kwargs,
+    ParsedArgs<N>& dst) {
+  TORCH_CHECK_VALUE(
+      N >= max_args,
+      "PythonArgParser: dst ParsedArgs buffer does not have enough capacity, expected ",
+      max_args,
+      " (got ",
+      N,
+      ")");
+  return raw_parse(self, args, kwargs, dst.args);
+}
+
+template <int N>
+inline PythonArgs PythonArgParser::parse(
+    PyObject* args,
+    PyObject* kwargs,
+    ParsedArgs<N>& dst) {
+  return parse(nullptr, args, kwargs, dst);
+}
+
+inline PythonArgs PythonArgParser::parse(PyObject* self, ParsedArgs<0>& dst) {
+  return parse(self, nullptr, nullptr, dst);
+}
+
+inline bool PythonArgs::has_torch_function() {
+  return !overloaded_args.empty() || at::impl::torch_function_mode_enabled();
+}
+
+inline std::string PythonArgs::get_func_name() {
+  return signature.name;
+}
+
+// TODO: this can return MaybeOwned
+inline at::Tensor PythonArgs::tensor(int i) {
+  if (args[i] && THPVariable_CheckExact(args[i])) {
+    return THPVariable_Unpack(args[i]);
+  }
+  return tensor_slow(i);
+}
+
+inline std::optional<at::Tensor> PythonArgs::optionalTensor(int i) {
+  at::Tensor t = tensor(i);
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  if (t.defined()) {
+    return t;
+  } else {
+    return std::nullopt;
+  }
+}
+
+inline at::Scalar PythonArgs::scalar(int i) {
+  if (!args[i])
+    return signature.params[i].default_scalar;
+  return scalar_slow(i);
+}
+
+inline std::vector<at::Scalar> PythonArgs::scalarlist(int i) {
+  if (!args[i])
+    return std::vector<at::Scalar>();
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
+  std::vector<at::Scalar> res(size);
+  for (const auto idx : c10::irange(size)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
+    res[idx] = scalar_slow(obj);
+  }
+  return res;
+}
+
+inline at::Scalar PythonArgs::scalarWithDefault(
+    int i,
+    const at::Scalar& default_scalar) {
+  if (!args[i])
+    return default_scalar;
+  return scalar_slow(i);
+}
+
+inline std::optional<at::Scalar> PythonArgs::scalarOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return scalar_slow(i);
+}
+
+inline std::vector<at::Tensor> PythonArgs::tensorlist(int i) {
+  if (!args[i])
+    return std::vector<at::Tensor>();
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
+  std::vector<at::Tensor> res(size);
+  for (const auto idx : c10::irange(size)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
+    // This is checked by the argument parser so it's safe to cast without
+    // checking if this is a tensor first
+    res[idx] = THPVariable_Unpack(obj);
+  }
+  return res;
+}
+
+inline torch::List<std::optional<at::Tensor>> PythonArgs::
+    list_of_optional_tensors(int i) {
+  if (!args[i])
+    return torch::List<std::optional<at::Tensor>>();
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
+  torch::List<std::optional<at::Tensor>> res;
+  res.reserve(size);
+  for (const auto idx : c10::irange(size)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
+    // This is checked by the argument parser so it's safe to cast without
+    // checking if this is a tensor first
+    res.push_back(THPVariable_Unpack(obj));
+  }
+  return res;
+}
+
+template <int N>
+inline std::array<at::Tensor, N> PythonArgs::tensorlist_n(int i) {
+  auto res = std::array<at::Tensor, N>();
+  if (!args[i])
+    return res;
+  auto tuple = PyTuple_Check(args[i]);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i]));
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]);
+  if (size != N) {
+    TORCH_CHECK_TYPE(
+        false,
+        fmt::format("expected tuple of {} elements but got {}", N, size));
+  }
+  for (const auto idx : c10::irange(size)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx);
+    // This is checked by the argument parser so it's safe to cast without
+    // checking if this is a tensor first
+    res[idx] = THPVariable_Unpack(obj);
+  }
+  return res;
+}
+
+inline std::vector<int64_t> PythonArgs::intlist(int i) {
+  return intlistWithDefault(i, signature.params[i].default_intlist);
+}
+
+inline PyObject* toPyObject(const c10::SymInt& symint) {
+  if (symint.is_symbolic()) {
+    auto r = py::cast(symint).release().ptr();
+    TORCH_INTERNAL_ASSERT(r);
+    return r;
+  } else {
+    auto m = symint.maybe_as_int();
+    // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
+    return THPUtils_packInt64(m.value());
+  }
+}
+
+inline void throw_intlist_exception(
+    const torch::PythonArgs* args,
+    size_t i,
+    PyObject* obj,
+    size_t idx,
+    const std::exception& e = python_error()) {
+  std::string error = strlen(e.what())
+      ? e.what()
+      : std::string("type must be ") + args->signature.params[i].type_name() +
+          ",but got " + Py_TYPE(obj)->tp_name;
+  TORCH_CHECK_TYPE(
+      false,
+      fmt::format(
+          "{}(): argument '{}' failed to unpack the object at pos {} with error \"{}\"",
+          args->signature.name,
+          args->signature.params[i].name,
+          idx + 1,
+          error));
+}
+
+inline std::vector<c10::SymInt> PythonArgs::symintlist(int i) {
+  if (!args[i]) {
+    return c10::fmap(signature.params[i].default_intlist, [](int64_t di) {
+      return c10::SymInt(di);
+    });
+  }
+
+  const auto size1 = signature.params[i].size;
+  if (size1 > 0 && THPUtils_checkLong(args[i])) {
+    return std::vector<c10::SymInt>(
+        size1, c10::SymInt(THPUtils_unpackLong(args[i])));
+  }
+
+  if (size1 > 0 && torch::is_symint(py::handle(args[i]))) {
+    auto si = py::handle(args[i]).cast<c10::SymInt>();
+    return std::vector<c10::SymInt>(size1, si);
+  }
+
+  if (size1 > 0 && THPVariable_Check(args[i])) {
+    return std::vector<c10::SymInt>(
+        size1, THPVariable_Unpack(args[i]).item().toSymInt());
+  }
+
+  PyObject* arg = args[i];
+  auto tuple = PyTuple_Check(arg);
+  if (!tuple) {
+    TORCH_INTERNAL_ASSERT(PyList_Check(arg), "expected tuple or list");
+  }
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  const auto size2 = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg);
+  std::vector<c10::SymInt> res;
+  res.reserve(size2);
+  for (const auto idx : c10::irange(size2)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(arg, idx) : PyList_GET_ITEM(arg, idx);
+
+    // Elements of torch.Size are tensors during tracing, and we need to
+    // record extra information before they are turned into an IntArrayRef
+    if (traceable && jit::tracer::isTracing() && THPVariable_Check(obj)) {
+      auto& var = THPVariable_Unpack(obj);
+      jit::tracer::ArgumentStash::stashIntArrayRefElem(
+          signature.params[i].name, size2, idx, var);
+      try {
+        res.emplace_back(var.item<int64_t>());
+        continue;
+      } catch (std::exception& e) {
+        throw_intlist_exception(this, i, obj, idx, e);
+      }
+      continue;
+    } else {
+      // convert tensor to scalar outside of try / catch,
+      // so that Tensor subclass exceptions will not be caught.
+      if (THPUtils_checkLongExact(obj)) {
+        // Fast path for plain numbers
+        try {
+          res.emplace_back(THPUtils_unpackLong(obj));
+        } catch (std::exception& e) {
+          throw_intlist_exception(this, i, obj, idx, e);
+        }
+      } else if (THPVariable_Check(obj)) {
+        auto& var = THPVariable_Unpack(obj);
+        if (var.numel() != 1 ||
+            !at::isIntegralType(
+                var.dtype().toScalarType(), /*include_bool*/ true)) {
+          throw_intlist_exception(this, i, obj, idx);
+        }
+        auto scalar = var.item();
+        TORCH_CHECK(scalar.isIntegral(/*include bool*/ false));
+        res.push_back(scalar.toSymInt());
+      } else {
+        try {
+          if (is_symint(py::handle(obj))) {
+            res.push_back(py::handle(obj).cast<c10::SymInt>());
+          } else if (is_dynint(py::handle(obj))) {
+            res.emplace_back(py::handle(obj).cast<int>());
+          } else {
+            res.emplace_back(THPUtils_unpackIndex(obj));
+          }
+        } catch (std::exception& e) {
+          throw_intlist_exception(this, i, obj, idx, e);
+        }
+      }
+    }
+  }
+
+  return res;
+}
+
+inline std::vector<int64_t> PythonArgs::intlistWithDefault(
+    int i,
+    std::vector<int64_t> default_intlist) {
+  if (!args[i])
+    return default_intlist;
+  PyObject* arg = args[i];
+  const auto size1 = signature.params[i].size;
+  if (size1 > 0 && THPUtils_checkLong(arg)) {
+    return std::vector<int64_t>(size1, THPUtils_unpackLong(arg));
+  }
+  if (size1 > 0 && torch::is_symint(py::handle(arg))) {
+    return std::vector<int64_t>(
+        size1,
+        py::handle(arg).cast<c10::SymInt>().guard_int(__FILE__, __LINE__));
+  }
+  if (size1 > 0 && torch::is_dynint(py::handle(arg))) {
+    return std::vector<int64_t>(size1, py::handle(arg).cast<int>());
+  }
+  if (size1 > 0 && THPVariable_Check(arg)) {
+    return std::vector<int64_t>(size1, THPVariable_Unpack(arg).item<int64_t>());
+  }
+  auto tuple = PyTuple_Check(arg);
+  if (!tuple) {
+    TORCH_INTERNAL_ASSERT(PyList_Check(arg), "expected tuple or list");
+  }
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  const auto size2 = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg);
+  std::vector<int64_t> res(size2);
+  for (const auto idx : c10::irange(size2)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(arg, idx) : PyList_GET_ITEM(arg, idx);
+    // Elements of torch.Size are tensors during tracing, and we need to
+    // record extra information before they are turned into an IntArrayRef
+    if (traceable && jit::tracer::isTracing() && THPVariable_Check(obj)) {
+      auto& var = THPVariable_Unpack(obj);
+      jit::tracer::ArgumentStash::stashIntArrayRefElem(
+          signature.params[i].name, size2, idx, var);
+      try {
+        res[idx] = var.item<int64_t>();
+        continue;
+      } catch (std::exception& e) {
+        throw_intlist_exception(this, i, obj, idx, e);
+      }
+    } else {
+      // convert tensor to scalar outside of try / catch,
+      // so that Tensor subclass exceptions will not be caught.
+      if (THPUtils_checkLongExact(obj)) {
+        // Fast path for plain numbers
+        try {
+          res[idx] = THPUtils_unpackLong(obj);
+        } catch (std::exception& e) {
+          throw_intlist_exception(this, i, obj, idx, e);
+        }
+      } else if (torch::is_symint(py::handle(obj))) {
+        res[idx] = py::cast<c10::SymInt>(py::handle(obj))
+                       .guard_int(__FILE__, __LINE__);
+      } else if (torch::is_dynint(py::handle(obj))) {
+        res[idx] = py::handle(obj).cast<int>();
+      } else if (THPVariable_Check(obj)) {
+        auto& var = THPVariable_Unpack(obj);
+        if (var.numel() != 1 ||
+            !at::isIntegralType(
+                var.dtype().toScalarType(), /*include_bool*/ true)) {
+          throw_intlist_exception(this, i, obj, idx);
+        }
+        res[idx] = var.item<int64_t>();
+      } else {
+        try {
+          res[idx] = THPUtils_unpackIndex(obj);
+        } catch (std::exception& e) {
+          throw_intlist_exception(this, i, obj, idx, e);
+        }
+      }
+    }
+  }
+  return res;
+}
+
+inline c10::OptionalArray<int64_t> PythonArgs::intlistOptional(int i) {
+  if (!args[i]) {
+    return {};
+  }
+  return intlist(i);
+}
+
+inline c10::OptionalArray<c10::SymInt> PythonArgs::symintlistOptional(int i) {
+  if (!args[i]) {
+    return {};
+  }
+  return symintlist(i);
+}
+
+inline std::vector<double> PythonArgs::getDoublelist(int i) {
+  PyObject* arg = args[i];
+  auto tuple = PyTuple_Check(arg);
+  if (!tuple) {
+    TORCH_INTERNAL_ASSERT(PyList_Check(arg), "expected tuple or list");
+  }
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  auto size = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg);
+  std::vector<double> res(size);
+  for (const auto idx : c10::irange(size)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(arg, idx) : PyList_GET_ITEM(arg, idx);
+    try {
+      if (torch::is_symfloat(py::handle(obj))) {
+        res[idx] = py::cast<c10::SymFloat>(py::handle(obj))
+                       .guard_float(__FILE__, __LINE__);
+      } else {
+        res[idx] = THPUtils_unpackDouble(obj);
+      }
+    } catch (const std::exception&) {
+      TORCH_CHECK_TYPE(
+          false,
+          fmt::format(
+              "{}(): argument '{}' must be {}, but found element of type {} at pos {}",
+              signature.name,
+              signature.params[i].name,
+              signature.params[i].type_name(),
+              Py_TYPE(obj)->tp_name,
+              idx + 1));
+    }
+  }
+  return res;
+}
+
+inline c10::OptionalArray<double> PythonArgs::doublelistOptional(int i) {
+  if (!args[i]) {
+    return {};
+  }
+  return this->getDoublelist(i);
+}
+
+inline std::vector<double> PythonArgs::doublelist(int i) {
+  if (!args[i]) {
+    return {};
+  }
+  return this->getDoublelist(i);
+}
+
+inline std::optional<c10::DispatchKeySet> PythonArgs::toDispatchKeySetOptional(
+    int i) {
+  if (!args[i]) {
+    return {};
+  }
+  return py::cast<c10::DispatchKeySet>(py::handle(args[i]));
+}
+
+inline at::ScalarType PythonArgs::scalartypeWithDefault(
+    int i,
+    at::ScalarType default_scalartype) {
+  if (!args[i])
+    return default_scalartype;
+  return scalartype(i);
+}
+
+inline at::ScalarType toScalarType(PyObject* obj) {
+  if (obj == (PyObject*)&PyFloat_Type) {
+    return at::ScalarType::Double;
+  }
+  if (obj == (PyObject*)&PyBool_Type) {
+    return at::ScalarType::Bool;
+  }
+  if (obj == (PyObject*)&PyLong_Type) {
+    return at::ScalarType::Long;
+  }
+  if (obj == (PyObject*)&PyComplex_Type) {
+    return at::ScalarType::ComplexDouble;
+  }
+  return reinterpret_cast<THPDtype*>(obj)->scalar_type;
+}
+
+inline at::ScalarType PythonArgs::scalartype(int i) {
+  if (!args[i]) {
+    auto scalartype = signature.params[i].default_scalartype;
+    return (scalartype == at::ScalarType::Undefined)
+        ? torch::tensors::get_default_scalar_type()
+        : scalartype;
+  }
+  PyObject* obj = args[i];
+  return toScalarType(obj);
+}
+
+inline std::optional<at::ScalarType> PythonArgs::scalartypeOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return scalartype(i);
+}
+
+inline at::Layout toLayout(PyObject* obj) {
+  const auto layout = reinterpret_cast<THPLayout*>(obj);
+  return layout->layout;
+}
+
+inline at::Layout PythonArgs::layout(int i) {
+  if (!args[i])
+    return signature.params[i].default_layout;
+  return toLayout(args[i]);
+}
+
+inline at::Layout PythonArgs::layoutWithDefault(
+    int i,
+    at::Layout default_layout) {
+  if (!args[i])
+    return default_layout;
+  return layout(i);
+}
+
+inline std::optional<at::Layout> PythonArgs::layoutOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return layout(i);
+}
+
+inline at::Device deviceFromLong(int64_t device_index) {
+  TORCH_CHECK(device_index >= 0, "Device index must not be negative");
+  return at::Device(
+      // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
+      at::getAccelerator(true).value(),
+      static_cast<c10::DeviceIndex>(device_index));
+}
+
+inline at::Device toDevice(PyObject* obj) {
+  if (THPDevice_Check(obj)) {
+    const auto device = reinterpret_cast<THPDevice*>(obj);
+    return device->device;
+  }
+  if (THPUtils_checkLong(obj)) {
+    return deviceFromLong(THPUtils_unpackLong(obj));
+  }
+  if (torch::is_symint(py::handle(obj))) {
+    auto device_index =
+        py::cast<c10::SymInt>(py::handle(obj)).guard_int(__FILE__, __LINE__);
+    return deviceFromLong(device_index);
+  }
+  if (torch::is_dynint(py::handle(obj))) {
+    auto device_index = py::cast<int>(py::handle(obj));
+    return deviceFromLong(device_index);
+  }
+  const std::string& device_str = THPUtils_unpackString(obj);
+  return at::Device(device_str);
+}
+
+inline at::Device PythonArgs::device(int i) {
+  if (!args[i]) {
+    return torch::tensors::get_default_device();
+  }
+  return toDevice(args[i]);
+}
+
+inline at::Device PythonArgs::deviceWithDefault(
+    int i,
+    const at::Device& default_device) {
+  if (!args[i])
+    return default_device;
+  return device(i);
+}
+
+inline std::optional<at::Device> PythonArgs::deviceOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return device(i);
+}
+
+inline at::Dimname PythonArgs::dimname(int i) {
+  TORCH_INTERNAL_ASSERT(args[i] != nullptr);
+  return THPDimname_parse(args[i]);
+}
+
+inline std::vector<at::Dimname> parseDimnameList(PyObject* arg) {
+  auto tuple = PyTuple_Check(arg);
+  if (!tuple) {
+    TORCH_INTERNAL_ASSERT(PyList_Check(arg), "expected tuple or list");
+  }
+  // NOLINTNEXTLINE(bugprone-branch-clone)
+  auto size = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg);
+  std::vector<at::Dimname> res;
+  res.reserve(size);
+  for (const auto idx : c10::irange(size)) {
+    PyObject* obj =
+        tuple ? PyTuple_GET_ITEM(arg, idx) : PyList_GET_ITEM(arg, idx);
+    res.push_back(THPDimname_parse(obj));
+  }
+  return res;
+}
+
+inline std::optional<std::vector<at::Dimname>> PythonArgs::
+    toDimnameListOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return parseDimnameList(args[i]);
+}
+
+inline std::vector<at::Dimname> PythonArgs::dimnamelist(int i) {
+  TORCH_INTERNAL_ASSERT(args[i]);
+  PyObject* arg = args[i];
+  auto size = signature.params[i].size;
+  TORCH_INTERNAL_ASSERT(size == 0 || size == 1);
+  if (size == 1 && THPUtils_checkDimname(arg)) {
+    return {THPDimname_parse(arg)};
+  }
+  return parseDimnameList(arg);
+}
+
+inline at::MemoryFormat PythonArgs::memoryformat(int i) {
+  if (!args[i])
+    return at::MemoryFormat::Contiguous;
+  TORCH_CHECK(
+      THPMemoryFormat_Check(args[i]),
+      "memory_format arg must be an instance of the torch.memory_format");
+  const auto memory_format = reinterpret_cast<THPMemoryFormat*>(args[i]);
+  return memory_format->memory_format;
+}
+
+inline std::optional<at::MemoryFormat> PythonArgs::memoryformatOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return memoryformat(i);
+}
+
+inline at::QScheme PythonArgs::toQScheme(int i) {
+  if (!args[i])
+    return at::kPerTensorAffine;
+  TORCH_CHECK(
+      THPQScheme_Check(args[i]),
+      "qscheme arg must be an instance of the torch.qscheme");
+  const auto qscheme = reinterpret_cast<THPQScheme*>(args[i]);
+  return qscheme->qscheme;
+}
+
+inline std::string PythonArgs::string(int i) {
+  return stringWithDefault(i, signature.params[i].default_string);
+}
+
+inline std::string PythonArgs::stringWithDefault(
+    int i,
+    const std::string& default_str) {
+  if (!args[i])
+    return default_str;
+  return THPUtils_unpackString(args[i]);
+}
+
+inline std::optional<std::string> PythonArgs::stringOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return THPUtils_unpackString(args[i]);
+}
+
+inline std::string_view PythonArgs::stringView(int i) {
+  return stringViewWithDefault(i, signature.params[i].default_string);
+}
+
+inline std::string_view PythonArgs::stringViewWithDefault(
+    int i,
+    const std::string_view default_str) {
+  if (!args[i])
+    return default_str;
+  return THPUtils_unpackStringView(args[i]);
+}
+
+inline std::optional<std::string_view> PythonArgs::stringViewOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return THPUtils_unpackStringView(args[i]);
+}
+
+inline int64_t PythonArgs::toInt64(int i) {
+  if (!args[i])
+    return signature.params[i].default_int;
+  if (traceable && jit::tracer::isTracing() && THPVariable_Check(args[i])) {
+    auto& var = THPVariable_Unpack(args[i]);
+    jit::tracer::ArgumentStash::stashValue(
+        signature.params[i].name, idx, var, c10::IntType::get());
+  }
+  if (torch::is_symint(py::handle(args[i]))) {
+    return py::cast<c10::SymInt>(py::handle(args[i]))
+        .guard_int(__FILE__, __LINE__);
+  }
+  if (torch::is_dynint(py::handle(args[i]))) {
+    return py::cast<int>(py::handle(args[i]));
+  }
+  return THPUtils_unpackLong(args[i]);
+}
+
+inline c10::SymInt PythonArgs::toSymInt(int i) {
+  if (!args[i]) {
+    return c10::SymInt(signature.params[i].default_int);
+  }
+
+  if (traceable && jit::tracer::isTracing() && THPVariable_Check(args[i])) {
+    auto& var = THPVariable_Unpack(args[i]);
+    jit::tracer::ArgumentStash::stashValue(
+        signature.params[i].name, idx, var, c10::IntType::get());
+  }
+
+  return py::cast<c10::SymInt>(py::handle(args[i]));
+}
+
+inline c10::SymBool PythonArgs::toSymBool(int i) {
+  if (!args[i]) {
+    return c10::SymBool(signature.params[i].default_bool);
+  }
+  if (traceable && jit::tracer::isTracing() && THPVariable_Check(args[i])) {
+    auto& var = THPVariable_Unpack(args[i]);
+    jit::tracer::ArgumentStash::stashValue(
+        signature.params[i].name, idx, var, c10::BoolType::get());
+  }
+
+  return py::cast<c10::SymBool>(py::handle(args[i]));
+}
+
+inline int64_t PythonArgs::toInt64WithDefault(int i, int64_t default_int) {
+  if (!args[i])
+    return default_int;
+  return toInt64(i);
+}
+
+inline std::optional<int64_t> PythonArgs::toInt64Optional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return toInt64(i);
+}
+
+inline std::optional<c10::SymInt> PythonArgs::toSymIntOptional(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return toSymInt(i);
+}
+
+inline std::optional<bool> PythonArgs::toBoolOptional(int i) {
+  if (!args[i]) {
+    return std::nullopt;
+  }
+  return toBool(i);
+}
+
+inline std::optional<double> PythonArgs::toDoubleOptional(int i) {
+  if (!args[i]) {
+    return std::nullopt;
+  }
+  return toDouble(i);
+}
+
+inline double PythonArgs::toDouble(int i) {
+  if (!args[i])
+    return signature.params[i].default_double;
+  if (torch::is_symfloat(py::handle(args[i]))) {
+    return py::cast<c10::SymFloat>(py::handle(args[i]))
+        .guard_float(__FILE__, __LINE__);
+  }
+  if (torch::is_symint(py::handle(args[i]))) {
+    return static_cast<double>(py::cast<c10::SymInt>(py::handle(args[i]))
+                                   .guard_int(__FILE__, __LINE__));
+  }
+  if (torch::is_dynint(py::handle(args[i]))) {
+    return static_cast<double>(py::cast<int>(py::handle(args[i])));
+  }
+  return THPUtils_unpackDouble(args[i]);
+}
+
+inline bool PythonArgs::toBool(int i) {
+  if (!args[i]) {
+    return signature.params[i].default_bool;
+  }
+  if (args[i] == Py_True) {
+    return true;
+  }
+  if (args[i] == Py_False) {
+    return false;
+  }
+  if (torch::is_symbool(py::handle(args[i]))) {
+    return py::cast<c10::SymBool>(py::handle(args[i]))
+        .guard_bool(__FILE__, __LINE__);
+  }
+  return false;
+}
+
+inline double PythonArgs::toDoubleWithDefault(int i, double default_double) {
+  if (!args[i])
+    return default_double;
+  return toDouble(i);
+}
+
+inline c10::complex<double> PythonArgs::toComplex(int i) {
+  if (!args[i])
+    return *(reinterpret_cast<const c10::complex<double>*>(
+        signature.params[i].default_complex));
+  return THPUtils_unpackComplexDouble(args[i]);
+}
+
+inline c10::complex<double> PythonArgs::toComplexWithDefault(
+    int i,
+    c10::complex<double> default_complex) {
+  if (!args[i])
+    return default_complex;
+  return toComplex(i);
+}
+
+inline bool PythonArgs::toBoolWithDefault(int i, bool default_bool) {
+  if (!args[i])
+    return default_bool;
+  return toBool(i);
+}
+
+inline bool PythonArgs::isNone(int i) {
+  return args[i] == nullptr;
+}
+
+inline std::optional<at::Generator> PythonArgs::generator(int i) {
+  if (!args[i])
+    return std::nullopt;
+  return reinterpret_cast<THPGenerator*>(args[i])->cdata;
+}
+
+inline at::Storage PythonArgs::storage(int i) {
+  if (!args[i])
+    return at::Storage();
+  return createStorage(args[i]);
+}
+
+inline at::Storage PythonArgs::storage(
+    int i,
+    at::ScalarType& storage_scalar_type,
+    bool& is_typed_storage) {
+  at::Storage storage;
+  if (!args[i]) {
+    storage = at::Storage();
+    is_typed_storage = false;
+    storage_scalar_type = at::ScalarType::Undefined;
+  } else {
+    std::tie(storage, storage_scalar_type, is_typed_storage) =
+        createStorageGetType(args[i]);
+  }
+  return storage;
+}
+
+inline c10::Stream PythonArgs::stream(int i) {
+  if (!args[i])
+    return c10::Stream(
+        c10::Stream::Default::DEFAULT, c10::Device(c10::DeviceType::CPU, -1));
+  if (!THPStream_Check(args[i])) {
+    TORCH_CHECK_TYPE(
+        false,
+        fmt::format(
+            "expected Stream object. Got '{}'", Py_TYPE(args[i])->tp_name));
+  }
+  return c10::Stream::unpack3(
+      ((THPStream*)args[i])->stream_id,
+      static_cast<c10::DeviceIndex>(((THPStream*)args[i])->device_index),
+      static_cast<c10::DeviceType>(((THPStream*)args[i])->device_type));
+}
+
+inline PyObject* PythonArgs::pyobject(int i) {
+  if (!args[i])
+    return Py_None;
+  return args[i];
+}
+
+/*
+ *
+ * Handle __torch_function__ overrides if we know that there are overloaded
+ * arguments.  All objects stored in r.overloaded_args must have a
+ * __torch_function__ implementation and the arguments must be ordered in order
+ * of precedence. Precedence goes from left to right in the order of the
+ * signature of the function the overloaded arguments were passed to, except
+ * subclasses are always considered before superclasses.
+ *
+ * If the result of calling __torch_function__ is NotImplemented, the
+ * next implementation in the precedence order is called. If all
+ * arguments return NotImplemented from their __torch_function__
+ * implementation, a TypeError is raised in Python.
+ *
+ * Assumes overloaded_args has at least one entry. All entries must have
+ * a __torch_function__ attribute that resolves to a callable that
+ * accepts a torch API function, a tuple of arguments, and a dict of
+ * keyword arguments for the torch API function.
+ *
+ * It is sufficient to call PythonArgs::has_torch_function before
+ * calling this function to verify that there are valid arguments
+ * present. If that is not done then special care must be taken to
+ * ensure there are arguments that are overloaded with
+ * __torch_function__.
+ *
+ * See torch._overrides.handle_torch_function for the equivalent
+ * code in the pure-python implementation.
+ *
+ * 'r' is a parsed PythonArgs instance, returned from
+ * PythonArgParser::parse.
+ *
+ * 'args' is a reference to the python tuple of arguments to the torch
+ * API function.
+ *
+ * 'kwargs' is a reference to the python dict of keyword arguments to
+ * the torch API function.
+ *
+ * 'torch_api' is a reference to a python torch API namespace.
+ *
+ * 'torch_api_function' is the reference to the original torch method, usually,
+ * we can use torch_api and func_name to get torch_api_function. In some cases,
+ * e.g., torch custom op, we create the function in C++, if we still use
+ * torch_api and func_name to fetch original api, a cyclic call will happen.
+ *
+ * 'overloaded_args' is the args which have overloaded __torch_function__.
+ *
+ * 'func_name' is the named of the original torch method.
+ *
+ * TODO: we could use different names for the following 'handle_torch_function'
+ * instead of overloading.
+ *
+ */
+// Used for Tensor methods with arguments.
+auto handle_torch_function(
+    PythonArgs& r,
+    PyObject* self,
+    PyObject* args,
+    PyObject* kwargs,
+    PyObject* torch_api,
+    const char* module_name,
+    const char* func_name_override = nullptr) -> PyObject*;
+
+// Used for functions which needs to parse python args.
+auto handle_torch_function(
+    PythonArgs& r,
+    PyObject* args,
+    PyObject* kwargs,
+    PyObject* torch_api,
+    const char* module_name,
+    const char* func_name_override = nullptr) -> PyObject*;
+
+// Used for functions that have no argument parsing.
+auto handle_torch_function(
+    PyObject* self,
+    const std::string& func_name,
+    PyObject* args = nullptr,
+    PyObject* kwargs = nullptr,
+    PyObject* torch_api = THPVariableClass,
+    const std::string& module_name = "torch.Tensor") -> PyObject*;
+
+// Used for functions created in C++, e.g., C++ custom op, which doesn't use
+// PythonArgParser to get overloaded_args.
+enum class TorchFunctionName { TorchFunction, TorchDispatch };
+
+auto TORCH_PYTHON_API handle_torch_function_no_python_arg_parser(
+    at::ArrayRef<PyObject*> overloaded_args,
+    PyObject* args,
+    PyObject* kwargs,
+    const char* func_name,
+    PyObject* torch_api_function,
+    const char* module_name,
+    TorchFunctionName torch_function_name = TorchFunctionName::TorchFunction)
+    -> PyObject*;
+
+auto handle_torch_function_no_python_arg_parser(
+    at::ArrayRef<PyObject*> overloaded_args,
+    PyObject* args,
+    PyObject* kwargs,
+    const char* func_name,
+    PyObject* torch_api_function,
+    const char* module_name,
+    const c10::OperatorHandle* opt_op,
+    torch::jit::Stack* opt_stack,
+    TorchFunctionName torch_function_name = TorchFunctionName::TorchFunction)
+    -> PyObject*;
+
+// Used for getters of Tensor properties
+auto handle_torch_function_getter(
+    THPVariable* self,
+    const std::string& property_name) -> PyObject*;
+
+// Used for setters of Tensor properties.
+auto handle_torch_function_setter(
+    THPVariable* self,
+    const std::string& property_name,
+    PyObject* value) -> int;
+
+// Used for __getitem__ and __setitem__
+auto handle_torch_function_indexing(
+    PyObject* self,
+    PyObject* index,
+    PyObject* val = nullptr) -> PyObject*;
+
+/*
+ * Check if the input obj is Tensor type, including its subclass, or overloaded
+ * type. If the type defines __torch_function__, it also returns true.
+ * Otherwise returns false. If the class is not torch.Tensor, and it defines
+ * __torch_function__, we append obj to overloaded_args.
+ *
+ * 'obj': the input argument to be checked
+ * 'overloaded_args': the vector to append the overloaded args.
+ */
+bool is_tensor_and_append_overloaded(
+    PyObject* obj,
+    std::vector<PyObject*>* overloaded_args);
+
+/*
+ * Check if the input obj is Tensor List or Tensor Tuple type. First check
+ * whether obj is Tuple or List type, if true, iterate over each element and
+ * check whether it is Tensor type, including its subclass or overloaded type.
+ * At the same time, the overloaded arg is appended to the overloaded_args.
+ *
+ * 'obj': the input argument to be checked
+ * 'overloaded_args': the vector to append the overloaded args.
+ * 'argnum': the number of total arguments of the function being checked.
+ * 'throw_error': whether throw error if any element in the list or tuple is
+ *                not tensor type or overloaded.
+ */
+bool is_tensor_list_and_append_overloaded(
+    PyObject* obj,
+    std::vector<PyObject*>* overloaded_args,
+    size_t argnum,
+    bool throw_error);
+
+/* Given an argument that is definitely a tensor and is definitely overloaded,
+ * append it to the overloaded arguments list.  Use this instead of
+ * is_tensor_and_append_overloaded in situations where you have a PyObject
+ * and you know it definitely is a Tensor and it is definitely overloaded.
+ *
+ * 'overloaded_args': the vector to append the overloaded args
+ * 'obj': the input tensor that is overloaded
+ */
+void append_overloaded_tensor(
+    std::vector<PyObject*>* overloaded_args,
+    PyObject* obj);
+
+/* Given an argument that is definitely a type and is definitely overloaded,
+ * append it to the overloaded arguments list. Use this only with
+ * __torch_dispatch__, where we operate on classes that have a
+ * __torch_dispatch__ classmethod.
+ *
+ * 'overloaded_args': the vector to append the overloaded type
+ * 'obj': the input class that has a __torch_dispatch__ classmethod.
+ */
+void append_overloaded_type(
+    std::vector<PyObject*>* overloaded_args,
+    PyObject* obj);
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_compat.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_compat.h
new file mode 100644
index 0000000000000000000000000000000000000000..fe047faa9c15bf044dd80f1045a845b36c1878de
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_compat.h
@@ -0,0 +1,44 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef PYTHON_COMPAT
+#define PYTHON_COMPAT
+
+#include <torch/csrc/utils/pythoncapi_compat.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// PyTorch-only compat functions
+
+#define IS_PYTHON_3_11_PLUS (PY_VERSION_HEX >= 0x030B00C1)
+#define IS_PYTHON_3_12_PLUS (PY_VERSION_HEX >= 0x030C0000)
+#define IS_PYTHON_3_13_PLUS (PY_VERSION_HEX >= 0x030D0000)
+#define IS_PYTHON_3_14_PLUS (PY_VERSION_HEX >= 0x030E0000)
+#define IS_PYTHON_3_15_PLUS (PY_VERSION_HEX >= 0x030F0000)
+
+static inline int PyCode_GetNCellvars(PyCodeObject* code) {
+// gh-26364 added co_ncellvars to Python 3.11.0rc1
+#if IS_PYTHON_3_11_PLUS
+  return code->co_ncellvars;
+#else
+  return PyTuple_GET_SIZE(code->co_cellvars);
+#endif
+}
+
+static inline int PyCode_GetNFreevars(PyCodeObject* code) {
+// gh-26364 added co_nfreevars to Python 3.11.0rc1
+#if IS_PYTHON_3_11_PLUS
+  return code->co_nfreevars;
+#else
+  return PyTuple_GET_SIZE(code->co_freevars);
+#endif
+}
+
+#ifdef __cplusplus
+}
+#endif
+#endif // PYTHON_COMPAT
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_dispatch.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_dispatch.h
new file mode 100644
index 0000000000000000000000000000000000000000..3106c43acf60806568929f9e56338d7493ea32a7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_dispatch.h
@@ -0,0 +1,21 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <pybind11/pybind11.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch::impl::dispatch {
+
+void initDispatchBindings(PyObject* module);
+
+void python_op_registration_trampoline_impl(
+    const c10::OperatorHandle& op,
+    c10::DispatchKey key,
+    c10::DispatchKeySet keyset,
+    torch::jit::Stack* stack,
+    bool with_keyset,
+    bool with_op);
+
+} // namespace torch::impl::dispatch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_numbers.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_numbers.h
new file mode 100644
index 0000000000000000000000000000000000000000..371e793c508109be654ece27230f0603fab4f266
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_numbers.h
@@ -0,0 +1,232 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/Device.h>
+#include <torch/csrc/Exceptions.h>
+#include <torch/csrc/jit/frontend/tracer.h>
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/object_ptr.h>
+#include <torch/csrc/utils/tensor_numpy.h>
+#include <cstdint>
+#include <limits>
+#include <stdexcept>
+
+// largest integer that can be represented consecutively in a double
+const int64_t DOUBLE_INT_MAX = 9007199254740992;
+
+inline PyObject* THPUtils_packDeviceIndex(c10::DeviceIndex value) {
+  return PyLong_FromLong(value);
+}
+
+inline PyObject* THPUtils_packInt32(int32_t value) {
+  return PyLong_FromLong(value);
+}
+
+inline PyObject* THPUtils_packInt64(int64_t value) {
+  return PyLong_FromLongLong(value);
+}
+
+inline PyObject* THPUtils_packUInt32(uint32_t value) {
+  return PyLong_FromUnsignedLong(value);
+}
+
+inline PyObject* THPUtils_packUInt64(uint64_t value) {
+  return PyLong_FromUnsignedLongLong(value);
+}
+
+inline PyObject* THPUtils_packDoubleAsInt(double value) {
+  return PyLong_FromDouble(value);
+}
+
+inline bool THPUtils_checkLongExact(PyObject* obj) {
+  return PyLong_CheckExact(obj) && !PyBool_Check(obj);
+}
+
+inline bool THPUtils_checkLong(PyObject* obj) {
+  // Fast path
+  if (THPUtils_checkLongExact(obj)) {
+    return true;
+  }
+
+#ifdef USE_NUMPY
+  if (torch::utils::is_numpy_int(obj)) {
+    return true;
+  }
+#endif
+
+  return PyLong_Check(obj) && !PyBool_Check(obj);
+}
+
+inline int32_t THPUtils_unpackInt(PyObject* obj) {
+  int overflow = 0;
+  long value = PyLong_AsLongAndOverflow(obj, &overflow);
+  if (value == -1 && PyErr_Occurred()) {
+    throw python_error();
+  }
+  TORCH_CHECK_VALUE(overflow == 0, "Overflow when unpacking long long");
+  TORCH_CHECK_VALUE(
+      value <= std::numeric_limits<int32_t>::max() &&
+          value >= std::numeric_limits<int32_t>::min(),
+      "Overflow when unpacking long");
+  return (int32_t)value;
+}
+
+inline int64_t THPUtils_unpackLong(PyObject* obj) {
+  int overflow = 0;
+  long long value = PyLong_AsLongLongAndOverflow(obj, &overflow);
+  if (value == -1 && PyErr_Occurred()) {
+    throw python_error();
+  }
+  TORCH_CHECK_VALUE(overflow == 0, "Overflow when unpacking long long");
+  return (int64_t)value;
+}
+
+inline uint32_t THPUtils_unpackUInt32(PyObject* obj) {
+  unsigned long value = PyLong_AsUnsignedLong(obj);
+  if (PyErr_Occurred()) {
+    throw python_error();
+  }
+  TORCH_CHECK_VALUE(
+      value <= std::numeric_limits<uint32_t>::max(),
+      "Overflow when unpacking long long");
+  return (uint32_t)value;
+}
+
+inline uint64_t THPUtils_unpackUInt64(PyObject* obj) {
+  unsigned long long value = PyLong_AsUnsignedLongLong(obj);
+  if (PyErr_Occurred()) {
+    throw python_error();
+  }
+  return (uint64_t)value;
+}
+
+bool THPUtils_checkIndex(PyObject* obj);
+
+inline int64_t THPUtils_unpackIndex(PyObject* obj) {
+  if (!THPUtils_checkLong(obj)) {
+    auto index = THPObjectPtr(PyNumber_Index(obj));
+    if (index == nullptr) {
+      throw python_error();
+    }
+    // NB: This needs to be called before `index` goes out of scope and the
+    // underlying object's refcount is decremented
+    return THPUtils_unpackLong(index.get());
+  }
+  return THPUtils_unpackLong(obj);
+}
+
+inline bool THPUtils_unpackBool(PyObject* obj) {
+  if (obj == Py_True) {
+    return true;
+  } else if (obj == Py_False) {
+    return false;
+  } else {
+    TORCH_CHECK(false, "couldn't convert python object to boolean");
+  }
+}
+
+inline bool THPUtils_checkBool(PyObject* obj) {
+#ifdef USE_NUMPY
+  if (torch::utils::is_numpy_bool(obj)) {
+    return true;
+  }
+#endif
+  return PyBool_Check(obj);
+}
+
+inline bool THPUtils_checkDouble(PyObject* obj) {
+#ifdef USE_NUMPY
+  if (torch::utils::is_numpy_scalar(obj)) {
+    return true;
+  }
+#endif
+  return PyFloat_Check(obj) || PyLong_Check(obj);
+}
+
+inline double THPUtils_unpackDouble(PyObject* obj) {
+  if (PyFloat_Check(obj)) {
+    return PyFloat_AS_DOUBLE(obj);
+  }
+  double value = PyFloat_AsDouble(obj);
+  if (value == -1 && PyErr_Occurred()) {
+    throw python_error();
+  }
+  return value;
+}
+
+inline c10::complex<double> THPUtils_unpackComplexDouble(PyObject* obj) {
+  Py_complex value = PyComplex_AsCComplex(obj);
+  if (value.real == -1.0 && PyErr_Occurred()) {
+    throw python_error();
+  }
+
+  return c10::complex<double>(value.real, value.imag);
+}
+
+inline bool THPUtils_unpackNumberAsBool(PyObject* obj) {
+#ifdef USE_NUMPY
+  // Handle NumPy boolean scalars (np.bool_)
+  if (torch::utils::is_numpy_bool(obj)) {
+    int truth = PyObject_IsTrue(obj);
+    if (truth == -1) {
+      throw python_error();
+    }
+    return truth != 0;
+  }
+#endif
+  if (PyFloat_Check(obj)) {
+    return (bool)PyFloat_AS_DOUBLE(obj);
+  }
+
+  if (PyComplex_Check(obj)) {
+    double real_val = PyComplex_RealAsDouble(obj);
+    double imag_val = PyComplex_ImagAsDouble(obj);
+    return !(real_val == 0 && imag_val == 0);
+  }
+
+  int overflow = 0;
+  long long value = PyLong_AsLongLongAndOverflow(obj, &overflow);
+  if (value == -1 && PyErr_Occurred()) {
+    throw python_error();
+  }
+  // No need to check overflow, because when overflow occurred, it should
+  // return true in order to keep the same behavior of numpy.
+  return (bool)value;
+}
+
+inline c10::DeviceIndex THPUtils_unpackDeviceIndex(PyObject* obj) {
+  int overflow = 0;
+  long value = PyLong_AsLongAndOverflow(obj, &overflow);
+  if (value == -1 && PyErr_Occurred()) {
+    throw python_error();
+  }
+  TORCH_CHECK(overflow == 0, "Overflow when unpacking DeviceIndex");
+  TORCH_CHECK(
+      value <= std::numeric_limits<c10::DeviceIndex>::max() &&
+          value >= std::numeric_limits<c10::DeviceIndex>::min(),
+      "Overflow when unpacking DeviceIndex");
+  return (c10::DeviceIndex)value;
+}
+
+template <typename T>
+inline T THPUtils_unpackInteger(PyObject* obj) {
+  int overflow = -1;
+  const auto value = PyLong_AsLongLongAndOverflow(obj, &overflow);
+  if (value == -1 && PyErr_Occurred()) {
+    throw python_error();
+  }
+  if (!overflow) {
+    return static_cast<int64_t>(value);
+  }
+  // try unsigned
+  const auto uvalue = PyLong_AsUnsignedLongLong(obj);
+  if (uvalue == static_cast<std::decay_t<decltype(uvalue)>>(-1) &&
+      PyErr_Occurred()) {
+    throw python_error();
+  }
+  return static_cast<uint64_t>(uvalue);
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_raii.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_raii.h
new file mode 100644
index 0000000000000000000000000000000000000000..25d0b62bf6fa77adc2c19cfe7cb921de7830a90c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_raii.h
@@ -0,0 +1,89 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#include <torch/csrc/utils/pybind.h>
+#include <optional>
+#include <tuple>
+
+namespace torch::impl {
+
+template <typename GuardT, typename... Args>
+struct RAIIContextManager {
+  explicit RAIIContextManager(Args&&... args)
+      : args_(std::forward<Args>(args)...) {}
+
+  void enter() {
+    auto emplace = [&](Args... args) {
+      guard_.emplace(std::forward<Args>(args)...);
+    };
+    std::apply(std::move(emplace), args_);
+  }
+
+  void exit() {
+    guard_ = std::nullopt;
+  }
+
+ private:
+  std::optional<GuardT> guard_;
+  std::tuple<Args...> args_;
+};
+
+// Turns a C++ RAII guard into a Python context manager.
+// See _ExcludeDispatchKeyGuard in python_dispatch.cpp for example.
+template <typename GuardT, typename... GuardArgs>
+void py_context_manager(const py::module& m, const char* name) {
+  using ContextManagerT = RAIIContextManager<GuardT, GuardArgs...>;
+  py::class_<ContextManagerT>(m, name)
+      .def(py::init<GuardArgs...>())
+      .def("__enter__", [](ContextManagerT& guard) { guard.enter(); })
+      .def(
+          "__exit__",
+          [](ContextManagerT& guard,
+             const py::object& exc_type,
+             const py::object& exc_value,
+             const py::object& traceback) { guard.exit(); });
+}
+
+template <typename GuardT, typename... Args>
+struct DeprecatedRAIIContextManager {
+  explicit DeprecatedRAIIContextManager(Args&&... args) {
+    guard_.emplace(std::forward<Args>(args)...);
+  }
+
+  void enter() {}
+
+  void exit() {
+    guard_ = std::nullopt;
+  }
+
+ private:
+  std::optional<GuardT> guard_;
+  std::tuple<Args...> args_;
+};
+
+// Definition: a "Python RAII guard" is an object in Python that acquires
+// a resource on init and releases the resource on deletion.
+//
+// This API turns a C++ RAII guard into an object can be used either as a
+// Python context manager or as a "Python RAII guard".
+//
+// Please prefer `py_context_manager` to this API if you are binding a new
+// RAII guard into Python because "Python RAII guards" don't work as expected
+// in Python (Python makes no guarantees about when an object gets deleted)
+template <typename GuardT, typename... GuardArgs>
+void py_context_manager_DEPRECATED(const py::module& m, const char* name) {
+  using ContextManagerT = DeprecatedRAIIContextManager<GuardT, GuardArgs...>;
+  py::class_<ContextManagerT>(m, name)
+      .def(py::init<GuardArgs...>())
+      .def("__enter__", [](ContextManagerT& guard) { guard.enter(); })
+      .def(
+          "__exit__",
+          [](ContextManagerT& guard,
+             const py::object& exc_type,
+             const py::object& exc_value,
+             const py::object& traceback) { guard.exit(); });
+}
+
+} // namespace torch::impl
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_scalars.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_scalars.h
new file mode 100644
index 0000000000000000000000000000000000000000..29ac9daf6016667b4143153dd637163bafa19f31
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_scalars.h
@@ -0,0 +1,173 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/util/TypeCast.h>
+#include <torch/csrc/python_headers.h>
+
+#include <torch/csrc/Exceptions.h>
+#include <torch/csrc/utils/python_numbers.h>
+
+namespace torch::utils {
+
+template <typename T>
+inline T unpackIntegral(PyObject* obj, const char* type) {
+  // In Python-3.10 floats can no longer be silently converted to integers
+  // Keep backward compatible behavior for now
+  if (PyFloat_Check(obj)) {
+    return c10::checked_convert<T>(THPUtils_unpackDouble(obj), type);
+  }
+  return c10::checked_convert<T>(THPUtils_unpackLong(obj), type);
+}
+
+inline void store_scalar(void* data, at::ScalarType scalarType, PyObject* obj) {
+  switch (scalarType) {
+    case at::kByte:
+      *(uint8_t*)data = unpackIntegral<uint8_t>(obj, "uint8");
+      break;
+    case at::kUInt16:
+      *(uint16_t*)data = unpackIntegral<uint16_t>(obj, "uint16");
+      break;
+    case at::kUInt32:
+      *(uint32_t*)data = unpackIntegral<uint32_t>(obj, "uint32");
+      break;
+    case at::kUInt64:
+      // NB: This doesn't allow implicit conversion of float to int
+      *(uint64_t*)data = THPUtils_unpackUInt64(obj);
+      break;
+    case at::kChar:
+      *(int8_t*)data = unpackIntegral<int8_t>(obj, "int8");
+      break;
+    case at::kShort:
+      *(int16_t*)data = unpackIntegral<int16_t>(obj, "int16");
+      break;
+    case at::kInt:
+      *(int32_t*)data = unpackIntegral<int32_t>(obj, "int32");
+      break;
+    case at::kLong:
+      *(int64_t*)data = unpackIntegral<int64_t>(obj, "int64");
+      break;
+    case at::kHalf:
+      *(at::Half*)data =
+          at::convert<at::Half, double>(THPUtils_unpackDouble(obj));
+      break;
+    case at::kFloat:
+      *(float*)data = (float)THPUtils_unpackDouble(obj);
+      break;
+    case at::kDouble:
+      *(double*)data = THPUtils_unpackDouble(obj);
+      break;
+    case at::kComplexHalf:
+      *(c10::complex<at::Half>*)data =
+          (c10::complex<at::Half>)static_cast<c10::complex<float>>(
+              THPUtils_unpackComplexDouble(obj));
+      break;
+    case at::kComplexFloat:
+      *(c10::complex<float>*)data =
+          (c10::complex<float>)THPUtils_unpackComplexDouble(obj);
+      break;
+    case at::kComplexDouble:
+      *(c10::complex<double>*)data = THPUtils_unpackComplexDouble(obj);
+      break;
+    case at::kBool:
+      *(bool*)data = THPUtils_unpackNumberAsBool(obj);
+      break;
+    case at::kBFloat16:
+      *(at::BFloat16*)data =
+          at::convert<at::BFloat16, double>(THPUtils_unpackDouble(obj));
+      break;
+    // TODO(#146647): simplify below with macros
+    case at::kFloat8_e5m2:
+      *(at::Float8_e5m2*)data =
+          at::convert<at::Float8_e5m2, double>(THPUtils_unpackDouble(obj));
+      break;
+    case at::kFloat8_e5m2fnuz:
+      *(at::Float8_e5m2fnuz*)data =
+          at::convert<at::Float8_e5m2fnuz, double>(THPUtils_unpackDouble(obj));
+      break;
+    case at::kFloat8_e4m3fn:
+      *(at::Float8_e4m3fn*)data =
+          at::convert<at::Float8_e4m3fn, double>(THPUtils_unpackDouble(obj));
+      break;
+    case at::kFloat8_e4m3fnuz:
+      *(at::Float8_e4m3fnuz*)data =
+          at::convert<at::Float8_e4m3fnuz, double>(THPUtils_unpackDouble(obj));
+      break;
+    case at::kFloat8_e8m0fnu:
+      *(at::Float8_e8m0fnu*)data =
+          at::convert<at::Float8_e8m0fnu, double>(THPUtils_unpackDouble(obj));
+      break;
+    default:
+      TORCH_CHECK(false, "store_scalar: invalid type");
+  }
+}
+
+inline PyObject* load_scalar(const void* data, at::ScalarType scalarType) {
+  switch (scalarType) {
+    case at::kByte:
+      return THPUtils_packInt64(*(uint8_t*)data);
+    case at::kUInt16:
+      return THPUtils_packInt64(*(uint16_t*)data);
+    case at::kUInt32:
+      return THPUtils_packUInt32(*(uint32_t*)data);
+    case at::kUInt64:
+      return THPUtils_packUInt64(*(uint64_t*)data);
+    case at::kChar:
+      return THPUtils_packInt64(*(int8_t*)data);
+    case at::kShort:
+      return THPUtils_packInt64(*(int16_t*)data);
+    case at::kInt:
+      return THPUtils_packInt64(*(int32_t*)data);
+    case at::kLong:
+      return THPUtils_packInt64(*(int64_t*)data);
+    case at::kHalf:
+      return PyFloat_FromDouble(
+          at::convert<double, at::Half>(*(at::Half*)data));
+    case at::kFloat:
+      return PyFloat_FromDouble(*(float*)data);
+    case at::kDouble:
+      return PyFloat_FromDouble(*(double*)data);
+    case at::kComplexHalf: {
+      auto data_ = reinterpret_cast<const c10::complex<at::Half>*>(data);
+      return PyComplex_FromDoubles(data_->real(), data_->imag());
+    }
+    case at::kComplexFloat: {
+      auto data_ = reinterpret_cast<const c10::complex<float>*>(data);
+      return PyComplex_FromDoubles(data_->real(), data_->imag());
+    }
+    case at::kComplexDouble:
+      return PyComplex_FromCComplex(
+          *reinterpret_cast<Py_complex*>((c10::complex<double>*)data));
+    case at::kBool:
+      // Don't use bool*, since it may take out-of-range byte as bool.
+      // Instead, we cast explicitly to avoid ASAN error.
+      return PyBool_FromLong(static_cast<bool>(*(uint8_t*)data));
+    case at::kBFloat16:
+      return PyFloat_FromDouble(
+          at::convert<double, at::BFloat16>(*(at::BFloat16*)data));
+    // TODO(#146647): simplify below with macros
+    case at::kFloat8_e5m2:
+      return PyFloat_FromDouble(
+          at::convert<double, at::Float8_e5m2>(*(at::Float8_e5m2*)data));
+    case at::kFloat8_e4m3fn:
+      return PyFloat_FromDouble(
+          at::convert<double, at::Float8_e4m3fn>(*(at::Float8_e4m3fn*)data));
+    case at::kFloat8_e5m2fnuz:
+      return PyFloat_FromDouble(at::convert<double, at::Float8_e5m2fnuz>(
+          *(at::Float8_e5m2fnuz*)data));
+    case at::kFloat8_e4m3fnuz:
+      return PyFloat_FromDouble(at::convert<double, at::Float8_e4m3fnuz>(
+          *(at::Float8_e4m3fnuz*)data));
+    case at::kFloat8_e8m0fnu:
+      return PyFloat_FromDouble(
+          at::convert<double, at::Float8_e8m0fnu>(*(at::Float8_e8m0fnu*)data));
+    default:
+      TORCH_CHECK(false, "load_scalar: invalid type");
+  }
+}
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_strings.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_strings.h
new file mode 100644
index 0000000000000000000000000000000000000000..bd4f0c88607cc0ad66f35162efe4433a7eb2ee99
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_strings.h
@@ -0,0 +1,140 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/object_ptr.h>
+#include <torch/csrc/utils/pybind.h>
+#include <torch/csrc/utils/python_compat.h>
+#include <stdexcept>
+#include <string>
+
+// Utilities for handling Python strings. Note that PyString, when defined, is
+// the same as PyBytes.
+
+// Returns true if obj is a bytes/str or unicode object
+// As of Python 3.6, this does not require the GIL
+inline bool THPUtils_checkString(PyObject* obj) {
+  return PyBytes_Check(obj) || PyUnicode_Check(obj);
+}
+
+// Unpacks PyBytes (PyString) or PyUnicode as std::string
+// PyBytes are unpacked as-is. PyUnicode is unpacked as UTF-8.
+// NOTE: this method requires the GIL
+inline std::string THPUtils_unpackString(PyObject* obj) {
+  if (PyBytes_Check(obj)) {
+    size_t size = PyBytes_GET_SIZE(obj);
+    return std::string(PyBytes_AS_STRING(obj), size);
+  }
+  if (PyUnicode_Check(obj)) {
+    Py_ssize_t size = 0;
+    const char* data = PyUnicode_AsUTF8AndSize(obj, &size);
+    TORCH_CHECK(data, "error unpacking string as utf-8");
+    return std::string(data, (size_t)size);
+  }
+  TORCH_CHECK(false, "unpackString: expected bytes or unicode object");
+}
+
+// Unpacks PyBytes (PyString) or PyUnicode as std::string_view
+// PyBytes are unpacked as-is. PyUnicode is unpacked as UTF-8.
+// NOTE: If `obj` is destroyed, then the non-owning std::string_view will
+//   become invalid. If the string needs to be accessed at any point after
+//   `obj` is destroyed, then the std::string_view should be copied into
+//   a std::string, or another owning object, and kept alive. For an example,
+//   look at how IValue and autograd nodes handle std::string_view arguments.
+// NOTE: this method requires the GIL
+inline std::string_view THPUtils_unpackStringView(PyObject* obj) {
+  if (PyBytes_Check(obj)) {
+    size_t size = PyBytes_GET_SIZE(obj);
+    return std::string_view(PyBytes_AS_STRING(obj), size);
+  }
+  if (PyUnicode_Check(obj)) {
+    Py_ssize_t size = 0;
+    const char* data = PyUnicode_AsUTF8AndSize(obj, &size);
+    TORCH_CHECK(data, "error unpacking string as utf-8");
+    return std::string_view(data, (size_t)size);
+  }
+  TORCH_CHECK(false, "unpackString: expected bytes or unicode object");
+}
+
+inline PyObject* THPUtils_packString(const char* str) {
+  return PyUnicode_FromString(str);
+}
+
+inline PyObject* THPUtils_packString(const std::string& str) {
+  return PyUnicode_FromStringAndSize(
+      str.c_str(), static_cast<Py_ssize_t>(str.size()));
+}
+
+inline PyObject* THPUtils_internString(const std::string& str) {
+  return PyUnicode_InternFromString(str.c_str());
+}
+
+// Precondition: THPUtils_checkString(obj) must be true
+inline bool THPUtils_isInterned(PyObject* obj) {
+  return PyUnicode_CHECK_INTERNED(obj);
+}
+
+// Precondition: THPUtils_checkString(obj) must be true
+inline void THPUtils_internStringInPlace(PyObject** obj) {
+  PyUnicode_InternInPlace(obj);
+}
+
+/*
+ * Reference:
+ * https://github.com/numpy/numpy/blob/f4c497c768e0646df740b647782df463825bfd27/numpy/core/src/common/get_attr_string.h#L42
+ *
+ * Stripped down version of PyObject_GetAttrString,
+ * avoids lookups for None, tuple, and List objects,
+ * and doesn't create a PyErr since this code ignores it.
+ *
+ * This can be much faster then PyObject_GetAttrString where
+ * exceptions are not used by caller.
+ *
+ * 'obj' is the object to search for attribute.
+ *
+ * 'name' is the attribute to search for.
+ *
+ * Returns a py::object wrapping the return value. If the attribute lookup
+ * failed the value will be NULL.
+ *
+ */
+
+inline py::object PyObject_FastGetAttrString(PyObject* obj, const char* name) {
+#if IS_PYTHON_3_13_PLUS
+  PyObject* res = (PyObject*)nullptr;
+  int result_code = PyObject_GetOptionalAttrString(obj, name, &res);
+  if (result_code == -1) {
+    PyErr_Clear();
+  }
+  return py::reinterpret_steal<py::object>(res);
+#else
+  PyTypeObject* tp = Py_TYPE(obj);
+  PyObject* res = (PyObject*)nullptr;
+
+  /* Attribute referenced by (char *)name */
+  if (tp->tp_getattr != nullptr) {
+    // This is OK per https://bugs.python.org/issue39620
+    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
+    res = (*tp->tp_getattr)(obj, const_cast<char*>(name));
+    if (res == nullptr) {
+      PyErr_Clear();
+    }
+  }
+  /* Attribute referenced by (PyObject *)name */
+  else if (tp->tp_getattro != nullptr) {
+    auto w = py::reinterpret_steal<py::object>(PyUnicode_FromString(name));
+    if (w.ptr() == nullptr) {
+      return py::object();
+    }
+    res = (*tp->tp_getattro)(obj, w.ptr());
+    if (res == nullptr) {
+      PyErr_Clear();
+    }
+  }
+  return py::reinterpret_steal<py::object>(res);
+#endif
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_stub.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_stub.h
new file mode 100644
index 0000000000000000000000000000000000000000..f457be5949a775e9ce3f4b8b39d8c4bbe95985b8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_stub.h
@@ -0,0 +1,9 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+struct _object;
+using PyObject = _object;
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_symnode.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_symnode.h
new file mode 100644
index 0000000000000000000000000000000000000000..793f050d3c1ec3faaab1ee97c4e8b2c5e8f38192
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_symnode.h
@@ -0,0 +1,332 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/SafePyObject.h>
+#include <c10/core/SymNodeImpl.h>
+
+#include <torch/csrc/PyInterpreter.h>
+#include <torch/csrc/autograd/python_variable.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace torch {
+
+TORCH_PYTHON_API py::handle get_symint_class();
+TORCH_PYTHON_API py::handle get_symfloat_class();
+TORCH_PYTHON_API py::handle get_symbool_class();
+TORCH_PYTHON_API py::handle get_dynint_class();
+
+// NB: These functions must not be called too early, otherwise torch not setup.
+// Alternate design is to have torch "register" the object to us
+inline bool is_symint(py::handle obj) {
+  return py::isinstance(obj, get_symint_class());
+}
+inline bool is_symfloat(py::handle obj) {
+  return py::isinstance(obj, get_symfloat_class());
+}
+inline bool is_symbool(py::handle obj) {
+  return py::isinstance(obj, get_symbool_class());
+}
+inline bool is_dynint(py::handle obj) {
+  return py::isinstance(obj, get_dynint_class());
+}
+
+namespace impl {
+
+// This c10::SymNodeImpl simply backends to a Python object that
+// implements the API.   The Python object is the source of truth,
+// this is just an adapter so C++ calls can get to the object.
+class PythonSymNodeImpl : public c10::SymNodeImpl {
+ public:
+  PythonSymNodeImpl(py::object pyobj) : c10::SymNodeImpl() {
+    pyobj_ = std::make_shared<c10::SafePyObject>(
+        pyobj.release().ptr(), getPyInterpreter());
+  }
+
+  c10::SymNode wrap_int(int64_t num) override {
+    py::gil_scoped_acquire acquire;
+    auto r = getPyObj().attr("wrap_int")(num);
+    return c10::make_intrusive<PythonSymNodeImpl>(std::move(r));
+  }
+
+  c10::SymNode wrap_float(double num) override {
+    py::gil_scoped_acquire acquire;
+    auto r = getPyObj().attr("wrap_float")(num);
+    return c10::make_intrusive<PythonSymNodeImpl>(std::move(r));
+  }
+
+  c10::SymNode wrap_bool(bool num) override {
+    py::gil_scoped_acquire acquire;
+    auto r = getPyObj().attr("wrap_bool")(num);
+    return c10::make_intrusive<PythonSymNodeImpl>(std::move(r));
+  }
+
+#define TORCH_SYMNODE_SIZES_STRIDES(n)                                        \
+  c10::SymNode n(                                                             \
+      c10::ArrayRef<c10::SymNode> sizes, c10::ArrayRef<c10::SymNode> strides) \
+      override {                                                              \
+    py::gil_scoped_acquire acquire;                                           \
+    auto r = getPyObj().attr(#n)(sizes, strides);                             \
+    return c10::make_intrusive<PythonSymNodeImpl>(std::move(r));              \
+  }
+
+  // clang-format off
+    TORCH_SYMNODE_SIZES_STRIDES(is_contiguous)
+    TORCH_SYMNODE_SIZES_STRIDES(is_channels_last_contiguous_2d)
+    TORCH_SYMNODE_SIZES_STRIDES(is_channels_last_contiguous_3d)
+    TORCH_SYMNODE_SIZES_STRIDES(is_channels_last_strides_2d)
+    TORCH_SYMNODE_SIZES_STRIDES(is_channels_last_strides_3d)
+    TORCH_SYMNODE_SIZES_STRIDES(is_non_overlapping_and_dense)
+  // clang-format on
+
+#undef TORCH_SYMNODE_SIZES_STRIDES
+
+  bool bool_() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("bool_")().is(py::handle(Py_True));
+  }
+
+  bool is_int() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("is_int")().is(py::handle(Py_True));
+  }
+
+  bool is_float() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("is_float")().is(py::handle(Py_True));
+  }
+
+  bool is_bool() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("is_bool")().is(py::handle(Py_True));
+  }
+
+  bool is_nested_int() const override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("is_nested_int")().is(py::handle(Py_True));
+  }
+
+  bool has_hint() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("has_hint")().is(py::handle(Py_True));
+  }
+
+  int64_t guard_int(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("guard_int")(file, line).cast<int64_t>();
+  }
+
+  double guard_float(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("guard_float")(file, line).cast<double>();
+  }
+
+  bool guard_bool(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("guard_bool")(file, line).cast<bool>();
+  }
+
+  bool expect_true(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("expect_true")(file, line).cast<bool>();
+  }
+
+  bool guard_size_oblivious(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("guard_size_oblivious")(file, line).cast<bool>();
+  }
+
+  bool guard_or_false(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("guard_or_false")(file, line).cast<bool>();
+  }
+
+  bool statically_known_true(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("statically_known_true")(file, line).cast<bool>();
+  }
+
+  bool guard_or_true(const char* file, int64_t line) override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("guard_or_true")(file, line).cast<bool>();
+  }
+
+  int64_t int_() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("int_")().cast<int64_t>();
+  }
+
+  std::optional<int64_t> maybe_as_int() override {
+    py::gil_scoped_acquire acquire;
+    const auto& r = getPyObj().attr("maybe_as_int")();
+    if (r.is_none()) {
+      return std::nullopt;
+    } else {
+      return r.cast<int64_t>();
+    }
+  }
+
+  std::string str() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("str")().cast<std::string>();
+  }
+
+  std::string _graph_repr() override {
+    py::gil_scoped_acquire acquire;
+    return getPyObj().attr("_graph_repr")().cast<std::string>();
+  }
+
+  c10::SymNode dispatch_sym_ite_(
+      const char* fname,
+      const c10::SymNode& other,
+      const c10::SymNode& third) {
+    auto pother = dynamic_cast<PythonSymNodeImpl*>(other.get());
+    auto pthird = dynamic_cast<PythonSymNodeImpl*>(third.get());
+    TORCH_CHECK(pother);
+    TORCH_CHECK(pthird);
+    py::gil_scoped_acquire acquire;
+    auto r = getPyObj().attr(fname)(pother->getPyObj(), pthird->getPyObj());
+    return c10::make_intrusive<PythonSymNodeImpl>(r);
+  }
+
+  c10::SymNode dispatch_common_(const char* fname, const c10::SymNode& other) {
+    auto pother = dynamic_cast<PythonSymNodeImpl*>(other.get());
+    TORCH_CHECK(pother);
+    py::gil_scoped_acquire acquire;
+    auto r = getPyObj().attr(fname)(pother->getPyObj());
+    return c10::make_intrusive<PythonSymNodeImpl>(r);
+  }
+
+  c10::SymNode dispatch_common_(const char* fname) {
+    py::gil_scoped_acquire acquire;
+    auto r = getPyObj().attr(fname)();
+    return c10::make_intrusive<PythonSymNodeImpl>(r);
+  }
+
+  c10::SymNode add(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode sub(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode mul(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode truediv(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode float_truediv(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode int_truediv(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode pow(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode float_pow(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode pow_by_natural(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode floordiv(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode int_floordiv(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode mod(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode eq(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode ne(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode gt(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode lt(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode le(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode ge(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode sym_min(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+  c10::SymNode sym_max(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode sym_and(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode sym_or(const c10::SymNode& other) override {
+    return dispatch_common_(__func__, other);
+  }
+
+  c10::SymNode sym_ite(const c10::SymNode& other, const c10::SymNode& third)
+      override {
+    return dispatch_sym_ite_(__func__, other, third);
+  }
+
+  c10::SymNode sym_not() override {
+    return dispatch_common_(__func__);
+  }
+
+  c10::SymNode ceil() override {
+    return dispatch_common_(__func__);
+  }
+
+  c10::SymNode floor() override {
+    return dispatch_common_(__func__);
+  }
+
+  c10::SymNode neg() override {
+    return dispatch_common_(__func__);
+  }
+
+  c10::SymNode clone() override {
+    return dispatch_common_(__func__);
+  }
+
+  c10::SymNode sym_float() override {
+    return dispatch_common_(__func__);
+  }
+
+  py::handle getPyObj() const {
+    return py::handle(pyobj_->ptr(getPyInterpreter()));
+  }
+  std::shared_ptr<c10::SafePyObject> pyobj_ = nullptr;
+};
+
+} // namespace impl
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_torch_function_mode.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_torch_function_mode.h
new file mode 100644
index 0000000000000000000000000000000000000000..32fbbeb84b814e0c13d1203efa80d4d5362d14eb
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_torch_function_mode.h
@@ -0,0 +1,34 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/PythonTorchFunctionTLS.h>
+
+namespace torch::overrides {
+
+struct StashTorchFunctionModeGuard {
+  StashTorchFunctionModeGuard() {
+    cur_mode_ = at::impl::PythonTorchFunctionTLS::pop_stack();
+  }
+  ~StashTorchFunctionModeGuard() {
+    at::impl::PythonTorchFunctionTLS::push_onto_stack(cur_mode_);
+  }
+  StashTorchFunctionModeGuard(const StashTorchFunctionModeGuard&) = delete;
+  StashTorchFunctionModeGuard(StashTorchFunctionModeGuard&&) = delete;
+  StashTorchFunctionModeGuard& operator=(const StashTorchFunctionModeGuard&) =
+      delete;
+  StashTorchFunctionModeGuard& operator=(StashTorchFunctionModeGuard&&) =
+      delete;
+
+  const std::shared_ptr<c10::SafePyObject>& get_cur_mode() {
+    return cur_mode_;
+  }
+
+ private:
+  std::shared_ptr<c10::SafePyObject> cur_mode_;
+};
+
+} // namespace torch::overrides
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_tuples.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_tuples.h
new file mode 100644
index 0000000000000000000000000000000000000000..8a35c34d5d9ed11af34d2871b906a1f2582ff47c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/python_tuples.h
@@ -0,0 +1,32 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Exceptions.h>
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/object_ptr.h>
+#include <torch/csrc/utils/python_numbers.h>
+
+inline void THPUtils_packInt64Array(
+    PyObject* tuple,
+    size_t size,
+    const int64_t* sizes) {
+  for (size_t i = 0; i != size; ++i) {
+    PyObject* i64 = THPUtils_packInt64(sizes[i]);
+    if (!i64) {
+      throw python_error();
+    }
+    PyTuple_SET_ITEM(tuple, i, i64);
+  }
+}
+
+inline PyObject* THPUtils_packInt64Array(size_t size, const int64_t* sizes) {
+  THPObjectPtr tuple(PyTuple_New(static_cast<Py_ssize_t>(size)));
+  if (!tuple)
+    throw python_error();
+  THPUtils_packInt64Array(tuple.get(), size, sizes);
+  return tuple.release();
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pythoncapi_compat.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pythoncapi_compat.h
new file mode 100644
index 0000000000000000000000000000000000000000..2d107ad578a7850c2a7ceb5448687ce8361743a1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/pythoncapi_compat.h
@@ -0,0 +1,2671 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// Header file providing new C API functions to old Python versions.
+//
+// File distributed under the Zero Clause BSD (0BSD) license.
+// Copyright Contributors to the pythoncapi_compat project.
+//
+// Homepage:
+// https://github.com/python/pythoncapi_compat
+//
+// Latest version:
+// https://raw.githubusercontent.com/python/pythoncapi-compat/main/pythoncapi_compat.h
+//
+// SPDX-License-Identifier: 0BSD
+
+#ifndef PYTHONCAPI_COMPAT
+#define PYTHONCAPI_COMPAT
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <Python.h>
+#include <stddef.h>               // offsetof()
+
+// Python 3.11.0b4 added PyFrame_Back() to Python.h
+#if PY_VERSION_HEX < 0x030b00B4 && !defined(PYPY_VERSION)
+#  include "frameobject.h"        // PyFrameObject, PyFrame_GetBack()
+#endif
+
+
+#ifndef _Py_CAST
+#  define _Py_CAST(type, expr) ((type)(expr))
+#endif
+
+// Static inline functions should use _Py_NULL rather than using directly NULL
+// to prevent C++ compiler warnings. On C23 and newer and on C++11 and newer,
+// _Py_NULL is defined as nullptr.
+#ifndef _Py_NULL
+#  if (defined (__STDC_VERSION__) && __STDC_VERSION__ > 201710L) \
+          || (defined(__cplusplus) && __cplusplus >= 201103)
+#    define _Py_NULL nullptr
+#  else
+#    define _Py_NULL NULL
+#  endif
+#endif
+
+// Cast argument to PyObject* type.
+#ifndef _PyObject_CAST
+#  define _PyObject_CAST(op) _Py_CAST(PyObject*, op)
+#endif
+
+#ifndef Py_BUILD_ASSERT
+#  define Py_BUILD_ASSERT(cond) \
+        do { \
+            (void)sizeof(char [1 - 2 * !(cond)]); \
+        } while(0)
+#endif
+
+
+// bpo-42262 added Py_NewRef() to Python 3.10.0a3
+#if PY_VERSION_HEX < 0x030A00A3 && !defined(Py_NewRef)
+static inline PyObject* _Py_NewRef(PyObject *obj)
+{
+    Py_INCREF(obj);
+    return obj;
+}
+#define Py_NewRef(obj) _Py_NewRef(_PyObject_CAST(obj))
+#endif
+
+
+// bpo-42262 added Py_XNewRef() to Python 3.10.0a3
+#if PY_VERSION_HEX < 0x030A00A3 && !defined(Py_XNewRef)
+static inline PyObject* _Py_XNewRef(PyObject *obj)
+{
+    Py_XINCREF(obj);
+    return obj;
+}
+#define Py_XNewRef(obj) _Py_XNewRef(_PyObject_CAST(obj))
+#endif
+
+
+// bpo-39573 added Py_SET_REFCNT() to Python 3.9.0a4
+#if PY_VERSION_HEX < 0x030900A4 && !defined(Py_SET_REFCNT)
+static inline void _Py_SET_REFCNT(PyObject *ob, Py_ssize_t refcnt)
+{
+    ob->ob_refcnt = refcnt;
+}
+#define Py_SET_REFCNT(ob, refcnt) _Py_SET_REFCNT(_PyObject_CAST(ob), refcnt)
+#endif
+
+
+// Py_SETREF() and Py_XSETREF() were added to Python 3.5.2.
+// It is excluded from the limited C API.
+#if (PY_VERSION_HEX < 0x03050200 && !defined(Py_SETREF)) && !defined(Py_LIMITED_API)
+#define Py_SETREF(dst, src)                                     \
+    do {                                                        \
+        PyObject **_tmp_dst_ptr = _Py_CAST(PyObject**, &(dst)); \
+        PyObject *_tmp_dst = (*_tmp_dst_ptr);                   \
+        *_tmp_dst_ptr = _PyObject_CAST(src);                    \
+        Py_DECREF(_tmp_dst);                                    \
+    } while (0)
+
+#define Py_XSETREF(dst, src)                                    \
+    do {                                                        \
+        PyObject **_tmp_dst_ptr = _Py_CAST(PyObject**, &(dst)); \
+        PyObject *_tmp_dst = (*_tmp_dst_ptr);                   \
+        *_tmp_dst_ptr = _PyObject_CAST(src);                    \
+        Py_XDECREF(_tmp_dst);                                   \
+    } while (0)
+#endif
+
+
+// bpo-43753 added Py_Is(), Py_IsNone(), Py_IsTrue() and Py_IsFalse()
+// to Python 3.10.0b1.
+#if PY_VERSION_HEX < 0x030A00B1 && !defined(Py_Is)
+#  define Py_Is(x, y) ((x) == (y))
+#endif
+#if PY_VERSION_HEX < 0x030A00B1 && !defined(Py_IsNone)
+#  define Py_IsNone(x) Py_Is(x, Py_None)
+#endif
+#if (PY_VERSION_HEX < 0x030A00B1 || defined(PYPY_VERSION)) && !defined(Py_IsTrue)
+#  define Py_IsTrue(x) Py_Is(x, Py_True)
+#endif
+#if (PY_VERSION_HEX < 0x030A00B1 || defined(PYPY_VERSION)) && !defined(Py_IsFalse)
+#  define Py_IsFalse(x) Py_Is(x, Py_False)
+#endif
+
+
+// bpo-39573 added Py_SET_TYPE() to Python 3.9.0a4
+#if PY_VERSION_HEX < 0x030900A4 && !defined(Py_SET_TYPE)
+static inline void _Py_SET_TYPE(PyObject *ob, PyTypeObject *type)
+{
+    ob->ob_type = type;
+}
+#define Py_SET_TYPE(ob, type) _Py_SET_TYPE(_PyObject_CAST(ob), type)
+#endif
+
+
+// bpo-39573 added Py_SET_SIZE() to Python 3.9.0a4
+#if PY_VERSION_HEX < 0x030900A4 && !defined(Py_SET_SIZE)
+static inline void _Py_SET_SIZE(PyVarObject *ob, Py_ssize_t size)
+{
+    ob->ob_size = size;
+}
+#define Py_SET_SIZE(ob, size) _Py_SET_SIZE((PyVarObject*)(ob), size)
+#endif
+
+
+// bpo-40421 added PyFrame_GetCode() to Python 3.9.0b1
+#if PY_VERSION_HEX < 0x030900B1 || defined(PYPY_VERSION)
+static inline PyCodeObject* PyFrame_GetCode(PyFrameObject *frame)
+{
+    assert(frame != _Py_NULL);
+    assert(frame->f_code != _Py_NULL);
+    return _Py_CAST(PyCodeObject*, Py_NewRef(frame->f_code));
+}
+#endif
+
+static inline PyCodeObject* _PyFrame_GetCodeBorrow(PyFrameObject *frame)
+{
+    PyCodeObject *code = PyFrame_GetCode(frame);
+    Py_DECREF(code);
+    return code;
+}
+
+
+// bpo-40421 added PyFrame_GetBack() to Python 3.9.0b1
+#if PY_VERSION_HEX < 0x030900B1 && !defined(PYPY_VERSION)
+static inline PyFrameObject* PyFrame_GetBack(PyFrameObject *frame)
+{
+    assert(frame != _Py_NULL);
+    return _Py_CAST(PyFrameObject*, Py_XNewRef(frame->f_back));
+}
+#endif
+
+#if !defined(PYPY_VERSION)
+static inline PyFrameObject* _PyFrame_GetBackBorrow(PyFrameObject *frame)
+{
+    PyFrameObject *back = PyFrame_GetBack(frame);
+    Py_XDECREF(back);
+    return back;
+}
+#endif
+
+
+// bpo-40421 added PyFrame_GetLocals() to Python 3.11.0a7
+#if PY_VERSION_HEX < 0x030B00A7 && !defined(PYPY_VERSION)
+static inline PyObject* PyFrame_GetLocals(PyFrameObject *frame)
+{
+#if PY_VERSION_HEX >= 0x030400B1
+    if (PyFrame_FastToLocalsWithError(frame) < 0) {
+        return NULL;
+    }
+#else
+    PyFrame_FastToLocals(frame);
+#endif
+    return Py_NewRef(frame->f_locals);
+}
+#endif
+
+
+// bpo-40421 added PyFrame_GetGlobals() to Python 3.11.0a7
+#if PY_VERSION_HEX < 0x030B00A7 && !defined(PYPY_VERSION)
+static inline PyObject* PyFrame_GetGlobals(PyFrameObject *frame)
+{
+    return Py_NewRef(frame->f_globals);
+}
+#endif
+
+
+// bpo-40421 added PyFrame_GetBuiltins() to Python 3.11.0a7
+#if PY_VERSION_HEX < 0x030B00A7 && !defined(PYPY_VERSION)
+static inline PyObject* PyFrame_GetBuiltins(PyFrameObject *frame)
+{
+    return Py_NewRef(frame->f_builtins);
+}
+#endif
+
+
+// bpo-40421 added PyFrame_GetLasti() to Python 3.11.0b1
+#if PY_VERSION_HEX < 0x030B00B1 && !defined(PYPY_VERSION)
+static inline int PyFrame_GetLasti(PyFrameObject *frame)
+{
+#if PY_VERSION_HEX >= 0x030A00A7
+    // bpo-27129: Since Python 3.10.0a7, f_lasti is an instruction offset,
+    // not a bytes offset anymore. Python uses 16-bit "wordcode" (2 bytes)
+    // instructions.
+    if (frame->f_lasti < 0) {
+        return -1;
+    }
+    return frame->f_lasti * 2;
+#else
+    return frame->f_lasti;
+#endif
+}
+#endif
+
+
+// gh-91248 added PyFrame_GetVar() to Python 3.12.0a2
+#if PY_VERSION_HEX < 0x030C00A2 && !defined(PYPY_VERSION)
+static inline PyObject* PyFrame_GetVar(PyFrameObject *frame, PyObject *name)
+{
+    PyObject *locals, *value;
+
+    locals = PyFrame_GetLocals(frame);
+    if (locals == NULL) {
+        return NULL;
+    }
+#if PY_VERSION_HEX >= 0x03000000
+    value = PyDict_GetItemWithError(locals, name);
+#else
+    value = _PyDict_GetItemWithError(locals, name);
+#endif
+    Py_DECREF(locals);
+
+    if (value == NULL) {
+        if (PyErr_Occurred()) {
+            return NULL;
+        }
+#if PY_VERSION_HEX >= 0x03000000
+        PyErr_Format(PyExc_NameError, "variable %R does not exist", name);
+#else
+        PyErr_SetString(PyExc_NameError, "variable does not exist");
+#endif
+        return NULL;
+    }
+    return Py_NewRef(value);
+}
+#endif
+
+
+// gh-91248 added PyFrame_GetVarString() to Python 3.12.0a2
+#if PY_VERSION_HEX < 0x030C00A2 && !defined(PYPY_VERSION)
+static inline PyObject*
+PyFrame_GetVarString(PyFrameObject *frame, const char *name)
+{
+    PyObject *name_obj, *value;
+#if PY_VERSION_HEX >= 0x03000000
+    name_obj = PyUnicode_FromString(name);
+#else
+    name_obj = PyString_FromString(name);
+#endif
+    if (name_obj == NULL) {
+        return NULL;
+    }
+    value = PyFrame_GetVar(frame, name_obj);
+    Py_DECREF(name_obj);
+    return value;
+}
+#endif
+
+
+// bpo-39947 added PyThreadState_GetInterpreter() to Python 3.9.0a5
+#if PY_VERSION_HEX < 0x030900A5 || (defined(PYPY_VERSION) && PY_VERSION_HEX < 0x030B0000)
+static inline PyInterpreterState *
+PyThreadState_GetInterpreter(PyThreadState *tstate)
+{
+    assert(tstate != _Py_NULL);
+    return tstate->interp;
+}
+#endif
+
+
+// bpo-40429 added PyThreadState_GetFrame() to Python 3.9.0b1
+#if PY_VERSION_HEX < 0x030900B1 && !defined(PYPY_VERSION)
+static inline PyFrameObject* PyThreadState_GetFrame(PyThreadState *tstate)
+{
+    assert(tstate != _Py_NULL);
+    return _Py_CAST(PyFrameObject *, Py_XNewRef(tstate->frame));
+}
+#endif
+
+#if !defined(PYPY_VERSION)
+static inline PyFrameObject*
+_PyThreadState_GetFrameBorrow(PyThreadState *tstate)
+{
+    PyFrameObject *frame = PyThreadState_GetFrame(tstate);
+    Py_XDECREF(frame);
+    return frame;
+}
+#endif
+
+
+// bpo-39947 added PyInterpreterState_Get() to Python 3.9.0a5
+#if PY_VERSION_HEX < 0x030900A5 || defined(PYPY_VERSION)
+static inline PyInterpreterState* PyInterpreterState_Get(void)
+{
+    PyThreadState *tstate;
+    PyInterpreterState *interp;
+
+    tstate = PyThreadState_GET();
+    if (tstate == _Py_NULL) {
+        Py_FatalError("GIL released (tstate is NULL)");
+    }
+    interp = tstate->interp;
+    if (interp == _Py_NULL) {
+        Py_FatalError("no current interpreter");
+    }
+    return interp;
+}
+#endif
+
+
+// bpo-39947 added PyInterpreterState_Get() to Python 3.9.0a6
+#if 0x030700A1 <= PY_VERSION_HEX && PY_VERSION_HEX < 0x030900A6 && !defined(PYPY_VERSION)
+static inline uint64_t PyThreadState_GetID(PyThreadState *tstate)
+{
+    assert(tstate != _Py_NULL);
+    return tstate->id;
+}
+#endif
+
+// bpo-43760 added PyThreadState_EnterTracing() to Python 3.11.0a2
+#if PY_VERSION_HEX < 0x030B00A2 && !defined(PYPY_VERSION)
+static inline void PyThreadState_EnterTracing(PyThreadState *tstate)
+{
+    tstate->tracing++;
+#if PY_VERSION_HEX >= 0x030A00A1
+    tstate->cframe->use_tracing = 0;
+#else
+    tstate->use_tracing = 0;
+#endif
+}
+#endif
+
+// bpo-43760 added PyThreadState_LeaveTracing() to Python 3.11.0a2
+#if PY_VERSION_HEX < 0x030B00A2 && !defined(PYPY_VERSION)
+static inline void PyThreadState_LeaveTracing(PyThreadState *tstate)
+{
+    int use_tracing = (tstate->c_tracefunc != _Py_NULL
+                       || tstate->c_profilefunc != _Py_NULL);
+    tstate->tracing--;
+#if PY_VERSION_HEX >= 0x030A00A1
+    tstate->cframe->use_tracing = use_tracing;
+#else
+    tstate->use_tracing = use_tracing;
+#endif
+}
+#endif
+
+
+// bpo-37194 added PyObject_CallNoArgs() to Python 3.9.0a1
+// PyObject_CallNoArgs() added to PyPy 3.9.16-v7.3.11
+#if !defined(PyObject_CallNoArgs) && PY_VERSION_HEX < 0x030900A1
+static inline PyObject* PyObject_CallNoArgs(PyObject *func)
+{
+    return PyObject_CallFunctionObjArgs(func, NULL);
+}
+#endif
+
+
+// bpo-39245 made PyObject_CallOneArg() public (previously called
+// _PyObject_CallOneArg) in Python 3.9.0a4
+// PyObject_CallOneArg() added to PyPy 3.9.16-v7.3.11
+#if !defined(PyObject_CallOneArg) && PY_VERSION_HEX < 0x030900A4
+static inline PyObject* PyObject_CallOneArg(PyObject *func, PyObject *arg)
+{
+    return PyObject_CallFunctionObjArgs(func, arg, NULL);
+}
+#endif
+
+
+// bpo-1635741 added PyModule_AddObjectRef() to Python 3.10.0a3
+#if PY_VERSION_HEX < 0x030A00A3
+static inline int
+PyModule_AddObjectRef(PyObject *module, const char *name, PyObject *value)
+{
+    int res;
+
+    if (!value && !PyErr_Occurred()) {
+        // PyModule_AddObject() raises TypeError in this case
+        PyErr_SetString(PyExc_SystemError,
+                        "PyModule_AddObjectRef() must be called "
+                        "with an exception raised if value is NULL");
+        return -1;
+    }
+
+    Py_XINCREF(value);
+    res = PyModule_AddObject(module, name, value);
+    if (res < 0) {
+        Py_XDECREF(value);
+    }
+    return res;
+}
+#endif
+
+
+// bpo-40024 added PyModule_AddType() to Python 3.9.0a5
+#if PY_VERSION_HEX < 0x030900A5
+static inline int PyModule_AddType(PyObject *module, PyTypeObject *type)
+{
+    const char *name, *dot;
+
+    if (PyType_Ready(type) < 0) {
+        return -1;
+    }
+
+    // inline _PyType_Name()
+    name = type->tp_name;
+    assert(name != _Py_NULL);
+    dot = strrchr(name, '.');
+    if (dot != _Py_NULL) {
+        name = dot + 1;
+    }
+
+    return PyModule_AddObjectRef(module, name, _PyObject_CAST(type));
+}
+#endif
+
+
+// bpo-40241 added PyObject_GC_IsTracked() to Python 3.9.0a6.
+// bpo-4688 added _PyObject_GC_IS_TRACKED() to Python 2.7.0a2.
+#if PY_VERSION_HEX < 0x030900A6 && !defined(PYPY_VERSION)
+static inline int PyObject_GC_IsTracked(PyObject* obj)
+{
+    return (PyObject_IS_GC(obj) && _PyObject_GC_IS_TRACKED(obj));
+}
+#endif
+
+// bpo-40241 added PyObject_GC_IsFinalized() to Python 3.9.0a6.
+// bpo-18112 added _PyGCHead_FINALIZED() to Python 3.4.0 final.
+#if PY_VERSION_HEX < 0x030900A6 && PY_VERSION_HEX >= 0x030400F0 && !defined(PYPY_VERSION)
+static inline int PyObject_GC_IsFinalized(PyObject *obj)
+{
+    PyGC_Head *gc = _Py_CAST(PyGC_Head*, obj) - 1;
+    return (PyObject_IS_GC(obj) && _PyGCHead_FINALIZED(gc));
+}
+#endif
+
+
+// bpo-39573 added Py_IS_TYPE() to Python 3.9.0a4
+#if PY_VERSION_HEX < 0x030900A4 && !defined(Py_IS_TYPE)
+static inline int _Py_IS_TYPE(PyObject *ob, PyTypeObject *type) {
+    return Py_TYPE(ob) == type;
+}
+#define Py_IS_TYPE(ob, type) _Py_IS_TYPE(_PyObject_CAST(ob), type)
+#endif
+
+
+// bpo-46906 added PyFloat_Pack2() and PyFloat_Unpack2() to Python 3.11a7.
+// bpo-11734 added _PyFloat_Pack2() and _PyFloat_Unpack2() to Python 3.6.0b1.
+// Python 3.11a2 moved _PyFloat_Pack2() and _PyFloat_Unpack2() to the internal
+// C API: Python 3.11a2-3.11a6 versions are not supported.
+#if 0x030600B1 <= PY_VERSION_HEX && PY_VERSION_HEX <= 0x030B00A1 && !defined(PYPY_VERSION)
+static inline int PyFloat_Pack2(double x, char *p, int le)
+{ return _PyFloat_Pack2(x, (unsigned char*)p, le); }
+
+static inline double PyFloat_Unpack2(const char *p, int le)
+{ return _PyFloat_Unpack2((const unsigned char *)p, le); }
+#endif
+
+
+// bpo-46906 added PyFloat_Pack4(), PyFloat_Pack8(), PyFloat_Unpack4() and
+// PyFloat_Unpack8() to Python 3.11a7.
+// Python 3.11a2 moved _PyFloat_Pack4(), _PyFloat_Pack8(), _PyFloat_Unpack4()
+// and _PyFloat_Unpack8() to the internal C API: Python 3.11a2-3.11a6 versions
+// are not supported.
+#if PY_VERSION_HEX <= 0x030B00A1 && !defined(PYPY_VERSION)
+static inline int PyFloat_Pack4(double x, char *p, int le)
+{ return _PyFloat_Pack4(x, (unsigned char*)p, le); }
+
+static inline int PyFloat_Pack8(double x, char *p, int le)
+{ return _PyFloat_Pack8(x, (unsigned char*)p, le); }
+
+static inline double PyFloat_Unpack4(const char *p, int le)
+{ return _PyFloat_Unpack4((const unsigned char *)p, le); }
+
+static inline double PyFloat_Unpack8(const char *p, int le)
+{ return _PyFloat_Unpack8((const unsigned char *)p, le); }
+#endif
+
+
+// gh-92154 added PyCode_GetCode() to Python 3.11.0b1
+#if PY_VERSION_HEX < 0x030B00B1 && !defined(PYPY_VERSION)
+static inline PyObject* PyCode_GetCode(PyCodeObject *code)
+{
+    return Py_NewRef(code->co_code);
+}
+#endif
+
+
+// gh-95008 added PyCode_GetVarnames() to Python 3.11.0rc1
+#if PY_VERSION_HEX < 0x030B00C1 && !defined(PYPY_VERSION)
+static inline PyObject* PyCode_GetVarnames(PyCodeObject *code)
+{
+    return Py_NewRef(code->co_varnames);
+}
+#endif
+
+// gh-95008 added PyCode_GetFreevars() to Python 3.11.0rc1
+#if PY_VERSION_HEX < 0x030B00C1 && !defined(PYPY_VERSION)
+static inline PyObject* PyCode_GetFreevars(PyCodeObject *code)
+{
+    return Py_NewRef(code->co_freevars);
+}
+#endif
+
+// gh-95008 added PyCode_GetCellvars() to Python 3.11.0rc1
+#if PY_VERSION_HEX < 0x030B00C1 && !defined(PYPY_VERSION)
+static inline PyObject* PyCode_GetCellvars(PyCodeObject *code)
+{
+    return Py_NewRef(code->co_cellvars);
+}
+#endif
+
+
+// Py_UNUSED() was added to Python 3.4.0b2.
+#if PY_VERSION_HEX < 0x030400B2 && !defined(Py_UNUSED)
+#  if defined(__GNUC__) || defined(__clang__)
+#    define Py_UNUSED(name) _unused_ ## name __attribute__((unused))
+#  else
+#    define Py_UNUSED(name) _unused_ ## name
+#  endif
+#endif
+
+
+// gh-105922 added PyImport_AddModuleRef() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A0
+static inline PyObject* PyImport_AddModuleRef(const char *name)
+{
+    return Py_XNewRef(PyImport_AddModule(name));
+}
+#endif
+
+
+// gh-105927 added PyWeakref_GetRef() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D0000
+static inline int PyWeakref_GetRef(PyObject *ref, PyObject **pobj)
+{
+    PyObject *obj;
+    if (ref != NULL && !PyWeakref_Check(ref)) {
+        *pobj = NULL;
+        PyErr_SetString(PyExc_TypeError, "expected a weakref");
+        return -1;
+    }
+    obj = PyWeakref_GetObject(ref);
+    if (obj == NULL) {
+        // SystemError if ref is NULL
+        *pobj = NULL;
+        return -1;
+    }
+    if (obj == Py_None) {
+        *pobj = NULL;
+        return 0;
+    }
+    *pobj = Py_NewRef(obj);
+    return 1;
+}
+#endif
+
+
+// bpo-36974 added PY_VECTORCALL_ARGUMENTS_OFFSET to Python 3.8b1
+#ifndef PY_VECTORCALL_ARGUMENTS_OFFSET
+#  define PY_VECTORCALL_ARGUMENTS_OFFSET (_Py_CAST(size_t, 1) << (8 * sizeof(size_t) - 1))
+#endif
+
+// bpo-36974 added PyVectorcall_NARGS() to Python 3.8b1
+#if PY_VERSION_HEX < 0x030800B1
+static inline Py_ssize_t PyVectorcall_NARGS(size_t n)
+{
+    return n & ~PY_VECTORCALL_ARGUMENTS_OFFSET;
+}
+#endif
+
+
+// gh-105922 added PyObject_Vectorcall() to Python 3.9.0a4
+#if PY_VERSION_HEX < 0x030900A4
+static inline PyObject*
+PyObject_Vectorcall(PyObject *callable, PyObject *const *args,
+                     size_t nargsf, PyObject *kwnames)
+{
+#if PY_VERSION_HEX >= 0x030800B1 && !defined(PYPY_VERSION)
+    // bpo-36974 added _PyObject_Vectorcall() to Python 3.8.0b1
+    return _PyObject_Vectorcall(callable, args, nargsf, kwnames);
+#else
+    PyObject *posargs = NULL, *kwargs = NULL;
+    PyObject *res;
+    Py_ssize_t nposargs, nkwargs, i;
+
+    if (nargsf != 0 && args == NULL) {
+        PyErr_BadInternalCall();
+        goto error;
+    }
+    if (kwnames != NULL && !PyTuple_Check(kwnames)) {
+        PyErr_BadInternalCall();
+        goto error;
+    }
+
+    nposargs = (Py_ssize_t)PyVectorcall_NARGS(nargsf);
+    if (kwnames) {
+        nkwargs = PyTuple_GET_SIZE(kwnames);
+    }
+    else {
+        nkwargs = 0;
+    }
+
+    posargs = PyTuple_New(nposargs);
+    if (posargs == NULL) {
+        goto error;
+    }
+    if (nposargs) {
+        for (i=0; i < nposargs; i++) {
+            PyTuple_SET_ITEM(posargs, i, Py_NewRef(*args));
+            args++;
+        }
+    }
+
+    if (nkwargs) {
+        kwargs = PyDict_New();
+        if (kwargs == NULL) {
+            goto error;
+        }
+
+        for (i = 0; i < nkwargs; i++) {
+            PyObject *key = PyTuple_GET_ITEM(kwnames, i);
+            PyObject *value = *args;
+            args++;
+            if (PyDict_SetItem(kwargs, key, value) < 0) {
+                goto error;
+            }
+        }
+    }
+    else {
+        kwargs = NULL;
+    }
+
+    res = PyObject_Call(callable, posargs, kwargs);
+    Py_DECREF(posargs);
+    Py_XDECREF(kwargs);
+    return res;
+
+error:
+    Py_DECREF(posargs);
+    Py_XDECREF(kwargs);
+    return NULL;
+#endif
+}
+#endif
+
+
+// gh-106521 added PyObject_GetOptionalAttr() and
+// PyObject_GetOptionalAttrString() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyObject_GetOptionalAttr(PyObject *obj, PyObject *attr_name, PyObject **result)
+{
+    // bpo-32571 added _PyObject_LookupAttr() to Python 3.7.0b1
+#if PY_VERSION_HEX >= 0x030700B1 && !defined(PYPY_VERSION)
+    return _PyObject_LookupAttr(obj, attr_name, result);
+#else
+    *result = PyObject_GetAttr(obj, attr_name);
+    if (*result != NULL) {
+        return 1;
+    }
+    if (!PyErr_Occurred()) {
+        return 0;
+    }
+    if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
+        PyErr_Clear();
+        return 0;
+    }
+    return -1;
+#endif
+}
+
+static inline int
+PyObject_GetOptionalAttrString(PyObject *obj, const char *attr_name, PyObject **result)
+{
+    PyObject *name_obj;
+    int rc;
+#if PY_VERSION_HEX >= 0x03000000
+    name_obj = PyUnicode_FromString(attr_name);
+#else
+    name_obj = PyString_FromString(attr_name);
+#endif
+    if (name_obj == NULL) {
+        *result = NULL;
+        return -1;
+    }
+    rc = PyObject_GetOptionalAttr(obj, name_obj, result);
+    Py_DECREF(name_obj);
+    return rc;
+}
+#endif
+
+
+// gh-106307 added PyObject_GetOptionalAttr() and
+// PyMapping_GetOptionalItemString() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyMapping_GetOptionalItem(PyObject *obj, PyObject *key, PyObject **result)
+{
+    *result = PyObject_GetItem(obj, key);
+    if (*result) {
+        return 1;
+    }
+    if (!PyErr_ExceptionMatches(PyExc_KeyError)) {
+        return -1;
+    }
+    PyErr_Clear();
+    return 0;
+}
+
+static inline int
+PyMapping_GetOptionalItemString(PyObject *obj, const char *key, PyObject **result)
+{
+    PyObject *key_obj;
+    int rc;
+#if PY_VERSION_HEX >= 0x03000000
+    key_obj = PyUnicode_FromString(key);
+#else
+    key_obj = PyString_FromString(key);
+#endif
+    if (key_obj == NULL) {
+        *result = NULL;
+        return -1;
+    }
+    rc = PyMapping_GetOptionalItem(obj, key_obj, result);
+    Py_DECREF(key_obj);
+    return rc;
+}
+#endif
+
+// gh-108511 added PyMapping_HasKeyWithError() and
+// PyMapping_HasKeyStringWithError() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyMapping_HasKeyWithError(PyObject *obj, PyObject *key)
+{
+    PyObject *res;
+    int rc = PyMapping_GetOptionalItem(obj, key, &res);
+    Py_XDECREF(res);
+    return rc;
+}
+
+static inline int
+PyMapping_HasKeyStringWithError(PyObject *obj, const char *key)
+{
+    PyObject *res;
+    int rc = PyMapping_GetOptionalItemString(obj, key, &res);
+    Py_XDECREF(res);
+    return rc;
+}
+#endif
+
+
+// gh-108511 added PyObject_HasAttrWithError() and
+// PyObject_HasAttrStringWithError() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyObject_HasAttrWithError(PyObject *obj, PyObject *attr)
+{
+    PyObject *res;
+    int rc = PyObject_GetOptionalAttr(obj, attr, &res);
+    Py_XDECREF(res);
+    return rc;
+}
+
+static inline int
+PyObject_HasAttrStringWithError(PyObject *obj, const char *attr)
+{
+    PyObject *res;
+    int rc = PyObject_GetOptionalAttrString(obj, attr, &res);
+    Py_XDECREF(res);
+    return rc;
+}
+#endif
+
+
+// gh-106004 added PyDict_GetItemRef() and PyDict_GetItemStringRef()
+// to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyDict_GetItemRef(PyObject *mp, PyObject *key, PyObject **result)
+{
+#if PY_VERSION_HEX >= 0x03000000
+    PyObject *item = PyDict_GetItemWithError(mp, key);
+#else
+    PyObject *item = _PyDict_GetItemWithError(mp, key);
+#endif
+    if (item != NULL) {
+        *result = Py_NewRef(item);
+        return 1;  // found
+    }
+    if (!PyErr_Occurred()) {
+        *result = NULL;
+        return 0;  // not found
+    }
+    *result = NULL;
+    return -1;
+}
+
+static inline int
+PyDict_GetItemStringRef(PyObject *mp, const char *key, PyObject **result)
+{
+    int res;
+#if PY_VERSION_HEX >= 0x03000000
+    PyObject *key_obj = PyUnicode_FromString(key);
+#else
+    PyObject *key_obj = PyString_FromString(key);
+#endif
+    if (key_obj == NULL) {
+        *result = NULL;
+        return -1;
+    }
+    res = PyDict_GetItemRef(mp, key_obj, result);
+    Py_DECREF(key_obj);
+    return res;
+}
+#endif
+
+
+// gh-106307 added PyModule_Add() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyModule_Add(PyObject *mod, const char *name, PyObject *value)
+{
+    int res = PyModule_AddObjectRef(mod, name, value);
+    Py_XDECREF(value);
+    return res;
+}
+#endif
+
+
+// gh-108014 added Py_IsFinalizing() to Python 3.13.0a1
+// bpo-1856 added _Py_Finalizing to Python 3.2.1b1.
+// _Py_IsFinalizing() was added to PyPy 7.3.0.
+#if (0x030201B1 <= PY_VERSION_HEX && PY_VERSION_HEX < 0x030D00A1) \
+        && (!defined(PYPY_VERSION_NUM) || PYPY_VERSION_NUM >= 0x7030000)
+static inline int Py_IsFinalizing(void)
+{
+#if PY_VERSION_HEX >= 0x030700A1
+    // _Py_IsFinalizing() was added to Python 3.7.0a1.
+    return _Py_IsFinalizing();
+#else
+    return (_Py_Finalizing != NULL);
+#endif
+}
+#endif
+
+
+// gh-108323 added PyDict_ContainsString() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int PyDict_ContainsString(PyObject *op, const char *key)
+{
+    PyObject *key_obj = PyUnicode_FromString(key);
+    if (key_obj == NULL) {
+        return -1;
+    }
+    int res = PyDict_Contains(op, key_obj);
+    Py_DECREF(key_obj);
+    return res;
+}
+#endif
+
+
+// gh-108445 added PyLong_AsInt() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int PyLong_AsInt(PyObject *obj)
+{
+#ifdef PYPY_VERSION
+    long value = PyLong_AsLong(obj);
+    if (value == -1 && PyErr_Occurred()) {
+        return -1;
+    }
+    if (value < (long)INT_MIN || (long)INT_MAX < value) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "Python int too large to convert to C int");
+        return -1;
+    }
+    return (int)value;
+#else
+    return _PyLong_AsInt(obj);
+#endif
+}
+#endif
+
+
+// gh-107073 added PyObject_VisitManagedDict() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyObject_VisitManagedDict(PyObject *obj, visitproc visit, void *arg)
+{
+    PyObject **dict = _PyObject_GetDictPtr(obj);
+    if (dict == NULL || *dict == NULL) {
+        return -1;
+    }
+    Py_VISIT(*dict);
+    return 0;
+}
+
+static inline void
+PyObject_ClearManagedDict(PyObject *obj)
+{
+    PyObject **dict = _PyObject_GetDictPtr(obj);
+    if (dict == NULL || *dict == NULL) {
+        return;
+    }
+    Py_CLEAR(*dict);
+}
+#endif
+
+// gh-108867 added PyThreadState_GetUnchecked() to Python 3.13.0a1
+// Python 3.5.2 added _PyThreadState_UncheckedGet().
+#if PY_VERSION_HEX >= 0x03050200 && PY_VERSION_HEX < 0x030D00A1
+static inline PyThreadState*
+PyThreadState_GetUnchecked(void)
+{
+    return _PyThreadState_UncheckedGet();
+}
+#endif
+
+// gh-110289 added PyUnicode_EqualToUTF8() and PyUnicode_EqualToUTF8AndSize()
+// to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyUnicode_EqualToUTF8AndSize(PyObject *unicode, const char *str, Py_ssize_t str_len)
+{
+    Py_ssize_t len;
+    const void *utf8;
+    PyObject *exc_type, *exc_value, *exc_tb;
+    int res;
+
+    // API cannot report errors so save/restore the exception
+    PyErr_Fetch(&exc_type, &exc_value, &exc_tb);
+
+    // Python 3.3.0a1 added PyUnicode_AsUTF8AndSize()
+#if PY_VERSION_HEX >= 0x030300A1
+    if (PyUnicode_IS_ASCII(unicode)) {
+        utf8 = PyUnicode_DATA(unicode);
+        len = PyUnicode_GET_LENGTH(unicode);
+    }
+    else {
+        utf8 = PyUnicode_AsUTF8AndSize(unicode, &len);
+        if (utf8 == NULL) {
+            // Memory allocation failure. The API cannot report error,
+            // so ignore the exception and return 0.
+            res = 0;
+            goto done;
+        }
+    }
+
+    if (len != str_len) {
+        res = 0;
+        goto done;
+    }
+    res = (memcmp(utf8, str, (size_t)len) == 0);
+#else
+    PyObject *bytes = PyUnicode_AsUTF8String(unicode);
+    if (bytes == NULL) {
+        // Memory allocation failure. The API cannot report error,
+        // so ignore the exception and return 0.
+        res = 0;
+        goto done;
+    }
+
+#if PY_VERSION_HEX >= 0x03000000
+    len = PyBytes_GET_SIZE(bytes);
+    utf8 = PyBytes_AS_STRING(bytes);
+#else
+    len = PyString_GET_SIZE(bytes);
+    utf8 = PyString_AS_STRING(bytes);
+#endif
+    if (len != str_len) {
+        Py_DECREF(bytes);
+        res = 0;
+        goto done;
+    }
+
+    res = (memcmp(utf8, str, (size_t)len) == 0);
+    Py_DECREF(bytes);
+#endif
+
+done:
+    PyErr_Restore(exc_type, exc_value, exc_tb);
+    return res;
+}
+
+static inline int
+PyUnicode_EqualToUTF8(PyObject *unicode, const char *str)
+{
+    return PyUnicode_EqualToUTF8AndSize(unicode, str, (Py_ssize_t)strlen(str));
+}
+#endif
+
+
+// gh-111138 added PyList_Extend() and PyList_Clear() to Python 3.13.0a2
+#if PY_VERSION_HEX < 0x030D00A2
+static inline int
+PyList_Extend(PyObject *list, PyObject *iterable)
+{
+    return PyList_SetSlice(list, PY_SSIZE_T_MAX, PY_SSIZE_T_MAX, iterable);
+}
+
+static inline int
+PyList_Clear(PyObject *list)
+{
+    return PyList_SetSlice(list, 0, PY_SSIZE_T_MAX, NULL);
+}
+#endif
+
+// gh-111262 added PyDict_Pop() and PyDict_PopString() to Python 3.13.0a2
+#if PY_VERSION_HEX < 0x030D00A2
+static inline int
+PyDict_Pop(PyObject *dict, PyObject *key, PyObject **result)
+{
+    PyObject *value;
+
+    if (!PyDict_Check(dict)) {
+        PyErr_BadInternalCall();
+        if (result) {
+            *result = NULL;
+        }
+        return -1;
+    }
+
+    // bpo-16991 added _PyDict_Pop() to Python 3.5.0b2.
+    // Python 3.6.0b3 changed _PyDict_Pop() first argument type to PyObject*.
+    // Python 3.13.0a1 removed _PyDict_Pop().
+#if defined(PYPY_VERSION) || PY_VERSION_HEX < 0x030500b2 || PY_VERSION_HEX >= 0x030D0000
+    value = PyObject_CallMethod(dict, "pop", "O", key);
+#elif PY_VERSION_HEX < 0x030600b3
+    value = _PyDict_Pop(_Py_CAST(PyDictObject*, dict), key, NULL);
+#else
+    value = _PyDict_Pop(dict, key, NULL);
+#endif
+    if (value == NULL) {
+        if (result) {
+            *result = NULL;
+        }
+        if (PyErr_Occurred() && !PyErr_ExceptionMatches(PyExc_KeyError)) {
+            return -1;
+        }
+        PyErr_Clear();
+        return 0;
+    }
+    if (result) {
+        *result = value;
+    }
+    else {
+        Py_DECREF(value);
+    }
+    return 1;
+}
+
+static inline int
+PyDict_PopString(PyObject *dict, const char *key, PyObject **result)
+{
+    PyObject *key_obj = PyUnicode_FromString(key);
+    if (key_obj == NULL) {
+        if (result != NULL) {
+            *result = NULL;
+        }
+        return -1;
+    }
+
+    int res = PyDict_Pop(dict, key_obj, result);
+    Py_DECREF(key_obj);
+    return res;
+}
+#endif
+
+
+#if PY_VERSION_HEX < 0x030200A4
+// Python 3.2.0a4 added Py_hash_t type
+typedef Py_ssize_t Py_hash_t;
+#endif
+
+
+// gh-111545 added Py_HashPointer() to Python 3.13.0a3
+#if PY_VERSION_HEX < 0x030D00A3
+static inline Py_hash_t Py_HashPointer(const void *ptr)
+{
+#if PY_VERSION_HEX >= 0x030900A4 && !defined(PYPY_VERSION)
+    return _Py_HashPointer(ptr);
+#else
+    return _Py_HashPointer(_Py_CAST(void*, ptr));
+#endif
+}
+#endif
+
+
+// Python 3.13a4 added a PyTime API.
+// Use the private API added to Python 3.5.
+#if PY_VERSION_HEX < 0x030D00A4 && PY_VERSION_HEX  >= 0x03050000
+typedef _PyTime_t PyTime_t;
+#define PyTime_MIN _PyTime_MIN
+#define PyTime_MAX _PyTime_MAX
+
+static inline double PyTime_AsSecondsDouble(PyTime_t t)
+{ return _PyTime_AsSecondsDouble(t); }
+
+static inline int PyTime_Monotonic(PyTime_t *result)
+{ return _PyTime_GetMonotonicClockWithInfo(result, NULL); }
+
+static inline int PyTime_Time(PyTime_t *result)
+{ return _PyTime_GetSystemClockWithInfo(result, NULL); }
+
+static inline int PyTime_PerfCounter(PyTime_t *result)
+{
+#if PY_VERSION_HEX >= 0x03070000 && !defined(PYPY_VERSION)
+    return _PyTime_GetPerfCounterWithInfo(result, NULL);
+#elif PY_VERSION_HEX >= 0x03070000
+    // Call time.perf_counter_ns() and convert Python int object to PyTime_t.
+    // Cache time.perf_counter_ns() function for best performance.
+    static PyObject *func = NULL;
+    if (func == NULL) {
+        PyObject *mod = PyImport_ImportModule("time");
+        if (mod == NULL) {
+            return -1;
+        }
+
+        func = PyObject_GetAttrString(mod, "perf_counter_ns");
+        Py_DECREF(mod);
+        if (func == NULL) {
+            return -1;
+        }
+    }
+
+    PyObject *res = PyObject_CallNoArgs(func);
+    if (res == NULL) {
+        return -1;
+    }
+    long long value = PyLong_AsLongLong(res);
+    Py_DECREF(res);
+
+    if (value == -1 && PyErr_Occurred()) {
+        return -1;
+    }
+
+    Py_BUILD_ASSERT(sizeof(value) >= sizeof(PyTime_t));
+    *result = (PyTime_t)value;
+    return 0;
+#else
+    // Call time.perf_counter() and convert C double to PyTime_t.
+    // Cache time.perf_counter() function for best performance.
+    static PyObject *func = NULL;
+    if (func == NULL) {
+        PyObject *mod = PyImport_ImportModule("time");
+        if (mod == NULL) {
+            return -1;
+        }
+
+        func = PyObject_GetAttrString(mod, "perf_counter");
+        Py_DECREF(mod);
+        if (func == NULL) {
+            return -1;
+        }
+    }
+
+    PyObject *res = PyObject_CallNoArgs(func);
+    if (res == NULL) {
+        return -1;
+    }
+    double d = PyFloat_AsDouble(res);
+    Py_DECREF(res);
+
+    if (d == -1.0 && PyErr_Occurred()) {
+        return -1;
+    }
+
+    // Avoid floor() to avoid having to link to libm
+    *result = (PyTime_t)(d * 1e9);
+    return 0;
+#endif
+}
+
+#endif
+
+// gh-111389 added hash constants to Python 3.13.0a5. These constants were
+// added first as private macros to Python 3.4.0b1 and PyPy 7.3.8.
+#if (!defined(PyHASH_BITS) \
+     && ((!defined(PYPY_VERSION) && PY_VERSION_HEX >= 0x030400B1) \
+         || (defined(PYPY_VERSION) && PY_VERSION_HEX >= 0x03070000 \
+             && PYPY_VERSION_NUM >= 0x07030800)))
+#  define PyHASH_BITS _PyHASH_BITS
+#  define PyHASH_MODULUS _PyHASH_MODULUS
+#  define PyHASH_INF _PyHASH_INF
+#  define PyHASH_IMAG _PyHASH_IMAG
+#endif
+
+
+// gh-111545 added Py_GetConstant() and Py_GetConstantBorrowed()
+// to Python 3.13.0a6
+#if PY_VERSION_HEX < 0x030D00A6 && !defined(Py_CONSTANT_NONE)
+
+#define Py_CONSTANT_NONE 0
+#define Py_CONSTANT_FALSE 1
+#define Py_CONSTANT_TRUE 2
+#define Py_CONSTANT_ELLIPSIS 3
+#define Py_CONSTANT_NOT_IMPLEMENTED 4
+#define Py_CONSTANT_ZERO 5
+#define Py_CONSTANT_ONE 6
+#define Py_CONSTANT_EMPTY_STR 7
+#define Py_CONSTANT_EMPTY_BYTES 8
+#define Py_CONSTANT_EMPTY_TUPLE 9
+
+static inline PyObject* Py_GetConstant(unsigned int constant_id)
+{
+    static PyObject* constants[Py_CONSTANT_EMPTY_TUPLE + 1] = {NULL};
+
+    if (constants[Py_CONSTANT_NONE] == NULL) {
+        constants[Py_CONSTANT_NONE] = Py_None;
+        constants[Py_CONSTANT_FALSE] = Py_False;
+        constants[Py_CONSTANT_TRUE] = Py_True;
+        constants[Py_CONSTANT_ELLIPSIS] = Py_Ellipsis;
+        constants[Py_CONSTANT_NOT_IMPLEMENTED] = Py_NotImplemented;
+
+        constants[Py_CONSTANT_ZERO] = PyLong_FromLong(0);
+        if (constants[Py_CONSTANT_ZERO] == NULL) {
+            goto fatal_error;
+        }
+
+        constants[Py_CONSTANT_ONE] = PyLong_FromLong(1);
+        if (constants[Py_CONSTANT_ONE] == NULL) {
+            goto fatal_error;
+        }
+
+        constants[Py_CONSTANT_EMPTY_STR] = PyUnicode_FromStringAndSize("", 0);
+        if (constants[Py_CONSTANT_EMPTY_STR] == NULL) {
+            goto fatal_error;
+        }
+
+        constants[Py_CONSTANT_EMPTY_BYTES] = PyBytes_FromStringAndSize("", 0);
+        if (constants[Py_CONSTANT_EMPTY_BYTES] == NULL) {
+            goto fatal_error;
+        }
+
+        constants[Py_CONSTANT_EMPTY_TUPLE] = PyTuple_New(0);
+        if (constants[Py_CONSTANT_EMPTY_TUPLE] == NULL) {
+            goto fatal_error;
+        }
+        // goto dance to avoid compiler warnings about Py_FatalError()
+        goto init_done;
+
+fatal_error:
+        // This case should never happen
+        Py_FatalError("Py_GetConstant() failed to get constants");
+    }
+
+init_done:
+    if (constant_id <= Py_CONSTANT_EMPTY_TUPLE) {
+        return Py_NewRef(constants[constant_id]);
+    }
+    else {
+        PyErr_BadInternalCall();
+        return NULL;
+    }
+}
+
+static inline PyObject* Py_GetConstantBorrowed(unsigned int constant_id)
+{
+    PyObject *obj = Py_GetConstant(constant_id);
+    Py_XDECREF(obj);
+    return obj;
+}
+#endif
+
+
+// gh-114329 added PyList_GetItemRef() to Python 3.13.0a4
+#if PY_VERSION_HEX < 0x030D00A4
+static inline PyObject *
+PyList_GetItemRef(PyObject *op, Py_ssize_t index)
+{
+    PyObject *item = PyList_GetItem(op, index);
+    Py_XINCREF(item);
+    return item;
+}
+#endif
+
+
+// gh-114329 added PyList_GetItemRef() to Python 3.13.0a4
+#if PY_VERSION_HEX < 0x030D00A4
+static inline int
+PyDict_SetDefaultRef(PyObject *d, PyObject *key, PyObject *default_value,
+                     PyObject **result)
+{
+    PyObject *value;
+    if (PyDict_GetItemRef(d, key, &value) < 0) {
+        // get error
+        if (result) {
+            *result = NULL;
+        }
+        return -1;
+    }
+    if (value != NULL) {
+        // present
+        if (result) {
+            *result = value;
+        }
+        else {
+            Py_DECREF(value);
+        }
+        return 1;
+    }
+
+    // missing: set the item
+    if (PyDict_SetItem(d, key, default_value) < 0) {
+        // set error
+        if (result) {
+            *result = NULL;
+        }
+        return -1;
+    }
+    if (result) {
+        *result = Py_NewRef(default_value);
+    }
+    return 0;
+}
+#endif
+
+#if PY_VERSION_HEX < 0x030D00B3
+#  define Py_BEGIN_CRITICAL_SECTION(op) {
+#  define Py_END_CRITICAL_SECTION() }
+#  define Py_BEGIN_CRITICAL_SECTION2(a, b) {
+#  define Py_END_CRITICAL_SECTION2() }
+#endif
+
+#if PY_VERSION_HEX < 0x030E0000 && PY_VERSION_HEX >= 0x03060000 && !defined(PYPY_VERSION)
+typedef struct PyUnicodeWriter PyUnicodeWriter;
+
+static inline void PyUnicodeWriter_Discard(PyUnicodeWriter *writer)
+{
+    _PyUnicodeWriter_Dealloc((_PyUnicodeWriter*)writer);
+    PyMem_Free(writer);
+}
+
+static inline PyUnicodeWriter* PyUnicodeWriter_Create(Py_ssize_t length)
+{
+    if (length < 0) {
+        PyErr_SetString(PyExc_ValueError,
+                        "length must be positive");
+        return NULL;
+    }
+
+    const size_t size = sizeof(_PyUnicodeWriter);
+    PyUnicodeWriter *pub_writer = (PyUnicodeWriter *)PyMem_Malloc(size);
+    if (pub_writer == _Py_NULL) {
+        PyErr_NoMemory();
+        return _Py_NULL;
+    }
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter *)pub_writer;
+
+    _PyUnicodeWriter_Init(writer);
+    if (_PyUnicodeWriter_Prepare(writer, length, 127) < 0) {
+        PyUnicodeWriter_Discard(pub_writer);
+        return NULL;
+    }
+    writer->overallocate = 1;
+    return pub_writer;
+}
+
+static inline PyObject* PyUnicodeWriter_Finish(PyUnicodeWriter *writer)
+{
+    PyObject *str = _PyUnicodeWriter_Finish((_PyUnicodeWriter*)writer);
+    assert(((_PyUnicodeWriter*)writer)->buffer == NULL);
+    PyMem_Free(writer);
+    return str;
+}
+
+static inline int
+PyUnicodeWriter_WriteChar(PyUnicodeWriter *writer, Py_UCS4 ch)
+{
+    if (ch > 0x10ffff) {
+        PyErr_SetString(PyExc_ValueError,
+                        "character must be in range(0x110000)");
+        return -1;
+    }
+
+    return _PyUnicodeWriter_WriteChar((_PyUnicodeWriter*)writer, ch);
+}
+
+static inline int
+PyUnicodeWriter_WriteStr(PyUnicodeWriter *writer, PyObject *obj)
+{
+    PyObject *str = PyObject_Str(obj);
+    if (str == NULL) {
+        return -1;
+    }
+
+    int res = _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, str);
+    Py_DECREF(str);
+    return res;
+}
+
+static inline int
+PyUnicodeWriter_WriteRepr(PyUnicodeWriter *writer, PyObject *obj)
+{
+    PyObject *str = PyObject_Repr(obj);
+    if (str == NULL) {
+        return -1;
+    }
+
+    int res = _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, str);
+    Py_DECREF(str);
+    return res;
+}
+
+static inline int
+PyUnicodeWriter_WriteUTF8(PyUnicodeWriter *writer,
+                          const char *str, Py_ssize_t size)
+{
+    if (size < 0) {
+        size = (Py_ssize_t)strlen(str);
+    }
+
+    PyObject *str_obj = PyUnicode_FromStringAndSize(str, size);
+    if (str_obj == _Py_NULL) {
+        return -1;
+    }
+
+    int res = _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, str_obj);
+    Py_DECREF(str_obj);
+    return res;
+}
+
+static inline int
+PyUnicodeWriter_WriteASCII(PyUnicodeWriter *writer,
+                           const char *str, Py_ssize_t size)
+{
+    if (size < 0) {
+        size = (Py_ssize_t)strlen(str);
+    }
+
+    return _PyUnicodeWriter_WriteASCIIString((_PyUnicodeWriter*)writer,
+                                             str, size);
+}
+
+static inline int
+PyUnicodeWriter_WriteWideChar(PyUnicodeWriter *writer,
+                              const wchar_t *str, Py_ssize_t size)
+{
+    if (size < 0) {
+        size = (Py_ssize_t)wcslen(str);
+    }
+
+    PyObject *str_obj = PyUnicode_FromWideChar(str, size);
+    if (str_obj == _Py_NULL) {
+        return -1;
+    }
+
+    int res = _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, str_obj);
+    Py_DECREF(str_obj);
+    return res;
+}
+
+static inline int
+PyUnicodeWriter_WriteSubstring(PyUnicodeWriter *writer, PyObject *str,
+                               Py_ssize_t start, Py_ssize_t end)
+{
+    if (!PyUnicode_Check(str)) {
+        PyErr_Format(PyExc_TypeError, "expect str, not %s",
+                     Py_TYPE(str)->tp_name);
+        return -1;
+    }
+    if (start < 0 || start > end) {
+        PyErr_Format(PyExc_ValueError, "invalid start argument");
+        return -1;
+    }
+    if (end > PyUnicode_GET_LENGTH(str)) {
+        PyErr_Format(PyExc_ValueError, "invalid end argument");
+        return -1;
+    }
+
+    return _PyUnicodeWriter_WriteSubstring((_PyUnicodeWriter*)writer, str,
+                                           start, end);
+}
+
+static inline int
+PyUnicodeWriter_Format(PyUnicodeWriter *writer, const char *format, ...)
+{
+    va_list vargs;
+    va_start(vargs, format);
+    PyObject *str = PyUnicode_FromFormatV(format, vargs);
+    va_end(vargs);
+    if (str == _Py_NULL) {
+        return -1;
+    }
+
+    int res = _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, str);
+    Py_DECREF(str);
+    return res;
+}
+#endif  // PY_VERSION_HEX < 0x030E0000
+
+// gh-116560 added PyLong_GetSign() to Python 3.14.0a0
+#if PY_VERSION_HEX < 0x030E00A0
+static inline int PyLong_GetSign(PyObject *obj, int *sign)
+{
+    if (!PyLong_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "expect int, got %s", Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+
+    *sign = _PyLong_Sign(obj);
+    return 0;
+}
+#endif
+
+// gh-126061 added PyLong_IsPositive/Negative/Zero() to Python in 3.14.0a2
+#if PY_VERSION_HEX < 0x030E00A2
+static inline int PyLong_IsPositive(PyObject *obj)
+{
+    if (!PyLong_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "expected int, got %s", Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+    return _PyLong_Sign(obj) == 1;
+}
+
+static inline int PyLong_IsNegative(PyObject *obj)
+{
+    if (!PyLong_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "expected int, got %s", Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+    return _PyLong_Sign(obj) == -1;
+}
+
+static inline int PyLong_IsZero(PyObject *obj)
+{
+    if (!PyLong_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "expected int, got %s", Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+    return _PyLong_Sign(obj) == 0;
+}
+#endif
+
+
+// gh-124502 added PyUnicode_Equal() to Python 3.14.0a0
+#if PY_VERSION_HEX < 0x030E00A0
+static inline int PyUnicode_Equal(PyObject *str1, PyObject *str2)
+{
+    if (!PyUnicode_Check(str1)) {
+        PyErr_Format(PyExc_TypeError, "first argument must be str, not %s",
+                     Py_TYPE(str1)->tp_name);
+        return -1;
+    }
+    if (!PyUnicode_Check(str2)) {
+        PyErr_Format(PyExc_TypeError, "second argument must be str, not %s",
+                     Py_TYPE(str2)->tp_name);
+        return -1;
+    }
+
+#if PY_VERSION_HEX >= 0x030d0000 && !defined(PYPY_VERSION)
+    PyAPI_FUNC(int) _PyUnicode_Equal(PyObject *str1, PyObject *str2);
+
+    return _PyUnicode_Equal(str1, str2);
+#elif PY_VERSION_HEX >= 0x03060000 && !defined(PYPY_VERSION)
+    return _PyUnicode_EQ(str1, str2);
+#elif PY_VERSION_HEX >= 0x03090000 && defined(PYPY_VERSION)
+    return _PyUnicode_EQ(str1, str2);
+#else
+    return (PyUnicode_Compare(str1, str2) == 0);
+#endif
+}
+#endif
+
+
+// gh-121645 added PyBytes_Join() to Python 3.14.0a0
+#if PY_VERSION_HEX < 0x030E00A0
+static inline PyObject* PyBytes_Join(PyObject *sep, PyObject *iterable)
+{
+    return _PyBytes_Join(sep, iterable);
+}
+#endif
+
+
+#if PY_VERSION_HEX < 0x030E00A0
+static inline Py_hash_t Py_HashBuffer(const void *ptr, Py_ssize_t len)
+{
+#if PY_VERSION_HEX >= 0x03000000 && !defined(PYPY_VERSION)
+    PyAPI_FUNC(Py_hash_t) _Py_HashBytes(const void *src, Py_ssize_t len);
+
+    return _Py_HashBytes(ptr, len);
+#else
+    Py_hash_t hash;
+    PyObject *bytes = PyBytes_FromStringAndSize((const char*)ptr, len);
+    if (bytes == NULL) {
+        return -1;
+    }
+    hash = PyObject_Hash(bytes);
+    Py_DECREF(bytes);
+    return hash;
+#endif
+}
+#endif
+
+
+#if PY_VERSION_HEX < 0x030E00A0
+static inline int PyIter_NextItem(PyObject *iter, PyObject **item)
+{
+    iternextfunc tp_iternext;
+
+    assert(iter != NULL);
+    assert(item != NULL);
+
+    tp_iternext = Py_TYPE(iter)->tp_iternext;
+    if (tp_iternext == NULL) {
+        *item = NULL;
+        PyErr_Format(PyExc_TypeError, "expected an iterator, got '%s'",
+                     Py_TYPE(iter)->tp_name);
+        return -1;
+    }
+
+    if ((*item = tp_iternext(iter))) {
+        return 1;
+    }
+    if (!PyErr_Occurred()) {
+        return 0;
+    }
+    if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
+        PyErr_Clear();
+        return 0;
+    }
+    return -1;
+}
+#endif
+
+
+#if PY_VERSION_HEX < 0x030E00A0
+static inline PyObject* PyLong_FromInt32(int32_t value)
+{
+    Py_BUILD_ASSERT(sizeof(long) >= 4);
+    return PyLong_FromLong(value);
+}
+
+static inline PyObject* PyLong_FromInt64(int64_t value)
+{
+    Py_BUILD_ASSERT(sizeof(long long) >= 8);
+    return PyLong_FromLongLong(value);
+}
+
+static inline PyObject* PyLong_FromUInt32(uint32_t value)
+{
+    Py_BUILD_ASSERT(sizeof(unsigned long) >= 4);
+    return PyLong_FromUnsignedLong(value);
+}
+
+static inline PyObject* PyLong_FromUInt64(uint64_t value)
+{
+    Py_BUILD_ASSERT(sizeof(unsigned long long) >= 8);
+    return PyLong_FromUnsignedLongLong(value);
+}
+
+static inline int PyLong_AsInt32(PyObject *obj, int32_t *pvalue)
+{
+    Py_BUILD_ASSERT(sizeof(int) == 4);
+    int value = PyLong_AsInt(obj);
+    if (value == -1 && PyErr_Occurred()) {
+        return -1;
+    }
+    *pvalue = (int32_t)value;
+    return 0;
+}
+
+static inline int PyLong_AsInt64(PyObject *obj, int64_t *pvalue)
+{
+    Py_BUILD_ASSERT(sizeof(long long) == 8);
+    long long value = PyLong_AsLongLong(obj);
+    if (value == -1 && PyErr_Occurred()) {
+        return -1;
+    }
+    *pvalue = (int64_t)value;
+    return 0;
+}
+
+static inline int PyLong_AsUInt32(PyObject *obj, uint32_t *pvalue)
+{
+    Py_BUILD_ASSERT(sizeof(long) >= 4);
+    unsigned long value = PyLong_AsUnsignedLong(obj);
+    if (value == (unsigned long)-1 && PyErr_Occurred()) {
+        return -1;
+    }
+#if SIZEOF_LONG > 4
+    if ((unsigned long)UINT32_MAX < value) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "Python int too large to convert to C uint32_t");
+        return -1;
+    }
+#endif
+    *pvalue = (uint32_t)value;
+    return 0;
+}
+
+static inline int PyLong_AsUInt64(PyObject *obj, uint64_t *pvalue)
+{
+    Py_BUILD_ASSERT(sizeof(long long) == 8);
+    unsigned long long value = PyLong_AsUnsignedLongLong(obj);
+    if (value == (unsigned long long)-1 && PyErr_Occurred()) {
+        return -1;
+    }
+    *pvalue = (uint64_t)value;
+    return 0;
+}
+#endif
+
+
+// gh-102471 added import and export API for integers to 3.14.0a2.
+#if PY_VERSION_HEX < 0x030E00A2 && PY_VERSION_HEX >= 0x03000000 && !defined(PYPY_VERSION)
+// Helpers to access PyLongObject internals.
+static inline void
+_PyLong_SetSignAndDigitCount(PyLongObject *op, int sign, Py_ssize_t size)
+{
+#if PY_VERSION_HEX >= 0x030C0000
+    op->long_value.lv_tag = (uintptr_t)(1 - sign) | ((uintptr_t)size << 3);
+#elif PY_VERSION_HEX >= 0x030900A4
+    Py_SET_SIZE(op, sign * size);
+#else
+    Py_SIZE(op) = sign * size;
+#endif
+}
+
+static inline Py_ssize_t
+_PyLong_DigitCount(const PyLongObject *op)
+{
+#if PY_VERSION_HEX >= 0x030C0000
+    return (Py_ssize_t)(op->long_value.lv_tag >> 3);
+#else
+    return _PyLong_Sign((PyObject*)op) < 0 ? -Py_SIZE(op) : Py_SIZE(op);
+#endif
+}
+
+static inline digit*
+_PyLong_GetDigits(const PyLongObject *op)
+{
+#if PY_VERSION_HEX >= 0x030C0000
+    return (digit*)(op->long_value.ob_digit);
+#else
+    return (digit*)(op->ob_digit);
+#endif
+}
+
+typedef struct PyLongLayout {
+    uint8_t bits_per_digit;
+    uint8_t digit_size;
+    int8_t digits_order;
+    int8_t digit_endianness;
+} PyLongLayout;
+
+typedef struct PyLongExport {
+    int64_t value;
+    uint8_t negative;
+    Py_ssize_t ndigits;
+    const void *digits;
+    Py_uintptr_t _reserved;
+} PyLongExport;
+
+typedef struct PyLongWriter PyLongWriter;
+
+static inline const PyLongLayout*
+PyLong_GetNativeLayout(void)
+{
+    static const PyLongLayout PyLong_LAYOUT = {
+        PyLong_SHIFT,
+        sizeof(digit),
+        -1,  // least significant first
+        PY_LITTLE_ENDIAN ? -1 : 1,
+    };
+
+    return &PyLong_LAYOUT;
+}
+
+static inline int
+PyLong_Export(PyObject *obj, PyLongExport *export_long)
+{
+    if (!PyLong_Check(obj)) {
+        memset(export_long, 0, sizeof(*export_long));
+        PyErr_Format(PyExc_TypeError, "expected int, got %s",
+                     Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+
+    // Fast-path: try to convert to a int64_t
+    PyLongObject *self = (PyLongObject*)obj;
+    int overflow;
+#if SIZEOF_LONG == 8
+    long value = PyLong_AsLongAndOverflow(obj, &overflow);
+#else
+    // Windows has 32-bit long, so use 64-bit long long instead
+    long long value = PyLong_AsLongLongAndOverflow(obj, &overflow);
+#endif
+    Py_BUILD_ASSERT(sizeof(value) == sizeof(int64_t));
+    // the function cannot fail since obj is a PyLongObject
+    assert(!(value == -1 && PyErr_Occurred()));
+
+    if (!overflow) {
+        export_long->value = value;
+        export_long->negative = 0;
+        export_long->ndigits = 0;
+        export_long->digits = 0;
+        export_long->_reserved = 0;
+    }
+    else {
+        export_long->value = 0;
+        export_long->negative = _PyLong_Sign(obj) < 0;
+        export_long->ndigits = _PyLong_DigitCount(self);
+        if (export_long->ndigits == 0) {
+            export_long->ndigits = 1;
+        }
+        export_long->digits = _PyLong_GetDigits(self);
+        export_long->_reserved = (Py_uintptr_t)Py_NewRef(obj);
+    }
+    return 0;
+}
+
+static inline void
+PyLong_FreeExport(PyLongExport *export_long)
+{
+    PyObject *obj = (PyObject*)export_long->_reserved;
+
+    if (obj) {
+        export_long->_reserved = 0;
+        Py_DECREF(obj);
+    }
+}
+
+static inline PyLongWriter*
+PyLongWriter_Create(int negative, Py_ssize_t ndigits, void **digits)
+{
+    if (ndigits <= 0) {
+        PyErr_SetString(PyExc_ValueError, "ndigits must be positive");
+        return NULL;
+    }
+    assert(digits != NULL);
+
+    PyLongObject *obj = _PyLong_New(ndigits);
+    if (obj == NULL) {
+        return NULL;
+    }
+    _PyLong_SetSignAndDigitCount(obj, negative?-1:1, ndigits);
+
+    *digits = _PyLong_GetDigits(obj);
+    return (PyLongWriter*)obj;
+}
+
+static inline void
+PyLongWriter_Discard(PyLongWriter *writer)
+{
+    PyLongObject *obj = (PyLongObject *)writer;
+
+    assert(Py_REFCNT(obj) == 1);
+    Py_DECREF(obj);
+}
+
+static inline PyObject*
+PyLongWriter_Finish(PyLongWriter *writer)
+{
+    PyObject *obj = (PyObject *)writer;
+    PyLongObject *self = (PyLongObject*)obj;
+    Py_ssize_t j = _PyLong_DigitCount(self);
+    Py_ssize_t i = j;
+    int sign = _PyLong_Sign(obj);
+
+    assert(Py_REFCNT(obj) == 1);
+
+    // Normalize and get singleton if possible
+    while (i > 0 && _PyLong_GetDigits(self)[i-1] == 0) {
+        --i;
+    }
+    if (i != j) {
+        if (i == 0) {
+            sign = 0;
+        }
+        _PyLong_SetSignAndDigitCount(self, sign, i);
+    }
+    if (i <= 1) {
+        long val = sign * (long)(_PyLong_GetDigits(self)[0]);
+        Py_DECREF(obj);
+        return PyLong_FromLong(val);
+    }
+
+    return obj;
+}
+#endif
+
+
+#if PY_VERSION_HEX < 0x030C00A3
+#  define Py_T_SHORT      0
+#  define Py_T_INT        1
+#  define Py_T_LONG       2
+#  define Py_T_FLOAT      3
+#  define Py_T_DOUBLE     4
+#  define Py_T_STRING     5
+#  define _Py_T_OBJECT    6
+#  define Py_T_CHAR       7
+#  define Py_T_BYTE       8
+#  define Py_T_UBYTE      9
+#  define Py_T_USHORT     10
+#  define Py_T_UINT       11
+#  define Py_T_ULONG      12
+#  define Py_T_STRING_INPLACE  13
+#  define Py_T_BOOL       14
+#  define Py_T_OBJECT_EX  16
+#  define Py_T_LONGLONG   17
+#  define Py_T_ULONGLONG  18
+#  define Py_T_PYSSIZET   19
+
+#  if PY_VERSION_HEX >= 0x03000000 && !defined(PYPY_VERSION)
+#    define _Py_T_NONE      20
+#  endif
+
+#  define Py_READONLY            1
+#  define Py_AUDIT_READ          2
+#  define _Py_WRITE_RESTRICTED   4
+#endif
+
+
+// gh-127350 added Py_fopen() and Py_fclose() to Python 3.14a4
+#if PY_VERSION_HEX < 0x030E00A4
+static inline FILE* Py_fopen(PyObject *path, const char *mode)
+{
+#if 0x030400A2 <= PY_VERSION_HEX && !defined(PYPY_VERSION)
+    PyAPI_FUNC(FILE*) _Py_fopen_obj(PyObject *path, const char *mode);
+
+    return _Py_fopen_obj(path, mode);
+#else
+    FILE *f;
+    PyObject *bytes;
+#if PY_VERSION_HEX >= 0x03000000
+    if (!PyUnicode_FSConverter(path, &bytes)) {
+        return NULL;
+    }
+#else
+    if (!PyString_Check(path)) {
+        PyErr_SetString(PyExc_TypeError, "except str");
+        return NULL;
+    }
+    bytes = Py_NewRef(path);
+#endif
+    const char *path_bytes = PyBytes_AS_STRING(bytes);
+
+    f = fopen(path_bytes, mode);
+    Py_DECREF(bytes);
+
+    if (f == NULL) {
+        PyErr_SetFromErrnoWithFilenameObject(PyExc_OSError, path);
+        return NULL;
+    }
+    return f;
+#endif
+}
+
+static inline int Py_fclose(FILE *file)
+{
+    return fclose(file);
+}
+#endif
+
+
+#if 0x03080000 <= PY_VERSION_HEX && PY_VERSION_HEX < 0x030E0000 && !defined(PYPY_VERSION)
+PyAPI_FUNC(const PyConfig*) _Py_GetConfig(void);
+
+static inline PyObject*
+PyConfig_Get(const char *name)
+{
+    typedef enum {
+        _PyConfig_MEMBER_INT,
+        _PyConfig_MEMBER_UINT,
+        _PyConfig_MEMBER_ULONG,
+        _PyConfig_MEMBER_BOOL,
+        _PyConfig_MEMBER_WSTR,
+        _PyConfig_MEMBER_WSTR_OPT,
+        _PyConfig_MEMBER_WSTR_LIST,
+    } PyConfigMemberType;
+
+    typedef struct {
+        const char *name;
+        size_t offset;
+        PyConfigMemberType type;
+        const char *sys_attr;
+    } PyConfigSpec;
+
+#define PYTHONCAPI_COMPAT_SPEC(MEMBER, TYPE, sys_attr) \
+    {#MEMBER, offsetof(PyConfig, MEMBER), \
+     _PyConfig_MEMBER_##TYPE, sys_attr}
+
+    static const PyConfigSpec config_spec[] = {
+        PYTHONCAPI_COMPAT_SPEC(argv, WSTR_LIST, "argv"),
+        PYTHONCAPI_COMPAT_SPEC(base_exec_prefix, WSTR_OPT, "base_exec_prefix"),
+        PYTHONCAPI_COMPAT_SPEC(base_executable, WSTR_OPT, "_base_executable"),
+        PYTHONCAPI_COMPAT_SPEC(base_prefix, WSTR_OPT, "base_prefix"),
+        PYTHONCAPI_COMPAT_SPEC(bytes_warning, UINT, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(exec_prefix, WSTR_OPT, "exec_prefix"),
+        PYTHONCAPI_COMPAT_SPEC(executable, WSTR_OPT, "executable"),
+        PYTHONCAPI_COMPAT_SPEC(inspect, BOOL, _Py_NULL),
+#if 0x030C0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(int_max_str_digits, UINT, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(interactive, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(module_search_paths, WSTR_LIST, "path"),
+        PYTHONCAPI_COMPAT_SPEC(optimization_level, UINT, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(parser_debug, BOOL, _Py_NULL),
+#if 0x03090000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(platlibdir, WSTR, "platlibdir"),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(prefix, WSTR_OPT, "prefix"),
+        PYTHONCAPI_COMPAT_SPEC(pycache_prefix, WSTR_OPT, "pycache_prefix"),
+        PYTHONCAPI_COMPAT_SPEC(quiet, BOOL, _Py_NULL),
+#if 0x030B0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(stdlib_dir, WSTR_OPT, "_stdlib_dir"),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(use_environment, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(verbose, UINT, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(warnoptions, WSTR_LIST, "warnoptions"),
+        PYTHONCAPI_COMPAT_SPEC(write_bytecode, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(xoptions, WSTR_LIST, "_xoptions"),
+        PYTHONCAPI_COMPAT_SPEC(buffered_stdio, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(check_hash_pycs_mode, WSTR, _Py_NULL),
+#if 0x030B0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(code_debug_ranges, BOOL, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(configure_c_stdio, BOOL, _Py_NULL),
+#if 0x030D0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(cpu_count, INT, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(dev_mode, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(dump_refs, BOOL, _Py_NULL),
+#if 0x030B0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(dump_refs_file, WSTR_OPT, _Py_NULL),
+#endif
+#ifdef Py_GIL_DISABLED
+        PYTHONCAPI_COMPAT_SPEC(enable_gil, INT, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(faulthandler, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(filesystem_encoding, WSTR, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(filesystem_errors, WSTR, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(hash_seed, ULONG, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(home, WSTR_OPT, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(import_time, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(install_signal_handlers, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(isolated, BOOL, _Py_NULL),
+#ifdef MS_WINDOWS
+        PYTHONCAPI_COMPAT_SPEC(legacy_windows_stdio, BOOL, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(malloc_stats, BOOL, _Py_NULL),
+#if 0x030A0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(orig_argv, WSTR_LIST, "orig_argv"),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(parse_argv, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(pathconfig_warnings, BOOL, _Py_NULL),
+#if 0x030C0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(perf_profiling, UINT, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(program_name, WSTR, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(run_command, WSTR_OPT, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(run_filename, WSTR_OPT, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(run_module, WSTR_OPT, _Py_NULL),
+#if 0x030B0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(safe_path, BOOL, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(show_ref_count, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(site_import, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(skip_source_first_line, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(stdio_encoding, WSTR, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(stdio_errors, WSTR, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(tracemalloc, UINT, _Py_NULL),
+#if 0x030B0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(use_frozen_modules, BOOL, _Py_NULL),
+#endif
+        PYTHONCAPI_COMPAT_SPEC(use_hash_seed, BOOL, _Py_NULL),
+        PYTHONCAPI_COMPAT_SPEC(user_site_directory, BOOL, _Py_NULL),
+#if 0x030A0000 <= PY_VERSION_HEX
+        PYTHONCAPI_COMPAT_SPEC(warn_default_encoding, BOOL, _Py_NULL),
+#endif
+    };
+
+#undef PYTHONCAPI_COMPAT_SPEC
+
+    const PyConfigSpec *spec;
+    int found = 0;
+    for (size_t i=0; i < sizeof(config_spec) / sizeof(config_spec[0]); i++) {
+        spec = &config_spec[i];
+        if (strcmp(spec->name, name) == 0) {
+            found = 1;
+            break;
+        }
+    }
+    if (found) {
+        if (spec->sys_attr != NULL) {
+            PyObject *value = PySys_GetObject(spec->sys_attr);
+            if (value == NULL) {
+                PyErr_Format(PyExc_RuntimeError, "lost sys.%s", spec->sys_attr);
+                return NULL;
+            }
+            return Py_NewRef(value);
+        }
+
+        const PyConfig *config = _Py_GetConfig();
+        void *member = (char *)config + spec->offset;
+        switch (spec->type) {
+        case _PyConfig_MEMBER_INT:
+        case _PyConfig_MEMBER_UINT:
+        {
+            int value = *(int *)member;
+            return PyLong_FromLong(value);
+        }
+        case _PyConfig_MEMBER_BOOL:
+        {
+            int value = *(int *)member;
+            return PyBool_FromLong(value != 0);
+        }
+        case _PyConfig_MEMBER_ULONG:
+        {
+            unsigned long value = *(unsigned long *)member;
+            return PyLong_FromUnsignedLong(value);
+        }
+        case _PyConfig_MEMBER_WSTR:
+        case _PyConfig_MEMBER_WSTR_OPT:
+        {
+            wchar_t *wstr = *(wchar_t **)member;
+            if (wstr != NULL) {
+                return PyUnicode_FromWideChar(wstr, -1);
+            }
+            else {
+                return Py_NewRef(Py_None);
+            }
+        }
+        case _PyConfig_MEMBER_WSTR_LIST:
+        {
+            const PyWideStringList *list = (const PyWideStringList *)member;
+            PyObject *tuple = PyTuple_New(list->length);
+            if (tuple == NULL) {
+                return NULL;
+            }
+
+            for (Py_ssize_t i = 0; i < list->length; i++) {
+                PyObject *item = PyUnicode_FromWideChar(list->items[i], -1);
+                if (item == NULL) {
+                    Py_DECREF(tuple);
+                    return NULL;
+                }
+                PyTuple_SET_ITEM(tuple, i, item);
+            }
+            return tuple;
+        }
+        default:
+            Py_UNREACHABLE();
+        }
+    }
+
+    PyErr_Format(PyExc_ValueError, "unknown config option name: %s", name);
+    return NULL;
+}
+
+static inline int
+PyConfig_GetInt(const char *name, int *value)
+{
+    PyObject *obj = PyConfig_Get(name);
+    if (obj == NULL) {
+        return -1;
+    }
+
+    if (!PyLong_Check(obj)) {
+        Py_DECREF(obj);
+        PyErr_Format(PyExc_TypeError, "config option %s is not an int", name);
+        return -1;
+    }
+
+    int as_int = PyLong_AsInt(obj);
+    Py_DECREF(obj);
+    if (as_int == -1 && PyErr_Occurred()) {
+        PyErr_Format(PyExc_OverflowError,
+                     "config option %s value does not fit into a C int", name);
+        return -1;
+    }
+
+    *value = as_int;
+    return 0;
+}
+#endif  // PY_VERSION_HEX > 0x03090000 && !defined(PYPY_VERSION)
+
+// gh-133144 added PyUnstable_Object_IsUniquelyReferenced() to Python 3.14.0b1.
+// Adapted from  _PyObject_IsUniquelyReferenced() implementation.
+#if PY_VERSION_HEX < 0x030E00B0
+static inline int PyUnstable_Object_IsUniquelyReferenced(PyObject *obj)
+{
+#if !defined(Py_GIL_DISABLED)
+    return Py_REFCNT(obj) == 1;
+#else
+    // NOTE: the entire ob_ref_shared field must be zero, including flags, to
+    // ensure that other threads cannot concurrently create new references to
+    // this object.
+    return (_Py_IsOwnedByCurrentThread(obj) &&
+            _Py_atomic_load_uint32_relaxed(&obj->ob_ref_local) == 1 &&
+            _Py_atomic_load_ssize_relaxed(&obj->ob_ref_shared) == 0);
+#endif
+}
+#endif
+
+// gh-128926 added PyUnstable_TryIncRef() and PyUnstable_EnableTryIncRef() to
+// Python 3.14.0a5. Adapted from _Py_TryIncref() and _PyObject_SetMaybeWeakref().
+#if PY_VERSION_HEX < 0x030E00A5
+static inline int PyUnstable_TryIncRef(PyObject *op)
+{
+#ifndef Py_GIL_DISABLED
+    if (Py_REFCNT(op) > 0) {
+        Py_INCREF(op);
+        return 1;
+    }
+    return 0;
+#else
+    // _Py_TryIncrefFast()
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
+    local += 1;
+    if (local == 0) {
+        // immortal
+        return 1;
+    }
+    if (_Py_IsOwnedByCurrentThread(op)) {
+        _Py_INCREF_STAT_INC();
+        _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, local);
+#ifdef Py_REF_DEBUG
+        _Py_INCREF_IncRefTotal();
+#endif
+        return 1;
+    }
+
+    // _Py_TryIncRefShared()
+    Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->ob_ref_shared);
+    for (;;) {
+        // If the shared refcount is zero and the object is either merged
+        // or may not have weak references, then we cannot incref it.
+        if (shared == 0 || shared == _Py_REF_MERGED) {
+            return 0;
+        }
+
+        if (_Py_atomic_compare_exchange_ssize(
+                &op->ob_ref_shared,
+                &shared,
+                shared + (1 << _Py_REF_SHARED_SHIFT))) {
+#ifdef Py_REF_DEBUG
+            _Py_INCREF_IncRefTotal();
+#endif
+            _Py_INCREF_STAT_INC();
+            return 1;
+        }
+    }
+#endif
+}
+
+static inline void PyUnstable_EnableTryIncRef(PyObject *op)
+{
+#ifdef Py_GIL_DISABLED
+    // _PyObject_SetMaybeWeakref()
+    if (_Py_IsImmortal(op)) {
+        return;
+    }
+    for (;;) {
+        Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->ob_ref_shared);
+        if ((shared & _Py_REF_SHARED_FLAG_MASK) != 0) {
+            // Nothing to do if it's in WEAKREFS, QUEUED, or MERGED states.
+            return;
+        }
+        if (_Py_atomic_compare_exchange_ssize(
+                &op->ob_ref_shared, &shared, shared | _Py_REF_MAYBE_WEAKREF)) {
+            return;
+        }
+    }
+#else
+    (void)op;  // unused argument
+#endif
+}
+#endif
+
+
+#if PY_VERSION_HEX < 0x030F0000
+static inline PyObject*
+PySys_GetAttrString(const char *name)
+{
+#if PY_VERSION_HEX >= 0x03000000
+    PyObject *value = Py_XNewRef(PySys_GetObject(name));
+#else
+    PyObject *value = Py_XNewRef(PySys_GetObject((char*)name));
+#endif
+    if (value != NULL) {
+        return value;
+    }
+    if (!PyErr_Occurred()) {
+        PyErr_Format(PyExc_RuntimeError, "lost sys.%s", name);
+    }
+    return NULL;
+}
+
+static inline PyObject*
+PySys_GetAttr(PyObject *name)
+{
+#if PY_VERSION_HEX >= 0x03000000
+    const char *name_str = PyUnicode_AsUTF8(name);
+#else
+    const char *name_str = PyString_AsString(name);
+#endif
+    if (name_str == NULL) {
+        return NULL;
+    }
+
+    return PySys_GetAttrString(name_str);
+}
+
+static inline int
+PySys_GetOptionalAttrString(const char *name, PyObject **value)
+{
+#if PY_VERSION_HEX >= 0x03000000
+    *value = Py_XNewRef(PySys_GetObject(name));
+#else
+    *value = Py_XNewRef(PySys_GetObject((char*)name));
+#endif
+    if (*value != NULL) {
+        return 1;
+    }
+    return 0;
+}
+
+static inline int
+PySys_GetOptionalAttr(PyObject *name, PyObject **value)
+{
+#if PY_VERSION_HEX >= 0x03000000
+    const char *name_str = PyUnicode_AsUTF8(name);
+#else
+    const char *name_str = PyString_AsString(name);
+#endif
+    if (name_str == NULL) {
+        *value = NULL;
+        return -1;
+    }
+
+    return PySys_GetOptionalAttrString(name_str, value);
+}
+#endif  // PY_VERSION_HEX < 0x030F00A1
+
+
+#if PY_VERSION_HEX < 0x030F00A1
+typedef struct PyBytesWriter {
+    char small_buffer[256];
+    PyObject *obj;
+    Py_ssize_t size;
+} PyBytesWriter;
+
+static inline Py_ssize_t
+_PyBytesWriter_GetAllocated(PyBytesWriter *writer)
+{
+    if (writer->obj == NULL) {
+        return sizeof(writer->small_buffer);
+    }
+    else {
+        return PyBytes_GET_SIZE(writer->obj);
+    }
+}
+
+
+static inline int
+_PyBytesWriter_Resize_impl(PyBytesWriter *writer, Py_ssize_t size,
+                           int resize)
+{
+    int overallocate = resize;
+    assert(size >= 0);
+
+    if (size <= _PyBytesWriter_GetAllocated(writer)) {
+        return 0;
+    }
+
+    if (overallocate) {
+#ifdef MS_WINDOWS
+        /* On Windows, overallocate by 50% is the best factor */
+        if (size <= (PY_SSIZE_T_MAX - size / 2)) {
+            size += size / 2;
+        }
+#else
+        /* On Linux, overallocate by 25% is the best factor */
+        if (size <= (PY_SSIZE_T_MAX - size / 4)) {
+            size += size / 4;
+        }
+#endif
+    }
+
+    if (writer->obj != NULL) {
+        if (_PyBytes_Resize(&writer->obj, size)) {
+            return -1;
+        }
+        assert(writer->obj != NULL);
+    }
+    else {
+        writer->obj = PyBytes_FromStringAndSize(NULL, size);
+        if (writer->obj == NULL) {
+            return -1;
+        }
+
+        if (resize) {
+            assert((size_t)size > sizeof(writer->small_buffer));
+            memcpy(PyBytes_AS_STRING(writer->obj),
+                   writer->small_buffer,
+                   sizeof(writer->small_buffer));
+        }
+    }
+    return 0;
+}
+
+static inline void*
+PyBytesWriter_GetData(PyBytesWriter *writer)
+{
+    if (writer->obj == NULL) {
+        return writer->small_buffer;
+    }
+    else {
+        return PyBytes_AS_STRING(writer->obj);
+    }
+}
+
+static inline Py_ssize_t
+PyBytesWriter_GetSize(PyBytesWriter *writer)
+{
+    return writer->size;
+}
+
+static inline void
+PyBytesWriter_Discard(PyBytesWriter *writer)
+{
+    if (writer == NULL) {
+        return;
+    }
+
+    Py_XDECREF(writer->obj);
+    PyMem_Free(writer);
+}
+
+static inline PyBytesWriter*
+PyBytesWriter_Create(Py_ssize_t size)
+{
+    if (size < 0) {
+        PyErr_SetString(PyExc_ValueError, "size must be >= 0");
+        return NULL;
+    }
+
+    PyBytesWriter *writer = (PyBytesWriter*)PyMem_Malloc(sizeof(PyBytesWriter));
+    if (writer == NULL) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+
+    writer->obj = NULL;
+    writer->size = 0;
+
+    if (size >= 1) {
+        if (_PyBytesWriter_Resize_impl(writer, size, 0) < 0) {
+            PyBytesWriter_Discard(writer);
+            return NULL;
+        }
+        writer->size = size;
+    }
+    return writer;
+}
+
+static inline PyObject*
+PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
+{
+    PyObject *result;
+    if (size == 0) {
+        result = PyBytes_FromStringAndSize("", 0);
+    }
+    else if (writer->obj != NULL) {
+        if (size != PyBytes_GET_SIZE(writer->obj)) {
+            if (_PyBytes_Resize(&writer->obj, size)) {
+                goto error;
+            }
+        }
+        result = writer->obj;
+        writer->obj = NULL;
+    }
+    else {
+        result = PyBytes_FromStringAndSize(writer->small_buffer, size);
+    }
+    PyBytesWriter_Discard(writer);
+    return result;
+
+error:
+    PyBytesWriter_Discard(writer);
+    return NULL;
+}
+
+static inline PyObject*
+PyBytesWriter_Finish(PyBytesWriter *writer)
+{
+    return PyBytesWriter_FinishWithSize(writer, writer->size);
+}
+
+static inline PyObject*
+PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf)
+{
+    Py_ssize_t size = (char*)buf - (char*)PyBytesWriter_GetData(writer);
+    if (size < 0 || size > _PyBytesWriter_GetAllocated(writer)) {
+        PyBytesWriter_Discard(writer);
+        PyErr_SetString(PyExc_ValueError, "invalid end pointer");
+        return NULL;
+    }
+
+    return PyBytesWriter_FinishWithSize(writer, size);
+}
+
+static inline int
+PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
+{
+    if (size < 0) {
+        PyErr_SetString(PyExc_ValueError, "size must be >= 0");
+        return -1;
+    }
+    if (_PyBytesWriter_Resize_impl(writer, size, 1) < 0) {
+        return -1;
+    }
+    writer->size = size;
+    return 0;
+}
+
+static inline int
+PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t size)
+{
+    if (size < 0 && writer->size + size < 0) {
+        PyErr_SetString(PyExc_ValueError, "invalid size");
+        return -1;
+    }
+    if (size > PY_SSIZE_T_MAX - writer->size) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    size = writer->size + size;
+
+    if (_PyBytesWriter_Resize_impl(writer, size, 1) < 0) {
+        return -1;
+    }
+    writer->size = size;
+    return 0;
+}
+
+static inline void*
+PyBytesWriter_GrowAndUpdatePointer(PyBytesWriter *writer,
+                                   Py_ssize_t size, void *buf)
+{
+    Py_ssize_t pos = (char*)buf - (char*)PyBytesWriter_GetData(writer);
+    if (PyBytesWriter_Grow(writer, size) < 0) {
+        return NULL;
+    }
+    return (char*)PyBytesWriter_GetData(writer) + pos;
+}
+
+static inline int
+PyBytesWriter_WriteBytes(PyBytesWriter *writer,
+                         const void *bytes, Py_ssize_t size)
+{
+    if (size < 0) {
+        size_t len = strlen((const char*)bytes);
+        if (len > (size_t)PY_SSIZE_T_MAX) {
+            PyErr_NoMemory();
+            return -1;
+        }
+        size = (Py_ssize_t)len;
+    }
+
+    Py_ssize_t pos = writer->size;
+    if (PyBytesWriter_Grow(writer, size) < 0) {
+        return -1;
+    }
+    char *buf = (char*)PyBytesWriter_GetData(writer);
+    memcpy(buf + pos, bytes, (size_t)size);
+    return 0;
+}
+
+static inline int
+PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
+                     Py_GCC_ATTRIBUTE((format(printf, 2, 3)));
+
+static inline int
+PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
+{
+    va_list vargs;
+    va_start(vargs, format);
+    PyObject *str = PyBytes_FromFormatV(format, vargs);
+    va_end(vargs);
+
+    if (str == NULL) {
+        return -1;
+    }
+    int res = PyBytesWriter_WriteBytes(writer,
+                                       PyBytes_AS_STRING(str),
+                                       PyBytes_GET_SIZE(str));
+    Py_DECREF(str);
+    return res;
+}
+#endif  // PY_VERSION_HEX < 0x030F00A1
+
+
+#if PY_VERSION_HEX < 0x030F00A1
+static inline PyObject*
+PyTuple_FromArray(PyObject *const *array, Py_ssize_t size)
+{
+    PyObject *tuple = PyTuple_New(size);
+    if (tuple == NULL) {
+        return NULL;
+    }
+    for (Py_ssize_t i=0; i < size; i++) {
+        PyObject *item = array[i];
+        PyTuple_SET_ITEM(tuple, i, Py_NewRef(item));
+    }
+    return tuple;
+}
+#endif
+
+
+#if PY_VERSION_HEX < 0x030F00A1
+static inline Py_hash_t
+PyUnstable_Unicode_GET_CACHED_HASH(PyObject *op)
+{
+#ifdef PYPY_VERSION
+    (void)op;  // unused argument
+    return -1;
+#elif PY_VERSION_HEX >= 0x03000000
+    return ((PyASCIIObject*)op)->hash;
+#else
+    return ((PyUnicodeObject*)op)->hash;
+#endif
+}
+#endif
+
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // PYTHONCAPI_COMPAT
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/schema_info.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/schema_info.h
new file mode 100644
index 0000000000000000000000000000000000000000..e022af2a2a585fabe2c8b8867f0e6034d22fbf84
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/schema_info.h
@@ -0,0 +1,121 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/jit/frontend/function_schema_parser.h>
+#include <unordered_set>
+
+namespace torch::utils {
+
+using SchemaSpecialCasePair =
+    std::pair<c10::FunctionSchema, std::unordered_set<std::string>>;
+/**
+ * class SchemaInfo
+ *
+ * FunctionSchema wrapper that publicizes argument value specific operator
+ * behavior (mutation, aliasing, special cases, etc...)
+ */
+
+struct TORCH_API SchemaInfo {
+ public:
+  explicit SchemaInfo(c10::FunctionSchema schema)
+      : schema_(std::move(schema)),
+        alias_maps_current_(false),
+        has_init_(false) {}
+  explicit SchemaInfo(const char* signature)
+      : schema_(torch::jit::parseSchema(signature)),
+        alias_maps_current_(false),
+        has_init_(false) {}
+
+  bool is_mutable();
+
+  bool is_mutable(const c10::SchemaArgument& argument);
+
+  bool is_mutable(std::string_view name);
+
+  bool has_argument(std::string_view name);
+
+  bool is_nondeterministic() const;
+
+  // Returns whether lhs and rhs may alias directly.
+  // This does not account for cases where lhs or rhs are a container that
+  // may contain elements that alias the other argument.
+  // Besides the checks already included in FunctionSchema::may_alias, this
+  // method also accounts special aliasing cases causes by aliasing argument
+  // values supplied from addArgumentValue.
+  bool may_alias(
+      const c10::SchemaArgument& lhs,
+      const c10::SchemaArgument& rhs);
+
+  // Returns whether lhs and rhs may alias directly or whether lhs/rhs are a
+  // container that may contain elements that alias the other argument. Besides
+  // the checks already included in FunctionSchema::may_contain_alias, this
+  // method also accounts for special aliasing cases causes by aliasing argument
+  // values supplied from addArgumentValue. bidirectional = false only returns
+  // whether lhs may contain an alias of rhs while bidirectional = true returns
+  // both directions.
+  bool may_contain_alias(
+      const c10::SchemaArgument& lhs,
+      const c10::SchemaArgument& rhs,
+      bool bidirectional = true);
+
+  void addArgumentValue(const std::string& name, const at::IValue& value);
+
+  void addArgumentValues(
+      const std::vector<std::optional<at::IValue>>& value_list);
+
+  void addArgumentValues(
+      const std::unordered_map<std::string, at::IValue>& values);
+
+  bool hasInputArgumentNamed(const std::string& name) const;
+
+ private:
+  // This function enforces more conservative results when the TORCH_WARN is
+  // triggered from above due to duplicates in an argument list
+  void ensureConservativity(
+      const std::unordered_set<at::Symbol>& duplicates,
+      const std::vector<c10::Argument>& arguments_list,
+      c10::SchemaArgType type);
+
+  void initSchemaInfo();
+
+  void generateAliasMaps();
+
+  bool mayContainAliasImpl(
+      const c10::SchemaArgument& lhs,
+      const c10::SchemaArgument& rhs);
+
+  static std::vector<c10::FunctionSchema> getNonDeterministicOps();
+
+  static std::vector<SchemaSpecialCasePair> getTrainingOps();
+
+  const std::unordered_set<c10::SchemaArgument>& wildcardSet();
+
+  const std::unordered_set<c10::SchemaArgument>& containerSet();
+
+  // Set of all wildcard arguments
+  std::unordered_set<c10::SchemaArgument> wildcard_set_;
+
+  // Set of all container arguments
+  std::unordered_set<c10::SchemaArgument> container_set_;
+
+  // Map of argument IValues
+  std::unordered_map<std::string, at::IValue> value_map_;
+
+  // Alias map of inputs with each other
+  std::vector<std::unordered_set<size_t>> input_alias_map_;
+
+  // Alias map of outputs to inputs
+  std::vector<std::unordered_set<size_t>> output_alias_map_;
+
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
+  const c10::FunctionSchema schema_;
+
+  bool alias_maps_current_;
+
+  bool has_init_;
+};
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/structseq.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/structseq.h
new file mode 100644
index 0000000000000000000000000000000000000000..8d40d6bf4adc54ded2638bf0c67062a733ce9f5e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/structseq.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace torch::utils {
+
+PyObject* returned_structseq_repr(PyStructSequence* obj);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_apply.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_apply.h
new file mode 100644
index 0000000000000000000000000000000000000000..062e093f123e11ee7492cf297b83c8087b336178
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_apply.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/Tensor.h>
+#include <torch/csrc/python_headers.h>
+
+namespace torch::utils {
+
+const at::Tensor& apply_(const at::Tensor& self, PyObject* fn);
+const at::Tensor& map_(
+    const at::Tensor& self,
+    const at::Tensor& other_,
+    PyObject* fn);
+const at::Tensor& map2_(
+    const at::Tensor& self,
+    const at::Tensor& x_,
+    const at::Tensor& y_,
+    PyObject* fn);
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_dtypes.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_dtypes.h
new file mode 100644
index 0000000000000000000000000000000000000000..978bc25e12052f95b40120ba8c097c0940a955f6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_dtypes.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/ScalarType.h>
+#include <string>
+#include <tuple>
+
+namespace torch::utils {
+
+std::pair<std::string, std::string> getDtypeNames(at::ScalarType scalarType);
+
+void initializeDtypes();
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_flatten.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_flatten.h
new file mode 100644
index 0000000000000000000000000000000000000000..ea7652eede79d5e40858224f096da4242b8ef41b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_flatten.h
@@ -0,0 +1,89 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/core/functional.h>
+#include <c10/core/TensorOptions.h>
+#include <torch/csrc/Export.h>
+#include <utility>
+
+namespace torch::utils {
+
+/// Generate an ID for a combination of tensor backend + scalar type to be used
+/// when ordering tensors ('like' tensors are grouped by pulling out their
+/// backend + scalar type, so this function combines that into a single number)
+inline size_t type_id(const at::Tensor& tensor) {
+  return static_cast<size_t>(tensor.options().backend()) *
+      static_cast<size_t>(at::ScalarType::NumOptions) +
+      static_cast<size_t>(tensor.scalar_type());
+}
+
+inline at::Tensor flatten_dense_tensors(at::TensorList tensors) {
+  return at::flatten_dense_tensors(tensors);
+}
+
+inline std::vector<at::Tensor> unflatten_dense_tensors(
+    const at::Tensor& flat,
+    at::TensorList tensors) {
+  return at::unflatten_dense_tensors(flat, tensors);
+}
+
+struct TensorGroup {
+  std::vector<at::Tensor> tensors;
+  size_t size = 0;
+
+  size_t type_id() {
+    AT_ASSERT(!tensors.empty());
+    return ::torch::utils::type_id(tensors[0]);
+  }
+
+  const at::TensorOptions options() {
+    AT_ASSERT(!tensors.empty());
+    return tensors[0].options();
+  }
+};
+
+// Helper function that takes a list of tensors and splits them into tensor
+// groups by the size limit and outputs these tensor groups. If the input
+// tensors are of different tensor types, they will be split into different
+// groups as well.
+//
+// Two options of splitting provided to the user,
+//
+// Imagine the size_limit is 256 and the list of input tensors are:
+// tensor_a(fp16 - 128 bytes),
+// tensor_b(fp32 - 256 bytes),
+// tensor_c(fp16 - 128 bytes),
+//
+// when fine_grained == false:
+// The function will read the list of tensors sequentially and accumulate
+// enough tensors for each data type until the size_limit, therefore:
+// it will output: {{tensor_a, tensor_c}, {tensor_b}}
+//
+// when fine_grained == true:
+// The function will read the list of tensors sequentially and  accumulate
+// enough tensors for all data types until the size_limit, and then split
+// the accumulated tensors into different groups by data types, therefore:
+// it will output: {{tensor_a}, {tensor_b}, {tensor_c}}
+TORCH_API std::vector<TensorGroup> take_tensors(
+    at::TensorList tensors,
+    size_t size_limit,
+    bool fine_grained = false);
+
+TORCH_API void reorder_tensors_like(
+    std::vector<at::Tensor>& tensors,
+    at::TensorList order);
+
+TORCH_API std::pair<at::Tensor, at::Tensor> flatten_sparse_tensors(
+    at::TensorList tensors);
+
+TORCH_API std::vector<at::Tensor> unflatten_sparse_tensors(
+    const at::Tensor& flat_indices,
+    const at::Tensor& flat_values,
+    at::TensorList tensors);
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_layouts.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_layouts.h
new file mode 100644
index 0000000000000000000000000000000000000000..5521d1c4e0b8875c063e1ee14dae951c58e32bba
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_layouts.h
@@ -0,0 +1,12 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+namespace torch::utils {
+
+void initializeLayouts();
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_list.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_list.h
new file mode 100644
index 0000000000000000000000000000000000000000..8705ab13d1483cb270536d660217527e0c21b308
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_list.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+namespace at {
+class Tensor;
+}
+
+namespace torch::utils {
+
+PyObject* tensor_to_list(const at::Tensor& tensor);
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_memoryformats.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_memoryformats.h
new file mode 100644
index 0000000000000000000000000000000000000000..489033081dd89b705f048aba9fbb7f228851868a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_memoryformats.h
@@ -0,0 +1,20 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/MemoryFormat.h>
+#include <torch/csrc/Export.h>
+#include <torch/csrc/utils/python_stub.h>
+
+namespace torch::utils {
+
+void initializeMemoryFormats();
+
+// This methods returns a borrowed reference!
+TORCH_PYTHON_API PyObject* getTHPMemoryFormat(
+    c10::MemoryFormat /*memory_format*/);
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_new.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_new.h
new file mode 100644
index 0000000000000000000000000000000000000000..690695179a5462c9b52d45659d738fb97b50eacf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_new.h
@@ -0,0 +1,141 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+#include <torch/csrc/utils/python_arg_parser.h>
+
+#include <ATen/core/Tensor.h>
+
+namespace torch::utils {
+
+// NOTE: [torch.tensor, lift_fresh, and device movement]
+//
+// The `only_lift_cpu_tensors` flag controls what happens on torch.tensor([1, 2,
+// 3], device="cuda") (or any non-CPU devices).
+//
+// If false (default):
+// - the data gets moved into a CPU Tensor
+// - then, it gets moved to cuda (via .to)
+// - finally, we call lift_fresh() on it.
+// Steps 1 and 2 happen with all modes disabled.
+//
+// If true:
+// - the data gets moved into a CPU Tensor (with correct dtype)
+// - we call lift_fresh() on it
+// - finally, we move it to cuda (via .to)
+// Step 1 happens with all modes disabled.
+//
+// `only_lift_cpu_tensors=true` is useful to prevent CUDA initialization under
+// FakeTensorMode because it avoids moving concrete data to CUDA.
+TORCH_API bool only_lift_cpu_tensors();
+TORCH_API void set_only_lift_cpu_tensors(bool value);
+
+at::Tensor base_tensor_ctor(PyObject* args, PyObject* kwargs);
+TORCH_PYTHON_API at::Tensor legacy_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+at::Tensor legacy_tensor_new(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+at::Tensor indexing_tensor_from_data(
+    c10::TensorOptions options,
+    at::ScalarType scalar_type,
+    std::optional<at::Device> device,
+    PyObject* data);
+at::Tensor sparse_coo_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+void _validate_sparse_coo_tensor_args(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+
+at::Tensor sparse_compressed_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+at::Tensor sparse_csr_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+at::Tensor sparse_csc_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+at::Tensor sparse_bsr_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+at::Tensor sparse_bsc_tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+
+void _validate_sparse_compressed_tensor_args(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+void _validate_sparse_csr_tensor_args(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+void _validate_sparse_csc_tensor_args(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+void _validate_sparse_bsr_tensor_args(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+void _validate_sparse_bsc_tensor_args(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+
+at::Tensor tensor_ctor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+at::Tensor as_tensor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PythonArgs& r);
+at::Tensor new_tensor(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+at::Tensor new_ones(
+    c10::DispatchKey dispatch_key,
+    at::ScalarType scalar_type,
+    PyObject* args,
+    PyObject* kwargs);
+at::Tensor tensor_frombuffer(
+    PyObject* buffer,
+    at::ScalarType dtype,
+    int64_t count,
+    int64_t offset,
+    bool requires_grad);
+at::Tensor tensor_fromDLPack(PyObject* data);
+at::Tensor asarray(
+    PyObject* obj,
+    std::optional<c10::ScalarType> dtype,
+    std::optional<c10::Device> device,
+    std::optional<bool> copy,
+    std::optional<bool> requires_grad);
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_numpy.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_numpy.h
new file mode 100644
index 0000000000000000000000000000000000000000..b43522b8708452c5847be2873908d5c0f57ae6ad
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_numpy.h
@@ -0,0 +1,37 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/Tensor.h>
+#include <torch/csrc/python_headers.h>
+
+namespace torch::utils {
+
+TORCH_API PyObject* tensor_to_numpy(
+    const at::Tensor& tensor,
+    bool force = false);
+
+TORCH_API at::Tensor tensor_from_numpy(
+    PyObject* obj,
+    bool warn_if_not_writeable = true);
+
+TORCH_API int aten_to_numpy_dtype(const at::ScalarType scalar_type);
+TORCH_API at::ScalarType numpy_dtype_to_aten(int dtype);
+
+TORCH_API bool is_numpy_available();
+TORCH_API bool is_numpy_int(PyObject* obj);
+TORCH_API bool is_numpy_bool(PyObject* obj);
+TORCH_API bool is_numpy_scalar(PyObject* obj);
+
+void warn_numpy_not_writeable();
+at::Tensor tensor_from_cuda_array_interface(
+    PyObject* obj,
+    std::optional<c10::Device> device_opt = std::nullopt);
+
+void validate_numpy_for_dlpack_deleter_bug();
+bool is_numpy_dlpack_deleter_bugged();
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_qschemes.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_qschemes.h
new file mode 100644
index 0000000000000000000000000000000000000000..d7b24333ad63eab535b31d529754b3a4d3b30b82
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_qschemes.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/QScheme.h>
+
+namespace torch::utils {
+
+PyObject* getTHPQScheme(at::QScheme qscheme);
+void initializeQSchemes();
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_types.h
new file mode 100644
index 0000000000000000000000000000000000000000..d7891fedb997e32d9481613e9cce8c1754c8bd11
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/tensor_types.h
@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/DeprecatedTypeProperties.h>
+#include <c10/core/TensorOptions.h>
+#include <utility>
+#include <vector>
+
+namespace torch::utils {
+
+std::string options_to_string(const at::TensorOptions& options);
+std::string type_to_string(const at::DeprecatedTypeProperties& type);
+at::TensorOptions options_from_string(const std::string& str);
+
+// return a vector of all "declared" types, even those that weren't compiled
+std::vector<std::pair<at::Backend, at::ScalarType>> all_declared_types();
+
+// return python module name of backend, like torch.cuda, torch.foo
+const char* backend_to_string(const at::Backend& backend);
+
+} // namespace torch::utils
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/throughput_benchmark-inl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/throughput_benchmark-inl.h
new file mode 100644
index 0000000000000000000000000000000000000000..fa267bc1fc3a8ef59b60e5cee3777f696bd6b509
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/throughput_benchmark-inl.h
@@ -0,0 +1,176 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <random>
+#include <thread>
+
+#include <torch/csrc/autograd/profiler.h>
+#include <torch/csrc/jit/python/pybind_utils.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include <ATen/Parallel.h>
+#include <ATen/autocast_mode.h>
+#include <c10/core/GradMode.h>
+#include <c10/core/impl/LocalDispatchKeySet.h>
+#include <c10/util/irange.h>
+
+namespace torch::throughput_benchmark::detail {
+
+template <class Input, class Output, class Model>
+BenchmarkExecutionStats BenchmarkHelper<Input, Output, Model>::benchmark(
+    const BenchmarkConfig& config) const {
+  CHECK(initialized_);
+  TORCH_CHECK(
+      config.num_worker_threads == 1,
+      "Only parallelization by callers is supported");
+
+  LOG(INFO) << at::get_parallel_info();
+
+  // We pre-generate inputs here for each of the threads. This allows us to
+  // safely move inputs out for each of the threads independently and thus avoid
+  // overhead from the benchmark runner itself
+  std::vector<std::vector<Input>> thread_inputs(config.num_calling_threads);
+  std::vector<size_t> input_iters(config.num_calling_threads);
+  {
+    std::random_device seeder;
+    std::mt19937 engine(seeder());
+    TORCH_CHECK(
+        !inputs_.empty(),
+        "Please provide benchmark inputs."
+        "Did you forget to call add_input()? ");
+    std::uniform_int_distribution<int> dist(0, inputs_.size() - 1);
+
+    for (const auto thread_id : c10::irange(config.num_calling_threads)) {
+      // Just in case we generate num_iters inputs for each of the threads
+      // This was if one thread does all the work we will be fine
+      for (const auto i [[maybe_unused]] :
+           c10::irange(config.num_iters + config.num_warmup_iters)) {
+        thread_inputs[thread_id].push_back(cloneInput(inputs_[dist(engine)]));
+      }
+      input_iters[thread_id] = 0;
+    }
+  }
+
+  std::mutex m;
+  std::condition_variable worker_main_cv;
+  std::condition_variable main_worker_cv;
+  // TODO: add GUARDED_BY once it is available
+  int64_t initialized{0};
+  int64_t finished{0};
+  bool start{false};
+  std::atomic<int64_t> num_attempted_iters{0};
+  std::vector<std::thread> callers;
+
+  callers.reserve(config.num_calling_threads);
+
+  static constexpr auto& DEVICES = at::autocast::_AUTOCAST_SUPPORTED_DEVICES;
+  std::array<bool, DEVICES.size()> autocast_enabled;
+  std::array<at::ScalarType, DEVICES.size()> autocast_dtype;
+  for (size_t i = 0; i < DEVICES.size(); i++) {
+    autocast_enabled[i] = at::autocast::is_autocast_enabled(DEVICES[i]);
+    autocast_dtype[i] = at::autocast::get_autocast_dtype(DEVICES[i]);
+  }
+  bool autocast_cache_enabled = at::autocast::is_autocast_cache_enabled();
+  bool tls_grad_enabled = c10::GradMode::is_enabled();
+  c10::impl::LocalDispatchKeySet tls_key_set =
+      c10::impl::tls_local_dispatch_key_set();
+
+  for (const auto thread_id : c10::irange(config.num_calling_threads)) {
+    callers.emplace_back([&, thread_id]() {
+      // We use conditional variable as a barrier to make sure each thread
+      // performs required warmeup iterations before we start measuring
+      c10::GradMode::set_enabled(tls_grad_enabled);
+      c10::impl::_force_tls_local_dispatch_key_set(tls_key_set);
+      for (size_t i = 0; i < DEVICES.size(); i++) {
+        at::autocast::set_autocast_enabled(DEVICES[i], autocast_enabled[i]);
+        at::autocast::set_autocast_dtype(DEVICES[i], autocast_dtype[i]);
+      }
+      at::autocast::set_autocast_cache_enabled(autocast_cache_enabled);
+
+      for (const auto j : c10::irange(config.num_warmup_iters)) {
+        (void)j;
+        runOnce(std::move(thread_inputs[thread_id][input_iters[thread_id]]));
+        ++input_iters[thread_id];
+      }
+      {
+        std::unique_lock<std::mutex> lock(m);
+        ++initialized;
+        worker_main_cv.notify_one();
+        // NOLINTNEXTLINE(bugprone-infinite-loop)
+        while (!start) {
+          main_worker_cv.wait(lock);
+        }
+      }
+      LOG(INFO) << "Starting forward thread " << thread_id;
+      while (num_attempted_iters.fetch_add(1) < config.num_iters) {
+        runOnce(std::move(thread_inputs[thread_id][input_iters[thread_id]]));
+        ++input_iters[thread_id];
+      }
+
+      {
+        std::unique_lock<std::mutex> lock(m);
+        ++finished;
+        worker_main_cv.notify_one();
+        LOG(INFO) << "Shutting down forward thread " << thread_id
+                  << ". Total number of finished threads: " << finished;
+      }
+    });
+  }
+
+  using Clock = std::chrono::high_resolution_clock;
+  using RecordProfile = torch::autograd::profiler::RecordProfile;
+  using TimePoint = std::chrono::time_point<Clock>;
+  TimePoint start_time;
+
+  std::unique_ptr<RecordProfile> profiler_guard;
+  {
+    std::unique_lock<std::mutex> lock(m);
+    while (initialized != config.num_calling_threads) {
+      worker_main_cv.wait(lock);
+    }
+    if (!config.profiler_output_path.empty()) {
+      LOG(INFO) << "Using Autograd profiler. Trace will be saved to "
+                << config.profiler_output_path;
+      profiler_guard =
+          std::make_unique<RecordProfile>(config.profiler_output_path);
+    }
+    LOG(INFO) << "Starting threads";
+    start = true;
+    start_time = Clock::now();
+  }
+
+  main_worker_cv.notify_all();
+  {
+    std::unique_lock<std::mutex> lock(m);
+    worker_main_cv.wait(
+        lock, [&]() { return finished == config.num_calling_threads; });
+  }
+  auto end_time = std::chrono::high_resolution_clock::now();
+  profiler_guard.reset();
+  LOG(INFO) << "Finished benchmark";
+
+  BenchmarkExecutionStats stats;
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  float total_time_ms = std::chrono::duration_cast<std::chrono::nanoseconds>(
+                            end_time - start_time)
+                            .count() /
+      1000.0 / 1000.0;
+  // We use config.num_iters instead of num_attempted_iters as it is
+  // repsesatative of the real work done. Last attempted iteration on each
+  // calling threads doesn't represent the real work (i.e. running the model)
+  stats.latency_avg_ms =
+      // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+      total_time_ms * config.num_calling_threads / config.num_iters;
+  stats.num_iters = config.num_iters;
+
+  for (auto& t : callers) {
+    t.join();
+  }
+  return stats;
+}
+
+} // namespace torch::throughput_benchmark::detail
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/throughput_benchmark.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/throughput_benchmark.h
new file mode 100644
index 0000000000000000000000000000000000000000..4206d38fa32a08ff661e4686ac6a41a822993770
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/throughput_benchmark.h
@@ -0,0 +1,204 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/ivalue.h>
+#include <pybind11/pybind11.h>
+#include <torch/csrc/jit/api/module.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include <torch/csrc/jit/python/pybind_utils.h>
+
+#include <iosfwd>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace py = pybind11;
+
+namespace torch::throughput_benchmark {
+
+/**
+ * The struct is used to provide results of a benchmark to the caller
+ * In the future all additional statistics should be added here.
+ */
+struct BenchmarkExecutionStats {
+  float latency_avg_ms{-1};
+  int64_t num_iters{-1};
+};
+
+std::ostream& operator<<(
+    std::ostream& os,
+    const BenchmarkExecutionStats& value);
+
+/**
+ * Use this struct in order to configure a throughput benchmark run.
+ * This struct should include parameters related to threading, batching, number
+ * of iterations, warm-up, etc. More configs can be added as needed.
+ * General rule here is that only things that c++ must(!) to be aware of should
+ * be here. If we can keep other parts in python, we should keep them there.
+ * This is typical for things that are not perf critical and don't affect
+ * execution statistics benchmark returns.
+ */
+struct BenchmarkConfig {
+ public:
+  // Calling threads are those threads that are calling into a module in
+  // parallel.
+  int num_calling_threads{1};
+  // Worker threads are not supported yet. This is just an example that we plan
+  // to support some sort of multi-threaded forward calls. We may change this
+  // setting in the future to support different intra and inter op parallelism
+  // which is not available in PyTorch yet
+  int num_worker_threads{1};
+  // Warmup iters are used to make sure we run a module a few times before
+  // actually measuring things. This way we avoid cold caches and any other
+  // similar problems
+  int num_warmup_iters{1};
+  // Number of iterations the benchmark should run with. This number is separate
+  // from the warmup iterations
+  int64_t num_iters{100};
+  // If set autograd profiler will be enabled. I.e. this variable would be
+  // created before the main benchmark loop (but after the warmup):
+  // RecordProfile guard(profiler_output_path);
+  std::string profiler_output_path;
+};
+
+namespace detail {
+
+/**
+ * A helper class to abstract out different models we test throughput of
+ */
+template <class Input, class Output, class Model>
+class BenchmarkHelper {
+ public:
+  BenchmarkHelper();
+  explicit BenchmarkHelper(Model model)
+      : model_(std::move(model)), initialized_(true) {}
+
+  // This method to be used in benchmark() method
+  // Note that there is no result. This way we don't have to call this under GIL
+  // even when running in the nn.Module mode. Otherwise destructor of the result
+  // would race with Python
+  void runOnce(Input&&) const;
+  // This method is to be used when calling from Python directly
+  Output runOnce(const py::args&, const py::kwargs&) const;
+  // Aggregate input in the format Model expects in order to avoid further
+  // conversions at the benchmark time
+  void addInput(py::args&&, py::kwargs&&);
+  void addInput(Input&&);
+  BenchmarkExecutionStats benchmark(const BenchmarkConfig& config) const;
+
+  bool initialized() const {
+    return initialized_;
+  }
+
+  // Destructor doesn't require the GIL because it is going to be executed on
+  // the PyThon thread
+  std::vector<Input> inputs_;
+  Model model_;
+  bool initialized_{false};
+};
+
+struct C10_HIDDEN ModuleInput {
+  ModuleInput(ModuleInput&& other) = default;
+
+  ModuleInput(const ModuleInput&) = delete;
+  ModuleInput& operator=(ModuleInput& other) = delete;
+  ModuleInput& operator=(ModuleInput&& other) = delete;
+  ~ModuleInput() = default;
+
+  ModuleInput(py::args&& args, py::kwargs&& kwargs)
+      : args(std::move(args)), kwargs(std::move(kwargs)) {}
+
+  py::args args;
+  py::kwargs kwargs;
+};
+typedef py::object ModuleOutput;
+typedef std::vector<at::IValue> ScriptModuleInput;
+typedef at::IValue ScriptModuleOutput;
+
+template <class Input>
+Input cloneInput(const Input& input);
+
+typedef BenchmarkHelper<ScriptModuleInput, at::IValue, jit::Module>
+    ScriptModuleBenchmark;
+template <>
+inline BenchmarkHelper<ScriptModuleInput, at::IValue, jit::Module>::
+    BenchmarkHelper()
+    : model_("Module", std::make_shared<jit::CompilationUnit>()),
+      initialized_(false) {}
+typedef BenchmarkHelper<ModuleInput, py::object, py::object> ModuleBenchmark;
+template <>
+inline BenchmarkHelper<ModuleInput, py::object, py::object>::BenchmarkHelper()
+    : initialized_(false) {}
+
+template <>
+void ScriptModuleBenchmark::runOnce(ScriptModuleInput&& input) const;
+
+template <>
+ScriptModuleOutput ScriptModuleBenchmark::runOnce(
+    const py::args& args,
+    const py::kwargs& kwargs) const;
+
+template <>
+void ModuleBenchmark::runOnce(ModuleInput&& input) const;
+
+template <>
+ModuleOutput ModuleBenchmark::runOnce(
+    const py::args& args,
+    const py::kwargs& kwargs) const;
+
+template <>
+void ScriptModuleBenchmark::addInput(py::args&& args, py::kwargs&& kwargs);
+template <>
+void ScriptModuleBenchmark::addInput(ScriptModuleInput&& input);
+
+template <>
+void ModuleBenchmark::addInput(py::args&& args, py::kwargs&& kwargs);
+
+} // namespace detail
+
+/**
+ * This class is a small c++ component responsible for executing a PyTorch
+ * module under an inference server like load. It can emulate multiple calling
+ * threads to a single module provided. In the future we plan to enhance this
+ * component to support inter and intra-op parallelism as well as multiple
+ * models running in a single process.
+ *
+ * For current available configurations refer to the BenchmarkConfig
+ * documentation
+ *
+ * The class supports working with either nn.Module or ScriptModule.
+ * Under the hood it just dispatches to corresponding specialization of
+ * class BenchmarkHelper<Input, Output, Model>
+ */
+class C10_HIDDEN ThroughputBenchmark {
+ public:
+  explicit ThroughputBenchmark(const jit::Module& module);
+  explicit ThroughputBenchmark(py::object module);
+
+  // Add one more input example. This input example should be in the exact
+  // format the module under test expects. It is responsibility of the module to
+  // perform any such format checks, the benchmark doesn't perform any
+  // validation of its own
+  void addInput(py::args args, py::kwargs kwargs);
+
+  // Equivalent to just running the model directly on the given input
+  py::object runOnce(const py::args& args, const py::kwargs& kwargs);
+
+  // The main method of the class allows to perform a multi-threaded benchmark
+  // It returns BenchmarkExecutionStats object with a lot of useful statistics
+  // about runtime execution. We can enhance this class in the future to provide
+  // more information to the user
+  BenchmarkExecutionStats benchmark(const BenchmarkConfig& config) const;
+
+ private:
+  detail::ScriptModuleBenchmark script_module_;
+  detail::ModuleBenchmark module_;
+};
+} // namespace torch::throughput_benchmark
+
+#include <torch/csrc/utils/throughput_benchmark-inl.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/torch_dispatch_mode.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/torch_dispatch_mode.h
new file mode 100644
index 0000000000000000000000000000000000000000..1fd554d248bd3b430bf4e54dc8dd74f14e78c64f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/torch_dispatch_mode.h
@@ -0,0 +1,73 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/core/impl/TorchDispatchModeTLS.h>
+
+namespace torch::torch_dispatch_mode {
+
+struct StashTorchDispatchModeGuard {
+ public:
+  StashTorchDispatchModeGuard() {
+    if (c10::impl::TorchDispatchModeTLS::any_modes_set(
+            /*skip_infra_modes=*/true)) {
+      saved_mode_ = c10::impl::TorchDispatchModeTLS::pop_stack();
+    } else {
+      auto mode_and_key =
+          c10::impl::TorchDispatchModeTLS::pop_highest_infra_mode();
+      saved_mode_ = std::move(std::get<0>(mode_and_key));
+      saved_mode_key_ = std::get<1>(mode_and_key);
+    }
+  }
+
+  ~StashTorchDispatchModeGuard() {
+    if (saved_mode_key_.has_value()) {
+      c10::impl::TorchDispatchModeTLS::set_mode(
+          saved_mode_, saved_mode_key_.value());
+    } else {
+      c10::impl::TorchDispatchModeTLS::push_non_infra_mode_onto_stack(
+          std::move(saved_mode_));
+    }
+  }
+  StashTorchDispatchModeGuard(const StashTorchDispatchModeGuard&) = delete;
+  StashTorchDispatchModeGuard(StashTorchDispatchModeGuard&&) = delete;
+  StashTorchDispatchModeGuard& operator=(const StashTorchDispatchModeGuard&) =
+      delete;
+  StashTorchDispatchModeGuard& operator=(StashTorchDispatchModeGuard&&) =
+      delete;
+
+  const std::shared_ptr<c10::impl::PyObject_TorchDispatchMode>& get_cur_mode() {
+    return saved_mode_;
+  }
+
+ private:
+  std::shared_ptr<c10::impl::PyObject_TorchDispatchMode> saved_mode_;
+  std::optional<c10::impl::TorchDispatchModeKey> saved_mode_key_;
+};
+
+struct StashTorchDispatchStackGuard {
+ public:
+  StashTorchDispatchStackGuard() {
+    auto old = c10::impl::TorchDispatchModeTLS::get_state();
+    c10::impl::TorchDispatchModeTLS::set_state(std::move(saved_state_));
+    saved_state_ = std::move(old);
+  }
+  StashTorchDispatchStackGuard(const StashTorchDispatchStackGuard&) = delete;
+  StashTorchDispatchStackGuard(StashTorchDispatchStackGuard&&) = delete;
+  StashTorchDispatchStackGuard& operator=(const StashTorchDispatchStackGuard&) =
+      delete;
+  StashTorchDispatchStackGuard& operator=(StashTorchDispatchStackGuard&&) =
+      delete;
+
+  ~StashTorchDispatchStackGuard() {
+    c10::impl::TorchDispatchModeTLS::set_state(std::move(saved_state_));
+  }
+
+ private:
+  c10::impl::TorchDispatchModeTLS saved_state_;
+};
+
+} // namespace torch::torch_dispatch_mode
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/variadic.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/variadic.h
new file mode 100644
index 0000000000000000000000000000000000000000..e080abdbeb64ab85591aad79431f683743f0aed9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/variadic.h
@@ -0,0 +1,116 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/Tensor.h>
+#include <ATen/core/Variadic.h>
+#include <torch/csrc/autograd/variable.h>
+
+#include <type_traits>
+#include <utility>
+
+namespace torch {
+
+using at::IterArgs;
+
+struct CountTensors : IterArgs<CountTensors> {
+  size_t out = 0;
+  void operator()(const at::Tensor& x) {
+    out += 1;
+  }
+  void operator()(const std::optional<at::Tensor>& x) {
+    out += x.has_value();
+  }
+  void operator()(at::ArrayRef<at::Tensor> xs) {
+    out += xs.size();
+  }
+};
+
+template <typename... Args>
+size_t count_tensors(Args&&... args) {
+  return CountTensors().apply(std::forward<Args>(args)...).out;
+}
+
+struct CountVariables : IterArgs<CountVariables> {
+  size_t out = 0;
+  void operator()(const autograd::Variable& x) {
+    out += 1;
+  }
+  void operator()(at::ArrayRef<autograd::Variable> xs) {
+    out += xs.size();
+  }
+};
+
+template <typename... Args>
+inline size_t count_variables(Args&&... args) {
+  return CountVariables().apply(std::forward<Args>(args)...).out;
+}
+
+//===----------------------------------------------------------------------===//
+//                std::index_sequence shim for C++11
+//===----------------------------------------------------------------------===//
+
+// A container of type-template parameter indices.
+template <size_t... Is>
+struct Indices {};
+
+// Decrements the index N, adds N-1 to the list of indices and forwards
+// whatever we already have.
+template <size_t N, size_t... Is>
+struct MakeIndices : MakeIndices<N - 1, N - 1, Is...> {};
+
+// Partial specialization that forms our base case. When N is zero, we stop
+// and define a typedef that will be visible to earlier classes due to
+// inheritance. The typedef we define is an index list containing the numbers
+// 0 through N-1.
+template <size_t... Is>
+struct MakeIndices<0, Is...> {
+  using indices = Indices<Is...>;
+};
+
+//===----------------------------------------------------------------------===//
+//                                 Utilities
+//===----------------------------------------------------------------------===//
+
+template <typename Function, typename... Ts>
+void apply(Function function, Ts&&... ts) {
+  // https://stackoverflow.com/questions/13978916/inserting-a-variadic-argument-list-into-a-vector
+  // Creates a dummy array, so that each function call is evaluated in order.
+  // `(function(), 0)` is because `function` should (!) return `void`, so
+  // according to the comma operator, it is evaluated and its result (`void`)
+  // is discarded. Then the zero is evaluated and used as an element in the
+  // array. The first zero ensures the array is not empty.
+  // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+  int _[]{0, (function(std::forward<Ts>(ts)), 0)...};
+  (void)_;
+}
+
+template <
+    typename ReturnType,
+    typename... Ts,
+    typename Function,
+    typename Accessor>
+ReturnType unpack(Function function, Accessor accessor) {
+  return ReturnType(unpack<ReturnType, Ts...>(
+      std::move(function),
+      std::move(accessor),
+      typename MakeIndices<sizeof...(Ts)>::indices()));
+}
+
+template <
+    typename ReturnType,
+    typename... Ts,
+    typename Function,
+    typename Accessor,
+    size_t... Is>
+ReturnType unpack(
+    Function function,
+    Accessor accessor,
+    Indices<Is...> /*unused*/) {
+  return ReturnType(function(accessor.template operator()<Ts>(Is)...));
+}
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/verbose.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/verbose.h
new file mode 100644
index 0000000000000000000000000000000000000000..54a879f7d456e18a73eade463c9b5b5188f19f33
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/utils/verbose.h
@@ -0,0 +1,13 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+#include <torch/csrc/python_headers.h>
+
+namespace torch {
+
+void initVerboseBindings(PyObject* module);
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Event.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Event.h
new file mode 100644
index 0000000000000000000000000000000000000000..3dc8e69a8758a86e38775d94db140eab22198d8d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Event.h
@@ -0,0 +1,21 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/xpu/XPUEvent.h>
+#include <torch/csrc/Event.h>
+#include <torch/csrc/python_headers.h>
+
+struct THXPEvent : THPEvent {
+  at::xpu::XPUEvent xpu_event;
+};
+extern PyObject* THXPEventClass;
+
+void THXPEvent_init(PyObject* module);
+
+inline bool THXPEvent_Check(PyObject* obj) {
+  return THXPEventClass && PyObject_IsInstance(obj, THXPEventClass);
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Module.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Module.h
new file mode 100644
index 0000000000000000000000000000000000000000..1f2eb36ed24981065ffdf7614eb9512f94bb2094
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Module.h
@@ -0,0 +1,16 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+PyMethodDef* THXPModule_methods();
+
+namespace torch::xpu {
+
+void initModule(PyObject* module);
+
+} // namespace torch::xpu
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Stream.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Stream.h
new file mode 100644
index 0000000000000000000000000000000000000000..de7e8366741469a70ceebb33215e0e13eb7119e3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/Stream.h
@@ -0,0 +1,22 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/xpu/XPUStream.h>
+#include <torch/csrc/Stream.h>
+#include <torch/csrc/python_headers.h>
+
+// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
+struct THXPStream : THPStream {
+  at::xpu::XPUStream xpu_stream;
+};
+extern PyObject* THXPStreamClass;
+
+void THXPStream_init(PyObject* module);
+
+inline bool THXPStream_Check(PyObject* obj) {
+  return THXPStreamClass && PyObject_IsInstance(obj, THXPStreamClass);
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/XPUPluggableAllocator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/XPUPluggableAllocator.h
new file mode 100644
index 0000000000000000000000000000000000000000..8a34baba3b47a204e44eb037e0125047487d3aba
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/XPUPluggableAllocator.h
@@ -0,0 +1,85 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/xpu/XPUCachingAllocator.h>
+#include <torch/csrc/Export.h>
+
+namespace torch::xpu::XPUPluggableAllocator {
+
+struct _AllocationMetadata {
+  _AllocationMetadata() {}
+  _AllocationMetadata(
+      size_t size,
+      c10::DeviceIndex device_idx,
+      sycl::queue* queue)
+      : size(size), device_idx(device_idx), queue(queue) {}
+  size_t size{0};
+  c10::DeviceIndex device_idx{-1};
+  sycl::queue* queue{};
+};
+
+struct TORCH_PYTHON_API XPUPluggableAllocator
+    : public c10::xpu::XPUCachingAllocator::XPUAllocator {
+  XPUPluggableAllocator(
+      std::function<void*(size_t, int, sycl::queue*)> alloc_fn,
+      std::function<void(void*, size_t, int, sycl::queue*)> free_fn)
+      : alloc_fn_(std::move(alloc_fn)), free_fn_(std::move(free_fn)) {}
+
+  C10_DISABLE_COPY_AND_ASSIGN(XPUPluggableAllocator);
+
+  ~XPUPluggableAllocator() override = default;
+
+  void* malloc(size_t size, c10::DeviceIndex device, sycl::queue* stream);
+
+  c10::DataPtr allocate(size_t size) override;
+  c10::DeleterFnPtr raw_deleter() const override;
+
+  void* raw_alloc(size_t nbytes) override;
+  void raw_delete(void* ptr) override;
+  void init(c10::DeviceIndex device_count) override;
+  bool initialized() override;
+  void copy_data(void* dest, const void* src, std::size_t count) const final;
+
+  void recordStream(const c10::DataPtr&, c10::Stream stream) override;
+  void emptyCache(c10::MempoolId_t mempool_id = {0, 0}) override;
+  c10::CachingDeviceAllocator::DeviceStats getDeviceStats(
+      c10::DeviceIndex device) override;
+  void resetAccumulatedStats(c10::DeviceIndex device) override;
+  void resetPeakStats(c10::DeviceIndex device) override;
+
+  void set_init_fn(std::function<void(int)> init_fn) {
+    init_fn_ = std::move(init_fn);
+  }
+  void set_record_stream_fn(
+      std::function<void(void* ptr, sycl::queue* queue)> record_stream_fn) {
+    record_stream_fn_ = std::move(record_stream_fn);
+  }
+
+ protected:
+  std::function<void*(size_t, int, sycl::queue*)> alloc_fn_;
+  std::function<void(void*, size_t, int, sycl::queue*)> free_fn_;
+  std::function<void(int)> init_fn_;
+  std::function<void(void* ptr, sycl::queue*)> record_stream_fn_;
+  std::mutex allocator_mutex_;
+  // We do the bookkeeping here in order to simplify custom allocators
+  std::unordered_map<void*, _AllocationMetadata> allocation_metadata_;
+  bool initialized_ = false;
+};
+
+TORCH_XPU_API std::shared_ptr<c10::xpu::XPUCachingAllocator::XPUAllocator>
+getCurrentAllocator();
+
+TORCH_XPU_API std::shared_ptr<c10::xpu::XPUCachingAllocator::XPUAllocator>
+createCustomAllocator(
+    std::function<void*(size_t, int, sycl::queue*)> alloc_fn,
+    std::function<void(void*, size_t, int, sycl::queue*)> free_fn);
+
+TORCH_XPU_API void changeCurrentAllocator(
+    const std::shared_ptr<c10::xpu::XPUCachingAllocator::XPUAllocator>&
+        allocator);
+
+} // namespace torch::xpu::XPUPluggableAllocator
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/memory_snapshot.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/memory_snapshot.h
new file mode 100644
index 0000000000000000000000000000000000000000..7966c07aa89b6dcc7dafd48032c52084ba6daa8f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/csrc/xpu/memory_snapshot.h
@@ -0,0 +1,24 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/Export.h>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace torch::xpu {
+
+TORCH_PYTHON_API void _record_memory_history(
+    std::optional<std::string> enabled = "all",
+    std::optional<std::string> context = "all",
+    const std::string& stacks = "all",
+    size_t max_entries = SIZE_MAX,
+    bool clear_history = false,
+    const std::vector<std::string>& skip_actions = {});
+
+} // namespace torch::xpu
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/custom_class.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/custom_class.h
new file mode 100644
index 0000000000000000000000000000000000000000..e7177ec77fac0e178e21021e5806e782f223c17c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/custom_class.h
@@ -0,0 +1,523 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/builtin_function.h>
+#include <ATen/core/function_schema.h>
+#include <ATen/core/ivalue.h>
+#include <ATen/core/class_type.h>
+#include <ATen/core/op_registration/infer_schema.h>
+#include <ATen/core/stack.h>
+#include <c10/util/Metaprogramming.h>
+#include <c10/util/TypeList.h>
+#include <c10/util/TypeTraits.h>
+#include <torch/custom_class_detail.h>
+#include <torch/library.h>
+
+#include <functional>
+#include <sstream>
+
+namespace torch {
+
+/// This function is used in conjunction with `class_::def()` to register
+/// a constructor for a given C++ class type. For example,
+/// `torch::init<int, std::string>()` would register a two-argument constructor
+/// taking an `int` and a `std::string` as argument.
+template <class... Types>
+detail::types<void, Types...> init() {
+  return detail::types<void, Types...>{};
+}
+
+template <typename Func, typename... ParameterTypeList>
+struct InitLambda {
+  Func f;
+};
+
+template <typename Func>
+decltype(auto) init(Func&& f) {
+  using InitTraits = c10::guts::infer_function_traits_t<std::decay_t<Func>>;
+  using ParameterTypeList = typename InitTraits::parameter_types;
+
+  InitLambda<Func, ParameterTypeList> init{std::forward<Func>(f)};
+  return init;
+}
+
+/// Entry point for custom C++ class registration. To register a C++ class
+/// in PyTorch, instantiate `torch::class_` with the desired class as the
+/// template parameter. Typically, this instantiation should be done in
+/// the initialization of a global variable, so that the class will be
+/// made available on dynamic library loading without any additional API
+/// calls needed. For example, to register a class named Foo, you might
+/// create a global variable like so:
+///
+///     static auto register_foo = torch::class_<Foo>("myclasses", "Foo")
+///       .def("myMethod", &Foo::myMethod)
+///       .def("lambdaMethod", [](const c10::intrusive_ptr<Foo>& self) {
+///         // Do something with `self`
+///       });
+///
+/// In addition to registering the class, this registration also chains
+/// `def()` calls to register methods. `myMethod()` is registered with
+/// a pointer to the Foo class's `myMethod()` method. `lambdaMethod()`
+/// is registered with a C++ lambda expression.
+template <class CurClass>
+class class_ : public ::torch::detail::class_base {
+  static_assert(
+      std::is_base_of_v<CustomClassHolder, CurClass>,
+      "torch::class_<T> requires T to inherit from CustomClassHolder");
+
+ public:
+  /// This constructor actually registers the class type.
+  /// String argument `namespaceName` is an identifier for the
+  /// namespace you would like this class to appear in.
+  /// String argument `className` is the name you would like to
+  /// see this class exposed as in Python and TorchScript. For example, if
+  /// you pass `foo` as the namespace name and `Bar` as the className, the
+  /// class will appear as `torch.classes.foo.Bar` in Python and TorchScript
+  explicit class_(
+      const std::string& namespaceName,
+      const std::string& className,
+      std::string doc_string = "")
+      : class_base(
+            namespaceName,
+            className,
+            std::move(doc_string),
+            typeid(c10::intrusive_ptr<CurClass>),
+            typeid(c10::tagged_capsule<CurClass>)) {}
+
+  /// def() can be used in conjunction with `torch::init()` to register
+  /// a constructor for a given C++ class type. For example, passing
+  /// `torch::init<int, std::string>()` would register a two-argument
+  /// constructor taking an `int` and a `std::string` as argument.
+  template <typename... Types>
+  class_& def(
+      torch::detail::types<void, Types...> /*unused*/,
+      std::string doc_string = "",
+      std::initializer_list<arg> default_args =
+          {}) { // Used in combination with
+    // torch::init<...>()
+    auto func = [](c10::tagged_capsule<CurClass> self, Types... args) {
+      auto classObj = c10::make_intrusive<CurClass>(args...);
+      auto object = self.ivalue.toObject();
+      object->setSlot(0, c10::IValue::make_capsule(std::move(classObj)));
+    };
+
+    defineMethod(
+        "__init__",
+        std::move(func),
+        std::move(doc_string),
+        default_args);
+    return *this;
+  }
+
+  // Used in combination with torch::init([]lambda(){......})
+  template <typename Func, typename... ParameterTypes>
+  class_& def(
+      InitLambda<Func, c10::guts::typelist::typelist<ParameterTypes...>> init,
+      std::string doc_string = "",
+      std::initializer_list<arg> default_args = {}) {
+    auto init_lambda_wrapper = [func = std::move(init.f)](
+                                   c10::tagged_capsule<CurClass> self,
+                                   ParameterTypes... arg) {
+      c10::intrusive_ptr<CurClass> classObj =
+          std::invoke(func, std::forward<ParameterTypes>(arg)...);
+      auto object = self.ivalue.toObject();
+      object->setSlot(0, c10::IValue::make_capsule(classObj));
+    };
+
+    defineMethod(
+        "__init__",
+        std::move(init_lambda_wrapper),
+        std::move(doc_string),
+        default_args);
+
+    return *this;
+  }
+
+  /// This is the normal method registration API. `name` is the name that
+  /// the method will be made accessible by in Python and TorchScript.
+  /// `f` is a callable object that defines the method. Typically `f`
+  /// will either be a pointer to a method on `CurClass`, or a lambda
+  /// expression that takes a `c10::intrusive_ptr<CurClass>` as the first
+  /// argument (emulating a `this` argument in a C++ method.)
+  ///
+  /// Examples:
+  ///
+  ///     // Exposes method `foo` on C++ class `Foo` as `call_foo()` in
+  ///     // Python and TorchScript
+  ///     .def("call_foo", &Foo::foo)
+  ///
+  ///     // Exposes the given lambda expression as method `call_lambda()`
+  ///     // in Python and TorchScript.
+  ///     .def("call_lambda", [](const c10::intrusive_ptr<Foo>& self) {
+  ///       // do something
+  ///     })
+  template <typename Func>
+  class_& def(
+      std::string name,
+      Func f,
+      std::string doc_string = "",
+      std::initializer_list<arg> default_args = {}) {
+    auto wrapped_f = detail::wrap_func<CurClass, Func>(std::move(f));
+    defineMethod(
+        std::move(name),
+        std::move(wrapped_f),
+        std::move(doc_string),
+        default_args);
+    return *this;
+  }
+
+  /// Method registration API for static methods.
+  template <typename Func>
+  class_& def_static(std::string name, Func func, std::string doc_string = "") {
+    auto qualMethodName = qualClassName + "." + name;
+    auto schema =
+        c10::inferFunctionSchemaSingleReturn<Func>(std::move(name), "");
+
+    auto wrapped_func =
+        [func = std::move(func)](jit::Stack& stack) mutable -> void {
+      using RetType =
+          typename c10::guts::infer_function_traits_t<Func>::return_type;
+      detail::BoxedProxy<RetType, Func>()(stack, func);
+    };
+    auto method = std::make_unique<jit::BuiltinOpFunction>(
+        std::move(qualMethodName),
+        std::move(schema),
+        std::move(wrapped_func),
+        std::move(doc_string));
+
+    classTypePtr->addStaticMethod(method.get());
+    registerCustomClassMethod(std::move(method));
+    return *this;
+  }
+
+  /// Property registration API for properties with both getter and setter
+  /// functions.
+  template <typename GetterFunc, typename SetterFunc>
+  class_& def_property(
+      const std::string& name,
+      GetterFunc getter_func,
+      SetterFunc setter_func,
+      std::string doc_string = "") {
+    torch::jit::Function* getter{};
+    torch::jit::Function* setter{};
+
+    auto wrapped_getter =
+        detail::wrap_func<CurClass, GetterFunc>(std::move(getter_func));
+    getter = defineMethod(name + "_getter", wrapped_getter, doc_string);
+
+    auto wrapped_setter =
+        detail::wrap_func<CurClass, SetterFunc>(std::move(setter_func));
+    setter = defineMethod(name + "_setter", wrapped_setter, doc_string);
+
+    classTypePtr->addProperty(name, getter, setter);
+    return *this;
+  }
+
+  /// Property registration API for properties with only getter function.
+  template <typename GetterFunc>
+  class_& def_property(
+      const std::string& name,
+      GetterFunc getter_func,
+      std::string doc_string = "") {
+    torch::jit::Function* getter{};
+
+    auto wrapped_getter =
+        detail::wrap_func<CurClass, GetterFunc>(std::move(getter_func));
+    getter = defineMethod(name + "_getter", wrapped_getter, doc_string);
+
+    classTypePtr->addProperty(name, getter, nullptr);
+    return *this;
+  }
+
+  /// Property registration API for properties with read-write access.
+  template <typename T>
+  class_& def_readwrite(const std::string& name, T CurClass::*field) {
+    auto getter_func = [field =
+                            field](const c10::intrusive_ptr<CurClass>& self) {
+      return self.get()->*field;
+    };
+
+    auto setter_func = [field = field](
+                           const c10::intrusive_ptr<CurClass>& self, T value) {
+      self.get()->*field = value;
+    };
+
+    return def_property(name, getter_func, setter_func);
+  }
+
+  /// Property registration API for properties with read-only access.
+  template <typename T>
+  class_& def_readonly(const std::string& name, T CurClass::*field) {
+    auto getter_func =
+        [field = std::move(field)](const c10::intrusive_ptr<CurClass>& self) {
+          return self.get()->*field;
+        };
+
+    return def_property(name, getter_func);
+  }
+
+  /// This is an unsafe method registration API added for adding custom JIT
+  /// backend support via custom C++ classes. It is not for general purpose use.
+  class_& _def_unboxed(
+      const std::string& name,
+      std::function<void(jit::Stack&)> func,
+      c10::FunctionSchema schema,
+      std::string doc_string = "") {
+    auto method = std::make_unique<jit::BuiltinOpFunction>(
+        qualClassName + "." + name,
+        std::move(schema),
+        std::move(func),
+        std::move(doc_string));
+    classTypePtr->addMethod(method.get());
+    registerCustomClassMethod(std::move(method));
+    return *this;
+  }
+
+  /// def_pickle() is used to define exactly what state gets serialized
+  /// or deserialized for a given instance of a custom C++ class in
+  /// Python or TorchScript. This protocol is equivalent to the Pickle
+  /// concept of `__getstate__` and `__setstate__` from Python
+  /// (https://docs.python.org/2/library/pickle.html#object.__getstate__)
+  ///
+  /// Currently, both the `get_state` and `set_state` callables must be
+  /// C++ lambda expressions. They should have the following signatures,
+  /// where `CurClass` is the class you're registering and `T1` is some object
+  /// that encapsulates the state of the object.
+  ///
+  ///     __getstate__(intrusive_ptr<CurClass>) -> T1
+  ///     __setstate__(T2) -> intrusive_ptr<CurClass>
+  ///
+  /// `T1` must be an object that is convertible to IValue by the same rules
+  /// for custom op/method registration.
+  ///
+  /// For the common case, T1 == T2. T1 can also be a subtype of T2. An
+  /// example where it makes sense for T1 and T2 to differ is if __setstate__
+  /// handles legacy formats in a backwards compatible way.
+  ///
+  /// Example:
+  ///
+  ///     .def_pickle(
+  ///         // __getstate__
+  ///         [](const c10::intrusive_ptr<MyStackClass<std::string>>& self) {
+  ///           return self->stack_;
+  ///         },
+  ///         [](std::vector<std::string> state) { // __setstate__
+  ///            return c10::make_intrusive<MyStackClass<std::string>>(
+  ///               std::vector<std::string>{"i", "was", "deserialized"});
+  ///         })
+  template <typename GetStateFn, typename SetStateFn>
+  class_& def_pickle(GetStateFn&& get_state, SetStateFn&& set_state) {
+    static_assert(
+        c10::guts::is_stateless_lambda<std::decay_t<GetStateFn>>::value &&
+            c10::guts::is_stateless_lambda<std::decay_t<SetStateFn>>::value,
+        "def_pickle() currently only supports lambdas as "
+        "__getstate__ and __setstate__ arguments.");
+    def("__getstate__", std::forward<GetStateFn>(get_state));
+
+    // __setstate__ needs to be registered with some custom handling:
+    // We need to wrap the invocation of the user-provided function
+    // such that we take the return value (i.e. c10::intrusive_ptr<CurrClass>)
+    // and assign it to the `capsule` attribute.
+    using SetStateTraits =
+        c10::guts::infer_function_traits_t<std::decay_t<SetStateFn>>;
+    using SetStateArg = typename c10::guts::typelist::head_t<
+        typename SetStateTraits::parameter_types>;
+    auto setstate_wrapper = [set_state = std::forward<SetStateFn>(set_state)](
+                                c10::tagged_capsule<CurClass> self,
+                                SetStateArg arg) {
+      c10::intrusive_ptr<CurClass> classObj =
+          std::invoke(set_state, std::move(arg));
+      auto object = self.ivalue.toObject();
+      object->setSlot(0, c10::IValue::make_capsule(classObj));
+    };
+    defineMethod(
+        "__setstate__",
+        detail::wrap_func<CurClass, decltype(setstate_wrapper)>(
+            std::move(setstate_wrapper)));
+
+    // type validation
+    auto getstate_schema = classTypePtr->getMethod("__getstate__").getSchema();
+#ifndef STRIP_ERROR_MESSAGES
+    auto format_getstate_schema = [&getstate_schema]() {
+      std::stringstream ss;
+      ss << getstate_schema;
+      return ss.str();
+    };
+#endif
+    TORCH_CHECK(
+        getstate_schema.arguments().size() == 1,
+        "__getstate__ should take exactly one argument: self. Got: ",
+        format_getstate_schema());
+    auto first_arg_type = getstate_schema.arguments().at(0).type();
+    TORCH_CHECK(
+        *first_arg_type == *classTypePtr,
+        "self argument of __getstate__ must be the custom class type. Got ",
+        first_arg_type->repr_str());
+    TORCH_CHECK(
+        getstate_schema.returns().size() == 1,
+        "__getstate__ should return exactly one value for serialization. Got: ",
+        format_getstate_schema());
+
+    auto ser_type = getstate_schema.returns().at(0).type();
+    auto setstate_schema = classTypePtr->getMethod("__setstate__").getSchema();
+    auto arg_type = setstate_schema.arguments().at(1).type();
+    TORCH_CHECK(
+        ser_type->isSubtypeOf(*arg_type),
+        "__getstate__'s return type should be a subtype of "
+        "input argument of __setstate__. Got ",
+        ser_type->repr_str(),
+        " but expected ",
+        arg_type->repr_str());
+
+    return *this;
+  }
+
+ private:
+  template <typename Func>
+  torch::jit::Function* defineMethod(
+      std::string name,
+      Func func,
+      std::string doc_string = "",
+      std::initializer_list<arg> default_args = {}) {
+    auto qualMethodName = qualClassName + "." + name;
+    auto schema =
+        c10::inferFunctionSchemaSingleReturn<Func>(std::move(name), "");
+
+    // If default values are provided for function arguments, there must be
+    // none (no default values) or default values for all function
+    // arguments, except for self. This is because argument names are not
+    // extracted by inferFunctionSchemaSingleReturn, and so there must be a
+    // torch::arg instance in default_args even for arguments that do not
+    // have an actual default value provided.
+    TORCH_CHECK(
+        default_args.size() == 0 ||
+            default_args.size() == schema.arguments().size() - 1,
+        "Default values must be specified for none or all arguments");
+
+    // If there are default args, copy the argument names and default values to
+    // the function schema.
+    if (default_args.size() > 0) {
+      schema = withNewArguments(schema, default_args);
+    }
+
+    auto wrapped_func =
+        [func = std::move(func)](jit::Stack& stack) mutable -> void {
+      // TODO: we need to figure out how to profile calls to custom functions
+      // like this! Currently can't do it because the profiler stuff is in
+      // libtorch and not ATen
+      using RetType =
+          typename c10::guts::infer_function_traits_t<Func>::return_type;
+      detail::BoxedProxy<RetType, Func>()(stack, func);
+    };
+    auto method = std::make_unique<jit::BuiltinOpFunction>(
+        qualMethodName,
+        std::move(schema),
+        std::move(wrapped_func),
+        std::move(doc_string));
+
+    // Register the method here to keep the Method alive.
+    // ClassTypes do not hold ownership of their methods (normally it
+    // those are held by the CompilationUnit), so we need a proxy for
+    // that behavior here.
+    auto method_val = method.get();
+    classTypePtr->addMethod(method_val);
+    registerCustomClassMethod(std::move(method));
+    return method_val;
+  }
+};
+
+/// make_custom_class() is a convenient way to create an instance of a
+/// registered custom class and wrap it in an IValue, for example when you want
+/// to pass the object to TorchScript. Its syntax is equivalent to APIs like
+/// `std::make_shared<>` or `c10::make_intrusive<>`.
+///
+/// For example, if you have a custom C++ class that can be constructed from an
+/// `int` and `std::string`, you might use this API like so:
+///
+///     IValue custom_class_iv = torch::make_custom_class<MyClass>(3,
+///     "foobarbaz");
+template <typename CurClass, typename... CtorArgs>
+c10::IValue make_custom_class(CtorArgs&&... args) {
+  auto userClassInstance =
+      c10::make_intrusive<CurClass>(std::forward<CtorArgs>(args)...);
+  return c10::IValue(std::move(userClassInstance));
+}
+
+// Alternative api for creating a torchbind class over torch::class_ this api is
+// preferred to prevent size regressions on Edge usecases. Must be used in
+// conjunction with TORCH_SELECTIVE_CLASS macro aka
+// selective_class<foo>("foo_namespace", TORCH_SELECTIVE_CLASS("foo"))
+template <class CurClass>
+inline class_<CurClass> selective_class_(
+    const std::string& namespace_name,
+    detail::SelectiveStr<true> className) {
+  auto class_name = std::string(className.operator const char*());
+  return torch::class_<CurClass>(namespace_name, class_name);
+}
+
+template <class CurClass>
+inline detail::ClassNotSelected selective_class_(
+    const std::string& /*unused*/,
+    detail::SelectiveStr<false> /*unused*/) {
+  return detail::ClassNotSelected();
+}
+
+// jit namespace for backward-compatibility
+// We previously defined everything in torch::jit but moved it out to
+// better reflect that these features are not limited only to TorchScript
+namespace jit {
+
+using ::torch::class_;
+using ::torch::getCustomClass;
+using ::torch::init;
+using ::torch::isCustomClass;
+
+} // namespace jit
+
+template <class CurClass>
+inline class_<CurClass> Library::class_(const std::string& className) {
+  TORCH_CHECK(
+      kind_ == DEF || kind_ == FRAGMENT,
+      "class_(\"",
+      className,
+      "\"): Cannot define a class inside of a TORCH_LIBRARY_IMPL block.  "
+      "All class_()s should be placed in the (unique) TORCH_LIBRARY block for their namespace.  "
+      "(Error occurred at ",
+      file_,
+      ":",
+      line_,
+      ")");
+  TORCH_INTERNAL_ASSERT(ns_.has_value(), file_, ":", line_);
+  return torch::class_<CurClass>(*ns_, className);
+}
+
+const std::unordered_set<std::string> getAllCustomClassesNames();
+
+template <class CurClass>
+inline class_<CurClass> Library::class_(detail::SelectiveStr<true> className) {
+  auto class_name = std::string(className.operator const char*());
+  TORCH_CHECK(
+      kind_ == DEF || kind_ == FRAGMENT,
+      "class_(\"",
+      class_name,
+      "\"): Cannot define a class inside of a TORCH_LIBRARY_IMPL block.  "
+      "All class_()s should be placed in the (unique) TORCH_LIBRARY block for their namespace.  "
+      "(Error occurred at ",
+      file_,
+      ":",
+      line_,
+      ")");
+  TORCH_INTERNAL_ASSERT(ns_.has_value(), file_, ":", line_);
+  return torch::class_<CurClass>(*ns_, class_name);
+}
+
+template <class CurClass>
+inline detail::ClassNotSelected Library::class_(detail::SelectiveStr<false> /*unused*/) {
+  return detail::ClassNotSelected();
+}
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/custom_class_detail.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/custom_class_detail.h
new file mode 100644
index 0000000000000000000000000000000000000000..9d1356b98fc20038723f0c4ff4cad66a1310382a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/custom_class_detail.h
@@ -0,0 +1,247 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h>
+#include <ATen/core/function.h>
+#include <c10/util/Metaprogramming.h>
+#include <c10/util/TypeTraits.h>
+#include <c10/util/irange.h>
+
+#include <functional>
+
+namespace torch {
+
+namespace detail {
+/**
+ * In the Facebook internal build (using BUCK), this macro is enabled by
+ * passing in -c pt.enable_record_kernel_dtype=1 when building the tracer
+ * binary.
+ */
+#if defined ENABLE_RECORD_KERNEL_FUNCTION_DTYPE
+TORCH_API void record_custom_class(std::string name);
+
+/**
+ * Record an instance of a custom class being loaded
+ * grab portion of string after final '.' from qualified name
+ * as this seemingly aligns with how users name their custom classes
+ * example: __torch__.torch.classes.xnnpack.Conv2dOpContext
+ */
+#define RECORD_CUSTOM_CLASS(NAME) \
+  auto name = std::string(NAME);  \
+  detail::record_custom_class(name.substr(name.find_last_of(".") + 1));
+#else
+#define RECORD_CUSTOM_CLASS(NAME)
+#endif
+} // namespace detail
+
+/// This struct is used to represent default values for arguments
+/// when registering methods for custom classes.
+///     static auto register_foo = torch::class_<Foo>("myclasses", "Foo")
+///       .def("myMethod", &Foo::myMethod, {torch::arg("name") = name});
+struct arg {
+  // Static method for representing a default value of None. This is meant to
+  // be used like so:
+  //     torch::arg("name") = torch::arg::none
+  // and is identical to:
+  //     torch::arg("name") = IValue()
+  static c10::IValue none() {
+    return c10::IValue();
+  }
+
+  // Explicit constructor.
+  explicit arg(std::string name)
+      : name_(std::move(name)), value_(std::nullopt) {}
+  // Assignment operator. This enables the pybind-like syntax of
+  // torch::arg("name") = value.
+  arg& operator=(const c10::IValue& rhs) {
+    value_ = rhs;
+    return *this;
+  }
+
+  // The name of the argument. This is copied to the schema; argument
+  // names cannot be extracted from the C++ declaration.
+  std::string name_;
+  // IValue's default constructor makes it None, which is not distinguishable
+  // from an actual, user-provided default value that is None. This boolean
+  // helps distinguish between the two cases.
+  std::optional<c10::IValue> value_;
+};
+
+namespace detail {
+
+// Argument type utilities
+template <class R, class...>
+struct types {
+  using type = types;
+};
+
+template <typename Method>
+struct WrapMethod;
+
+template <typename R, typename CurrClass, typename... Args>
+struct WrapMethod<R (CurrClass::*)(Args...)> {
+  WrapMethod(R (CurrClass::*m)(Args...)) : m(std::move(m)) {}
+
+  R operator()(c10::intrusive_ptr<CurrClass> cur, Args... args) {
+    return std::invoke(m, *cur, args...);
+  }
+
+  R (CurrClass::*m)(Args...);
+};
+
+template <typename R, typename CurrClass, typename... Args>
+struct WrapMethod<R (CurrClass::*)(Args...) const> {
+  WrapMethod(R (CurrClass::*m)(Args...) const) : m(std::move(m)) {}
+
+  R operator()(c10::intrusive_ptr<CurrClass> cur, Args... args) {
+    return std::invoke(m, *cur, args...);
+  }
+
+  R (CurrClass::*m)(Args...) const;
+};
+
+// Adapter for different callable types
+template <
+    typename CurClass,
+    typename Func,
+    std::enable_if_t<
+        std::is_member_function_pointer_v<std::decay_t<Func>>,
+        bool> = false>
+WrapMethod<Func> wrap_func(Func f) {
+  return WrapMethod<Func>(std::move(f));
+}
+
+template <
+    typename CurClass,
+    typename Func,
+    std::enable_if_t<
+        !std::is_member_function_pointer_v<std::decay_t<Func>>,
+        bool> = false>
+Func wrap_func(Func f) {
+  return f;
+}
+
+template <
+    class Functor,
+    bool AllowDeprecatedTypes,
+    size_t... ivalue_arg_indices>
+typename c10::guts::infer_function_traits_t<Functor>::return_type
+call_torchbind_method_from_stack(
+    Functor& functor,
+    jit::Stack& stack,
+    std::index_sequence<ivalue_arg_indices...> /*unused*/) {
+  (void)stack; // when sizeof...(ivalue_arg_indices) == 0, this argument would
+                 // be unused and we have to silence the compiler warning.
+
+  constexpr size_t num_ivalue_args = sizeof...(ivalue_arg_indices);
+
+  using IValueArgTypes =
+      typename c10::guts::infer_function_traits_t<Functor>::parameter_types;
+  // TODO We shouldn't use c10::impl stuff directly here. We should use the
+  // KernelFunction API instead.
+  return functor(c10::impl::ivalue_to_arg<
+                   typename c10::impl::decay_if_not_tensor<
+                       c10::guts::typelist::
+                           element_t<ivalue_arg_indices, IValueArgTypes>>::type,
+                   AllowDeprecatedTypes>::
+                       call(torch::jit::peek(
+                           stack, ivalue_arg_indices, num_ivalue_args))...);
+}
+
+template <class Functor, bool AllowDeprecatedTypes>
+typename c10::guts::infer_function_traits_t<Functor>::return_type
+call_torchbind_method_from_stack(Functor& functor, jit::Stack& stack) {
+  constexpr size_t num_ivalue_args =
+      c10::guts::infer_function_traits_t<Functor>::number_of_parameters;
+  return call_torchbind_method_from_stack<Functor, AllowDeprecatedTypes>(
+      functor, stack, std::make_index_sequence<num_ivalue_args>());
+}
+
+template <class RetType, class Func>
+struct BoxedProxy;
+
+template <class RetType, class Func>
+struct BoxedProxy {
+  void operator()(jit::Stack& stack, Func& func) {
+    auto retval = call_torchbind_method_from_stack<Func, false>(func, stack);
+    constexpr size_t num_ivalue_args =
+        c10::guts::infer_function_traits_t<Func>::number_of_parameters;
+    torch::jit::drop(stack, num_ivalue_args);
+    stack.emplace_back(c10::ivalue::from(std::move(retval)));
+  }
+};
+
+template <class Func>
+struct BoxedProxy<void, Func> {
+  void operator()(jit::Stack& stack, Func& func) {
+    call_torchbind_method_from_stack<Func, false>(func, stack);
+    constexpr size_t num_ivalue_args =
+        c10::guts::infer_function_traits_t<Func>::number_of_parameters;
+    torch::jit::drop(stack, num_ivalue_args);
+    stack.emplace_back();
+  }
+};
+
+inline bool validIdent(size_t i, char n) {
+  return isalpha(n) || n == '_' || (i > 0 && isdigit(n));
+}
+
+inline void checkValidIdent(const std::string& str, const char* type) {
+  for (const auto i : c10::irange(str.size())) {
+    TORCH_CHECK(
+        validIdent(i, str[i]),
+        type,
+        " must be a valid Python/C++ identifier."
+        " Character '",
+        str[i],
+        "' at index ",
+        i,
+        " is illegal.");
+  }
+}
+
+class TORCH_API class_base {
+ protected:
+  explicit class_base(
+      const std::string& namespaceName,
+      const std::string& className,
+      std::string doc_string,
+      const std::type_info& intrusivePtrClassTypeid,
+      const std::type_info& taggedCapsuleClass);
+
+  static c10::FunctionSchema withNewArguments(
+      const c10::FunctionSchema& schema,
+      std::initializer_list<arg> default_args);
+  std::string qualClassName;
+  at::ClassTypePtr classTypePtr;
+};
+
+} // namespace detail
+
+TORCH_API void registerCustomClass(at::ClassTypePtr class_type);
+TORCH_API void registerCustomClassMethod(std::unique_ptr<jit::Function> method);
+
+// Given a qualified name (e.g. __torch__.torch.classes.Foo), return
+// the ClassType pointer to the Type that describes that custom class,
+// or nullptr if no class by that name was found.
+TORCH_API at::ClassTypePtr getCustomClass(const std::string& name);
+
+// Given an IValue, return true if the object contained in that IValue
+// is a custom C++ class, otherwise return false.
+// NOLINTNEXTLINE(readability-redundant-declaration)
+TORCH_API bool isCustomClass(const c10::IValue& v);
+
+// This API is for testing purposes ONLY. It should not be used in
+// any load-bearing code.
+TORCH_API std::vector<c10::FunctionSchema> customClassSchemasForBCCheck();
+
+namespace jit {
+using ::torch::registerCustomClass;
+using ::torch::registerCustomClassMethod;
+} // namespace jit
+
+} // namespace torch
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/extension.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/extension.h
new file mode 100644
index 0000000000000000000000000000000000000000..ddce96571507358894d4a5232b0954c3461cc7ac
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/extension.h
@@ -0,0 +1,14 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifndef TORCH_INDUCTOR_CPP_WRAPPER
+// All pure C++ headers for the C++ frontend.
+#include <torch/all.h>
+#endif
+
+// Python bindings for the C++ frontend (includes Python.h).
+#include <torch/python.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/DeviceType.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/DeviceType.h
new file mode 100644
index 0000000000000000000000000000000000000000..9db3ef2568d341fb9a341a36c02ac91c4dcca30f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/DeviceType.h
@@ -0,0 +1,128 @@
+#pragma once
+
+// This is directly synchronized with caffe2/proto/caffe2.proto, but
+// doesn't require me to figure out how to get Protobuf headers into
+// ATen/core (which would require a lot more build system hacking.)
+// If you modify me, keep me synchronized with that file.
+
+#include <torch/headeronly/macros/Export.h>
+#include <torch/headeronly/macros/Macros.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+
+namespace c10 {
+
+// These contains all device types that also have a BackendComponent
+// and therefore participate in per-backend functionality dispatch keys.
+// This is most backends except PrivateUse2 and PrivateUse3
+#define C10_FORALL_BACKEND_DEVICE_TYPES(_, extra) \
+  _(CPU, extra)                                   \
+  _(CUDA, extra)                                  \
+  _(HIP, extra)                                   \
+  _(XLA, extra)                                   \
+  _(MPS, extra)                                   \
+  _(IPU, extra)                                   \
+  _(XPU, extra)                                   \
+  _(HPU, extra)                                   \
+  _(VE, extra)                                    \
+  _(Lazy, extra)                                  \
+  _(Meta, extra)                                  \
+  _(MTIA, extra)                                  \
+  _(PrivateUse1, extra)
+
+enum class DeviceType : int8_t {
+  CPU = 0,
+  CUDA = 1, // CUDA.
+  MKLDNN = 2, // Reserved for explicit MKLDNN
+  OPENGL = 3, // OpenGL
+  OPENCL = 4, // OpenCL
+  IDEEP = 5, // IDEEP.
+  HIP = 6, // AMD HIP
+  FPGA = 7, // FPGA
+  MAIA = 8, // ONNX Runtime / Microsoft
+  XLA = 9, // XLA / TPU
+  Vulkan = 10, // Vulkan
+  Metal = 11, // Metal
+  XPU = 12, // XPU
+  MPS = 13, // MPS
+  Meta = 14, // Meta (tensors with no data)
+  HPU = 15, // HPU / HABANA
+  VE = 16, // SX-Aurora / NEC
+  Lazy = 17, // Lazy Tensors
+  IPU = 18, // Graphcore IPU
+  MTIA = 19, // Meta training and inference devices
+  PrivateUse1 = 20, // PrivateUse1 device
+  // NB: If you add more devices:
+  //  - Change the implementations of DeviceTypeName and isValidDeviceType
+  //    in c10/core/DeviceType.cpp
+  //  - Change the number below
+  COMPILE_TIME_MAX_DEVICE_TYPES = 21,
+};
+
+constexpr DeviceType kCPU = DeviceType::CPU;
+constexpr DeviceType kCUDA = DeviceType::CUDA;
+constexpr DeviceType kHIP = DeviceType::HIP;
+constexpr DeviceType kFPGA = DeviceType::FPGA;
+constexpr DeviceType kMAIA = DeviceType::MAIA;
+constexpr DeviceType kXLA = DeviceType::XLA;
+constexpr DeviceType kMPS = DeviceType::MPS;
+constexpr DeviceType kMeta = DeviceType::Meta;
+constexpr DeviceType kVulkan = DeviceType::Vulkan;
+constexpr DeviceType kMetal = DeviceType::Metal;
+constexpr DeviceType kXPU = DeviceType::XPU;
+constexpr DeviceType kHPU = DeviceType::HPU;
+constexpr DeviceType kVE = DeviceType::VE;
+constexpr DeviceType kLazy = DeviceType::Lazy;
+constexpr DeviceType kIPU = DeviceType::IPU;
+constexpr DeviceType kMTIA = DeviceType::MTIA;
+constexpr DeviceType kPrivateUse1 = DeviceType::PrivateUse1;
+
+// define explicit int constant
+constexpr int COMPILE_TIME_MAX_DEVICE_TYPES =
+    static_cast<int>(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES);
+
+static_assert(
+    COMPILE_TIME_MAX_DEVICE_TYPES <= 21,
+    "Hey!  You seem to be adding a lot of new DeviceTypes.  The intent was "
+    "for this constant to reflect the actual number of DeviceTypes we support "
+    "in PyTorch; it's important that this number is not too large as we "
+    "use this to allocate stack arrays in some places in our code.  If you "
+    "are indeed just adding the 20th device type, feel free to change "
+    "the check to 32; but if you are adding some sort of extensible device "
+    "types registration, please be aware that you are affecting code that "
+    "this number is small.  Try auditing uses of this constant.");
+
+} // namespace c10
+
+namespace std {
+template <>
+struct hash<c10::DeviceType> {
+  std::size_t operator()(c10::DeviceType k) const {
+    return std::hash<int>()(static_cast<int>(k));
+  }
+};
+} // namespace std
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::COMPILE_TIME_MAX_DEVICE_TYPES;
+using c10::DeviceType;
+using c10::kCPU;
+using c10::kCUDA;
+using c10::kFPGA;
+using c10::kHIP;
+using c10::kHPU;
+using c10::kIPU;
+using c10::kLazy;
+using c10::kMAIA;
+using c10::kMeta;
+using c10::kMetal;
+using c10::kMPS;
+using c10::kMTIA;
+using c10::kPrivateUse1;
+using c10::kVE;
+using c10::kVulkan;
+using c10::kXLA;
+using c10::kXPU;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Dispatch.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Dispatch.h
new file mode 100644
index 0000000000000000000000000000000000000000..43293ef701ddac509c51abd9a4c1f923e61118f9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Dispatch.h
@@ -0,0 +1,73 @@
+#pragma once
+
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/macros/Macros.h>
+
+// THO_PRIVATE_CASE_TYPE_USING_HINT_TMPL is same as
+// AT_PRIVATE_CASE_TYPE_USING_HINT but with a custom PRELUDE macro:
+#define THO_PRIVATE_CASE_TYPE_USING_HINT_TMPL(PRELUDE, enum_type, HINT, ...) \
+  case enum_type: {                                                          \
+    PRELUDE(enum_type);                                                      \
+    using HINT [[maybe_unused]] =                                            \
+        torch::headeronly::impl::ScalarTypeToCPPTypeT<enum_type>;            \
+    return __VA_ARGS__();                                                    \
+  }
+
+// THO_DISPATCH_CASE_TMPL is same as AT_DISPATCH_CASE but with a
+// custom CASE_TYPE_USING_HINT macro:
+#define THO_DISPATCH_CASE_TMPL(CASE_TYPE_USING_HINT, enum_type, ...) \
+  CASE_TYPE_USING_HINT(enum_type, scalar_t, __VA_ARGS__)
+
+namespace detail {
+inline torch::headeronly::ScalarType scalar_type(
+    torch::headeronly::ScalarType s) {
+  return s;
+}
+} // namespace detail
+
+// THO_DISPATCH_SWITCH_TMPL is same as AT_DISPATCH_SWITCH but with
+// custom PRELUDE and CHECK_NOT_IMPLEMENTED macros:
+#define THO_DISPATCH_SWITCH_TMPL(                                           \
+    PRELUDE, CHECK_NOT_IMPLEMENTED, TYPE, NAME, ...)                        \
+  [&] {                                                                     \
+    const auto& the_type = TYPE;                                            \
+    constexpr const char* at_dispatch_name = NAME;                          \
+    /* don't use TYPE again in case it is an expensive or side-effect op */ \
+    torch::headeronly::ScalarType _st = ::detail::scalar_type(the_type);    \
+    PRELUDE(at_dispatch_name, _st);                                         \
+    C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")             \
+    switch (_st) {                                                          \
+      __VA_ARGS__                                                           \
+      default:                                                              \
+        CHECK_NOT_IMPLEMENTED(                                              \
+            false,                                                          \
+            '"',                                                            \
+            at_dispatch_name,                                               \
+            "\" not implemented for '",                                     \
+            torch::headeronly::toString(_st),                               \
+            "'");                                                           \
+    }                                                                       \
+    C10_DIAGNOSTIC_POP()                                                    \
+  }()
+
+// THO_EMPTY is a helper macro that discards its arguments.
+#define THO_EMPTY(...)
+
+// THO_PRIVATE_CASE_TYPE_USING_HINT is same as
+// AT_PRIVATE_CASE_TYPE_USING_HINT with call to macro
+// AT_PRIVATE_CHECK_SELECTIVE_BUILD removed.
+#define THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, HINT, ...) \
+  THO_PRIVATE_CASE_TYPE_USING_HINT_TMPL(THO_EMPTY, enum_type, HINT, __VA_ARGS__)
+
+// THO_DISPATCH_SWITCH is same as AT_DISPATCH_SWITCH with call to
+// macro RECORD_KERNEL_FUNCTION_DTYPE removed and using
+// STD_TORCH_CHECK instead of TORCH_CHECK_NOT_IMPLEMENTED.
+#define THO_DISPATCH_SWITCH(TYPE, NAME, ...) \
+  THO_DISPATCH_SWITCH_TMPL(THO_EMPTY, STD_TORCH_CHECK, TYPE, NAME, __VA_ARGS__)
+
+// THO_DISPATCH_CASE is same as AT_DISPATCH_CASE but using
+// THO_PRIVATE_CASE_TYPE_USING_HINT instead of
+// AT_PRIVATE_CASE_TYPE_USING_HINT.
+#define THO_DISPATCH_CASE(enum_type, ...) \
+  THO_DISPATCH_CASE_TMPL(                 \
+      THO_PRIVATE_CASE_TYPE_USING_HINT, enum_type, __VA_ARGS__)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Dispatch_v2.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Dispatch_v2.h
new file mode 100644
index 0000000000000000000000000000000000000000..aa1bc1d9e9f72b1f211c10a313f0b97d94eea0df
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Dispatch_v2.h
@@ -0,0 +1,157 @@
+#pragma once
+
+#include <torch/headeronly/core/Dispatch.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+// This file provides THO_DISPATCH_V2_TMPL macro that is a generalized
+// version of the original AT_DISPATCH_V2 (see ATen/Dispatch_v2.h for
+// documentation): THO_DISPATCH_V2_TMPL extends AT_DISPATCH_V2 with
+// extra DISPATCH_SWITCH and DISPATCH_CASE arguments for specifying
+// custom implementations of the original AT_DISPATCH_SWITCH and
+// AT_DISPATCH_CASE macros. Use the provided macros
+// THO_DISPATCH_SWITCH_TMPL and THO_DISPATCH_CASE_TMPL to define the
+// custom implementations of the switch and case macros, respectively.
+
+// Public API macros
+
+// THO_DISPATCH_V2_TMPL is same as AT_DISPATCH_V2 but with custom
+// DISPATCH_SWITCH and DISPATCH_CASE macro arguments:
+#define THO_DISPATCH_V2_TMPL(                              \
+    DISPATCH_SWITCH, DISPATCH_CASE, TYPE, NAME, BODY, ...) \
+  DISPATCH_SWITCH(                                         \
+      TYPE,                                                \
+      NAME,                                                \
+      THO_AP_VAR_TMPL(DISPATCH_CASE, AT_WRAP(BODY), TYPE, __VA_ARGS__))
+
+// THO_DISPATCH_V2 is same as AT_DISPATCH_V2 but using
+// THO_DISPATCH_SWITCH and THO_DISPATCH_CASE instead of
+// AT_DISPATCH_SWITCH and AT_DISPATCH_CASE, respectively.
+#define THO_DISPATCH_V2(TYPE, NAME, BODY, ...) \
+  THO_DISPATCH_V2_TMPL(                        \
+      THO_DISPATCH_SWITCH, THO_DISPATCH_CASE, TYPE, NAME, BODY, __VA_ARGS__)
+
+// Type collection macros
+
+// This macro lets you pass an arbitrary expression that may contain internal
+// commas to another macro without having the commas causing the expression
+// to be interpreted as being multiple arguments
+#define AT_WRAP(...) __VA_ARGS__
+
+#define AT_FLOAT8_TYPES                               \
+  torch::headeronly::ScalarType::Float8_e5m2,         \
+      torch::headeronly::ScalarType::Float8_e5m2fnuz, \
+      torch::headeronly::ScalarType::Float8_e4m3fn,   \
+      torch::headeronly::ScalarType::Float8_e4m3fnuz, \
+      torch::headeronly::ScalarType::Float8_e8m0fnu
+
+#define AT_INTEGRAL_TYPES                                                      \
+  torch::headeronly::ScalarType::Byte, torch::headeronly::ScalarType::Char,    \
+      torch::headeronly::ScalarType::Int, torch::headeronly::ScalarType::Long, \
+      torch::headeronly::ScalarType::Short
+#define AT_FLOATING_TYPES \
+  torch::headeronly::ScalarType::Double, torch::headeronly::ScalarType::Float
+#define AT_BAREBONES_UNSIGNED_TYPES          \
+  torch::headeronly::ScalarType::UInt16,     \
+      torch::headeronly::ScalarType::UInt32, \
+      torch::headeronly::ScalarType::UInt64
+#define AT_OPAQUE_TYPES                                                       \
+  torch::headeronly::ScalarType::Byte, torch::headeronly::ScalarType::UInt16, \
+      torch::headeronly::ScalarType::UInt32,                                  \
+      torch::headeronly::ScalarType::UInt64,                                  \
+      torch::headeronly::ScalarType::ComplexDouble
+#define AT_INTEGRAL_TYPES_V2 \
+  AT_EXPAND(AT_INTEGRAL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES)
+#define AT_COMPLEX_TYPES                        \
+  torch::headeronly::ScalarType::ComplexDouble, \
+      torch::headeronly::ScalarType::ComplexFloat
+#define AT_QINT_TYPES                                                          \
+  torch::headeronly::ScalarType::QInt8, torch::headeronly::ScalarType::QUInt8, \
+      torch::headeronly::ScalarType::QInt32
+// NB: not *actually* all types
+#define AT_ALL_TYPES AT_EXPAND(AT_INTEGRAL_TYPES), AT_EXPAND(AT_FLOATING_TYPES)
+#define AT_ALL_TYPES_AND_COMPLEX \
+  AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_COMPLEX_TYPES)
+
+// Helper macros
+
+// THO_AP_VAR_TMPL is same as AT_AP_VAR but with a custom
+// DISPATCH_CASE macro argument:
+#define THO_AP_VAR_TMPL(C, N, T, ...) \
+  AT_EXPAND(                          \
+      AT_CONCAT(THO_AP, AT_NUM_ARGS(__VA_ARGS__))(C, AT_WRAP(N), __VA_ARGS__))
+#define AT_CONCAT(a, b) AT_CONCAT_AUX(a, b)
+#define AT_CONCAT_AUX(a, b) a##b
+#define AT_EXPAND(X) X
+
+// Ensure we never have too many scalar types for the expansion here to
+// support.  To bump this, you must regenerate the macros below.
+static_assert(static_cast<int>(torch::headeronly::ScalarType::NumOptions) < 60);
+
+// Begin generated code
+// clang-format off
+
+#define AT_NUM_ARGS(...) AT_EXPAND(AT_NUM_ARGS_AUX(__VA_ARGS__, 60, 59, 58, 57, 56, 55, 54, 53, 52, 51, 50, 49, 48, 47, 46, 45, 44, 43, 42, 41, 40, 39, 38, 37, 36, 35, 34, 33, 32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0))
+#define AT_NUM_ARGS_AUX(_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54, _55, _56, _57, _58, _59, _60, N, ...) N
+#define THO_AP1(C, N, _1) C(_1, N)
+#define THO_AP2(C, N, _1, _2) C(_1, N) C(_2, N)
+#define THO_AP3(C, N, _1, _2, _3) C(_1, N) C(_2, N) C(_3, N)
+#define THO_AP4(C, N, _1, _2, _3, _4) C(_1, N) C(_2, N) C(_3, N) C(_4, N)
+#define THO_AP5(C, N, _1, _2, _3, _4, _5) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N)
+#define THO_AP6(C, N, _1, _2, _3, _4, _5, _6) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N)
+#define THO_AP7(C, N, _1, _2, _3, _4, _5, _6, _7) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N)
+#define THO_AP8(C, N, _1, _2, _3, _4, _5, _6, _7, _8) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N)
+#define THO_AP9(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N)
+#define THO_AP10(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N)
+#define THO_AP11(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N)
+#define THO_AP12(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N)
+#define THO_AP13(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N)
+#define THO_AP14(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N)
+#define THO_AP15(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N)
+#define THO_AP16(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N)
+#define THO_AP17(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N)
+#define THO_AP18(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N)
+#define THO_AP19(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N)
+#define THO_AP20(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N)
+#define THO_AP21(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N)
+#define THO_AP22(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N)
+#define THO_AP23(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N)
+#define THO_AP24(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N)
+#define THO_AP25(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N)
+#define THO_AP26(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N)
+#define THO_AP27(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N)
+#define THO_AP28(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N)
+#define THO_AP29(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N)
+#define THO_AP30(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N)
+#define THO_AP31(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N)
+#define THO_AP32(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N)
+#define THO_AP33(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N)
+#define THO_AP34(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N)
+#define THO_AP35(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N)
+#define THO_AP36(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N)
+#define THO_AP37(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N)
+#define THO_AP38(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N)
+#define THO_AP39(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N)
+#define THO_AP40(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N)
+#define THO_AP41(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N)
+#define THO_AP42(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N)
+#define THO_AP43(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N)
+#define THO_AP44(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N)
+#define THO_AP45(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N)
+#define THO_AP46(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N)
+#define THO_AP47(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N)
+#define THO_AP48(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N)
+#define THO_AP49(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N)
+#define THO_AP50(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N)
+#define THO_AP51(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N)
+#define THO_AP52(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N)
+#define THO_AP53(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N)
+#define THO_AP54(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N) C(_54, N)
+#define THO_AP55(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54, _55) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N) C(_54, N) C(_55, N)
+#define THO_AP56(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54, _55, _56) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N) C(_54, N) C(_55, N) C(_56, N)
+#define THO_AP57(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54, _55, _56, _57) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N) C(_54, N) C(_55, N) C(_56, N) C(_57, N)
+#define THO_AP58(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54, _55, _56, _57, _58) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N) C(_54, N) C(_55, N) C(_56, N) C(_57, N) C(_58, N)
+#define THO_AP59(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54, _55, _56, _57, _58, _59) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N) C(_54, N) C(_55, N) C(_56, N) C(_57, N) C(_58, N) C(_59, N)
+#define THO_AP60(C, N, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _50, _51, _52, _53, _54, _55, _56, _57, _58, _59, _60) C(_1, N) C(_2, N) C(_3, N) C(_4, N) C(_5, N) C(_6, N) C(_7, N) C(_8, N) C(_9, N) C(_10, N) C(_11, N) C(_12, N) C(_13, N) C(_14, N) C(_15, N) C(_16, N) C(_17, N) C(_18, N) C(_19, N) C(_20, N) C(_21, N) C(_22, N) C(_23, N) C(_24, N) C(_25, N) C(_26, N) C(_27, N) C(_28, N) C(_29, N) C(_30, N) C(_31, N) C(_32, N) C(_33, N) C(_34, N) C(_35, N) C(_36, N) C(_37, N) C(_38, N) C(_39, N) C(_40, N) C(_41, N) C(_42, N) C(_43, N) C(_44, N) C(_45, N) C(_46, N) C(_47, N) C(_48, N) C(_49, N) C(_50, N) C(_51, N) C(_52, N) C(_53, N) C(_54, N) C(_55, N) C(_56, N) C(_57, N) C(_58, N) C(_59, N) C(_60, N)
+
+// End generated code
+// clang-format on
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Layout.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Layout.h
new file mode 100644
index 0000000000000000000000000000000000000000..62e34ff67b457ae3b6fb5ff4d804d85b3da9cb7a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/Layout.h
@@ -0,0 +1,44 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include <cstdint>
+#include <ostream>
+
+namespace c10 {
+
+enum class Layout : int8_t {
+  Strided,
+  Sparse,
+  SparseCsr,
+  Mkldnn,
+  SparseCsc,
+  SparseBsr,
+  SparseBsc,
+  Jagged,
+  NumOptions
+};
+
+constexpr auto kStrided = Layout::Strided;
+constexpr auto kSparse = Layout::Sparse;
+constexpr auto kSparseCsr = Layout::SparseCsr;
+constexpr auto kMkldnn = Layout::Mkldnn;
+constexpr auto kSparseCsc = Layout::SparseCsc;
+constexpr auto kSparseBsr = Layout::SparseBsr;
+constexpr auto kSparseBsc = Layout::SparseBsc;
+constexpr auto kJagged = Layout::Jagged;
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::kJagged;
+using c10::kMkldnn;
+using c10::kSparse;
+using c10::kSparseBsc;
+using c10::kSparseBsr;
+using c10::kSparseCsc;
+using c10::kSparseCsr;
+using c10::kStrided;
+using c10::Layout;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/MemoryFormat.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/MemoryFormat.h
new file mode 100644
index 0000000000000000000000000000000000000000..ad02a901e0169ad6c2e2169a90c72c88fb609097
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/MemoryFormat.h
@@ -0,0 +1,46 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include <cstdint>
+#include <ostream>
+
+// Memory format is not the property of a Tensor. It is the way to tell an
+// operator how the result should be organized in memory and nothing more. That
+// means memory format should never be used as return value for any tensor state
+// interrogation functions (internally and externally).
+//
+// Possible options are:
+//  Preserve:
+//    If any of the input tensors is in channels_last format, operator output
+//    should be in channels_last format
+//
+//  Contiguous:
+//    Regardless of input tensors format, the output should be contiguous
+//    Tensor.
+//
+//  ChannelsLast:
+//    Regardless of input tensors format, the output should be in channels_last
+//    format.
+
+namespace c10 {
+
+enum class MemoryFormat : int8_t {
+  Contiguous,
+  Preserve,
+  ChannelsLast,
+  ChannelsLast3d,
+  NumOptions
+};
+
+inline MemoryFormat get_contiguous_memory_format() {
+  return MemoryFormat::Contiguous;
+}
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::get_contiguous_memory_format;
+using c10::MemoryFormat;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/ScalarType.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/ScalarType.h
new file mode 100644
index 0000000000000000000000000000000000000000..ce43ce6866cd954adc778cec54c3405d2d1569a1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/ScalarType.h
@@ -0,0 +1,381 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Float4_e2m1fn_x2.h>
+#include <torch/headeronly/util/Float8_e4m3fn.h>
+#include <torch/headeronly/util/Float8_e4m3fnuz.h>
+#include <torch/headeronly/util/Float8_e5m2.h>
+#include <torch/headeronly/util/Float8_e5m2fnuz.h>
+#include <torch/headeronly/util/Float8_e8m0fnu.h>
+#include <torch/headeronly/util/Half.h>
+#include <torch/headeronly/util/bits.h>
+#include <torch/headeronly/util/complex.h>
+#include <torch/headeronly/util/qint32.h>
+#include <torch/headeronly/util/qint8.h>
+#include <torch/headeronly/util/quint2x4.h>
+#include <torch/headeronly/util/quint4x2.h>
+#include <torch/headeronly/util/quint8.h>
+
+#include <cstdint>
+
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")
+
+namespace c10 {
+
+// dummy struct for uint1 to uint7, actual functionality
+// of these dtypes will be implemented in python with Tensor subclass
+template <unsigned int N>
+struct dummy_uint1_7_t {};
+
+// dummy struct for int1 to int7, actual functionality
+// of these dtypes will be implemented in python with Tensor subclass
+template <unsigned int N>
+struct dummy_int1_7_t {};
+
+// [dtype Macros note] For the macros below:
+//
+// For users: If you want to macro some code for all non-QInt scalar types
+// (i.e. types with complete information, you probably want one of the
+// AT_FORALL_SCALAR_TYPES / AT_FORALL_SCALAR_TYPES_AND macros below, which are
+// designed to behave similarly to the Dispatch macros with the same name.
+//
+// For adding a new dtype: In the beginning, we had an idea that there was a
+// list of all scalar types, and you could use AT_FORALL_SCALAR_TYPES to
+// iterate over them.  But over the years we added weird types which couldn't
+// be handled uniformly everywhere and so in the end we ended up with some
+// mish-mosh of some helper macros, but mostly use sites making a call about
+// what dtypes they can or can't support.  So if you want to add a new dtype,
+// the preferred resolution is to find a dtype similar to what you want,
+// grep for it and edit all the sites you find this way.  If you need to add
+// a completely new kind of dtype, you're going to have to laboriously audit
+// all of the sites everywhere to figure out how it should work.  Consulting
+// some old PRs where we added new dtypes (check history of this file) can
+// help give you an idea where to start.
+
+// If you want to support ComplexHalf for real, add ComplexHalf
+// into this macro (and change the name).  But beware: convert()
+// doesn't work for all the conversions you need...
+//
+// TODO: To add unsigned int types here, we must define accumulate type.
+// But uint8 currently accumulates into int64, so we would have to make
+// an inconsistent choice for the larger types.  Difficult.
+#define AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF_F8NZ(_) \
+  _(uint8_t, Byte)                                                      \
+  _(int8_t, Char)                                                       \
+  _(int16_t, Short)                                                     \
+  _(int, Int)                                                           \
+  _(int64_t, Long)                                                      \
+  _(c10::Half, Half)                                                    \
+  _(float, Float)                                                       \
+  _(double, Double)                                                     \
+  _(c10::complex<float>, ComplexFloat)                                  \
+  _(c10::complex<double>, ComplexDouble)                                \
+  _(bool, Bool)                                                         \
+  _(c10::BFloat16, BFloat16)                                            \
+  _(c10::Float8_e5m2, Float8_e5m2)                                      \
+  _(c10::Float8_e4m3fn, Float8_e4m3fn)
+
+// This macro controls many of our C++ APIs, including constructors
+// for Scalar as well as the data() and item() accessors on Tensor
+#define AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(_) \
+  _(uint8_t, Byte)                             \
+  _(int8_t, Char)                              \
+  _(int16_t, Short)                            \
+  _(int, Int)                                  \
+  _(int64_t, Long)                             \
+  _(c10::Half, Half)                           \
+  _(float, Float)                              \
+  _(double, Double)                            \
+  _(c10::complex<c10::Half>, ComplexHalf)      \
+  _(c10::complex<float>, ComplexFloat)         \
+  _(c10::complex<double>, ComplexDouble)       \
+  _(bool, Bool)                                \
+  _(c10::BFloat16, BFloat16)                   \
+  _(c10::Float8_e5m2, Float8_e5m2)             \
+  _(c10::Float8_e4m3fn, Float8_e4m3fn)         \
+  _(c10::Float8_e5m2fnuz, Float8_e5m2fnuz)     \
+  _(c10::Float8_e4m3fnuz, Float8_e4m3fnuz)     \
+  _(c10::Float8_e8m0fnu, Float8_e8m0fnu)
+
+// NB: Order matters for this macro; it is relied upon in
+// _promoteTypesLookup and the serialization format.
+#define AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS(_) \
+  _(uint8_t, Byte) /* 0 */                               \
+  _(int8_t, Char) /* 1 */                                \
+  _(int16_t, Short) /* 2 */                              \
+  _(int, Int) /* 3 */                                    \
+  _(int64_t, Long) /* 4 */                               \
+  _(c10::Half, Half) /* 5 */                             \
+  _(float, Float) /* 6 */                                \
+  _(double, Double) /* 7 */                              \
+  _(c10::complex<c10::Half>, ComplexHalf) /* 8 */        \
+  _(c10::complex<float>, ComplexFloat) /* 9 */           \
+  _(c10::complex<double>, ComplexDouble) /* 10 */        \
+  _(bool, Bool) /* 11 */                                 \
+  _(c10::qint8, QInt8) /* 12 */                          \
+  _(c10::quint8, QUInt8) /* 13 */                        \
+  _(c10::qint32, QInt32) /* 14 */                        \
+  _(c10::BFloat16, BFloat16) /* 15 */                    \
+  _(c10::quint4x2, QUInt4x2) /* 16 */                    \
+  _(c10::quint2x4, QUInt2x4) /* 17 */                    \
+  _(c10::bits1x8, Bits1x8) /* 18 */                      \
+  _(c10::bits2x4, Bits2x4) /* 19 */                      \
+  _(c10::bits4x2, Bits4x2) /* 20 */                      \
+  _(c10::bits8, Bits8) /* 21 */                          \
+  _(c10::bits16, Bits16) /* 22 */                        \
+  _(c10::Float8_e5m2, Float8_e5m2) /* 23 */              \
+  _(c10::Float8_e4m3fn, Float8_e4m3fn) /* 24 */          \
+  _(c10::Float8_e5m2fnuz, Float8_e5m2fnuz) /* 25 */      \
+  _(c10::Float8_e4m3fnuz, Float8_e4m3fnuz) /* 26 */      \
+  _(uint16_t, UInt16) /* 27 */                           \
+  _(uint32_t, UInt32) /* 28 */                           \
+  _(uint64_t, UInt64) /* 29 */                           \
+  _(c10::dummy_uint1_7_t<1>, UInt1) /* 30 */             \
+  _(c10::dummy_uint1_7_t<2>, UInt2) /* 31 */             \
+  _(c10::dummy_uint1_7_t<3>, UInt3) /* 32 */             \
+  _(c10::dummy_uint1_7_t<4>, UInt4) /* 33 */             \
+  _(c10::dummy_uint1_7_t<5>, UInt5) /* 34 */             \
+  _(c10::dummy_uint1_7_t<6>, UInt6) /* 35 */             \
+  _(c10::dummy_uint1_7_t<7>, UInt7) /* 36 */             \
+  _(c10::dummy_int1_7_t<1>, Int1) /* 37 */               \
+  _(c10::dummy_int1_7_t<2>, Int2) /* 38 */               \
+  _(c10::dummy_int1_7_t<3>, Int3) /* 39 */               \
+  _(c10::dummy_int1_7_t<4>, Int4) /* 40 */               \
+  _(c10::dummy_int1_7_t<5>, Int5) /* 41 */               \
+  _(c10::dummy_int1_7_t<6>, Int6) /* 42 */               \
+  _(c10::dummy_int1_7_t<7>, Int7) /* 43 */               \
+  _(c10::Float8_e8m0fnu, Float8_e8m0fnu) /* 44 */        \
+  _(c10::Float4_e2m1fn_x2, Float4_e2m1fn_x2) /* 45 */
+
+// NB: despite its generic sounding name, the macros that don't take _AND
+// are mostly only used by tensorexpr
+#define AT_FORALL_INT_TYPES(_) \
+  _(uint8_t, Byte)             \
+  _(int8_t, Char)              \
+  _(int16_t, Short)            \
+  _(int, Int)                  \
+  _(int64_t, Long)
+
+#define AT_FORALL_SCALAR_TYPES(_) \
+  _(uint8_t, Byte)                \
+  _(int8_t, Char)                 \
+  _(int16_t, Short)               \
+  _(int, Int)                     \
+  _(int64_t, Long)                \
+  _(float, Float)                 \
+  _(double, Double)
+
+// These macros are often controlling how many template instantiations we
+// create for kernels.  It is typically inappropriate to add new dtypes here,
+// instead, new types should be added to use sites on a case-by-case basis.
+// We generally are not accepting new dtypes due to binary size concerns.
+
+#define AT_FORALL_SCALAR_TYPES_AND(SCALARTYPE, _) \
+  _(uint8_t, Byte)                                \
+  _(int8_t, Char)                                 \
+  _(int16_t, Short)                               \
+  _(int, Int)                                     \
+  _(int64_t, Long)                                \
+  _(float, Float)                                 \
+  _(double, Double)                               \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE>, SCALARTYPE)
+
+#define AT_FORALL_SCALAR_TYPES_AND2(SCALARTYPE1, SCALARTYPE2, _)   \
+  _(uint8_t, Byte)                                                 \
+  _(int8_t, Char)                                                  \
+  _(int16_t, Short)                                                \
+  _(int, Int)                                                      \
+  _(int64_t, Long)                                                 \
+  _(float, Float)                                                  \
+  _(double, Double)                                                \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE1>, \
+    SCALARTYPE1)                                                   \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE2>, SCALARTYPE2)
+
+#define AT_FORALL_SCALAR_TYPES_AND3(SCALARTYPE1, SCALARTYPE2, SCALARTYPE3, _) \
+  _(uint8_t, Byte)                                                            \
+  _(int8_t, Char)                                                             \
+  _(int16_t, Short)                                                           \
+  _(int, Int)                                                                 \
+  _(int64_t, Long)                                                            \
+  _(float, Float)                                                             \
+  _(double, Double)                                                           \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE1>,            \
+    SCALARTYPE1)                                                              \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE2>,            \
+    SCALARTYPE2)                                                              \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE3>, SCALARTYPE3)
+
+#define AT_FORALL_SCALAR_TYPES_AND7(                               \
+    SCALARTYPE1,                                                   \
+    SCALARTYPE2,                                                   \
+    SCALARTYPE3,                                                   \
+    SCALARTYPE4,                                                   \
+    SCALARTYPE5,                                                   \
+    SCALARTYPE6,                                                   \
+    SCALARTYPE7,                                                   \
+    _)                                                             \
+  _(uint8_t, Byte)                                                 \
+  _(int8_t, Char)                                                  \
+  _(int16_t, Short)                                                \
+  _(int, Int)                                                      \
+  _(int64_t, Long)                                                 \
+  _(float, Float)                                                  \
+  _(double, Double)                                                \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE1>, \
+    SCALARTYPE1)                                                   \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE2>, \
+    SCALARTYPE2)                                                   \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE3>, \
+    SCALARTYPE3)                                                   \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE4>, \
+    SCALARTYPE4)                                                   \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE5>, \
+    SCALARTYPE5)                                                   \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE6>, \
+    SCALARTYPE6)                                                   \
+  _(c10::impl::ScalarTypeToCPPTypeT<c10::ScalarType::SCALARTYPE7>, SCALARTYPE7)
+
+#define AT_FORALL_QINT_TYPES(_) \
+  _(c10::qint8, QInt8)          \
+  _(c10::quint8, QUInt8)        \
+  _(c10::qint32, QInt32)        \
+  _(c10::quint4x2, QUInt4x2)    \
+  _(c10::quint2x4, QUInt2x4)
+
+#define AT_FORALL_FLOAT8_TYPES(_)          \
+  _(c10::Float8_e5m2, Float8_e5m2)         \
+  _(c10::Float8_e4m3fn, Float8_e4m3fn)     \
+  _(c10::Float8_e5m2fnuz, Float8_e5m2fnuz) \
+  _(c10::Float8_e4m3fnuz, Float8_e4m3fnuz) \
+  _(c10::Float8_e8m0fnu, Float8_e8m0fnu)
+
+#define AT_FORALL_COMPLEX_TYPES(_)     \
+  _(c10::complex<float>, ComplexFloat) \
+  _(c10::complex<double>, ComplexDouble)
+
+enum class ScalarType : int8_t {
+#define DEFINE_ST_ENUM_VAL_(_1, n) n,
+  AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS(DEFINE_ST_ENUM_VAL_)
+#undef DEFINE_ENUM_ST_ENUM_VAL_
+      Undefined,
+  NumOptions
+};
+
+constexpr uint16_t NumScalarTypes =
+    static_cast<uint16_t>(ScalarType::NumOptions);
+
+// Map from C++ type to ScalarType enum
+template <typename T>
+struct CppTypeToScalarType;
+
+#define SPECIALIZE_CppTypeToScalarType(cpp_type, scalar_type)                  \
+  template <>                                                                  \
+  struct CppTypeToScalarType<cpp_type>                                         \
+      : std::                                                                  \
+            integral_constant<c10::ScalarType, c10::ScalarType::scalar_type> { \
+  };
+
+AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS(SPECIALIZE_CppTypeToScalarType)
+
+#undef SPECIALIZE_CppTypeToScalarType
+
+namespace impl {
+
+// These are used to map ScalarTypes to C++ types.
+
+template <c10::ScalarType N>
+struct ScalarTypeToCPPType;
+
+#define SPECIALIZE_ScalarTypeToCPPType(cpp_type, scalar_type)                \
+  template <>                                                                \
+  struct ScalarTypeToCPPType<c10::ScalarType::scalar_type> {                 \
+    using type = cpp_type;                                                   \
+                                                                             \
+    /* This is a workaround for the CUDA bug which prevents */               \
+    /* ::detail::ScalarTypeToCType<T>::type being used directly due to */    \
+    /* ambiguous reference which can't to be resolved. For some reason it */ \
+    /* can't pick between at::detail and at::cuda::detail. */                \
+    /* For repro example, please see: */                                     \
+    /* https://gist.github.com/izdeby/952ae7cf256ddb740a73776d39a7e7ba */    \
+    /* UPDATE: while the CUDA bug is fixed, we cannot remove the  */         \
+    /* workaround as it is BC breaking. However, it is recommended to  */    \
+    /* update any code that contains */                                      \
+    /*   decltype(ScalarTypeToCPPType<T>::t) */                              \
+    /* with */                                                               \
+    /*   ScalarTypeToCPPTypeT<T> */                                          \
+    static type t;                                                           \
+  };
+
+AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS(SPECIALIZE_ScalarTypeToCPPType)
+
+#undef SPECIALIZE_ScalarTypeToCPPType
+
+template <c10::ScalarType N>
+using ScalarTypeToCPPTypeT = typename ScalarTypeToCPPType<N>::type;
+
+} // namespace impl
+
+inline const char* toString(ScalarType t) {
+#define DEFINE_CASE(_, name) \
+  case ScalarType::name:     \
+    return #name;
+
+  switch (t) {
+    AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS(DEFINE_CASE)
+    default:
+      return "UNKNOWN_SCALAR";
+  }
+#undef DEFINE_CASE
+}
+
+inline std::ostream& operator<<(
+    std::ostream& stream,
+    c10::ScalarType scalar_type) {
+  return stream << toString(scalar_type);
+}
+
+inline bool isQIntType(ScalarType t) {
+  // Don't forget to extend this when adding new QInt types
+  return t == ScalarType::QInt8 || t == ScalarType::QUInt8 ||
+      t == ScalarType::QInt32 || t == ScalarType::QUInt4x2 ||
+      t == ScalarType::QUInt2x4;
+}
+
+inline ScalarType toUnderlying(ScalarType t) {
+  switch (t) {
+    case ScalarType::QUInt8:
+    case ScalarType::QUInt4x2:
+      [[fallthrough]];
+    case ScalarType::QUInt2x4:
+      return ScalarType::Byte;
+    case ScalarType::QInt8:
+      return ScalarType::Char;
+    case ScalarType::QInt32:
+      return ScalarType::Int;
+    default:
+      return t;
+  }
+}
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::CppTypeToScalarType;
+using c10::dummy_int1_7_t;
+using c10::dummy_uint1_7_t;
+using c10::NumScalarTypes;
+using c10::ScalarType;
+using c10::toString;
+using c10::operator<<;
+using c10::isQIntType;
+using c10::toUnderlying;
+
+namespace impl {
+using c10::impl::ScalarTypeToCPPTypeT;
+} // namespace impl
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+C10_DIAGNOSTIC_POP()
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/TensorAccessor.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/TensorAccessor.h
new file mode 100644
index 0000000000000000000000000000000000000000..9019c7ac3104dd90e78b89631e4f011d07786ffd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/TensorAccessor.h
@@ -0,0 +1,462 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Exception.h>
+#include <torch/headeronly/util/HeaderOnlyArrayRef.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <type_traits>
+
+namespace torch::headeronly {
+
+// The PtrTraits argument to the TensorAccessor/GenericPackedTensorAccessor
+// is used to enable the __restrict__ keyword/modifier for the data
+// passed to cuda.
+template <typename T>
+struct DefaultPtrTraits {
+  typedef T* PtrType;
+};
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+template <typename T>
+struct RestrictPtrTraits {
+  typedef T* __restrict__ PtrType;
+};
+#endif
+
+namespace detail {
+// Template classes in torch::headeronly::detail namespace are used
+// to construct accessor template classes with custom ArrayRef and
+// index bound check implementations. For instance,
+// at::TensorAccessor and torch::headeronly::TensorAccessor template
+// classes use c10::IntArrayRef and
+// torch::headeronly::IntHeaderOnlyArrayRef classes, respectively,
+// as return value types of sizes() and strides() methods.
+
+// TensorAccessorBase and TensorAccessor are used for both CPU and CUDA tensors.
+// For CUDA tensors it is used in device code (only). This means that we
+// restrict ourselves to functions and types available there (e.g. IntArrayRef
+// isn't).
+
+// The PtrTraits argument is only relevant to cuda to support `__restrict__`
+// pointers.
+template <
+    class ArrayRefCls,
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+class TensorAccessorBase {
+ public:
+  typedef typename PtrTraits<T>::PtrType PtrType;
+
+  C10_HOST_DEVICE TensorAccessorBase(
+      PtrType data_,
+      const index_t* sizes_,
+      const index_t* strides_)
+      : data_(data_), sizes_(sizes_), strides_(strides_) {}
+  C10_HOST ArrayRefCls sizes() const {
+    return ArrayRefCls(sizes_, N);
+  }
+  C10_HOST ArrayRefCls strides() const {
+    return ArrayRefCls(strides_, N);
+  }
+  C10_HOST_DEVICE index_t stride(index_t i) const {
+    return strides_[i];
+  }
+  C10_HOST_DEVICE index_t size(index_t i) const {
+    return sizes_[i];
+  }
+  C10_HOST_DEVICE PtrType data() {
+    return data_;
+  }
+  C10_HOST_DEVICE const PtrType data() const {
+    return data_;
+  }
+
+ protected:
+  PtrType data_;
+  const index_t* sizes_;
+  const index_t* strides_;
+};
+
+// The `TensorAccessor` is typically instantiated for CPU `Tensor`s using
+// `Tensor.accessor<T, N>()`.
+// For CUDA `Tensor`s, `GenericPackedTensorAccessor` is used on the host and
+// only indexing on the device uses `TensorAccessor`s.
+template <
+    class ArrayRefCls,
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+class TensorAccessor
+    : public TensorAccessorBase<ArrayRefCls, T, N, PtrTraits, index_t> {
+ public:
+  typedef typename PtrTraits<T>::PtrType PtrType;
+
+  C10_HOST_DEVICE TensorAccessor(
+      PtrType data_,
+      const index_t* sizes_,
+      const index_t* strides_)
+      : TensorAccessorBase<ArrayRefCls, T, N, PtrTraits, index_t>(
+            data_,
+            sizes_,
+            strides_) {}
+
+  C10_HOST_DEVICE TensorAccessor<ArrayRefCls, T, N - 1, PtrTraits, index_t>
+  operator[](index_t i) {
+    return TensorAccessor<ArrayRefCls, T, N - 1, PtrTraits, index_t>(
+        this->data_ + this->strides_[0] * i,
+        this->sizes_ + 1,
+        this->strides_ + 1);
+  }
+
+  C10_HOST_DEVICE const TensorAccessor<
+      ArrayRefCls,
+      T,
+      N - 1,
+      PtrTraits,
+      index_t>
+  operator[](index_t i) const {
+    return TensorAccessor<ArrayRefCls, T, N - 1, PtrTraits, index_t>(
+        this->data_ + this->strides_[0] * i,
+        this->sizes_ + 1,
+        this->strides_ + 1);
+  }
+};
+
+template <
+    class ArrayRefCls,
+    typename T,
+    template <typename U> class PtrTraits,
+    typename index_t>
+class TensorAccessor<ArrayRefCls, T, 1, PtrTraits, index_t>
+    : public TensorAccessorBase<ArrayRefCls, T, 1, PtrTraits, index_t> {
+ public:
+  typedef typename PtrTraits<T>::PtrType PtrType;
+
+  C10_HOST_DEVICE TensorAccessor(
+      PtrType data_,
+      const index_t* sizes_,
+      const index_t* strides_)
+      : TensorAccessorBase<ArrayRefCls, T, 1, PtrTraits, index_t>(
+            data_,
+            sizes_,
+            strides_) {}
+  C10_HOST_DEVICE T& operator[](index_t i) {
+    // NOLINTNEXTLINE(clang-analyzer-core.NullDereference)
+    return this->data_[this->strides_[0] * i];
+  }
+  C10_HOST_DEVICE const T& operator[](index_t i) const {
+    return this->data_[this->strides_[0] * i];
+  }
+};
+
+// GenericPackedTensorAccessorBase and GenericPackedTensorAccessor are used on
+// for CUDA `Tensor`s on the host and as in contrast to `TensorAccessor`s, they
+// copy the strides and sizes on instantiation (on the host) in order to
+// transfer them on the device when calling kernels. On the device, indexing of
+// multidimensional tensors gives to `TensorAccessor`s. Use RestrictPtrTraits as
+// PtrTraits if you want the tensor's data pointer to be marked as __restrict__.
+// Instantiation from data, sizes, strides is only needed on the host and
+// std::copy isn't available on the device, so those functions are host only.
+template <
+    typename IndexBoundsCheck,
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+class GenericPackedTensorAccessorBase {
+ public:
+  typedef typename PtrTraits<T>::PtrType PtrType;
+  C10_HOST GenericPackedTensorAccessorBase(
+      PtrType data_,
+      const index_t* sizes_,
+      const index_t* strides_)
+      : data_(data_) {
+    std::copy(sizes_, sizes_ + N, std::begin(this->sizes_));
+    std::copy(strides_, strides_ + N, std::begin(this->strides_));
+  }
+
+  // if index_t is not int64_t, we want to have an int64_t constructor
+  template <
+      typename source_index_t,
+      class = std::enable_if_t<std::is_same_v<source_index_t, int64_t>>>
+  C10_HOST GenericPackedTensorAccessorBase(
+      PtrType data_,
+      const source_index_t* sizes_,
+      const source_index_t* strides_)
+      : data_(data_) {
+    for (size_t i = 0; i < N; ++i) {
+      this->sizes_[i] = sizes_[i];
+      this->strides_[i] = strides_[i];
+    }
+  }
+
+  C10_HOST_DEVICE index_t stride(index_t i) const {
+    return strides_[i];
+  }
+  C10_HOST_DEVICE index_t size(index_t i) const {
+    return sizes_[i];
+  }
+  C10_HOST_DEVICE PtrType data() {
+    return data_;
+  }
+  C10_HOST_DEVICE const PtrType data() const {
+    return data_;
+  }
+
+ protected:
+  PtrType data_;
+  // NOLINTNEXTLINE(*c-arrays*)
+  index_t sizes_[N];
+  // NOLINTNEXTLINE(*c-arrays*)
+  index_t strides_[N];
+  C10_HOST void bounds_check_(index_t i) const {
+    IndexBoundsCheck _(i);
+  }
+};
+
+template <
+    typename ItemAccessor,
+    typename IndexBoundsCheck,
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+class GenericPackedTensorAccessor : public GenericPackedTensorAccessorBase<
+                                        IndexBoundsCheck,
+                                        T,
+                                        N,
+                                        PtrTraits,
+                                        index_t> {
+ public:
+  typedef typename PtrTraits<T>::PtrType PtrType;
+
+  C10_HOST GenericPackedTensorAccessor(
+      PtrType data_,
+      const index_t* sizes_,
+      const index_t* strides_)
+      : GenericPackedTensorAccessorBase<
+            IndexBoundsCheck,
+            T,
+            N,
+            PtrTraits,
+            index_t>(data_, sizes_, strides_) {}
+
+  // if index_t is not int64_t, we want to have an int64_t constructor
+  template <
+      typename source_index_t,
+      class = std::enable_if_t<std::is_same_v<source_index_t, int64_t>>>
+  C10_HOST GenericPackedTensorAccessor(
+      PtrType data_,
+      const source_index_t* sizes_,
+      const source_index_t* strides_)
+      : GenericPackedTensorAccessorBase<
+            IndexBoundsCheck,
+            T,
+            N,
+            PtrTraits,
+            index_t>(data_, sizes_, strides_) {}
+
+  C10_DEVICE ItemAccessor operator[](index_t i) {
+    index_t* new_sizes = this->sizes_ + 1;
+    index_t* new_strides = this->strides_ + 1;
+    return ItemAccessor(
+        this->data_ + this->strides_[0] * i, new_sizes, new_strides);
+  }
+
+  C10_DEVICE const ItemAccessor operator[](index_t i) const {
+    const index_t* new_sizes = this->sizes_ + 1;
+    const index_t* new_strides = this->strides_ + 1;
+    return ItemAccessor(
+        this->data_ + this->strides_[0] * i, new_sizes, new_strides);
+  }
+
+  /// Returns a PackedTensorAccessor of the same dimension after transposing the
+  /// two dimensions given. Does not actually move elements; transposition is
+  /// made by permuting the size/stride arrays. If the dimensions are not valid,
+  /// asserts.
+  C10_HOST GenericPackedTensorAccessor<
+      ItemAccessor,
+      IndexBoundsCheck,
+      T,
+      N,
+      PtrTraits,
+      index_t>
+  transpose(index_t dim1, index_t dim2) const {
+    this->bounds_check_(dim1);
+    this->bounds_check_(dim2);
+    GenericPackedTensorAccessor<
+        ItemAccessor,
+        IndexBoundsCheck,
+        T,
+        N,
+        PtrTraits,
+        index_t>
+        result(this->data_, this->sizes_, this->strides_);
+    std::swap(result.strides_[dim1], result.strides_[dim2]);
+    std::swap(result.sizes_[dim1], result.sizes_[dim2]);
+    return result;
+  }
+};
+
+template <
+    typename ItemAccessor,
+    typename IndexBoundsCheck,
+    typename T,
+    template <typename U> class PtrTraits,
+    typename index_t>
+class GenericPackedTensorAccessor<
+    ItemAccessor,
+    IndexBoundsCheck,
+    T,
+    1,
+    PtrTraits,
+    index_t>
+    : public GenericPackedTensorAccessorBase<
+          IndexBoundsCheck,
+          T,
+          1,
+          PtrTraits,
+          index_t> {
+ public:
+  typedef typename PtrTraits<T>::PtrType PtrType;
+  C10_HOST GenericPackedTensorAccessor(
+      PtrType data_,
+      const index_t* sizes_,
+      const index_t* strides_)
+      : GenericPackedTensorAccessorBase<
+            IndexBoundsCheck,
+            T,
+            1,
+            PtrTraits,
+            index_t>(data_, sizes_, strides_) {}
+
+  // if index_t is not int64_t, we want to have an int64_t constructor
+  template <
+      typename source_index_t,
+      class = std::enable_if_t<std::is_same_v<source_index_t, int64_t>>>
+  C10_HOST GenericPackedTensorAccessor(
+      PtrType data_,
+      const source_index_t* sizes_,
+      const source_index_t* strides_)
+      : GenericPackedTensorAccessorBase<
+            IndexBoundsCheck,
+            T,
+            1,
+            PtrTraits,
+            index_t>(data_, sizes_, strides_) {}
+
+  C10_DEVICE T& operator[](index_t i) {
+    return this->data_[this->strides_[0] * i];
+  }
+  C10_DEVICE const T& operator[](index_t i) const {
+    return this->data_[this->strides_[0] * i];
+  }
+
+  // Same as in the general N-dimensional case, but note that in the
+  // 1-dimensional case the returned PackedTensorAccessor will always be an
+  // identical copy of the original
+  C10_HOST GenericPackedTensorAccessor<
+      ItemAccessor,
+      IndexBoundsCheck,
+      T,
+      1,
+      PtrTraits,
+      index_t>
+  transpose(index_t dim1, index_t dim2) const {
+    this->bounds_check_(dim1);
+    this->bounds_check_(dim2);
+    return GenericPackedTensorAccessor<
+        ItemAccessor,
+        IndexBoundsCheck,
+        T,
+        1,
+        PtrTraits,
+        index_t>(this->data_, this->sizes_, this->strides_);
+  }
+};
+
+template <size_t N, typename index_t>
+struct HeaderOnlyIndexBoundsCheck {
+  HeaderOnlyIndexBoundsCheck(index_t i) {
+    STD_TORCH_CHECK(
+        0 <= i && i < index_t{N},
+        "Index ",
+        i,
+        " is not within bounds of a tensor of dimension ",
+        N);
+  }
+};
+
+} // namespace detail
+
+// HeaderOnlyTensorAccessorBase is same as at::TensorAccessorBase
+// except sizes() and strides() return IntHeaderOnlyArrayRef instead
+// of IntArrayRef.
+template <
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+using HeaderOnlyTensorAccessorBase = detail::TensorAccessorBase<
+    torch::headeronly::IntHeaderOnlyArrayRef,
+    T,
+    N,
+    PtrTraits,
+    index_t>;
+
+// HeaderOnlyTensorAccessor is same as at::TensorAccessor except
+// sizes() and strides() return IntHeaderOnlyArrayRef instead of
+// IntArrayRef.
+template <
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+using HeaderOnlyTensorAccessor = detail::TensorAccessor<
+    torch::headeronly::IntHeaderOnlyArrayRef,
+    T,
+    N,
+    PtrTraits,
+    index_t>;
+
+// HeaderOnlyGenericPackedTensorAccessorBase is same as
+// at::GenericPackedTensorAccessorBase except sizes() and strides()
+// return IntHeaderOnlyArrayRef instead of IntArrayRef.
+template <
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+using HeaderOnlyGenericPackedTensorAccessorBase =
+    detail::GenericPackedTensorAccessorBase<
+        detail::HeaderOnlyIndexBoundsCheck<N, index_t>,
+        T,
+        N,
+        PtrTraits,
+        index_t>;
+
+// HeaderOnlyGenericPackedTensorAccessor is same as
+// at::GenericPackedTensorAccessor except sizes() and strides() return
+// IntHeaderOnlyArrayRef instead of IntArrayRef, and bounds check uses
+// STD_TORCH_CHECK instead of TORCH_CHECK_INDEX.
+template <
+    typename T,
+    size_t N,
+    template <typename U> class PtrTraits = DefaultPtrTraits,
+    typename index_t = int64_t>
+using HeaderOnlyGenericPackedTensorAccessor =
+    detail::GenericPackedTensorAccessor<
+        HeaderOnlyTensorAccessor<T, N - 1, PtrTraits, index_t>,
+        detail::HeaderOnlyIndexBoundsCheck<N, index_t>,
+        T,
+        N,
+        PtrTraits,
+        index_t>;
+
+} // namespace torch::headeronly
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/enum_tag.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/enum_tag.h
new file mode 100644
index 0000000000000000000000000000000000000000..f45f3f49e86fb315e6afd58c79d75141447e9395
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/core/enum_tag.h
@@ -0,0 +1,36 @@
+#pragma once
+
+// @generated by torchgen/gen.py from enum_tag.h
+
+#include <torch/headeronly/macros/Macros.h>
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+
+// Enum of valid tags obtained from the entries in tags.yaml
+enum class Tag {
+    core,
+    cudagraph_unsafe,
+    data_dependent_output,
+    dynamic_output_shape,
+    flexible_layout,
+    generated,
+    inplace_view,
+    maybe_aliasing_or_mutating,
+    needs_contiguous_strides,
+    needs_exact_strides,
+    needs_fixed_stride_order,
+    nondeterministic_bitwise,
+    nondeterministic_seeded,
+    out_variant,
+    pointwise,
+    pt2_compliant_tag,
+    reduction,
+    view_copy
+};
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+// Re-expose in the at:: namespace for backward compatibility
+namespace at {
+    using torch::headeronly::Tag;
+}
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/cpu/vec/intrinsics.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/cpu/vec/intrinsics.h
new file mode 100644
index 0000000000000000000000000000000000000000..3cf427dae64bce378f9c284e0db6dd452cb0e3d5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/cpu/vec/intrinsics.h
@@ -0,0 +1,50 @@
+#pragma once
+#if defined(__GNUC__) && (defined(__x86_64__) || defined(__i386__))
+/* GCC or clang-compatible compiler, targeting x86/x86-64 */
+#include <x86intrin.h>
+#elif defined(__clang__) && (defined(__ARM_NEON__) || defined(__aarch64__))
+/* Clang-compatible compiler, targeting arm neon */
+#include <arm_neon.h>
+#if defined(__ARM_FEATURE_SVE)
+/* CLANG-compatible compiler, targeting ARM with SVE */
+#include <arm_sve.h>
+#endif
+#elif defined(_MSC_VER)
+/* Microsoft C/C++-compatible compiler */
+#include <intrin.h>
+#if _MSC_VER <= 1900
+#define _mm256_extract_epi64(X, Y) \
+  (_mm_extract_epi64(_mm256_extractf128_si256(X, Y >> 1), Y % 2))
+#define _mm256_extract_epi32(X, Y) \
+  (_mm_extract_epi32(_mm256_extractf128_si256(X, Y >> 2), Y % 4))
+#define _mm256_extract_epi16(X, Y) \
+  (_mm_extract_epi16(_mm256_extractf128_si256(X, Y >> 3), Y % 8))
+#define _mm256_extract_epi8(X, Y) \
+  (_mm_extract_epi8(_mm256_extractf128_si256(X, Y >> 4), Y % 16))
+#endif
+#elif defined(__GNUC__) && (defined(__ARM_NEON__) || defined(__aarch64__))
+/* GCC-compatible compiler, targeting ARM with NEON */
+#include <arm_neon.h>
+#if defined(__ARM_FEATURE_SVE)
+/* GCC-compatible compiler, targeting ARM with SVE */
+#include <arm_sve.h>
+#endif
+#elif defined(__GNUC__) && defined(__IWMMXT__)
+/* GCC-compatible compiler, targeting ARM with WMMX */
+#include <mmintrin.h>
+#elif defined(__s390x__)
+// targets Z/architecture
+// we will include vecintrin later
+#elif (defined(__GNUC__) || defined(__xlC__)) && \
+    (defined(__VEC__) || defined(__ALTIVEC__))
+/* XLC or GCC-compatible compiler, targeting PowerPC with VMX/VSX */
+#include <altivec.h>
+/* We need to undef those tokens defined by <altivec.h> to avoid conflicts
+   with the C++ types. => Can still use __bool/__vector */
+#undef bool
+#undef vector
+#undef pixel
+#elif defined(__GNUC__) && defined(__SPE__)
+/* GCC-compatible compiler, targeting PowerPC with SPE */
+#include <spe.h>
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/cpu/vec/vec_half.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/cpu/vec/vec_half.h
new file mode 100644
index 0000000000000000000000000000000000000000..6ad37b1948306903ea4f92c64d8402ffa3423b23
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/cpu/vec/vec_half.h
@@ -0,0 +1,59 @@
+#pragma once
+
+#include <torch/headeronly/cpu/vec/intrinsics.h>
+#include <torch/headeronly/macros/Macros.h>
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly, vec)
+// See Note [CPU_CAPABILITY namespace]
+inline namespace CPU_CAPABILITY {
+
+#if (defined(CPU_CAPABILITY_AVX2) || defined(CPU_CAPABILITY_AVX512)) && \
+    !defined(__APPLE__)
+static inline uint16_t float2half_scalar(float val) {
+#if defined(CPU_CAPABILITY_AVX2)
+#if defined(_MSC_VER)
+  __m256 v = _mm256_set1_ps(val);
+  __m128i o =
+      _mm256_cvtps_ph(v, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
+  return static_cast<std::uint16_t>(_mm_cvtsi128_si32(o));
+#else
+  return _cvtss_sh(val, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+#endif
+#elif defined(CPU_CAPABILITY_AVX512)
+  __m512 v = _mm512_set1_ps(val);
+  __m256i o =
+      _mm512_cvtps_ph(v, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
+  return static_cast<std::uint16_t>(
+      _mm_cvtsi128_si32(_mm256_castsi256_si128(o)));
+#endif
+}
+
+static inline float half2float_scalar(uint16_t val) {
+#if defined(CPU_CAPABILITY_AVX2)
+#if defined(_MSC_VER)
+  __m128i v = _mm_cvtsi32_si128(val);
+  __m256 o = _mm256_cvtph_ps(v);
+  return _mm256_cvtss_f32(o);
+#else
+  return _cvtsh_ss(val);
+#endif
+#elif defined(CPU_CAPABILITY_AVX512)
+  __m256i v =
+      _mm256_setr_epi16(val, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
+  __m512 o = _mm512_cvtph_ps(v);
+  return _mm512_cvtss_f32(o);
+#endif
+}
+
+#endif
+
+} // namespace CPU_CAPABILITY
+HIDDEN_NAMESPACE_END(torch, headeronly, vec)
+
+namespace at::vec {
+#if (defined(CPU_CAPABILITY_AVX2) || defined(CPU_CAPABILITY_AVX512)) && \
+    !defined(__APPLE__)
+using torch::headeronly::vec::float2half_scalar;
+using torch::headeronly::vec::half2float_scalar;
+#endif
+} // namespace at::vec
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/Export.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/Export.h
new file mode 100644
index 0000000000000000000000000000000000000000..14b26c18e7b1b702d2ecd1522e95001aaeb18d60
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/Export.h
@@ -0,0 +1,145 @@
+#pragma once
+
+#ifndef C10_MACROS_EXPORT_H_
+#define C10_MACROS_EXPORT_H_
+
+#ifndef C10_USING_CUSTOM_GENERATED_MACROS
+#include <torch/headeronly/macros/cmake_macros.h>
+#endif // C10_USING_CUSTOM_GENERATED_MACROS
+
+/* Header file to define the common scaffolding for exported symbols.
+ *
+ * Export is by itself a quite tricky situation to deal with, and if you are
+ * hitting this file, make sure you start with the background here:
+ * - Linux: https://gcc.gnu.org/wiki/Visibility
+ * - Windows:
+ * https://docs.microsoft.com/en-us/cpp/cpp/dllexport-dllimport?view=vs-2017
+ *
+ * Do NOT include this file directly. Instead, use c10/macros/Macros.h
+ */
+
+// You do not need to edit this part of file unless you are changing the core
+// pytorch export abstractions.
+//
+// This part defines the C10 core export and import macros. This is controlled
+// by whether we are building shared libraries or not, which is determined
+// during build time and codified in c10/core/cmake_macros.h.
+// When the library is built as a shared lib, EXPORT and IMPORT will contain
+// visibility attributes. If it is being built as a static lib, then EXPORT
+// and IMPORT basically have no effect.
+
+// As a rule of thumb, you should almost NEVER mix static and shared builds for
+// libraries that depend on c10. AKA, if c10 is built as a static library, we
+// recommend everything dependent on c10 to be built statically. If c10 is built
+// as a shared library, everything dependent on it should be built as shared. In
+// the PyTorch project, all native libraries shall use the macro
+// C10_BUILD_SHARED_LIB to check whether pytorch is building shared or static
+// libraries.
+
+// For build systems that do not directly depend on CMake and directly build
+// from the source directory (such as Buck), one may not have a cmake_macros.h
+// file at all. In this case, the build system is responsible for providing
+// correct macro definitions corresponding to the cmake_macros.h.in file.
+//
+// In such scenarios, one should define the macro
+//     C10_USING_CUSTOM_GENERATED_MACROS
+// to inform this header that it does not need to include the cmake_macros.h
+// file.
+
+#ifdef _WIN32
+#define C10_HIDDEN
+#if defined(C10_BUILD_SHARED_LIBS)
+#define C10_EXPORT __declspec(dllexport)
+#define C10_IMPORT __declspec(dllimport)
+#else
+#define C10_EXPORT
+#define C10_IMPORT
+#endif
+#else // _WIN32
+#if defined(__GNUC__)
+#define C10_EXPORT __attribute__((__visibility__("default")))
+#define C10_HIDDEN __attribute__((__visibility__("hidden")))
+#else // defined(__GNUC__)
+#define C10_EXPORT
+#define C10_HIDDEN
+#endif // defined(__GNUC__)
+#define C10_IMPORT C10_EXPORT
+#endif // _WIN32
+
+#ifdef NO_EXPORT
+#undef C10_EXPORT
+#define C10_EXPORT
+#endif
+
+// Definition of an adaptive XX_API macro, that depends on whether you are
+// building the library itself or not, routes to XX_EXPORT and XX_IMPORT.
+// Basically, you will need to do this for each shared library that you are
+// building, and the instruction is as follows: assuming that you are building
+// a library called libawesome.so. You should:
+// (1) for your cmake target (usually done by "add_library(awesome, ...)"),
+//     define a macro called AWESOME_BUILD_MAIN_LIB using
+//     target_compile_options.
+// (2) define the AWESOME_API macro similar to the one below.
+// And in the source file of your awesome library, use AWESOME_API to
+// annotate public symbols.
+
+// Here, for the C10 library, we will define the macro C10_API for both import
+// and export.
+
+// This one is being used by libc10.so
+#ifdef C10_BUILD_MAIN_LIB
+#define C10_API C10_EXPORT
+#else
+#define C10_API C10_IMPORT
+#endif
+
+// This one is being used by libtorch.so
+#ifdef CAFFE2_BUILD_MAIN_LIB
+#define TORCH_API C10_EXPORT
+#else
+#define TORCH_API C10_IMPORT
+#endif
+
+// You may be wondering why we have TORCH_CUDA_CPP_API and TORCH_CUDA_CU_API
+// belonging to the same library instead of just one TORCH_CUDA_API. Well, it
+// can indeed just be one TORCH_CUDA_API (and used to be)! TORCH_CUDA_CPP_API
+// and TORCH_CUDA_CU_API are artifacts of when we needed a split build to
+// avoid relocation marker linking errors. The context is as follows:
+//
+// Once upon a time, there _was_ only TORCH_CUDA_API. All was happy until we
+// tried to compile PyTorch for CUDA 11.1, which ran into relocation marker
+// issues when linking big binaries.
+// (https://github.com/pytorch/pytorch/issues/39968) We had two choices:
+//    (1) Stop supporting so many GPU architectures
+//    (2) Do something else
+// We chose #2 and decided to split the behemoth that was torch_cuda into two
+// smaller libraries, one with most of the core kernel functions (torch_cuda_cu)
+// and the other that had..well..everything else (torch_cuda_cpp). The idea was
+// this: instead of linking our static libraries (like the hefty
+// libcudnn_static.a) with another huge library, torch_cuda, and run into pesky
+// relocation marker issues, we could link our static libraries to a smaller
+// part of torch_cuda (torch_cuda_cpp) and avoid the issues.
+
+// libtorch_cuda.so (where torch_cuda_cu and torch_cuda_cpp are a part of the
+// same api)
+#ifdef TORCH_CUDA_BUILD_MAIN_LIB
+#define TORCH_CUDA_CPP_API C10_EXPORT
+#define TORCH_CUDA_CU_API C10_EXPORT
+#else
+#define TORCH_CUDA_CPP_API C10_IMPORT
+#define TORCH_CUDA_CU_API C10_IMPORT
+#endif
+
+#if defined(TORCH_XPU_BUILD_MAIN_LIB)
+#define TORCH_XPU_API C10_EXPORT
+#else
+#define TORCH_XPU_API C10_IMPORT
+#endif
+
+// Enums only need to be exported on windows for non-CUDA files
+#if defined(_WIN32) && defined(__CUDACC__)
+#define C10_API_ENUM C10_API
+#else
+#define C10_API_ENUM
+#endif
+#endif // C10_MACROS_EXPORT_H_
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/Macros.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/Macros.h
new file mode 100644
index 0000000000000000000000000000000000000000..cef99df3f566fbbe091f9b5932ba00523fc3c127
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/Macros.h
@@ -0,0 +1,741 @@
+#ifndef C10_MACROS_MACROS_H_
+#define C10_MACROS_MACROS_H_
+
+#ifdef __cplusplus
+#include <cassert>
+#else
+#include <assert.h>
+#endif
+
+/* Main entry for torch/headeronly/macros (used to be c10/macros).
+ *
+ * In your code, include torch/headeronly/macros/Macros.h directly, instead of
+ * individual files in this folder.
+ */
+
+// For build systems that do not directly depend on CMake and directly build
+// from the source directory (such as Buck), one may not have a cmake_macros.h
+// file at all. In this case, the build system is responsible for providing
+// correct macro definitions corresponding to the cmake_macros.h.in file.
+//
+// In such scenarios, one should define the macro
+//     C10_USING_CUSTOM_GENERATED_MACROS
+// to inform this header that it does not need to include the cmake_macros.h
+// file.
+
+#ifndef C10_USING_CUSTOM_GENERATED_MACROS
+#include <torch/headeronly/macros/cmake_macros.h>
+#endif // C10_USING_CUSTOM_GENERATED_MACROS
+
+#include <torch/headeronly/macros/Export.h>
+
+#if defined(__clang__)
+#define __ubsan_ignore_float_divide_by_zero__ \
+  __attribute__((no_sanitize("float-divide-by-zero")))
+#define __ubsan_ignore_undefined__ __attribute__((no_sanitize("undefined")))
+#define __ubsan_ignore_signed_int_overflow__ \
+  __attribute__((no_sanitize("signed-integer-overflow")))
+#define __ubsan_ignore_pointer_overflow__ \
+  __attribute__((no_sanitize("pointer-overflow")))
+#define __ubsan_ignore_function__ __attribute__((no_sanitize("function")))
+#define __ubsan_ignore_float_cast_overflow__ \
+  __attribute__((no_sanitize("float-cast-overflow")))
+#else
+#define __ubsan_ignore_float_divide_by_zero__
+#define __ubsan_ignore_undefined__
+#define __ubsan_ignore_signed_int_overflow__
+#define __ubsan_ignore_pointer_overflow__
+#define __ubsan_ignore_function__
+#define __ubsan_ignore_float_cast_overflow__
+#endif
+
+// Detect address sanitizer as some stuff doesn't work with it
+#undef C10_ASAN_ENABLED
+
+// for clang
+#if defined(__has_feature)
+#if ((__has_feature(address_sanitizer)))
+#define C10_ASAN_ENABLED 1
+#endif
+#endif
+
+// for gcc
+#if defined(__SANITIZE_ADDRESS__)
+#if __SANITIZE_ADDRESS__
+#if !defined(C10_ASAN_ENABLED)
+#define C10_ASAN_ENABLED 1
+#endif
+#endif
+#endif
+
+#if !defined(C10_ASAN_ENABLED)
+#define C10_ASAN_ENABLED 0
+#endif
+
+// Detect undefined-behavior sanitizer (UBSAN)
+#undef C10_UBSAN_ENABLED
+
+// for clang or gcc >= 14
+// NB: gcc 14 adds support for Clang's __has_feature
+//   https://gcc.gnu.org/gcc-14/changes.html
+//   gcc < 14 doesn't have a macro for UBSAN
+//   (e.g. __SANITIZE_UNDEFINED__ does not exist in gcc)
+//   https://github.com/google/sanitizers/issues/765
+#if defined(__has_feature)
+#if ((__has_feature(undefined_behavior_sanitizer)))
+#define C10_UBSAN_ENABLED 1
+#endif
+#endif
+
+#if !defined(C10_UBSAN_ENABLED)
+#define C10_UBSAN_ENABLED 0
+#endif
+
+// Disable the copy and assignment operator for a class. Note that this will
+// disable the usage of the class in std containers.
+#define C10_DISABLE_COPY_AND_ASSIGN(classname) \
+  classname(const classname&) = delete;        \
+  classname& operator=(const classname&) = delete
+
+#define C10_CONCATENATE_IMPL(s1, s2) s1##s2
+#define C10_CONCATENATE(s1, s2) C10_CONCATENATE_IMPL(s1, s2)
+
+#define C10_MACRO_EXPAND(args) args
+
+#define C10_STRINGIZE_IMPL(x) #x
+#define C10_STRINGIZE(x) C10_STRINGIZE_IMPL(x)
+
+/**
+ * C10_ANONYMOUS_VARIABLE(str) introduces a new identifier which starts with
+ * str and ends with a unique number.
+ */
+#ifdef __COUNTER__
+#define C10_UID __COUNTER__
+#define C10_ANONYMOUS_VARIABLE(str) C10_CONCATENATE(str, __COUNTER__)
+#else
+#define C10_UID __LINE__
+#define C10_ANONYMOUS_VARIABLE(str) C10_CONCATENATE(str, __LINE__)
+#endif
+
+#ifdef __has_cpp_attribute
+#define C10_HAS_CPP_ATTRIBUTE(x) __has_cpp_attribute(x)
+#else
+#define C10_HAS_CPP_ATTRIBUTE(x) (0)
+#endif
+
+#ifndef FBCODE_CAFFE2
+/// DEPRECATED: Warn if a type or return value is discarded.
+#define C10_NODISCARD [[nodiscard]]
+
+/// DEPRECATED: Suppress an unused variable.
+#define C10_UNUSED [[maybe_unused]]
+#endif
+
+#if !defined(__has_attribute)
+#define __has_attribute(x) 0
+#endif
+
+// Direct port of LLVM_ATTRIBUTE_USED.
+#if __has_attribute(used)
+#define C10_USED __attribute__((__used__))
+#else
+#define C10_USED
+#endif
+
+#define C10_RESTRICT __restrict
+
+#ifdef __cplusplus
+
+// Simply define the namespace, in case a dependent library want to refer to
+// the c10 namespace but not any nontrivial files.
+namespace c10 {}
+namespace c10::cuda {}
+namespace c10::hip {}
+namespace c10::xpu {}
+
+// Since C10 is the core library for caffe2 (and aten), we will simply reroute
+// all abstractions defined in c10 to be available in caffe2 as well.
+// This is only for backwards compatibility. Please use the symbols from the
+// c10 namespace where possible.
+namespace caffe2 {
+using namespace c10;
+}
+namespace at {
+using namespace c10;
+}
+namespace at::cuda {
+using namespace c10::cuda;
+} // namespace at::cuda
+
+// WARNING!!! THIS IS A GIANT HACK!!!
+// This line means you cannot simultaneously include c10/hip
+// and c10/cuda and then use them from the at::cuda namespace.
+// This is true in practice, because HIPIFY works inplace on
+// files in ATen/cuda, so it assumes that c10::hip is available
+// from at::cuda.  This namespace makes that happen.  When
+// HIPIFY is no longer out-of-place, we can switch the cuda
+// here to hip and everyone is happy.
+namespace at::cuda {
+using namespace c10::hip;
+} // namespace at::cuda
+
+namespace at::xpu {
+using namespace c10::xpu;
+} // namespace at::xpu
+
+#endif // __cplusplus
+
+// C10_LIKELY/C10_UNLIKELY
+//
+// These macros provide parentheses, so you can use these macros as:
+//
+//    if C10_LIKELY(some_expr) {
+//      ...
+//    }
+//
+// NB: static_cast to boolean is mandatory in C++, because __builtin_expect
+// takes a long argument, which means you may trigger the wrong conversion
+// without it.
+//
+#if defined(__GNUC__) || defined(__ICL) || defined(__clang__)
+#define C10_LIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 1))
+#define C10_UNLIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 0))
+#else
+#define C10_LIKELY(expr) (expr)
+#define C10_UNLIKELY(expr) (expr)
+#endif
+
+/// C10_NOINLINE - Functions whose declaration is annotated with this will not
+/// be inlined.
+#ifdef __GNUC__
+#define C10_NOINLINE __attribute__((noinline))
+#elif _MSC_VER
+#define C10_NOINLINE __declspec(noinline)
+#else
+#define C10_NOINLINE
+#endif
+
+#if defined(_MSC_VER)
+#define C10_ALWAYS_INLINE __forceinline
+#elif __has_attribute(always_inline) || defined(__GNUC__)
+#define C10_ALWAYS_INLINE __attribute__((__always_inline__)) inline
+#else
+#define C10_ALWAYS_INLINE inline
+#endif
+
+// Unlike C10_ALWAYS_INLINE, C10_ALWAYS_INLINE_ATTRIBUTE can be used
+// on a lambda.
+#if defined(_MSC_VER)
+// MSVC 14.39 is reasonably recent and doesn't like
+// [[msvc::forceinline]] on a lambda, so don't try to use it.
+#define C10_ALWAYS_INLINE_ATTRIBUTE
+#elif __has_attribute(always_inline) || defined(__GNUC__)
+#define C10_ALWAYS_INLINE_ATTRIBUTE __attribute__((__always_inline__))
+#else
+#define C10_ALWAYS_INLINE_ATTRIBUTE
+#endif
+
+#if defined(_MSC_VER)
+#define C10_ATTR_VISIBILITY_HIDDEN
+#elif defined(__GNUC__)
+#define C10_ATTR_VISIBILITY_HIDDEN __attribute__((__visibility__("hidden")))
+#else
+#define C10_ATTR_VISIBILITY_HIDDEN
+#endif
+
+#define C10_ERASE C10_ALWAYS_INLINE C10_ATTR_VISIBILITY_HIDDEN
+
+#ifdef __cplusplus
+#include <cstdint>
+#else
+#include <stdint.h>
+#endif
+
+#ifdef __HIPCC__
+// Unlike CUDA, HIP requires a HIP header to be included for __host__ to work.
+// We do this #include here so that C10_HOST_DEVICE and friends will Just Work.
+// See https://github.com/ROCm/hip/issues/441
+#include <hip/hip_runtime.h>
+#endif
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+// Designates functions callable from the host (CPU) and the device (GPU)
+#define C10_HOST_DEVICE __host__ __device__
+#define C10_DEVICE __device__
+#define C10_HOST __host__
+// constants from
+// (https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications)
+// The maximum number of threads per multiprocessor is 1024 for Turing
+// architecture (7.5), 1536 for Geforce Ampere (8.6)/Jetson Orin (8.7), and
+// 2048 for all other architectures. You'll get warnings if you exceed these
+// constants. Hence, the following macros adjust the input values from the user
+// to resolve potential warnings.
+#if __CUDA_ARCH__ == 750
+constexpr uint32_t CUDA_MAX_THREADS_PER_SM = 1024;
+#elif __CUDA_ARCH__ == 860 || __CUDA_ARCH__ == 870 || __CUDA_ARCH__ == 890 || \
+    __CUDA_ARCH__ == 1200
+constexpr uint32_t CUDA_MAX_THREADS_PER_SM = 1536;
+#else
+constexpr uint32_t CUDA_MAX_THREADS_PER_SM = 2048;
+#endif
+// CUDA_MAX_THREADS_PER_BLOCK is same for all architectures currently
+constexpr uint32_t CUDA_MAX_THREADS_PER_BLOCK = 1024;
+// CUDA_THREADS_PER_BLOCK_FALLBACK is the "canonical fallback" choice of block
+// size. 256 is a good number for this fallback and should give good occupancy
+// and versatility across all architectures.
+constexpr uint32_t CUDA_THREADS_PER_BLOCK_FALLBACK = 256;
+// NOTE: if you are thinking of constexpr-ify the inputs to launch bounds, it
+//       turns out that although __launch_bounds__ can take constexpr, it
+//       can't take a constexpr that has anything to do with templates.
+//       Currently we use launch_bounds that depend on template arguments in
+//       Loops.cuh, Reduce.cuh and LossCTC.cuh. Hence, C10_MAX_THREADS_PER_BLOCK
+//       and C10_MIN_BLOCKS_PER_SM are kept as macros.
+// Suppose you were planning to write __launch_bounds__(a, b), based on your
+// performance tuning on a modern GPU. Instead, you should write
+// __launch_bounds__(C10_MAX_THREADS_PER_BLOCK(a), C10_MIN_BLOCKS_PER_SM(a, b)),
+// which will also properly respect limits on old architectures.
+#define C10_MAX_THREADS_PER_BLOCK(val)           \
+  (((val) <= CUDA_MAX_THREADS_PER_BLOCK) ? (val) \
+                                         : CUDA_THREADS_PER_BLOCK_FALLBACK)
+#define C10_MIN_BLOCKS_PER_SM(threads_per_block, blocks_per_sm)        \
+  ((((threads_per_block) * (blocks_per_sm) <= CUDA_MAX_THREADS_PER_SM) \
+        ? (blocks_per_sm)                                              \
+        : ((CUDA_MAX_THREADS_PER_SM + (threads_per_block) - 1) /       \
+           (threads_per_block))))
+// C10_LAUNCH_BOUNDS is analogous to __launch_bounds__
+#define C10_LAUNCH_BOUNDS_0 \
+  __launch_bounds__(        \
+      256, 4) // default launch bounds that should give good occupancy and
+              // versatility across all architectures.
+#define C10_LAUNCH_BOUNDS_1(max_threads_per_block) \
+  __launch_bounds__((C10_MAX_THREADS_PER_BLOCK((max_threads_per_block))))
+#define C10_LAUNCH_BOUNDS_2(max_threads_per_block, min_blocks_per_sm) \
+  __launch_bounds__(                                                  \
+      (C10_MAX_THREADS_PER_BLOCK((max_threads_per_block))),           \
+      (C10_MIN_BLOCKS_PER_SM((max_threads_per_block), (min_blocks_per_sm))))
+#else
+#define C10_HOST_DEVICE
+#define C10_HOST
+#define C10_DEVICE
+#endif
+
+#if defined(USE_ROCM)
+#define C10_HIP_HOST_DEVICE __host__ __device__
+#else
+#define C10_HIP_HOST_DEVICE
+#endif
+
+// C10_WARP_SIZE is only allowed for device code.
+// Host code dynamically-sized launch configs _must_ use at::cuda::warp_size().
+// Host or device statically-sized arrays _must_ use either
+// C10_WARP_SIZE_UPPER_BOUND or C10_WARP_SIZE_LOWER_BOUND, as needed.
+//
+// HIP header used to define warpSize as a constexpr that was either 32 or 64
+// depending on the target device, and then always set it to 64 for host code.
+// For a time, that allowed C10_WARP_SIZE to be defined like so:
+//
+// #ifdef USE_ROCM
+// #define C10_WARP_SIZE warpSize
+// #else
+// #define C10_WARP_SIZE 32
+// #endif
+//
+// In ROCm 7, warpSize is no longer constexpr, matching CUDA behavior.
+// We can now only use warpSize for C10_WARP_SIZE in device code and this is
+// enforced by using __device__ in its definition.  In host code where
+// C10_WARP_SIZE was previously used as a compile-time constant, this will now
+// cause a compile-time error.
+//
+// If an array was previously expected to be sized at compile-time using
+// C10_WARP_SIZE, users must now use either C10_WARP_SIZE_UPPER_BOUND or
+// C10_WARP_SIZE_LOWER_BOUND depending on the situation.
+//
+// If C10_WARP_SIZE was previously used to determine kernel launch sizes, users
+// must now use at::cuda::warp_size() for the dynamic runtime query.
+//
+// Unfortunately, C10_WARP_SIZE has been public and available for both host and
+// device since approximately 2019, so forcing it to be device-only would break
+// existing code in the wild.
+#if defined(USE_ROCM)
+namespace at::cuda {
+TORCH_CUDA_CPP_API int warp_size();
+}
+#if defined(__HIPCC__)
+static __host__ inline int C10_WARP_SIZE_INTERNAL() {
+  return at::cuda::warp_size();
+}
+// NOTE: __device__ C10_WARP_SIZE_INTERNAL
+// For __SPIRV__, we must use dynamic warpSize. When not targeting __SPIRV__,
+// we can use constexpr. This matches prior behavior. We preserve this for
+// backward compatibility instead of forcing old code to use dynamic warpSize
+// and losing constexpr. However, compiling for --offload-arch=amdgcnspirv
+// could expose where C10_WARP_SIZE was used incorrectly where the dynamic
+// warpSize is not allowed.
+#if defined(__SPIRV__)
+static __device__ inline int C10_WARP_SIZE_INTERNAL() {
+  return warpSize;
+}
+#else // __SPIRV__
+static __device__ inline constexpr int C10_WARP_SIZE_INTERNAL() {
+#if defined(__GFX9__)
+  return 64;
+#else // __GFX9__
+  return 32;
+#endif // __GFX9__
+}
+#endif // __SPIRV__
+#if defined(__SPIRV__)
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 64
+#elif defined(__GFX9__)
+#define C10_WARP_SIZE_LOWER_BOUND 64
+#define C10_WARP_SIZE_UPPER_BOUND 64
+#else
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 32
+#endif
+#else // !__HIPCC__
+static inline int C10_WARP_SIZE_INTERNAL() {
+  return at::cuda::warp_size();
+}
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 64
+#endif // __HIPCC__
+#define C10_WARP_SIZE (C10_WARP_SIZE_INTERNAL())
+#else // !USE_ROCM
+#define C10_WARP_SIZE 32
+#define C10_WARP_SIZE_LOWER_BOUND 32
+#define C10_WARP_SIZE_UPPER_BOUND 32
+#endif // USE_ROCM
+
+#if defined(_MSC_VER) && _MSC_VER <= 1900
+#define __func__ __FUNCTION__
+#endif
+
+// CUDA_KERNEL_ASSERT checks the assertion
+// even when NDEBUG is defined. This is useful for important assertions in CUDA
+// code that would otherwise be suppressed when building Release.
+#if defined(__ANDROID__) || defined(__APPLE__) || defined(__FreeBSD__)
+// Those platforms do not support assert()
+#define CUDA_KERNEL_ASSERT(cond)
+#define CUDA_KERNEL_ASSERT_MSG(cond, msg)
+#define CUDA_KERNEL_ASSERT_PRINTF(cond, msg, ...)
+#define SYCL_KERNEL_ASSERT(cond)
+#elif defined(_MSC_VER)
+#if defined(NDEBUG)
+extern "C" {
+C10_IMPORT
+#if defined(__SYCL_DEVICE_ONLY__)
+extern SYCL_EXTERNAL void _wassert(
+    const wchar_t* wexpr,
+    const wchar_t* wfile,
+    unsigned line);
+#else
+#if defined(__CUDA_ARCH__)
+__host__ __device__
+#endif // __CUDA_ARCH__
+    void
+    _wassert(wchar_t const* _Message, wchar_t const* _File, unsigned _Line);
+#endif // __SYCL_DEVICE_ONLY__
+}
+#endif // NDEBUG
+#define CUDA_KERNEL_ASSERT(cond)                 \
+  if (C10_UNLIKELY(!(cond))) {                   \
+    (void)(_wassert(                             \
+               _CRT_WIDE(#cond),                 \
+               _CRT_WIDE(__FILE__),              \
+               static_cast<unsigned>(__LINE__)), \
+           0);                                   \
+  }
+// TODO: This doesn't assert the message because I (chilli) couldn't figure out
+// a nice way to convert a char* to a wchar_t*
+#define CUDA_KERNEL_ASSERT_MSG(cond, msg)        \
+  if (C10_UNLIKELY(!(cond))) {                   \
+    (void)(_wassert(                             \
+               _CRT_WIDE(#cond),                 \
+               _CRT_WIDE(__FILE__),              \
+               static_cast<unsigned>(__LINE__)), \
+           0);                                   \
+  }
+#define CUDA_KERNEL_ASSERT_PRINTF(cond, msg, ...)                     \
+  if (C10_UNLIKELY(!(cond))) {                                        \
+    (void)(printf(                                                    \
+        "[CUDA_KERNEL_ASSERT] " __FILE__ ":" C10_STRINGIZE(           \
+            __LINE__) ": %s: block: [%d,%d,%d], thread: [%d,%d,%d]: " \
+                      "Assertion failed: `" #cond "`: " msg "\n",     \
+        __func__,                                                     \
+        blockIdx.x,                                                   \
+        blockIdx.y,                                                   \
+        blockIdx.z,                                                   \
+        threadIdx.x,                                                  \
+        threadIdx.y,                                                  \
+        threadIdx.z,                                                  \
+        ##__VA_ARGS__));                                              \
+    (void)(_wassert(                                                  \
+               _CRT_WIDE(#cond),                                      \
+               _CRT_WIDE(__FILE__),                                   \
+               static_cast<unsigned>(__LINE__)),                      \
+           0);                                                        \
+  }
+#define SYCL_KERNEL_ASSERT(cond)                 \
+  if (C10_UNLIKELY(!(cond))) {                   \
+    (void)(_wassert(                             \
+               _CRT_WIDE(#cond),                 \
+               _CRT_WIDE(__FILE__),              \
+               static_cast<unsigned>(__LINE__)), \
+           0);                                   \
+  }
+#else // __APPLE__, _MSC_VER
+#if defined(NDEBUG)
+extern "C" {
+#if defined(__SYCL_DEVICE_ONLY__)
+extern SYCL_EXTERNAL void __assert_fail(
+    const char* expr,
+    const char* file,
+    unsigned int line,
+    const char* func);
+#elif (defined(__EMSCRIPTEN__))
+// As defined in assert.h in the Emscripten stdlib
+_Noreturn void __assert_fail(
+    const char* expr,
+    const char* file,
+    int line,
+    const char* func);
+#else // __SYCL_DEVICE_ONLY__
+#if (defined(__CUDA_ARCH__) && !(defined(__clang__) && defined(__CUDA__)))
+// CUDA supports __assert_fail function which are common for both device
+// and host side code.
+__host__ __device__
+#endif
+
+    // This forward declaration matching the declaration of __assert_fail
+    // exactly how it is in glibc in case parts of the program are compiled with
+    // different NDEBUG settings. Otherwise we might get 'ambiguous declaration'
+    // error. Note: On ROCm - this declaration serves for host side compilation.
+    void
+    __assert_fail(
+        const char* assertion,
+        const char* file,
+        unsigned int line,
+        const char* function) noexcept __attribute__((__noreturn__));
+
+#endif // __SYCL_DEVICE_ONLY__
+}
+#endif // NDEBUG
+// ROCm disables kernel assert by default for performance considerations.
+// Though ROCm supports __assert_fail, it uses kernel printf which has
+// a non-negligible performance impact even if the assert condition is
+// never triggered. We choose to use abort() instead which will still
+// terminate the application but without a more useful error message.
+#if !defined(C10_USE_ROCM_KERNEL_ASSERT) && defined(USE_ROCM)
+#define CUDA_KERNEL_ASSERT(cond) \
+  if C10_UNLIKELY (!(cond)) {    \
+    abort();                     \
+  }
+#define CUDA_KERNEL_ASSERT_MSG(cond, msg) \
+  if C10_UNLIKELY (!(cond)) {             \
+    abort();                              \
+  }
+#define CUDA_KERNEL_ASSERT_PRINTF(cond, msg, ...) \
+  if C10_UNLIKELY (!(cond)) {                     \
+    abort();                                      \
+  }
+#define SYCL_KERNEL_ASSERT(cond) \
+  if C10_UNLIKELY (!(cond)) {    \
+    abort();                     \
+  }
+#else
+#define CUDA_KERNEL_ASSERT(cond)                                         \
+  if (C10_UNLIKELY(!(cond))) {                                           \
+    __assert_fail(                                                       \
+        #cond, __FILE__, static_cast<unsigned int>(__LINE__), __func__); \
+  }
+#define CUDA_KERNEL_ASSERT_MSG(cond, msg)                              \
+  if (C10_UNLIKELY(!(cond))) {                                         \
+    __assert_fail(                                                     \
+        msg, __FILE__, static_cast<unsigned int>(__LINE__), __func__); \
+  }
+#define CUDA_KERNEL_ASSERT_PRINTF(cond, msg, ...)                        \
+  if (C10_UNLIKELY(!(cond))) {                                           \
+    printf(                                                            \
+        "[CUDA_KERNEL_ASSERT] " __FILE__ ":" C10_STRINGIZE(            \
+            __LINE__) ": %s: block: [%d,%d,%d], thread: [%d,%d,%d]: "  \
+            "Assertion failed: `" #cond "`: " msg "\n",                \
+        __func__,                                                      \
+        blockIdx.x,                                                    \
+        blockIdx.y,                                                    \
+        blockIdx.z,                                                    \
+        threadIdx.x,                                                   \
+        threadIdx.y,                                                   \
+        threadIdx.z,                                                   \
+        ##__VA_ARGS__); \
+    __assert_fail(                                                       \
+        #cond, __FILE__, static_cast<unsigned int>(__LINE__), __func__); \
+  }
+#define SYCL_KERNEL_ASSERT(cond)                                         \
+  if (C10_UNLIKELY(!(cond))) {                                           \
+    __assert_fail(                                                       \
+        #cond, __FILE__, static_cast<unsigned int>(__LINE__), __func__); \
+  }
+#endif //  C10_USE_ROCM_KERNEL_ASSERT && USE_ROCM
+#endif // __APPLE__
+
+// Compile-time switch to control how assertions are logged inside CUDA kernels.
+// If C10_CUDA_VERBOSE_ASSERT is defined,  CUDA_KERNEL_ASSERT_VERBOSE will
+// take addition information passed to the macro and forward them to
+// CUDA_KERNEL_ASSERT_PRINTF If C10_CUDA_VERBOSE_ASSERT is not defined,
+// CUDA_KERNEL_ASSERT_VERBOSE will behave the same as CUDA_KERNEL_ASSERT.
+#ifdef C10_ENABLE_VERBOSE_ASSERT
+#define CUDA_KERNEL_ASSERT_VERBOSE(cond, ...) \
+  CUDA_KERNEL_ASSERT_PRINTF(cond, __VA_ARGS__)
+#else
+#define CUDA_KERNEL_ASSERT_VERBOSE(cond, ...) CUDA_KERNEL_ASSERT(cond)
+#endif
+
+#ifdef __APPLE__
+#include <TargetConditionals.h>
+#endif
+
+#if defined(__ANDROID__)
+#define C10_ANDROID 1
+#define C10_MOBILE 1
+#elif (                   \
+    defined(__APPLE__) && \
+    (TARGET_IPHONE_SIMULATOR || TARGET_OS_SIMULATOR || TARGET_OS_IPHONE))
+#define C10_IOS 1
+#define C10_MOBILE 1
+#endif // ANDROID / IOS
+
+#if defined(C10_MOBILE) && C10_MOBILE
+#define C10_ALWAYS_INLINE_UNLESS_MOBILE inline
+#else
+#define C10_ALWAYS_INLINE_UNLESS_MOBILE C10_ALWAYS_INLINE
+#endif
+
+#if !defined(FBCODE_CAFFE2) && !defined(C10_NODEPRECATED)
+#define CONSTEXPR_EXCEPT_WIN_CUDA constexpr
+#define C10_HOST_CONSTEXPR_EXCEPT_WIN_CUDA constexpr
+
+#define STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(field, val) \
+  static constexpr const char field[] = val;
+#define STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA(cls, field, val)
+#endif // !defined(FBCODE_CAFFE2) && !defined(C10_NODEPRECATED)
+
+#ifndef HAS_DEMANGLE
+#if defined(__ANDROID__) || defined(_WIN32) || defined(__EMSCRIPTEN__)
+#define HAS_DEMANGLE 0
+#elif defined(__APPLE__) && \
+    (TARGET_IPHONE_SIMULATOR || TARGET_OS_SIMULATOR || TARGET_OS_IPHONE)
+#define HAS_DEMANGLE 0
+#else
+#define HAS_DEMANGLE 1
+#endif
+#endif // HAS_DEMANGLE
+
+#define _C10_PRAGMA__(string) _Pragma(#string)
+#define _C10_PRAGMA_(string) _C10_PRAGMA__(string)
+
+#ifdef __clang__
+#define C10_CLANG_DIAGNOSTIC_PUSH() _Pragma("clang diagnostic push")
+#define C10_CLANG_DIAGNOSTIC_POP() _Pragma("clang diagnostic pop")
+#define C10_CLANG_DIAGNOSTIC_IGNORE(flag) \
+  _C10_PRAGMA_(clang diagnostic ignored flag)
+#define C10_CLANG_HAS_WARNING(flag) __has_warning(flag)
+#else
+#define C10_CLANG_DIAGNOSTIC_PUSH()
+#define C10_CLANG_DIAGNOSTIC_POP()
+#define C10_CLANG_DIAGNOSTIC_IGNORE(flag)
+#define C10_CLANG_HAS_WARNING(flag) 0
+#endif
+
+#ifdef __clang__
+
+#define C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED(warning)         \
+  _C10_PRAGMA_(clang diagnostic push)                               \
+  _C10_PRAGMA_(clang diagnostic ignored "-Wunknown-warning-option") \
+  _C10_PRAGMA_(clang diagnostic ignored warning)
+
+#define C10_DIAGNOSTIC_POP() _C10_PRAGMA_(clang diagnostic pop)
+
+#elif __GNUC__
+
+#define C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED(warning) \
+  _C10_PRAGMA_(GCC diagnostic push)                         \
+  _C10_PRAGMA_(GCC diagnostic ignored "-Wpragmas")          \
+  _C10_PRAGMA_(GCC diagnostic ignored warning)
+
+#define C10_DIAGNOSTIC_POP() _C10_PRAGMA_(GCC diagnostic pop)
+
+#else
+
+#define C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED(warning)
+#define C10_DIAGNOSTIC_POP()
+
+#endif
+
+// This macro is used to find older C++ compilers
+// that don't support move optimization for return values.
+
+#if (defined(__GNUC__) && __GNUC__ < 13 && __cplusplus < 202002L) || \
+    (defined(__clang_major__) && __clang_major__ < 13)
+#define C10_RETURN_MOVE_IF_OLD_COMPILER 1
+#else
+#define C10_RETURN_MOVE_IF_OLD_COMPILER 0
+#endif
+
+// The HIDDEN_NAMESPACE_BEGIN and HIDDEN_NAMESPACE_END below
+// are needed for maintaining robustness in our header APIs in
+// torch/headeronly and torch/csrc/stable under the namespaces
+// torch::headeronly and torch::stable respectively. We enforce
+// hidden visibility for these APIs because we want to enable
+// loading custom extensions compiled against different libtorch
+// versions where these APIs may have changed.
+
+// Helper macros to handle 1-3 hidden namespace levels when not windows
+#define _HIDDEN_NS_GET_MACRO(_1, _2, _3, NAME, ...) NAME
+#define _HIDDEN_NS_1(n1) namespace n1 __attribute__((visibility("hidden"))) {
+#define _HIDDEN_NS_2(n1, n2) \
+  namespace n1 {             \
+  namespace n2 __attribute__((visibility("hidden"))) {
+#define _HIDDEN_NS_3(n1, n2, n3) \
+  namespace n1::n2 {             \
+  namespace n3 __attribute__((visibility("hidden"))) {
+
+// Helper macros to close namespaces when not windows
+#define _HIDDEN_NS_END_1(n1) }
+#define _HIDDEN_NS_END_N(n1, ...) \
+  }                               \
+  }
+
+// Helper macros to join strs with :: (for win, where symbols are hidden by
+// default)
+#define _EXPAND(...) __VA_ARGS__
+#define _JOIN_GET_MACRO(_1, _2, _3, NAME, ...) NAME
+#define _JOIN_NS1(a) a
+#define _JOIN_NS2(a, b) a::b
+#define _JOIN_NS3(a, b, c) a::b::c
+
+#if !defined(HIDDEN_NAMESPACE_BEGIN)
+#if defined(__GNUG__) && !defined(_WIN32)
+#define HIDDEN_NAMESPACE_BEGIN(...) \
+  _HIDDEN_NS_GET_MACRO(             \
+      __VA_ARGS__, _HIDDEN_NS_3, _HIDDEN_NS_2, _HIDDEN_NS_1)(__VA_ARGS__)
+#else
+#define HIDDEN_NAMESPACE_BEGIN(...)  \
+  namespace _EXPAND(_JOIN_GET_MACRO( \
+      __VA_ARGS__, _JOIN_NS3, _JOIN_NS2, _JOIN_NS1)(__VA_ARGS__)) {
+#endif
+#endif
+
+#if !defined(HIDDEN_NAMESPACE_END)
+#if defined(__GNUG__) && !defined(_WIN32)
+#define HIDDEN_NAMESPACE_END(...)                                         \
+  _HIDDEN_NS_GET_MACRO(                                                   \
+      __VA_ARGS__, _HIDDEN_NS_END_N, _HIDDEN_NS_END_N, _HIDDEN_NS_END_1)( \
+      __VA_ARGS__)
+#else
+#define HIDDEN_NAMESPACE_END(...) }
+#endif
+#endif
+
+#endif // C10_MACROS_MACROS_H_
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/cmake_macros.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/cmake_macros.h
new file mode 100644
index 0000000000000000000000000000000000000000..aafbb2a9ae4381d694e9f12a395ed28c42d5386f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/macros/cmake_macros.h
@@ -0,0 +1,14 @@
+#ifndef C10_MACROS_CMAKE_MACROS_H_
+#define C10_MACROS_CMAKE_MACROS_H_
+
+// Automatically generated header file for the C10 library.
+// Do not include this file directly. Instead, include torch/headeronly/macros/Macros.h.
+
+#define C10_BUILD_SHARED_LIBS
+/* #undef C10_USE_GLOG */
+/* #undef C10_USE_GFLAGS */
+#define C10_USE_NUMA
+/* #undef C10_USE_MSVC_STATIC_RUNTIME */
+/* #undef C10_USE_ROCM_KERNEL_ASSERT */
+
+#endif // C10_MACROS_CMAKE_MACROS_H_
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/BFloat16.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/BFloat16.h
new file mode 100644
index 0000000000000000000000000000000000000000..9aa08c265bd2c9e67c34fab9838f5bcf010ab912
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/BFloat16.h
@@ -0,0 +1,481 @@
+#pragma once
+
+// Defines the bloat16 type (brain floating-point). This representation uses
+// 1 bit for the sign, 8 bits for the exponent and 7 bits for the mantissa.
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/bit_cast.h>
+
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <iosfwd>
+#include <ostream>
+
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
+#include <cuda_bf16.h>
+#endif
+
+#if defined(CL_SYCL_LANGUAGE_VERSION)
+#include <CL/sycl.hpp> // for SYCL 1.2.1
+#elif defined(SYCL_LANGUAGE_VERSION)
+#include <sycl/sycl.hpp> // for SYCL 2020
+#endif
+
+namespace c10 {
+
+struct alignas(2) BFloat16 {
+  uint16_t x;
+
+  // HIP wants __host__ __device__ tag, CUDA does not
+#if defined(USE_ROCM) && defined(__HIPCC__)
+  C10_HOST_DEVICE BFloat16() = default;
+#else
+  BFloat16() = default;
+#endif
+
+  struct from_bits_t {};
+  static constexpr C10_HOST_DEVICE from_bits_t from_bits() {
+    return from_bits_t();
+  }
+
+  constexpr C10_HOST_DEVICE BFloat16(
+      unsigned short bits,
+      from_bits_t /*unused*/)
+      : x(bits) {}
+  /* implicit */ inline C10_HOST_DEVICE BFloat16(float value);
+  inline C10_HOST_DEVICE operator float() const;
+
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
+  inline C10_HOST_DEVICE BFloat16(const __nv_bfloat16& value);
+  explicit inline C10_HOST_DEVICE operator __nv_bfloat16() const;
+#endif
+
+#if defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS)
+  inline C10_HOST_DEVICE BFloat16(const sycl::ext::oneapi::bfloat16& value);
+  explicit inline C10_HOST_DEVICE operator sycl::ext::oneapi::bfloat16() const;
+#endif
+};
+
+inline std::ostream& operator<<(std::ostream& out, const BFloat16& value) {
+  out << (float)value;
+  return out;
+}
+
+namespace detail {
+inline C10_HOST_DEVICE float f32_from_bits(uint16_t src) {
+  float res = 0;
+  uint32_t tmp = src;
+  tmp <<= 16;
+
+#if defined(USE_ROCM) && defined(__HIPCC__)
+  float* tempRes;
+
+  // We should be using memcpy in order to respect the strict aliasing rule
+  // but it fails in the HIP environment.
+  tempRes = reinterpret_cast<float*>(&tmp);
+  res = *tempRes;
+#else
+  std::memcpy(&res, &tmp, sizeof(tmp));
+#endif
+
+  return res;
+}
+
+inline C10_HOST_DEVICE uint16_t bits_from_f32(float src) {
+  uint32_t res = 0;
+
+#if defined(USE_ROCM) && defined(__HIPCC__)
+  // We should be using memcpy in order to respect the strict aliasing rule
+  // but it fails in the HIP environment.
+  uint32_t* tempRes = reinterpret_cast<uint32_t*>(&src);
+  res = *tempRes;
+#else
+  std::memcpy(&res, &src, sizeof(res));
+#endif
+
+  return res >> 16;
+}
+
+inline C10_HOST_DEVICE uint16_t round_to_nearest_even(float src) {
+#if defined(USE_ROCM) && defined(__HIPCC__)
+  if (src != src) {
+#elif defined(_MSC_VER)
+  if (isnan(src)) {
+#else
+  if (std::isnan(src)) {
+#endif
+    return UINT16_C(0x7FC0);
+  } else {
+    const uint32_t U32 = c10::bit_cast<uint32_t>(src);
+    uint32_t rounding_bias = ((U32 >> 16) & 1) + UINT32_C(0x7FFF);
+    return static_cast<uint16_t>((U32 + rounding_bias) >> 16);
+  }
+}
+
+} // namespace detail
+
+//-------- the following is copied from c10/util/BFloat16-inl.h ---------//
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+/// Constructors
+inline C10_HOST_DEVICE BFloat16::BFloat16(float value)
+    :
+#if defined(__CUDACC__) &&                                                   \
+    (!defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || \
+     defined(USE_ROCM) && (TORCH_HIP_VERSION >= 702))
+      x(__bfloat16_as_ushort(__float2bfloat16(value)))
+#elif defined(__SYCL_DEVICE_ONLY__) && \
+    defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS)
+      x(c10::bit_cast<uint16_t>(sycl::ext::oneapi::bfloat16(value)))
+#else
+      // RNE by default
+      x(detail::round_to_nearest_even(value))
+#endif
+{
+}
+
+/// Implicit conversions
+inline C10_HOST_DEVICE BFloat16::operator float() const {
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
+  return __bfloat162float(*reinterpret_cast<const __nv_bfloat16*>(&x));
+#elif defined(__SYCL_DEVICE_ONLY__) && \
+    defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS)
+  return float(*reinterpret_cast<const sycl::ext::oneapi::bfloat16*>(&x));
+#else
+  return detail::f32_from_bits(x);
+#endif
+}
+
+#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702))
+inline C10_HOST_DEVICE BFloat16::BFloat16(const __nv_bfloat16& value) {
+  x = *reinterpret_cast<const unsigned short*>(&value);
+}
+inline C10_HOST_DEVICE BFloat16::operator __nv_bfloat16() const {
+  return *reinterpret_cast<const __nv_bfloat16*>(&x);
+}
+#endif
+
+#if defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS)
+inline C10_HOST_DEVICE BFloat16::BFloat16(
+    const sycl::ext::oneapi::bfloat16& value) {
+  x = *reinterpret_cast<const unsigned short*>(&value);
+}
+inline C10_HOST_DEVICE BFloat16::operator sycl::ext::oneapi::bfloat16() const {
+  return *reinterpret_cast<const sycl::ext::oneapi::bfloat16*>(&x);
+}
+#endif
+
+// CUDA intrinsics
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+inline C10_DEVICE BFloat16 __ldg(const BFloat16* ptr) {
+#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  return __ldg(reinterpret_cast<const __nv_bfloat16*>(ptr));
+#else
+  return *ptr;
+#endif
+}
+#endif
+
+/// Arithmetic
+
+inline C10_HOST_DEVICE BFloat16
+operator+(const BFloat16& a, const BFloat16& b) {
+  return static_cast<float>(a) + static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE BFloat16
+operator-(const BFloat16& a, const BFloat16& b) {
+  return static_cast<float>(a) - static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE BFloat16
+operator*(const BFloat16& a, const BFloat16& b) {
+  return static_cast<float>(a) * static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE BFloat16 operator/(const BFloat16& a, const BFloat16& b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE BFloat16 operator-(const BFloat16& a) {
+  return -static_cast<float>(a);
+}
+
+inline C10_HOST_DEVICE BFloat16& operator+=(BFloat16& a, const BFloat16& b) {
+  a = a + b;
+  return a;
+}
+
+inline C10_HOST_DEVICE BFloat16& operator-=(BFloat16& a, const BFloat16& b) {
+  a = a - b;
+  return a;
+}
+
+inline C10_HOST_DEVICE BFloat16& operator*=(BFloat16& a, const BFloat16& b) {
+  a = a * b;
+  return a;
+}
+
+inline C10_HOST_DEVICE BFloat16& operator/=(BFloat16& a, const BFloat16& b) {
+  a = a / b;
+  return a;
+}
+
+inline C10_HOST_DEVICE BFloat16& operator|(BFloat16& a, const BFloat16& b) {
+  a.x = a.x | b.x;
+  return a;
+}
+
+inline C10_HOST_DEVICE BFloat16& operator^(BFloat16& a, const BFloat16& b) {
+  a.x = a.x ^ b.x;
+  return a;
+}
+
+inline C10_HOST_DEVICE BFloat16& operator&(BFloat16& a, const BFloat16& b) {
+  a.x = a.x & b.x;
+  return a;
+}
+
+/// Arithmetic with floats
+
+inline C10_HOST_DEVICE float operator+(BFloat16 a, float b) {
+  return static_cast<float>(a) + b;
+}
+inline C10_HOST_DEVICE float operator-(BFloat16 a, float b) {
+  return static_cast<float>(a) - b;
+}
+inline C10_HOST_DEVICE float operator*(BFloat16 a, float b) {
+  return static_cast<float>(a) * b;
+}
+inline C10_HOST_DEVICE float operator/(BFloat16 a, float b) {
+  return static_cast<float>(a) / b;
+}
+
+inline C10_HOST_DEVICE float operator+(float a, BFloat16 b) {
+  return a + static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator-(float a, BFloat16 b) {
+  return a - static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator*(float a, BFloat16 b) {
+  return a * static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator/(float a, BFloat16 b) {
+  return a / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE float& operator+=(float& a, const BFloat16& b) {
+  return a += static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator-=(float& a, const BFloat16& b) {
+  return a -= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator*=(float& a, const BFloat16& b) {
+  return a *= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator/=(float& a, const BFloat16& b) {
+  return a /= static_cast<float>(b);
+}
+
+/// Arithmetic with doubles
+
+inline C10_HOST_DEVICE double operator+(BFloat16 a, double b) {
+  return static_cast<double>(a) + b;
+}
+inline C10_HOST_DEVICE double operator-(BFloat16 a, double b) {
+  return static_cast<double>(a) - b;
+}
+inline C10_HOST_DEVICE double operator*(BFloat16 a, double b) {
+  return static_cast<double>(a) * b;
+}
+inline C10_HOST_DEVICE double operator/(BFloat16 a, double b) {
+  return static_cast<double>(a) / b;
+}
+
+inline C10_HOST_DEVICE double operator+(double a, BFloat16 b) {
+  return a + static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator-(double a, BFloat16 b) {
+  return a - static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator*(double a, BFloat16 b) {
+  return a * static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator/(double a, BFloat16 b) {
+  return a / static_cast<double>(b);
+}
+
+/// Arithmetic with ints
+
+inline C10_HOST_DEVICE BFloat16 operator+(BFloat16 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<BFloat16>(b);
+}
+inline C10_HOST_DEVICE BFloat16 operator-(BFloat16 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<BFloat16>(b);
+}
+inline C10_HOST_DEVICE BFloat16 operator*(BFloat16 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<BFloat16>(b);
+}
+inline C10_HOST_DEVICE BFloat16 operator/(BFloat16 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<BFloat16>(b);
+}
+
+inline C10_HOST_DEVICE BFloat16 operator+(int a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) + b;
+}
+inline C10_HOST_DEVICE BFloat16 operator-(int a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) - b;
+}
+inline C10_HOST_DEVICE BFloat16 operator*(int a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) * b;
+}
+inline C10_HOST_DEVICE BFloat16 operator/(int a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) / b;
+}
+
+//// Arithmetic with int64_t
+
+inline C10_HOST_DEVICE BFloat16 operator+(BFloat16 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<BFloat16>(b);
+}
+inline C10_HOST_DEVICE BFloat16 operator-(BFloat16 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<BFloat16>(b);
+}
+inline C10_HOST_DEVICE BFloat16 operator*(BFloat16 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<BFloat16>(b);
+}
+inline C10_HOST_DEVICE BFloat16 operator/(BFloat16 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<BFloat16>(b);
+}
+
+inline C10_HOST_DEVICE BFloat16 operator+(int64_t a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) + b;
+}
+inline C10_HOST_DEVICE BFloat16 operator-(int64_t a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) - b;
+}
+inline C10_HOST_DEVICE BFloat16 operator*(int64_t a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) * b;
+}
+inline C10_HOST_DEVICE BFloat16 operator/(int64_t a, BFloat16 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<BFloat16>(a) / b;
+}
+
+// Overloading < and > operators, because std::max and std::min use them.
+
+inline C10_HOST_DEVICE bool operator>(BFloat16& lhs, BFloat16& rhs) {
+  return float(lhs) > float(rhs);
+}
+
+inline C10_HOST_DEVICE bool operator<(BFloat16& lhs, BFloat16& rhs) {
+  return float(lhs) < float(rhs);
+}
+
+C10_CLANG_DIAGNOSTIC_POP()
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+
+namespace detail {
+using c10::detail::bits_from_f32;
+using c10::detail::f32_from_bits;
+using c10::detail::round_to_nearest_even;
+} // namespace detail
+
+using c10::BFloat16;
+using c10::operator+;
+using c10::operator-;
+using c10::operator*;
+using c10::operator/;
+using c10::operator+=;
+using c10::operator-=;
+using c10::operator*=;
+using c10::operator/=;
+using c10::operator<;
+using c10::operator>;
+using c10::operator<<;
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace std {
+
+template <>
+class numeric_limits<c10::BFloat16> {
+ public:
+  static constexpr bool is_signed = true;
+  static constexpr bool is_specialized = true;
+  static constexpr bool is_integer = false;
+  static constexpr bool is_exact = false;
+  static constexpr bool has_infinity = true;
+  static constexpr bool has_quiet_NaN = true;
+  static constexpr bool has_signaling_NaN = true;
+  static constexpr auto has_denorm = numeric_limits<float>::has_denorm;
+  static constexpr auto has_denorm_loss =
+      numeric_limits<float>::has_denorm_loss;
+  static constexpr auto round_style = numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = false;
+  static constexpr bool is_bounded = true;
+  static constexpr bool is_modulo = false;
+  static constexpr int digits = 8;
+  static constexpr int digits10 = 2;
+  static constexpr int max_digits10 = 4;
+  static constexpr int radix = 2;
+  static constexpr int min_exponent = -125;
+  static constexpr int min_exponent10 = -37;
+  static constexpr int max_exponent = 128;
+  static constexpr int max_exponent10 = 38;
+  static constexpr auto traps = numeric_limits<float>::traps;
+  static constexpr auto tinyness_before =
+      numeric_limits<float>::tinyness_before;
+
+  static constexpr c10::BFloat16 min() {
+    return c10::BFloat16(0x0080, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 lowest() {
+    return c10::BFloat16(0xFF7F, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 max() {
+    return c10::BFloat16(0x7F7F, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 epsilon() {
+    return c10::BFloat16(0x3C00, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 round_error() {
+    return c10::BFloat16(0x3F00, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 infinity() {
+    return c10::BFloat16(0x7F80, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 quiet_NaN() {
+    return c10::BFloat16(0x7FC0, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 signaling_NaN() {
+    return c10::BFloat16(0x7F80, c10::BFloat16::from_bits());
+  }
+  static constexpr c10::BFloat16 denorm_min() {
+    return c10::BFloat16(0x0001, c10::BFloat16::from_bits());
+  }
+};
+
+} // namespace std
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Deprecated.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Deprecated.h
new file mode 100644
index 0000000000000000000000000000000000000000..88440a0242eb4e9e87433278006863fd38c5450d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Deprecated.h
@@ -0,0 +1,102 @@
+#pragma once
+
+/**
+ * This file provides portable macros for marking declarations
+ * as deprecated.  You should generally use C10_DEPRECATED,
+ * except when marking 'using' declarations as deprecated,
+ * in which case you should use C10_DEFINE_DEPRECATED_USING
+ * (due to portability concerns).
+ */
+
+// Sample usage:
+//
+//    C10_DEPRECATED void bad_func();
+//    struct C10_DEPRECATED BadStruct {
+//      ...
+//    };
+
+// NB: __cplusplus doesn't work for MSVC, so for now MSVC always uses
+// the "__declspec(deprecated)" implementation and not the C++14
+// "[[deprecated]]" attribute. We tried enabling "[[deprecated]]" for C++14 on
+// MSVC, but ran into issues with some older MSVC versions.
+#if (defined(__cplusplus) && __cplusplus >= 201402L)
+#define C10_DEPRECATED [[deprecated]]
+#define C10_DEPRECATED_MESSAGE(message) [[deprecated(message)]]
+#elif defined(__GNUC__)
+#define C10_DEPRECATED __attribute__((deprecated))
+// TODO Is there some way to implement this?
+#define C10_DEPRECATED_MESSAGE(message) __attribute__((deprecated))
+
+#elif defined(_MSC_VER)
+#define C10_DEPRECATED __declspec(deprecated)
+#define C10_DEPRECATED_MESSAGE(message) __declspec(deprecated(message))
+#else
+#warning "You need to implement C10_DEPRECATED for this compiler"
+#define C10_DEPRECATED
+#endif
+
+// Sample usage:
+//
+//    C10_DEFINE_DEPRECATED_USING(BadType, int)
+//
+//   which is the portable version of
+//
+//    using BadType [[deprecated]] = int;
+
+// technically [[deprecated]] syntax is from c++14 standard, but it works in
+// many compilers.
+#if defined(__has_cpp_attribute)
+#if __has_cpp_attribute(deprecated) && !defined(__CUDACC__)
+#define C10_DEFINE_DEPRECATED_USING(TypeName, TypeThingy) \
+  using TypeName [[deprecated]] = TypeThingy;
+#endif
+#endif
+
+#if defined(_MSC_VER)
+#if defined(__CUDACC__)
+// neither [[deprecated]] nor __declspec(deprecated) work on nvcc on Windows;
+// you get the error:
+//
+//    error: attribute does not apply to any entity
+//
+// So we just turn the macro off in this case.
+#if defined(C10_DEFINE_DEPRECATED_USING)
+#undef C10_DEFINE_DEPRECATED_USING
+#endif
+#define C10_DEFINE_DEPRECATED_USING(TypeName, TypeThingy) \
+  using TypeName = TypeThingy;
+#else
+// [[deprecated]] does work in windows without nvcc, though msc doesn't support
+// `__has_cpp_attribute` when c++14 is supported, otherwise
+// __declspec(deprecated) is used as the alternative.
+#ifndef C10_DEFINE_DEPRECATED_USING
+#if defined(_MSVC_LANG) && _MSVC_LANG >= 201402L
+#define C10_DEFINE_DEPRECATED_USING(TypeName, TypeThingy) \
+  using TypeName [[deprecated]] = TypeThingy;
+#else
+#define C10_DEFINE_DEPRECATED_USING(TypeName, TypeThingy) \
+  using TypeName = __declspec(deprecated) TypeThingy;
+#endif
+#endif
+#endif
+#endif
+
+#if !defined(C10_DEFINE_DEPRECATED_USING) && defined(__GNUC__)
+// nvcc has a bug where it doesn't understand __attribute__((deprecated))
+// declarations even when the host compiler supports it. We'll only use this gcc
+// attribute when not cuda, and when using a GCC compiler that doesn't support
+// the c++14 syntax we checked for above (available in __GNUC__ >= 5)
+#if !defined(__CUDACC__)
+#define C10_DEFINE_DEPRECATED_USING(TypeName, TypeThingy) \
+  using TypeName __attribute__((deprecated)) = TypeThingy;
+#else
+// using cuda + gcc < 5, neither deprecated syntax is available so turning off.
+#define C10_DEFINE_DEPRECATED_USING(TypeName, TypeThingy) \
+  using TypeName = TypeThingy;
+#endif
+#endif
+
+#if !defined(C10_DEFINE_DEPRECATED_USING)
+#warning "You need to implement C10_DEFINE_DEPRECATED_USING for this compiler"
+#define C10_DEFINE_DEPRECATED_USING
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Exception.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Exception.h
new file mode 100644
index 0000000000000000000000000000000000000000..0e067244a50e15db01d65ecfeb5908c29aecaeba
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Exception.h
@@ -0,0 +1,83 @@
+#pragma once
+
+#include <torch/headeronly/macros/Export.h>
+#include <torch/headeronly/macros/Macros.h>
+
+#include <sstream>
+#include <string>
+
+namespace c10 {
+// On nvcc, C10_UNLIKELY thwarts missing return statement analysis.  In cases
+// where the unlikely expression may be a constant, use this macro to ensure
+// return statement analysis keeps working (at the cost of not getting the
+// likely/unlikely annotation on nvcc).
+// https://github.com/pytorch/pytorch/issues/21418
+//
+// Currently, this is only used in the error reporting macros below.  If you
+// want to use it more generally, move me to Macros.h
+//
+// TODO: Brian Vaughan observed that we might be able to get this to work on
+// nvcc by writing some sort of C++ overload that distinguishes constexpr inputs
+// from non-constexpr.  Since there isn't any evidence that losing C10_UNLIKELY
+// in nvcc is causing us perf problems, this is not yet implemented, but this
+// might be an interesting piece of C++ code for an intrepid bootcamper to
+// write.
+#if defined(__CUDACC__)
+#define C10_UNLIKELY_OR_CONST(e) e
+#else
+#define C10_UNLIKELY_OR_CONST(e) C10_UNLIKELY(e)
+#endif
+
+} // namespace c10
+
+// STD_TORCH_CHECK throws std::runtime_error instead of c10::Error which is
+// useful when certain headers are used in a libtorch-independent way,
+// e.g. when Vectorized<T> is used in AOTInductor generated code, or
+// for custom ops to have an ABI stable dependency on libtorch.
+#ifdef STRIP_ERROR_MESSAGES
+#define STD_TORCH_CHECK_MSG(cond, type, ...) \
+  (#cond #type " CHECK FAILED at " C10_STRINGIZE(__FILE__))
+#else // so STRIP_ERROR_MESSAGES is not defined
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly, detail)
+template <typename... Args>
+std::string stdTorchCheckMsgImpl(const char* /*msg*/, const Args&... args) {
+  // This is similar to the one in c10/util/Exception.h, but does
+  // not depend on the more complex c10::str() function. ostringstream
+  // supports fewer data types than c10::str(), but should be sufficient
+  // in the headeronly world.
+  std::ostringstream oss;
+  ((oss << args), ...);
+  return oss.str();
+}
+
+inline const char* stdTorchCheckMsgImpl(const char* msg) {
+  return msg;
+}
+// If there is just 1 user-provided C-string argument, use it.
+inline const char* stdTorchCheckMsgImpl(const char* /*msg*/, const char* args) {
+  return args;
+}
+HIDDEN_NAMESPACE_END(torch, headeronly, detail)
+
+#define STD_TORCH_CHECK_MSG(cond, type, ...)               \
+  (torch::headeronly::detail::stdTorchCheckMsgImpl(        \
+      "Expected " #cond                                    \
+      " to be true, but got false.  "                      \
+      "(Could this error message be improved?  If so, "    \
+      "please report an enhancement request to PyTorch.)", \
+      ##__VA_ARGS__))
+#endif // STRIP_ERROR_MESSAGES
+
+#define STD_TORCH_CHECK(cond, ...)                \
+  if (C10_UNLIKELY_OR_CONST(!(cond))) {           \
+    throw std::runtime_error(STD_TORCH_CHECK_MSG( \
+        cond,                                     \
+        "",                                       \
+        __func__,                                 \
+        ", ",                                     \
+        __FILE__,                                 \
+        ":",                                      \
+        __LINE__,                                 \
+        ", ",                                     \
+        ##__VA_ARGS__));                          \
+  }
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float4_e2m1fn_x2.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float4_e2m1fn_x2.h
new file mode 100644
index 0000000000000000000000000000000000000000..00075914cdc346af4d212f089731c5fea6943873
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float4_e2m1fn_x2.h
@@ -0,0 +1,47 @@
+#pragma once
+#include <cstdint>
+
+#include <torch/headeronly/macros/Macros.h>
+
+/// Defines the Float4_e2m1fn_x2 type (4-bit floating-point, two elements packed
+/// into one byte). This is the FP4 dtype from the OCP MX format spec
+/// (https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf,
+/// Section 5.3.3)
+///
+/// Given two high precision values val0 and val1, here is the
+/// binary configuration of their packed representation, from MSB to LSB:
+///
+///   original value             | val1 : val0
+///   ========================================
+///   bit index (MSB==7, LSB==0) | 7654 : 3210
+///   sign/exponent/mantissa     | seem : seem
+///
+
+namespace c10 {
+
+struct alignas(1) Float4_e2m1fn_x2 {
+  uint8_t val_;
+  Float4_e2m1fn_x2() = default;
+  C10_HOST_DEVICE explicit Float4_e2m1fn_x2(uint8_t val) : val_(val) {}
+};
+
+/// Comparison operators
+inline C10_HOST_DEVICE bool operator==(
+    const Float4_e2m1fn_x2& a,
+    const Float4_e2m1fn_x2& b) {
+  return a.val_ == b.val_;
+}
+
+inline C10_HOST_DEVICE bool operator!=(
+    const Float4_e2m1fn_x2& a,
+    const Float4_e2m1fn_x2& b) {
+  return a.val_ != b.val_;
+}
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::Float4_e2m1fn_x2;
+using c10::operator==;
+using c10::operator!=;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e4m3fn.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e4m3fn.h
new file mode 100644
index 0000000000000000000000000000000000000000..06d70e6b4982c908fd13b54a3530d11381453c72
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e4m3fn.h
@@ -0,0 +1,537 @@
+#pragma once
+
+/// Defines the Float8_e4m3fn type (8-bit floating-point) including conversions
+/// to standard C types and basic arithmetic operations. Note that arithmetic
+/// operations are implemented by converting to floating point and
+/// performing the operation in float32.
+/// Binary configuration:
+/// s eeee mmm
+/// 1 sign bit
+/// 4 exponent bits
+/// 3 mantissa bits
+/// bias = 7
+///
+/// Implementation based on the paper https://arxiv.org/pdf/2209.05433.pdf
+/// and inspired by Half implementation from pytorch/c10/util/Half.h
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/floating_point_utils.h>
+
+#if defined(__cplusplus)
+#include <cmath>
+#include <cstdint>
+#elif !defined(__OPENCL_VERSION__)
+#include <math.h>
+#include <stdint.h>
+#endif
+
+#ifdef _MSC_VER
+#include <intrin.h>
+#endif
+
+#include <climits>
+#include <iostream>
+
+namespace c10 {
+
+struct alignas(1) Float8_e4m3fn {
+  uint8_t x;
+
+  struct from_bits_t {};
+  C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
+    return from_bits_t();
+  }
+
+  Float8_e4m3fn() = default;
+
+  constexpr C10_HOST_DEVICE Float8_e4m3fn(uint8_t bits, from_bits_t /*unused*/)
+      : x(bits) {}
+  inline C10_HOST_DEVICE Float8_e4m3fn(float value);
+  inline C10_HOST_DEVICE operator float() const;
+  inline C10_HOST_DEVICE bool isnan() const;
+  inline C10_HOST_DEVICE bool isinf() const;
+};
+
+inline std::ostream& operator<<(std::ostream& out, const Float8_e4m3fn& value) {
+  out << (float)value;
+  return out;
+}
+
+namespace detail {
+
+/*
+ * Convert a 8-bit floating-point number in fp8 E4M3FN format, in bit
+ * representation, to a 32-bit floating-point number in IEEE single-precision
+ * format, in bit representation.
+ *
+ * @note The implementation doesn't use any floating-point operations.
+ */
+inline C10_HOST_DEVICE float fp8e4m3fn_to_fp32_value(uint8_t input) {
+  /*
+   * Extend the fp8 E4M3FN number to 32 bits and shift to the
+   * upper part of the 32-bit word:
+   *      +---+----+---+-----------------------------+
+   *      | S |EEEE|MMM|0000 0000 0000 0000 0000 0000|
+   *      +---+----+---+-----------------------------+
+   * Bits  31 27-30 24-26          0-23
+   *
+   * S - sign bit, E - bits of the biased exponent, M - bits of the mantissa, 0
+   * - zero bits.
+   */
+  const uint32_t w = (uint32_t)input << 24;
+  /*
+   * Extract the sign of the input number into the high bit of the 32-bit word:
+   *
+   *      +---+----------------------------------+
+   *      | S |0000000 00000000 00000000 00000000|
+   *      +---+----------------------------------+
+   * Bits  31                 0-31
+   */
+  const uint32_t sign = w & UINT32_C(0x80000000);
+  /*
+   * Extract mantissa and biased exponent of the input number into the bits 0-30
+   * of the 32-bit word:
+   *
+   *      +---+----+---+-----------------------------+
+   *      | S |EEEE|MMM|0000 0000 0000 0000 0000 0000|
+   *      +---+----+---+-----------------------------+
+   * Bits  31  27-30 24-26      0-23
+   */
+  const uint32_t nonsign = w & UINT32_C(0x7FFFFFFF);
+  /*
+   * Renorm shift is the number of bits to shift mantissa left to make the
+   * half-precision number normalized. If the initial number is normalized, some
+   * of its high 5 bits (sign == 0 and 4-bit exponent) equals one. In this case
+   * renorm_shift == 0. If the number is denormalize, renorm_shift > 0. Note
+   * that if we shift denormalized nonsign by renorm_shift, the unit bit of
+   * mantissa will shift into exponent, turning the biased exponent into 1, and
+   * making mantissa normalized (i.e. without leading 1).
+   */
+#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
+  uint32_t renorm_shift = __clz(nonsign);
+#elif defined(__SYCL_DEVICE_ONLY__)
+  // Note: zero is not a supported input into `__builtin_clz`
+  uint32_t renorm_shift =
+      nonsign != 0 ? __builtin_clz(nonsign) : sizeof(uint32_t) * CHAR_BIT;
+#elif defined(_MSC_VER) && !defined(__clang__)
+  unsigned long nonsign_bsr;
+  _BitScanReverse(&nonsign_bsr, (unsigned long)nonsign);
+  uint32_t renorm_shift = (uint32_t)nonsign_bsr ^ 31;
+#else
+  // Note: zero is not a supported input into `__builtin_clz`
+  uint32_t renorm_shift =
+      nonsign != 0 ? __builtin_clz(nonsign) : sizeof(uint32_t) * CHAR_BIT;
+#endif
+  renorm_shift = renorm_shift > 4 ? renorm_shift - 4 : 0;
+  /*
+   * Iff fp8e4m3fn number has all exponent and mantissa bits set to 1,
+   * the addition overflows it into bit 31, and the subsequent shift turns the
+   * high 9 bits into 1. Thus inf_nan_mask == 0x7F800000 if the fp8e4m3fn number
+   * is Nan, 0x00000000 otherwise
+   */
+  const int32_t inf_nan_mask =
+      ((int32_t)(nonsign + 0x01000000) >> 8) & INT32_C(0x7F800000);
+  /*
+   * Iff nonsign is 0, it overflows into 0xFFFFFFFF, turning bit 31
+   * into 1. Otherwise, bit 31 remains 0. The signed shift right by 31
+   * broadcasts bit 31 into all bits of the zero_mask. Thus zero_mask ==
+   * 0xFFFFFFFF if the half-precision number was zero (+0.0h or -0.0h)
+   * 0x00000000 otherwise
+   */
+  const int32_t zero_mask = (int32_t)(nonsign - 1) >> 31;
+  /*
+   * 1. Shift nonsign left by renorm_shift to normalize it (if the input
+   * was denormal)
+   * 2. Shift nonsign right by 4 so the exponent (4 bits originally)
+   * becomes an 8-bit field and 3-bit mantissa shifts into the 3 high
+   * bits of the 23-bit mantissa of IEEE single-precision number.
+   * 3. Add 0x78 to the exponent (starting at bit 23) to compensate the
+   * different in exponent bias (0x7F for single-precision number less 0x07
+   * for fp8e4m3fn number).
+   * 4. Subtract renorm_shift from the exponent (starting at bit 23) to
+   * account for renormalization. As renorm_shift is less than 0x78, this
+   * can be combined with step 3.
+   * 5. Binary OR with inf_nan_mask to turn the exponent into 0xFF if the
+   * input was NaN or infinity.
+   * 6. Binary ANDNOT with zero_mask to turn the mantissa and exponent
+   * into zero if the input was zero.
+   * 7. Combine with the sign of the input number.
+   */
+  uint32_t result = sign |
+      ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) |
+        inf_nan_mask) &
+       ~zero_mask);
+  return fp32_from_bits(result);
+}
+
+/*
+ * Convert a 32-bit floating-point number in IEEE single-precision format to a
+ * 8-bit floating-point number in fp8 E4M3FN format, in bit representation.
+ */
+inline C10_HOST_DEVICE uint8_t fp8e4m3fn_from_fp32_value(float f) {
+  /*
+   * Binary representation of 480.0f, which is the first value
+   * not representable in fp8e4m3fn range:
+   * 0 1111 111 - fp8e4m3fn
+   * 0 10000111 11100000000000000000000 - fp32
+   */
+  constexpr uint32_t fp8_max = UINT32_C(1087) << 20;
+
+  /*
+   * A mask for converting fp32 numbers lower than fp8e4m3fn normal range
+   * into denorm representation
+   * magic number: ((127 - 7) + (23 - 3) + 1)
+   */
+  constexpr uint32_t denorm_mask = UINT32_C(141) << 23;
+
+  uint32_t f_bits = fp32_to_bits(f);
+
+  uint8_t result = 0u;
+
+  /*
+   * Extract the sign of the input number into the high bit of the 32-bit word:
+   *
+   *      +---+----------------------------------+
+   *      | S |0000000 00000000 00000000 00000000|
+   *      +---+----------------------------------+
+   * Bits  31                 0-31
+   */
+  const uint32_t sign = f_bits & UINT32_C(0x80000000);
+
+  /*
+   * Set sign bit to 0
+   */
+  f_bits ^= sign;
+
+  if (f_bits >= fp8_max) {
+    // NaN - all exponent and mantissa bits set to 1
+    result = 0x7f;
+  } else {
+    if (f_bits < (UINT32_C(121) << 23)) {
+      // Input number is smaller than 2^(-6), which is the smallest
+      // fp8e4m3fn normal number
+      f_bits =
+          fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
+      result = static_cast<uint8_t>(f_bits - denorm_mask);
+    } else {
+      // resulting mantissa is odd
+      uint8_t mant_odd = (f_bits >> 20) & 1;
+
+      // update exponent, rounding bias part 1
+      f_bits += ((uint32_t)(7 - 127) << 23) + 0x7FFFF;
+
+      // rounding bias part 2
+      f_bits += mant_odd;
+
+      // take the bits!
+      result = static_cast<uint8_t>(f_bits >> 20);
+    }
+  }
+
+  result |= static_cast<uint8_t>(sign >> 24);
+  return result;
+}
+
+} // namespace detail
+
+// -------- below is copied from c10/util/Float8_e4m3fn-inl.h --------//
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+/// Constructors
+
+inline C10_HOST_DEVICE Float8_e4m3fn::Float8_e4m3fn(float value)
+    : x(detail::fp8e4m3fn_from_fp32_value(value)) {}
+
+/// Implicit conversions
+
+inline C10_HOST_DEVICE Float8_e4m3fn::operator float() const {
+  return detail::fp8e4m3fn_to_fp32_value(x);
+}
+
+/// Special values helper
+
+inline C10_HOST_DEVICE bool Float8_e4m3fn::isnan() const {
+  return (x & 0b01111111) == 0b01111111;
+}
+
+inline C10_HOST_DEVICE bool Float8_e4m3fn::isinf() const {
+  // Note: fp8e4m3fn does not have infinity, so this always returns false.
+  return false;
+}
+
+/// Arithmetic
+
+inline C10_HOST_DEVICE Float8_e4m3fn
+operator+(const Float8_e4m3fn& a, const Float8_e4m3fn& b) {
+  return static_cast<float>(a) + static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn
+operator-(const Float8_e4m3fn& a, const Float8_e4m3fn& b) {
+  return static_cast<float>(a) - static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn
+operator*(const Float8_e4m3fn& a, const Float8_e4m3fn& b) {
+  return static_cast<float>(a) * static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn operator/(
+    const Float8_e4m3fn& a,
+    const Float8_e4m3fn& b) __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn operator-(const Float8_e4m3fn& a) {
+  return -static_cast<float>(a);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn& operator+=(
+    Float8_e4m3fn& a,
+    const Float8_e4m3fn& b) {
+  a = a + b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn& operator-=(
+    Float8_e4m3fn& a,
+    const Float8_e4m3fn& b) {
+  a = a - b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn& operator*=(
+    Float8_e4m3fn& a,
+    const Float8_e4m3fn& b) {
+  a = a * b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn& operator/=(
+    Float8_e4m3fn& a,
+    const Float8_e4m3fn& b) {
+  a = a / b;
+  return a;
+}
+
+/// Arithmetic with floats
+
+inline C10_HOST_DEVICE float operator+(Float8_e4m3fn a, float b) {
+  return static_cast<float>(a) + b;
+}
+inline C10_HOST_DEVICE float operator-(Float8_e4m3fn a, float b) {
+  return static_cast<float>(a) - b;
+}
+inline C10_HOST_DEVICE float operator*(Float8_e4m3fn a, float b) {
+  return static_cast<float>(a) * b;
+}
+inline C10_HOST_DEVICE float operator/(Float8_e4m3fn a, float b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / b;
+}
+
+inline C10_HOST_DEVICE float operator+(float a, Float8_e4m3fn b) {
+  return a + static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator-(float a, Float8_e4m3fn b) {
+  return a - static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator*(float a, Float8_e4m3fn b) {
+  return a * static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator/(float a, Float8_e4m3fn b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE float& operator+=(float& a, const Float8_e4m3fn& b) {
+  return a += static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator-=(float& a, const Float8_e4m3fn& b) {
+  return a -= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator*=(float& a, const Float8_e4m3fn& b) {
+  return a *= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator/=(float& a, const Float8_e4m3fn& b) {
+  return a /= static_cast<float>(b);
+}
+
+/// Arithmetic with doubles
+
+inline C10_HOST_DEVICE double operator+(Float8_e4m3fn a, double b) {
+  return static_cast<double>(a) + b;
+}
+inline C10_HOST_DEVICE double operator-(Float8_e4m3fn a, double b) {
+  return static_cast<double>(a) - b;
+}
+inline C10_HOST_DEVICE double operator*(Float8_e4m3fn a, double b) {
+  return static_cast<double>(a) * b;
+}
+inline C10_HOST_DEVICE double operator/(Float8_e4m3fn a, double b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<double>(a) / b;
+}
+
+inline C10_HOST_DEVICE double operator+(double a, Float8_e4m3fn b) {
+  return a + static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator-(double a, Float8_e4m3fn b) {
+  return a - static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator*(double a, Float8_e4m3fn b) {
+  return a * static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator/(double a, Float8_e4m3fn b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<double>(b);
+}
+
+/// Arithmetic with ints
+
+inline C10_HOST_DEVICE Float8_e4m3fn operator+(Float8_e4m3fn a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e4m3fn>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator-(Float8_e4m3fn a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e4m3fn>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator*(Float8_e4m3fn a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e4m3fn>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator/(Float8_e4m3fn a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e4m3fn>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn operator+(int a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator-(int a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator*(int a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator/(int a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) / b;
+}
+
+//// Arithmetic with int64_t
+
+inline C10_HOST_DEVICE Float8_e4m3fn operator+(Float8_e4m3fn a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e4m3fn>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator-(Float8_e4m3fn a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e4m3fn>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator*(Float8_e4m3fn a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e4m3fn>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator/(Float8_e4m3fn a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e4m3fn>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fn operator+(int64_t a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator-(int64_t a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator*(int64_t a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fn operator/(int64_t a, Float8_e4m3fn b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fn>(a) / b;
+}
+
+/// NOTE: we do not define comparisons directly and instead rely on the implicit
+/// conversion from c10::Float8_e4m3fn to float.
+
+C10_CLANG_DIAGNOSTIC_POP()
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::Float8_e4m3fn;
+using c10::operator<<;
+using c10::operator+;
+using c10::operator-;
+using c10::operator*;
+using c10::operator/;
+using c10::operator+=;
+using c10::operator-=;
+using c10::operator*=;
+using c10::operator/=;
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace std {
+
+template <>
+class numeric_limits<c10::Float8_e4m3fn> {
+ public:
+  static constexpr bool is_specialized = true;
+  static constexpr bool is_signed = true;
+  static constexpr bool is_integer = false;
+  static constexpr bool is_exact = false;
+  static constexpr bool has_infinity = false;
+  static constexpr bool has_quiet_NaN = true;
+  static constexpr bool has_signaling_NaN = false;
+  static constexpr auto has_denorm = true;
+  static constexpr auto has_denorm_loss = true;
+  static constexpr auto round_style = numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = false;
+  static constexpr bool is_bounded = true;
+  static constexpr bool is_modulo = false;
+  static constexpr int digits = 4;
+  static constexpr int digits10 = 0;
+  static constexpr int max_digits10 = 3;
+  static constexpr int radix = 2;
+  static constexpr int min_exponent = -5;
+  static constexpr int min_exponent10 = -1;
+  static constexpr int max_exponent = 8;
+  static constexpr int max_exponent10 = 2;
+  static constexpr auto traps = numeric_limits<float>::traps;
+  static constexpr auto tinyness_before = false;
+
+  static constexpr c10::Float8_e4m3fn min() {
+    return c10::Float8_e4m3fn(0x08, c10::Float8_e4m3fn::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fn lowest() {
+    return c10::Float8_e4m3fn(0xFE, c10::Float8_e4m3fn::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fn max() {
+    return c10::Float8_e4m3fn(0x7E, c10::Float8_e4m3fn::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fn epsilon() {
+    return c10::Float8_e4m3fn(0x20, c10::Float8_e4m3fn::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fn round_error() {
+    return c10::Float8_e4m3fn(0x30, c10::Float8_e4m3fn::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fn quiet_NaN() {
+    return c10::Float8_e4m3fn(0x7F, c10::Float8_e4m3fn::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fn denorm_min() {
+    return c10::Float8_e4m3fn(0x01, c10::Float8_e4m3fn::from_bits());
+  }
+};
+
+} // namespace std
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e4m3fnuz.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e4m3fnuz.h
new file mode 100644
index 0000000000000000000000000000000000000000..4c11b3be05593fe3f9a2d8787750f546bd875bff
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e4m3fnuz.h
@@ -0,0 +1,450 @@
+#pragma once
+
+/// Defines the Float8_e4m3fnuz type (8-bit floating-point) including
+/// conversions to standard C types and basic arithmetic operations. Note that
+/// arithmetic operations are implemented by converting to floating point and
+/// performing the operation in float32.
+/// Binary configuration remains the same as Float8_e4m3fn:
+/// s eeee mmm
+/// 1 sign bit
+/// 4 exponent bits
+/// 3 mantissa bits
+/// The key differences versus Float8_e4m3fn are:
+/// bias = 8
+/// no infinities or negative zero
+/// NaN only when sign bit is 1, rest all 0s
+///
+/// Implementation based on the paper https://arxiv.org/pdf/2206.02915.pdf and
+/// the existing Float8_e4m3fn implementation.
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Float8_fnuz_cvt.h>
+#include <torch/headeronly/util/floating_point_utils.h>
+
+#include <limits>
+
+#if defined(__cplusplus)
+#include <cstdint>
+#elif !defined(__OPENCL_VERSION__)
+#include <math.h>
+#include <stdint.h>
+#endif
+
+#include <iosfwd>
+#include <ostream>
+
+namespace c10 {
+
+struct alignas(1) Float8_e4m3fnuz {
+  uint8_t x;
+
+  struct from_bits_t {};
+  C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
+    return from_bits_t();
+  }
+
+  Float8_e4m3fnuz() = default;
+
+  constexpr C10_HOST_DEVICE Float8_e4m3fnuz(
+      uint8_t bits,
+      from_bits_t /*unused*/)
+      : x(bits) {}
+  inline C10_HOST_DEVICE Float8_e4m3fnuz(float value);
+  inline C10_HOST_DEVICE operator float() const;
+  inline C10_HOST_DEVICE bool isnan() const;
+  inline C10_HOST_DEVICE bool isinf() const;
+};
+
+inline std::ostream& operator<<(
+    std::ostream& out,
+    const Float8_e4m3fnuz& value) {
+  out << (float)value;
+  return out;
+}
+
+namespace detail {
+
+/*
+ * Convert a 32-bit floating-point number in IEEE single-precision format to a
+ * 8-bit floating-point number in fp8 E4M3FNUZ format, in bit representation.
+ */
+inline C10_HOST_DEVICE uint8_t fp8e4m3fnuz_from_fp32_value(float f) {
+  /*
+   * Binary representation of 256.0f, which is the first value not representable
+   * (i.e. the first value which would overflow in to the sign bit, resulting in
+   * a NaN) in fp8e4m3fnuz range:
+   * 1 0000 000 - fp8e4m3fnuz
+   * 0 10000111 00000000000000000000000 - fp32
+   */
+  constexpr uint32_t fnuz_max = UINT32_C(0x87) << 23;
+
+  /*
+   * A mask for converting fp32 numbers lower than fp8e4m3fnuz normal range
+   * into denorm representation
+   * magic number: ((127 - 8) + (23 - 3) + 1)
+   */
+  constexpr uint32_t denorm_mask = UINT32_C(0x8C) << 23;
+
+  uint32_t f_bits = fp32_to_bits(f);
+
+  uint32_t result = 0u;
+
+  /*
+   * Extract the sign of the input number into the high bit of the 32-bit word:
+   *
+   *      +---+----------------------------------+
+   *      | S |0000000 00000000 00000000 00000000|
+   *      +---+----------------------------------+
+   * Bits  31                 0-31
+   */
+  const uint32_t sign = f_bits & UINT32_C(0x80000000);
+
+  /*
+   * Set sign bit to 0
+   */
+  f_bits ^= sign;
+
+  if (f_bits >= fnuz_max) {
+    // NaN -- sign bit set to 1, rest 0s.
+    return 0x80;
+  }
+
+  if (f_bits < (UINT32_C(0x78) << 23) /* 2^-7 in float32 */) {
+    // Input exponent is less than -7, the smallest e4m3fnuz exponent, so the
+    // number will become subnormal.
+    f_bits = fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
+    result = static_cast<uint8_t>(f_bits - denorm_mask);
+    if (result == 0) {
+      // fnuz types don't have negative zero.
+      return 0;
+    }
+  } else {
+    // resulting mantissa is odd
+    uint8_t mant_odd = (f_bits >> 20) & 1;
+
+    // update exponent, rounding bias part 1
+    f_bits += ((uint32_t)(8 - 127) << 23) + 0x7FFFF;
+
+    // rounding bias part 2
+    f_bits += mant_odd;
+
+    // take the bits!
+    result = static_cast<uint8_t>(f_bits >> 20);
+  }
+
+  result |= sign >> 24;
+  return result;
+}
+
+} // namespace detail
+
+//------ below is copied from c10/util/Float8_e4m3fnuz-inl.h ------//
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+/// Constructors
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz::Float8_e4m3fnuz(float value)
+    : x(detail::fp8e4m3fnuz_from_fp32_value(value)) {}
+
+/// Implicit conversions
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz::operator float() const {
+  return torch::headeronly::detail::fp8_fnuz_to_fp32_value<4, 3>(x);
+}
+
+/// Special values helper
+
+inline C10_HOST_DEVICE bool Float8_e4m3fnuz::isnan() const {
+  return x == 0b10000000;
+}
+
+inline C10_HOST_DEVICE bool Float8_e4m3fnuz::isinf() const {
+  // Note: fp8e4m3fnuz does not have infinity, so this always returns false.
+  return false;
+}
+
+/// Arithmetic
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz
+operator+(const Float8_e4m3fnuz& a, const Float8_e4m3fnuz& b) {
+  return static_cast<float>(a) + static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz
+operator-(const Float8_e4m3fnuz& a, const Float8_e4m3fnuz& b) {
+  return static_cast<float>(a) - static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz
+operator*(const Float8_e4m3fnuz& a, const Float8_e4m3fnuz& b) {
+  return static_cast<float>(a) * static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator/(
+    const Float8_e4m3fnuz& a,
+    const Float8_e4m3fnuz& b) __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator-(const Float8_e4m3fnuz& a) {
+  return -static_cast<float>(a);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz& operator+=(
+    Float8_e4m3fnuz& a,
+    const Float8_e4m3fnuz& b) {
+  a = a + b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz& operator-=(
+    Float8_e4m3fnuz& a,
+    const Float8_e4m3fnuz& b) {
+  a = a - b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz& operator*=(
+    Float8_e4m3fnuz& a,
+    const Float8_e4m3fnuz& b) {
+  a = a * b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz& operator/=(
+    Float8_e4m3fnuz& a,
+    const Float8_e4m3fnuz& b) {
+  a = a / b;
+  return a;
+}
+
+/// Arithmetic with floats
+
+inline C10_HOST_DEVICE float operator+(Float8_e4m3fnuz a, float b) {
+  return static_cast<float>(a) + b;
+}
+inline C10_HOST_DEVICE float operator-(Float8_e4m3fnuz a, float b) {
+  return static_cast<float>(a) - b;
+}
+inline C10_HOST_DEVICE float operator*(Float8_e4m3fnuz a, float b) {
+  return static_cast<float>(a) * b;
+}
+inline C10_HOST_DEVICE float operator/(Float8_e4m3fnuz a, float b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / b;
+}
+
+inline C10_HOST_DEVICE float operator+(float a, Float8_e4m3fnuz b) {
+  return a + static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator-(float a, Float8_e4m3fnuz b) {
+  return a - static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator*(float a, Float8_e4m3fnuz b) {
+  return a * static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator/(float a, Float8_e4m3fnuz b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE float& operator+=(float& a, const Float8_e4m3fnuz& b) {
+  return a += static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator-=(float& a, const Float8_e4m3fnuz& b) {
+  return a -= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator*=(float& a, const Float8_e4m3fnuz& b) {
+  return a *= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator/=(float& a, const Float8_e4m3fnuz& b) {
+  return a /= static_cast<float>(b);
+}
+
+/// Arithmetic with doubles
+
+inline C10_HOST_DEVICE double operator+(Float8_e4m3fnuz a, double b) {
+  return static_cast<double>(a) + b;
+}
+inline C10_HOST_DEVICE double operator-(Float8_e4m3fnuz a, double b) {
+  return static_cast<double>(a) - b;
+}
+inline C10_HOST_DEVICE double operator*(Float8_e4m3fnuz a, double b) {
+  return static_cast<double>(a) * b;
+}
+inline C10_HOST_DEVICE double operator/(Float8_e4m3fnuz a, double b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<double>(a) / b;
+}
+
+inline C10_HOST_DEVICE double operator+(double a, Float8_e4m3fnuz b) {
+  return a + static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator-(double a, Float8_e4m3fnuz b) {
+  return a - static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator*(double a, Float8_e4m3fnuz b) {
+  return a * static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator/(double a, Float8_e4m3fnuz b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<double>(b);
+}
+
+/// Arithmetic with ints
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator+(Float8_e4m3fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e4m3fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator-(Float8_e4m3fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e4m3fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator*(Float8_e4m3fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e4m3fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator/(Float8_e4m3fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e4m3fnuz>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator+(int a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator-(int a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator*(int a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator/(int a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) / b;
+}
+
+//// Arithmetic with int64_t
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator+(Float8_e4m3fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e4m3fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator-(Float8_e4m3fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e4m3fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator*(Float8_e4m3fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e4m3fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator/(Float8_e4m3fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e4m3fnuz>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator+(int64_t a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator-(int64_t a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator*(int64_t a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e4m3fnuz operator/(int64_t a, Float8_e4m3fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e4m3fnuz>(a) / b;
+}
+
+/// NOTE: we do not define comparisons directly and instead rely on the implicit
+/// conversion from c10::Float8_e4m3fnuz to float.
+
+C10_CLANG_DIAGNOSTIC_POP()
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::Float8_e4m3fnuz;
+using c10::operator+;
+using c10::operator-;
+using c10::operator*;
+using c10::operator/;
+using c10::operator+=;
+using c10::operator-=;
+using c10::operator*=;
+using c10::operator/=;
+using c10::operator<<;
+
+namespace detail {
+using c10::detail::fp8e4m3fnuz_from_fp32_value;
+} // namespace detail
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace std {
+
+template <>
+class numeric_limits<c10::Float8_e4m3fnuz> {
+ public:
+  static constexpr bool is_specialized = true;
+  static constexpr bool is_signed = true;
+  static constexpr bool is_integer = false;
+  static constexpr bool is_exact = false;
+  static constexpr bool has_infinity = false;
+  static constexpr bool has_quiet_NaN = true;
+  static constexpr bool has_signaling_NaN = false;
+  static constexpr auto has_denorm = true;
+  static constexpr auto has_denorm_loss = true;
+  static constexpr auto round_style = numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = false;
+  static constexpr bool is_bounded = true;
+  static constexpr bool is_modulo = false;
+  static constexpr int digits = 4;
+  static constexpr int digits10 = 0;
+  static constexpr int max_digits10 = 3;
+  static constexpr int radix = 2;
+  static constexpr int min_exponent = -6;
+  static constexpr int min_exponent10 = -1;
+  static constexpr int max_exponent = 8;
+  static constexpr int max_exponent10 = 2;
+  static constexpr auto traps = numeric_limits<float>::traps;
+  static constexpr auto tinyness_before = false;
+
+  static constexpr c10::Float8_e4m3fnuz min() {
+    return c10::Float8_e4m3fnuz(0x08, c10::Float8_e4m3fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fnuz lowest() {
+    return c10::Float8_e4m3fnuz(0xFF, c10::Float8_e4m3fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fnuz max() {
+    return c10::Float8_e4m3fnuz(0x7F, c10::Float8_e4m3fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fnuz epsilon() {
+    return c10::Float8_e4m3fnuz(0x28, c10::Float8_e4m3fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fnuz round_error() {
+    return c10::Float8_e4m3fnuz(0x38, c10::Float8_e4m3fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fnuz infinity() {
+    // NaN (no infinities)
+    return c10::Float8_e4m3fnuz(0x80, c10::Float8_e4m3fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fnuz quiet_NaN() {
+    return c10::Float8_e4m3fnuz(0x80, c10::Float8_e4m3fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e4m3fnuz denorm_min() {
+    return c10::Float8_e4m3fnuz(0x01, c10::Float8_e4m3fnuz::from_bits());
+  }
+};
+
+} // namespace std
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e5m2.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e5m2.h
new file mode 100644
index 0000000000000000000000000000000000000000..0aa856d0d5546fc53b7bf7148fefef72c563b262
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e5m2.h
@@ -0,0 +1,458 @@
+#pragma once
+
+/// Defines the Float8_e5m2 type (8-bit floating-point) including conversions
+/// to standard C types and basic arithmetic operations. Note that arithmetic
+/// operations are implemented by converting to floating point and
+/// performing the operation in float32.
+/// Binary configuration:
+/// s eeeee mm
+/// 1 sign bit
+/// 5 exponent bits
+/// 2 mantissa bits
+/// bias = 15
+///
+/// Implementation based on the paper https://arxiv.org/pdf/2209.05433.pdf
+/// and inspired by Half implementation from pytorch/c10/util/Half.h
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Half.h>
+
+#include <limits>
+
+namespace c10 {
+
+struct alignas(1) Float8_e5m2 {
+  uint8_t x;
+
+  struct from_bits_t {};
+  C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
+    return from_bits_t();
+  }
+
+  Float8_e5m2() = default;
+
+  constexpr C10_HOST_DEVICE Float8_e5m2(uint8_t bits, from_bits_t /*unused*/)
+      : x(bits) {}
+  inline C10_HOST_DEVICE Float8_e5m2(float value);
+  inline C10_HOST_DEVICE operator float() const;
+  inline C10_HOST_DEVICE bool isnan() const;
+  inline C10_HOST_DEVICE bool isinf() const;
+};
+
+inline std::ostream& operator<<(std::ostream& out, const Float8_e5m2& value) {
+  out << (float)value;
+  return out;
+}
+
+namespace detail {
+
+/*
+ * Convert a 8-bit floating-point number in fp8 E5M2 format, in bit
+ * representation, to a 32-bit floating-point number in IEEE single-precision
+ * format, in bit representation.
+ *
+ * @note The implementation doesn't use any floating-point operations.
+ */
+inline C10_HOST_DEVICE float fp8e5m2_to_fp32_value(uint8_t input) {
+  /*
+   * Extend the fp8 E5M2 number to 32 bits and shift to the
+   * upper part of the 32-bit word:
+   *      +---+----+---+-----------------------------+
+   *      | S |EEEEE|MM|0000 0000 0000 0000 0000 0000|
+   *      +---+----+---+-----------------------------+
+   * Bits  31 26-30 24-25          0-23
+   *
+   * S - sign bit, E - bits of the biased exponent, M - bits of the mantissa, 0
+   * - zero bits.
+   */
+  uint16_t half_representation = input;
+  half_representation <<= 8;
+  return fp16_ieee_to_fp32_value(half_representation);
+}
+
+/*
+ * Convert a 32-bit floating-point number in IEEE single-precision format to a
+ * 8-bit floating-point number in fp8 E5M2 format, in bit representation.
+ */
+inline C10_HOST_DEVICE uint8_t fp8e5m2_from_fp32_value(float f) {
+  /*
+   * Binary representation of fp32 infinity
+   * 0 11111111 00000000000000000000000
+   */
+  constexpr uint32_t fp32_inf = UINT32_C(255) << 23;
+
+  /*
+   * Binary representation of 65536.0f, which is the first value
+   * not representable in fp8e5m2 range:
+   * 0 11111 00 - fp8e5m2
+   * 0 10001111 00000000000000000000000 - fp32
+   */
+  constexpr uint32_t fp8_max = UINT32_C(143) << 23;
+
+  /*
+   * A mask for converting fp32 numbers lower than fp8e5m2 normal range
+   * into denorm representation
+   * magic number: ((127 - 15) + (23 - 2) + 1)
+   */
+  constexpr uint32_t denorm_mask = UINT32_C(134) << 23;
+
+  uint32_t f_bits = fp32_to_bits(f);
+  uint8_t result = 0u;
+
+  /*
+   * Extract the sign of the input number into the high bit of the 32-bit word:
+   *
+   *      +---+----------------------------------+
+   *      | S |0000000 00000000 00000000 00000000|
+   *      +---+----------------------------------+
+   * Bits  31                 0-31
+   */
+  const uint32_t sign = f_bits & UINT32_C(0x80000000);
+
+  /*
+   * Set sign bit to 0
+   */
+  f_bits ^= sign;
+
+  if (f_bits >= fp8_max) {
+    // NaN - all exponent and mantissa bits set to 1
+    result = f_bits > fp32_inf ? UINT8_C(0x7F) : UINT8_C(0x7C);
+  } else {
+    if (f_bits < (UINT32_C(113) << 23)) {
+      // Input number is smaller than 2^(-14), which is the smallest
+      // fp8e5m2 normal number
+      f_bits =
+          fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
+      result = static_cast<uint8_t>(f_bits - denorm_mask);
+    } else {
+      // resulting mantissa is odd
+      uint32_t mant_odd = (f_bits >> 21) & 1;
+
+      // update exponent, rounding bias part 1
+      f_bits += ((uint32_t)(15 - 127) << 23) + 0xFFFFF;
+
+      // rounding bias part 2
+      f_bits += mant_odd;
+
+      // take the bits!
+      result = static_cast<uint8_t>(f_bits >> 21);
+    }
+  }
+
+  result |= static_cast<uint8_t>(sign >> 24);
+  return result;
+}
+
+} // namespace detail
+
+// -------- below is copied from c10/util/Float8_e5m2-inl.h --------//
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+#define EXP_WIDTH_FP8 5
+#define MAN_WIDTH_FP8 2
+#define EXP_BIAS_FP8 15
+
+/// Constructors
+
+inline C10_HOST_DEVICE Float8_e5m2::Float8_e5m2(float value)
+    : x(detail::fp8e5m2_from_fp32_value(value)) {}
+
+/// Implicit conversions
+
+inline C10_HOST_DEVICE Float8_e5m2::operator float() const {
+  return detail::fp8e5m2_to_fp32_value(x);
+}
+
+/// Special values helpers
+
+inline C10_HOST_DEVICE bool Float8_e5m2::isnan() const {
+  return (x & 0b01111111) > 0b01111100;
+}
+
+inline C10_HOST_DEVICE bool Float8_e5m2::isinf() const {
+  return (x & 0b01111111) == 0b01111100;
+}
+
+/// Arithmetic
+
+inline C10_HOST_DEVICE Float8_e5m2
+operator+(const Float8_e5m2& a, const Float8_e5m2& b) {
+  return static_cast<float>(a) + static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2
+operator-(const Float8_e5m2& a, const Float8_e5m2& b) {
+  return static_cast<float>(a) - static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2
+operator*(const Float8_e5m2& a, const Float8_e5m2& b) {
+  return static_cast<float>(a) * static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2 operator/(
+    const Float8_e5m2& a,
+    const Float8_e5m2& b) __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2 operator-(const Float8_e5m2& a) {
+  return -static_cast<float>(a);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2& operator+=(
+    Float8_e5m2& a,
+    const Float8_e5m2& b) {
+  a = a + b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e5m2& operator-=(
+    Float8_e5m2& a,
+    const Float8_e5m2& b) {
+  a = a - b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e5m2& operator*=(
+    Float8_e5m2& a,
+    const Float8_e5m2& b) {
+  a = a * b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e5m2& operator/=(
+    Float8_e5m2& a,
+    const Float8_e5m2& b) {
+  a = a / b;
+  return a;
+}
+
+/// Arithmetic with floats
+
+inline C10_HOST_DEVICE float operator+(Float8_e5m2 a, float b) {
+  return static_cast<float>(a) + b;
+}
+inline C10_HOST_DEVICE float operator-(Float8_e5m2 a, float b) {
+  return static_cast<float>(a) - b;
+}
+inline C10_HOST_DEVICE float operator*(Float8_e5m2 a, float b) {
+  return static_cast<float>(a) * b;
+}
+inline C10_HOST_DEVICE float operator/(Float8_e5m2 a, float b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / b;
+}
+
+inline C10_HOST_DEVICE float operator+(float a, Float8_e5m2 b) {
+  return a + static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator-(float a, Float8_e5m2 b) {
+  return a - static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator*(float a, Float8_e5m2 b) {
+  return a * static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator/(float a, Float8_e5m2 b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE float& operator+=(float& a, const Float8_e5m2& b) {
+  return a += static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator-=(float& a, const Float8_e5m2& b) {
+  return a -= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator*=(float& a, const Float8_e5m2& b) {
+  return a *= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator/=(float& a, const Float8_e5m2& b) {
+  return a /= static_cast<float>(b);
+}
+
+/// Arithmetic with doubles
+
+inline C10_HOST_DEVICE double operator+(Float8_e5m2 a, double b) {
+  return static_cast<double>(a) + b;
+}
+inline C10_HOST_DEVICE double operator-(Float8_e5m2 a, double b) {
+  return static_cast<double>(a) - b;
+}
+inline C10_HOST_DEVICE double operator*(Float8_e5m2 a, double b) {
+  return static_cast<double>(a) * b;
+}
+inline C10_HOST_DEVICE double operator/(Float8_e5m2 a, double b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<double>(a) / b;
+}
+
+inline C10_HOST_DEVICE double operator+(double a, Float8_e5m2 b) {
+  return a + static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator-(double a, Float8_e5m2 b) {
+  return a - static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator*(double a, Float8_e5m2 b) {
+  return a * static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator/(double a, Float8_e5m2 b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<double>(b);
+}
+
+/// Arithmetic with ints
+
+inline C10_HOST_DEVICE Float8_e5m2 operator+(Float8_e5m2 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e5m2>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator-(Float8_e5m2 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e5m2>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator*(Float8_e5m2 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e5m2>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator/(Float8_e5m2 a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e5m2>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2 operator+(int a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator-(int a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator*(int a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator/(int a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) / b;
+}
+
+//// Arithmetic with int64_t
+
+inline C10_HOST_DEVICE Float8_e5m2 operator+(Float8_e5m2 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e5m2>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator-(Float8_e5m2 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e5m2>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator*(Float8_e5m2 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e5m2>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator/(Float8_e5m2 a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e5m2>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2 operator+(int64_t a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator-(int64_t a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator*(int64_t a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e5m2 operator/(int64_t a, Float8_e5m2 b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2>(a) / b;
+}
+
+/// NOTE: we do not define comparisons directly and instead rely on the implicit
+/// conversion from c10::Float8_e5m2 to float.
+C10_CLANG_DIAGNOSTIC_POP()
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::Float8_e5m2;
+using c10::operator<<;
+using c10::operator+;
+using c10::operator-;
+using c10::operator*;
+using c10::operator/;
+using c10::operator+=;
+using c10::operator-=;
+using c10::operator*=;
+using c10::operator/=;
+
+namespace detail {
+using c10::detail::fp8e5m2_from_fp32_value;
+using c10::detail::fp8e5m2_to_fp32_value;
+} // namespace detail
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace std {
+
+template <>
+class numeric_limits<c10::Float8_e5m2> {
+ public:
+  static constexpr bool is_signed = true;
+  static constexpr bool is_integer = false;
+  static constexpr bool is_specialized = true;
+  static constexpr bool is_exact = false;
+  static constexpr bool has_infinity = true;
+  static constexpr bool has_quiet_NaN = true;
+  static constexpr bool has_signaling_NaN = false;
+  static constexpr auto has_denorm = true;
+  static constexpr auto has_denorm_loss = true;
+  static constexpr auto round_style = numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = false;
+  static constexpr bool is_bounded = true;
+  static constexpr bool is_modulo = false;
+  static constexpr int digits = 3;
+  static constexpr int digits10 = 0;
+  static constexpr int max_digits10 = 2;
+  static constexpr int radix = 2;
+  static constexpr int min_exponent = -13;
+  static constexpr int min_exponent10 = -4;
+  static constexpr int max_exponent = 16;
+  static constexpr int max_exponent10 = 4;
+  static constexpr auto traps = numeric_limits<float>::traps;
+  static constexpr auto tinyness_before =
+      numeric_limits<float>::tinyness_before;
+
+  static constexpr c10::Float8_e5m2 min() {
+    return c10::Float8_e5m2(0x4, c10::Float8_e5m2::from_bits());
+  }
+  static constexpr c10::Float8_e5m2 max() {
+    return c10::Float8_e5m2(0x7B, c10::Float8_e5m2::from_bits());
+  }
+  static constexpr c10::Float8_e5m2 lowest() {
+    return c10::Float8_e5m2(0xFB, c10::Float8_e5m2::from_bits());
+  }
+  static constexpr c10::Float8_e5m2 epsilon() {
+    return c10::Float8_e5m2(0x34, c10::Float8_e5m2::from_bits());
+  }
+  static constexpr c10::Float8_e5m2 round_error() {
+    return c10::Float8_e5m2(0x38, c10::Float8_e5m2::from_bits());
+  }
+  static constexpr c10::Float8_e5m2 infinity() {
+    return c10::Float8_e5m2(0x7C, c10::Float8_e5m2::from_bits());
+  }
+  static constexpr c10::Float8_e5m2 quiet_NaN() {
+    return c10::Float8_e5m2(0x7F, c10::Float8_e5m2::from_bits());
+  }
+  static constexpr c10::Float8_e5m2 denorm_min() {
+    return c10::Float8_e5m2(0x01, c10::Float8_e5m2::from_bits());
+  }
+};
+
+} // namespace std
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e5m2fnuz.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e5m2fnuz.h
new file mode 100644
index 0000000000000000000000000000000000000000..6b0f79b75adea4a994a3c44e666b73767b68c66d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e5m2fnuz.h
@@ -0,0 +1,448 @@
+#pragma once
+
+/// Defines the Float8_e5m2fnuz type (8-bit floating-point) including
+/// conversions to standard C types and basic arithmetic operations. Note that
+/// arithmetic operations are implemented by converting to floating point and
+/// performing the operation in float32.
+/// Binary configuration remains the same as e5m2:
+/// s eeeee mm
+/// 1 sign bit
+/// 5 exponent bits
+/// 2 mantissa bits
+/// The key differences that e5m2fnuz brings are:
+/// bias = 16
+/// no infinities or negative zero
+/// NaN only when sign bit is 1, rest all 0s
+///
+/// Implementation based on the paper https://arxiv.org/pdf/2206.02915.pdf and
+/// the existing Float8_e4m3fn implementation.
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Float8_fnuz_cvt.h>
+#include <torch/headeronly/util/TypeSafeSignMath.h>
+#include <torch/headeronly/util/floating_point_utils.h>
+
+#if defined(__cplusplus)
+#include <cstdint>
+#elif !defined(__OPENCL_VERSION__)
+#include <math.h>
+#include <stdint.h>
+#endif
+
+#include <iosfwd>
+#include <ostream>
+
+namespace c10 {
+
+struct alignas(1) Float8_e5m2fnuz {
+  uint8_t x;
+
+  struct from_bits_t {};
+  C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
+    return from_bits_t();
+  }
+
+  Float8_e5m2fnuz() = default;
+
+  constexpr C10_HOST_DEVICE Float8_e5m2fnuz(
+      uint8_t bits,
+      from_bits_t /*unused*/)
+      : x(bits) {}
+  inline C10_HOST_DEVICE Float8_e5m2fnuz(float value);
+  inline C10_HOST_DEVICE operator float() const;
+  inline C10_HOST_DEVICE bool isnan() const;
+  inline C10_HOST_DEVICE bool isinf() const;
+};
+
+inline std::ostream& operator<<(
+    std::ostream& out,
+    const Float8_e5m2fnuz& value) {
+  out << (float)value;
+  return out;
+}
+
+namespace detail {
+
+/*
+ * Convert a 32-bit floating-point number in IEEE single-precision format to a
+ * 8-bit floating-point number in fp8 E5M2 format, in bit representation.
+ */
+inline C10_HOST_DEVICE uint8_t fp8e5m2fnuz_from_fp32_value(float f) {
+  /*
+   * Binary representation of 65536.0f, which is the first value not
+   * representable (i.e. the first value which would overflow in to the sign
+   * bit, resulting in a NaN) in fp8e4m3fnuz range:
+   * 1 00000 00 - fp8e5m2fnuz
+   * 0 10001111 00000000000000000000000 - fp32
+   */
+  constexpr uint32_t fnuz_max = UINT32_C(0x8F) << 23;
+
+  /*
+   * A mask for converting fp32 numbers lower than fp8e5m2fnuz normal range
+   * into denormalized representation.
+   * magic number: ((127 - 16) + (23 - 2) + 1)
+   */
+  constexpr uint32_t denorm_mask = UINT32_C(0x85) << 23;
+
+  uint32_t f_bits = fp32_to_bits(f);
+  uint32_t result = 0u;
+
+  /*
+   * Extract the sign of the input number into the high bit of the 32-bit word:
+   *
+   *      +---+----------------------------------+
+   *      | S |0000000 00000000 00000000 00000000|
+   *      +---+----------------------------------+
+   * Bits  31                 0-31
+   */
+  const uint32_t sign = f_bits & UINT32_C(0x80000000);
+
+  /*
+   * Set sign bit to 0
+   */
+  f_bits ^= sign;
+
+  if (f_bits >= fnuz_max) {
+    // NaN -- sign bit set to 1, rest 0s
+    return 0x80;
+  }
+
+  if (f_bits < (UINT32_C(0x70) << 23) /* 2^-15 in float32 */) {
+    // Input exponent is less than -15, the smallest e5m2fnuz exponent, so the
+    // number will become subnormal.
+    f_bits = fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
+    result = static_cast<uint8_t>(f_bits - denorm_mask);
+    if (result == 0) {
+      // fnuz types don't have negative zero.
+      return 0;
+    }
+  } else {
+    // resulting mantissa is odd
+    uint8_t mant_odd = (f_bits >> 21) & 1;
+
+    // update exponent, rounding bias part 1
+    f_bits += ((uint32_t)(16 - 127) << 23) + 0xFFFFF;
+
+    // rounding bias part 2
+    f_bits += mant_odd;
+
+    // take the bits!
+    result = static_cast<uint8_t>(f_bits >> 21);
+  }
+
+  result |= sign >> 24;
+  return result;
+}
+
+} // namespace detail
+
+//------ below is copied from c10/util/Float8_e5m2fnuz-inl.h ------//
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+/// Constructors
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz::Float8_e5m2fnuz(float value)
+    : x(detail::fp8e5m2fnuz_from_fp32_value(value)) {}
+
+/// Implicit conversions
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz::operator float() const {
+  return torch::headeronly::detail::fp8_fnuz_to_fp32_value<5, 2>(x);
+}
+
+/// Special values helpers
+
+inline C10_HOST_DEVICE bool Float8_e5m2fnuz::isnan() const {
+  return x == 0b10000000;
+}
+
+inline C10_HOST_DEVICE bool Float8_e5m2fnuz::isinf() const {
+  return false;
+}
+
+/// Arithmetic
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz
+operator+(const Float8_e5m2fnuz& a, const Float8_e5m2fnuz& b) {
+  return static_cast<float>(a) + static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz
+operator-(const Float8_e5m2fnuz& a, const Float8_e5m2fnuz& b) {
+  return static_cast<float>(a) - static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz
+operator*(const Float8_e5m2fnuz& a, const Float8_e5m2fnuz& b) {
+  return static_cast<float>(a) * static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator/(
+    const Float8_e5m2fnuz& a,
+    const Float8_e5m2fnuz& b) __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator-(const Float8_e5m2fnuz& a) {
+  return -static_cast<float>(a);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz& operator+=(
+    Float8_e5m2fnuz& a,
+    const Float8_e5m2fnuz& b) {
+  a = a + b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz& operator-=(
+    Float8_e5m2fnuz& a,
+    const Float8_e5m2fnuz& b) {
+  a = a - b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz& operator*=(
+    Float8_e5m2fnuz& a,
+    const Float8_e5m2fnuz& b) {
+  a = a * b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz& operator/=(
+    Float8_e5m2fnuz& a,
+    const Float8_e5m2fnuz& b) {
+  a = a / b;
+  return a;
+}
+
+/// Arithmetic with floats
+
+inline C10_HOST_DEVICE float operator+(Float8_e5m2fnuz a, float b) {
+  return static_cast<float>(a) + b;
+}
+inline C10_HOST_DEVICE float operator-(Float8_e5m2fnuz a, float b) {
+  return static_cast<float>(a) - b;
+}
+inline C10_HOST_DEVICE float operator*(Float8_e5m2fnuz a, float b) {
+  return static_cast<float>(a) * b;
+}
+inline C10_HOST_DEVICE float operator/(Float8_e5m2fnuz a, float b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / b;
+}
+
+inline C10_HOST_DEVICE float operator+(float a, Float8_e5m2fnuz b) {
+  return a + static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator-(float a, Float8_e5m2fnuz b) {
+  return a - static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator*(float a, Float8_e5m2fnuz b) {
+  return a * static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator/(float a, Float8_e5m2fnuz b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE float& operator+=(float& a, const Float8_e5m2fnuz& b) {
+  return a += static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator-=(float& a, const Float8_e5m2fnuz& b) {
+  return a -= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator*=(float& a, const Float8_e5m2fnuz& b) {
+  return a *= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator/=(float& a, const Float8_e5m2fnuz& b) {
+  return a /= static_cast<float>(b);
+}
+
+/// Arithmetic with doubles
+
+inline C10_HOST_DEVICE double operator+(Float8_e5m2fnuz a, double b) {
+  return static_cast<double>(a) + b;
+}
+inline C10_HOST_DEVICE double operator-(Float8_e5m2fnuz a, double b) {
+  return static_cast<double>(a) - b;
+}
+inline C10_HOST_DEVICE double operator*(Float8_e5m2fnuz a, double b) {
+  return static_cast<double>(a) * b;
+}
+inline C10_HOST_DEVICE double operator/(Float8_e5m2fnuz a, double b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<double>(a) / b;
+}
+
+inline C10_HOST_DEVICE double operator+(double a, Float8_e5m2fnuz b) {
+  return a + static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator-(double a, Float8_e5m2fnuz b) {
+  return a - static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator*(double a, Float8_e5m2fnuz b) {
+  return a * static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator/(double a, Float8_e5m2fnuz b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<double>(b);
+}
+
+/// Arithmetic with ints
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator+(Float8_e5m2fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e5m2fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator-(Float8_e5m2fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e5m2fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator*(Float8_e5m2fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e5m2fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator/(Float8_e5m2fnuz a, int b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e5m2fnuz>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator+(int a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator-(int a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator*(int a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator/(int a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) / b;
+}
+
+//// Arithmetic with int64_t
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator+(Float8_e5m2fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a + static_cast<Float8_e5m2fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator-(Float8_e5m2fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a - static_cast<Float8_e5m2fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator*(Float8_e5m2fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a * static_cast<Float8_e5m2fnuz>(b);
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator/(Float8_e5m2fnuz a, int64_t b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return a / static_cast<Float8_e5m2fnuz>(b);
+}
+
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator+(int64_t a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) + b;
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator-(int64_t a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) - b;
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator*(int64_t a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) * b;
+}
+inline C10_HOST_DEVICE Float8_e5m2fnuz operator/(int64_t a, Float8_e5m2fnuz b) {
+  // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions)
+  return static_cast<Float8_e5m2fnuz>(a) / b;
+}
+
+/// NOTE: we do not define comparisons directly and instead rely on the implicit
+/// conversion from c10::Float8_e5m2fnuz to float.
+
+C10_CLANG_DIAGNOSTIC_POP()
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::Float8_e5m2fnuz;
+using c10::operator<<;
+using c10::operator+;
+using c10::operator-;
+using c10::operator*;
+using c10::operator/;
+using c10::operator+=;
+using c10::operator-=;
+using c10::operator*=;
+using c10::operator/=;
+
+namespace detail {
+using c10::detail::fp8e5m2fnuz_from_fp32_value;
+}
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace std {
+
+template <>
+class numeric_limits<c10::Float8_e5m2fnuz> {
+ public:
+  static constexpr bool is_signed = true;
+  static constexpr bool is_integer = false;
+  static constexpr bool is_specialized = true;
+  static constexpr bool is_exact = false;
+  static constexpr bool has_infinity = false;
+  static constexpr bool has_quiet_NaN = true;
+  static constexpr bool has_signaling_NaN = false;
+  static constexpr auto has_denorm = true;
+  static constexpr auto has_denorm_loss = true;
+  static constexpr auto round_style = numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = false;
+  static constexpr bool is_bounded = true;
+  static constexpr bool is_modulo = false;
+  static constexpr int digits = 3;
+  static constexpr int digits10 = 0;
+  static constexpr int max_digits10 = 2;
+  static constexpr int radix = 2;
+  static constexpr int min_exponent = -14;
+  static constexpr int min_exponent10 = -4;
+  static constexpr int max_exponent = 16;
+  static constexpr int max_exponent10 = 4;
+  static constexpr auto traps = numeric_limits<float>::traps;
+  static constexpr auto tinyness_before =
+      numeric_limits<float>::tinyness_before;
+
+  static constexpr c10::Float8_e5m2fnuz min() {
+    return c10::Float8_e5m2fnuz(0x04, c10::Float8_e5m2fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e5m2fnuz max() {
+    return c10::Float8_e5m2fnuz(0x7F, c10::Float8_e5m2fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e5m2fnuz lowest() {
+    return c10::Float8_e5m2fnuz(0xFF, c10::Float8_e5m2fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e5m2fnuz epsilon() {
+    return c10::Float8_e5m2fnuz(0x34, c10::Float8_e5m2fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e5m2fnuz round_error() {
+    return c10::Float8_e5m2fnuz(0x38, c10::Float8_e5m2fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e5m2fnuz infinity() {
+    return c10::Float8_e5m2fnuz(0x80, c10::Float8_e5m2fnuz::from_bits());
+  }
+  // TODO(future): we are mapping neg_zero to both inf and NaN, this is
+  // surprising and we should figure out what to do about it.
+  static constexpr c10::Float8_e5m2fnuz quiet_NaN() {
+    return c10::Float8_e5m2fnuz(0x80, c10::Float8_e5m2fnuz::from_bits());
+  }
+  static constexpr c10::Float8_e5m2fnuz denorm_min() {
+    return c10::Float8_e5m2fnuz(0x01, c10::Float8_e5m2fnuz::from_bits());
+  }
+};
+
+} // namespace std
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e8m0fnu.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e8m0fnu.h
new file mode 100644
index 0000000000000000000000000000000000000000..80d89439c90b939259e1ec20b0b0c0256d3ced84
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_e8m0fnu.h
@@ -0,0 +1,226 @@
+#pragma once
+
+/// Defines the Float8_e8m0fnu type (8-bit floating-point) including
+/// conversions to standard C types
+/// Binary configuration :
+/// eeeeeeee
+/// no sign bits
+/// 8 exponent bits
+/// no mantissa bits
+///
+/// This is the E8M0 dtype from the OCP MX format spec
+/// (https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf,
+/// Section 5.4.1)
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/floating_point_utils.h>
+
+// TODO(#146647): do we need to special case OPENCL?
+#if defined(__cplusplus)
+#include <cstdint>
+#elif !defined(__OPENCL_VERSION__)
+#include <math.h>
+#include <stdint.h>
+#endif
+
+#include <iosfwd>
+#include <limits>
+#include <ostream>
+
+namespace c10 {
+
+struct alignas(1) Float8_e8m0fnu {
+  uint8_t x;
+
+  struct from_bits_t {};
+  C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
+    return from_bits_t();
+  }
+
+  Float8_e8m0fnu() = default;
+
+  constexpr C10_HOST_DEVICE Float8_e8m0fnu(uint8_t bits, from_bits_t /*unused*/)
+      : x(bits) {}
+  inline C10_HOST_DEVICE Float8_e8m0fnu(float value);
+  inline C10_HOST_DEVICE operator float() const;
+  inline C10_HOST_DEVICE bool isnan() const;
+};
+
+inline std::ostream& operator<<(
+    std::ostream& out,
+    const Float8_e8m0fnu& value) {
+  out << (float)value;
+  return out;
+}
+
+namespace detail {
+/*
+ * Convert a 32-bit floating-point number in IEEE single-precision format to a
+ * 8-bit floating-point number in fp8 e8m0fnu format, in bit representation.
+ */
+inline C10_HOST_DEVICE uint8_t fp8e8m0fnu_from_fp32_value(float f) {
+  // TODO(#146647): maybe rewrite without control flow
+
+  uint32_t f_bits = c10::detail::fp32_to_bits(f);
+
+  // extract the exponent
+  uint32_t exponent = (f_bits >> 23) & 0b11111111;
+
+  // special case float32 NaN and +-inf to map to e8m0 nan
+  if (exponent == 0b11111111) {
+    return exponent;
+  }
+
+  // next, we use guard, round, sticky bits and the LSB to implement round to
+  // nearest, with ties to even
+
+  // guard bit - bit 23, or 22 zero-indexed
+  uint8_t g = (f_bits & 0x400000) > 0;
+  // round bit - bit 22, or 21 zero-indexed
+  uint8_t r = (f_bits & 0x200000) > 0;
+  // sticky bit - bits 21 to 1, or 20 to 0 zero-indexed
+  uint8_t s = (f_bits & 0x1FFFFF) > 0;
+  // in casting to e8m0, LSB is the implied mantissa bit. It equals to 0 if the
+  // original float32 is denormal, and to 1 if the original float32 is normal.
+  uint8_t lsb = exponent > 0;
+
+  // implement the RNE logic
+  bool round_up = false;
+
+  // if g == 0, round down (no-op)
+  if (g == 1) {
+    if ((r == 1) || (s == 1)) {
+      // round up
+      round_up = true;
+    } else {
+      if (lsb == 1) {
+        // round up
+        round_up = true;
+      }
+      // if lsb == 0, round down (no-op)
+    }
+  }
+
+  if (round_up) {
+    // adjust exponent
+    // note that if exponent was 255 we would have already returned earlier, so
+    // we know we can add one safely without running out of bounds
+    exponent++;
+  }
+
+  return exponent;
+}
+
+} // namespace detail
+
+//------- the below is from c10/util/Float8_e8m0fnu-inl.h  ------//
+// TODO(#146647): Can we remove the below warning?
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+/// Constructors
+inline C10_HOST_DEVICE Float8_e8m0fnu::Float8_e8m0fnu(float value)
+    : x(detail::fp8e8m0fnu_from_fp32_value(value)) {}
+
+/// Implicit conversions
+
+inline C10_HOST_DEVICE Float8_e8m0fnu::operator float() const {
+  // TODO(#146647): maybe rewrite without control flow
+
+  // if exponent is zero, need to special case to return 2^-127 instead of zero
+  if (x == 0) {
+    return c10::detail::fp32_from_bits(0x00400000);
+  }
+
+  // if exponent is NaN, need to special case to return properly encoded NaN
+  if (isnan()) {
+    return c10::detail::fp32_from_bits(0x7f800001);
+  }
+
+  // leave sign at 0, set the exponent bits, leave stored mantissa at 0
+  uint32_t res = x << 23;
+
+  return c10::detail::fp32_from_bits(res);
+}
+
+/// Special values helper
+
+inline C10_HOST_DEVICE bool Float8_e8m0fnu::isnan() const {
+  return x == 0b11111111;
+}
+
+/// NOTE: we do not define comparisons directly and instead rely on the implicit
+/// conversion from c10::Float8_e8m0fnu to float.
+C10_CLANG_DIAGNOSTIC_POP()
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::Float8_e8m0fnu;
+using c10::operator<<;
+
+namespace detail {
+using c10::detail::fp8e8m0fnu_from_fp32_value;
+} // namespace detail
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace std {
+
+template <>
+class numeric_limits<c10::Float8_e8m0fnu> {
+ public:
+  static constexpr bool is_specialized = true;
+  static constexpr bool is_signed = false;
+  static constexpr bool is_integer = false;
+  static constexpr bool is_exact = false;
+  static constexpr bool has_infinity = false;
+  static constexpr bool has_quiet_NaN = true;
+  static constexpr bool has_signaling_NaN = false;
+  static constexpr auto has_denorm = false;
+  static constexpr auto has_denorm_loss = false;
+  static constexpr auto round_style = numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = false;
+  static constexpr bool is_bounded = true;
+  static constexpr bool is_modulo = false;
+  static constexpr int digits = 1;
+  static constexpr int digits10 = 0;
+  static constexpr int max_digits10 = 1; // just a 2!
+  static constexpr int radix = 2;
+  static constexpr int min_exponent = -126;
+  static constexpr int min_exponent10 = -38;
+  static constexpr int max_exponent = 128;
+  static constexpr int max_exponent10 = 38;
+  static constexpr auto traps = numeric_limits<float>::traps;
+  static constexpr auto tinyness_before = false;
+
+  static constexpr c10::Float8_e8m0fnu min() {
+    // 2^-127
+    return c10::Float8_e8m0fnu(0b00000000, c10::Float8_e8m0fnu::from_bits());
+  }
+  static constexpr c10::Float8_e8m0fnu lowest() {
+    // 2^-127
+    return c10::Float8_e8m0fnu(0b00000000, c10::Float8_e8m0fnu::from_bits());
+  }
+  static constexpr c10::Float8_e8m0fnu max() {
+    // 254 biased, which is 127 unbiased, so 2^127
+    return c10::Float8_e8m0fnu(0b11111110, c10::Float8_e8m0fnu::from_bits());
+  }
+  static constexpr c10::Float8_e8m0fnu epsilon() {
+    // according to https://en.cppreference.com/w/cpp/types/numeric_limits, this
+    // is "the difference between 1.0 and the next representable value of the
+    // given floating-point type". The next representable value is 2.0, so the
+    // difference is 1.0 which is 2^0. 0 unbiased is 127 biased.
+    return c10::Float8_e8m0fnu(0b01111111, c10::Float8_e8m0fnu::from_bits());
+  }
+  static constexpr c10::Float8_e8m0fnu round_error() {
+    // 0.5 in float, which is 2^-1, and -1 + 127 = 126
+    return c10::Float8_e8m0fnu(0b01111110, c10::Float8_e8m0fnu::from_bits());
+  }
+  static constexpr c10::Float8_e8m0fnu quiet_NaN() {
+    return c10::Float8_e8m0fnu(0b11111111, c10::Float8_e8m0fnu::from_bits());
+  }
+};
+
+} // namespace std
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_fnuz_cvt.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_fnuz_cvt.h
new file mode 100644
index 0000000000000000000000000000000000000000..59fab8cc28dbfb2c9c8cb4951deaf2bb06dd06ec
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Float8_fnuz_cvt.h
@@ -0,0 +1,69 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/floating_point_utils.h>
+
+#include <cstdint>
+
+#if defined(SYCL_LANGUAGE_VERSION)
+#include <sycl/sycl.hpp>
+#endif
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly, detail)
+
+/*
+ * Convert a 8-bit floating-point number in either f8 E4M3FNUZ or bf8 E5M2FNUZ
+ * format, in bit representation, to a 32-bit floating-point number.
+ */
+template <uint32_t we, uint32_t wm>
+inline C10_HOST_DEVICE float fp8_fnuz_to_fp32_value(uint8_t x) {
+  static_assert((we == 4 && wm == 3) || (we == 5 && wm == 2));
+  constexpr uint32_t weo = 8;
+  constexpr uint32_t wmo = 23;
+
+  if (x == 0) {
+    return 0;
+  }
+
+  if (x == 0x80) {
+    constexpr uint32_t ifNaN = 0x7F800001;
+    return fp32_from_bits(ifNaN);
+  }
+
+  uint32_t mantissa = x & ((1 << wm) - 1);
+  uint32_t exponent = (x & 0x7F) >> wm;
+
+  // subnormal input
+  if (exponent == 0) {
+    // guaranteed mantissa!=0 since cases 0x0 and 0x80 are handled above
+#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
+    uint32_t renorm_shift = __clz(mantissa);
+#elif defined(__SYCL_DEVICE_ONLY__)
+    uint32_t renorm_shift = sycl::clz(mantissa);
+#elif defined(_MSC_VER)
+    unsigned long nonsign_bsr;
+    _BitScanReverse(&nonsign_bsr, (unsigned long)mantissa);
+    uint32_t renorm_shift = (uint32_t)nonsign_bsr ^ 31;
+#else
+    uint32_t renorm_shift = __builtin_clz(mantissa);
+#endif
+    uint32_t sh = 1 + renorm_shift - (32 - wm);
+    mantissa <<= sh;
+    exponent += 1 - sh;
+    mantissa &= ((1 << wm) - 1);
+  }
+
+  const uint32_t exp_low_cutoff = (1 << (weo - 1)) - (1 << (we - 1));
+  exponent += exp_low_cutoff - 1;
+  mantissa <<= wmo - wm;
+
+  uint32_t sign = x >> 7;
+  uint32_t retval = (sign << 31) | (exponent << 23) | mantissa;
+  return fp32_from_bits(retval);
+}
+
+HIDDEN_NAMESPACE_END(torch, headeronly, detail)
+
+namespace c10::detail {
+using torch::headeronly::detail::fp8_fnuz_to_fp32_value;
+}
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Half.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Half.h
new file mode 100644
index 0000000000000000000000000000000000000000..a9c0b166ba2ea825ee8c2933e519de40112dae6f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Half.h
@@ -0,0 +1,788 @@
+#pragma once
+
+/// Defines the Half type (half-precision floating-point) including conversions
+/// to standard C types and basic arithmetic operations. Note that arithmetic
+/// operations are implemented by converting to floating point and
+/// performing the operation in float32, instead of using CUDA half intrinsics.
+/// Most uses of this type within ATen are memory bound, including the
+/// element-wise kernels, and the half intrinsics aren't efficient on all GPUs.
+/// If you are writing a compute bound kernel, you can use the CUDA half
+/// intrinsics directly on the Half type from device code.
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/bit_cast.h>
+#include <torch/headeronly/util/floating_point_utils.h>
+
+#if defined(__cplusplus)
+#include <cmath>
+#elif !defined(__OPENCL_VERSION__)
+#include <math.h>
+#endif
+
+#ifdef _MSC_VER
+#include <intrin.h>
+#endif
+
+#include <cstdint>
+#include <cstring>
+#include <ostream>
+
+#ifdef __CUDACC__
+#include <cuda_fp16.h>
+#endif
+
+#ifdef __HIPCC__
+#include <hip/hip_fp16.h>
+#endif
+
+#if defined(CL_SYCL_LANGUAGE_VERSION)
+#include <CL/sycl.hpp> // for SYCL 1.2.1
+#elif defined(SYCL_LANGUAGE_VERSION)
+#include <sycl/sycl.hpp> // for SYCL 2020
+#endif
+
+#if (defined(CPU_CAPABILITY_AVX2) || defined(CPU_CAPABILITY_AVX512)) && \
+    !defined(__APPLE__)
+#include <torch/headeronly/cpu/vec/vec_half.h>
+#endif
+
+#if defined(__aarch64__) && !defined(__CUDACC__)
+#include <arm_neon.h>
+#endif
+
+#if defined(__GNUC__) || defined(__clang__)
+#if defined(__x86_64__) || defined(_M_X64) || defined(__i386) || \
+    defined(_M_IX86)
+#if defined(__F16C__) &&                               \
+    !(defined(__CUDA_ARCH__) || defined(__CUDACC__) || \
+      defined(__HIP_DEVICE_COMPILE__))
+#define C10_X86_F16 1
+#include <immintrin.h> // import conversion ops from f16cintrin.h
+#endif // defined(__F16C__) && !(defined(__CUDA_ARCH__) || defined(__CUDACC__)
+       // || defined(__HIP_DEVICE_COMPILE__))
+#endif // __x86_64__ || _M_X64 || __i386 || _M_IX86
+#endif // __GNUC__ || __clang__
+
+namespace c10 {
+
+struct alignas(2) Half {
+  unsigned short x;
+
+  struct from_bits_t {};
+  C10_HOST_DEVICE static constexpr from_bits_t from_bits() {
+    return from_bits_t();
+  }
+
+  // HIP wants __host__ __device__ tag, CUDA does not
+#if defined(USE_ROCM)
+  C10_HOST_DEVICE Half() = default;
+#else
+  Half() = default;
+#endif
+
+  constexpr C10_HOST_DEVICE Half(unsigned short bits, from_bits_t /*unused*/)
+      : x(bits) {}
+#if defined(__aarch64__) && !defined(__CUDACC__)
+  inline Half(float16_t value);
+  inline operator float16_t() const;
+#else
+  inline C10_HOST_DEVICE Half(float value);
+  inline C10_HOST_DEVICE operator float() const;
+#endif
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+  inline C10_HOST_DEVICE Half(const __half& value);
+  inline C10_HOST_DEVICE operator __half() const;
+#endif
+#ifdef SYCL_LANGUAGE_VERSION
+  inline C10_HOST_DEVICE Half(const sycl::half& value);
+  inline C10_HOST_DEVICE operator sycl::half() const;
+#endif
+};
+
+inline std::ostream& operator<<(std::ostream& out, const Half& value) {
+  out << (float)value;
+  return out;
+}
+
+namespace detail {
+/*
+ * Convert a 16-bit floating-point number in IEEE half-precision format, in bit
+ * representation, to a 32-bit floating-point number in IEEE single-precision
+ * format.
+ *
+ * @note The implementation relies on IEEE-like (no assumption about rounding
+ * mode and no operations on denormals) floating-point operations and bitcasts
+ * between integer and floating-point variables.
+ */
+C10_HOST_DEVICE inline float fp16_ieee_to_fp32_value(uint16_t h) {
+#ifdef C10_X86_F16
+  return _cvtsh_ss(h);
+#else
+  /*
+   * Extend the half-precision floating-point number to 32 bits and shift to the
+   * upper part of the 32-bit word:
+   *      +---+-----+------------+-------------------+
+   *      | S |EEEEE|MM MMMM MMMM|0000 0000 0000 0000|
+   *      +---+-----+------------+-------------------+
+   * Bits  31  26-30    16-25            0-15
+   *
+   * S - sign bit, E - bits of the biased exponent, M - bits of the mantissa, 0
+   * - zero bits.
+   */
+  const uint32_t w = (uint32_t)h << 16;
+  /*
+   * Extract the sign of the input number into the high bit of the 32-bit word:
+   *
+   *      +---+----------------------------------+
+   *      | S |0000000 00000000 00000000 00000000|
+   *      +---+----------------------------------+
+   * Bits  31                 0-31
+   */
+  const uint32_t sign = w & UINT32_C(0x80000000);
+  /*
+   * Extract mantissa and biased exponent of the input number into the high bits
+   * of the 32-bit word:
+   *
+   *      +-----+------------+---------------------+
+   *      |EEEEE|MM MMMM MMMM|0 0000 0000 0000 0000|
+   *      +-----+------------+---------------------+
+   * Bits  27-31    17-26            0-16
+   */
+  const uint32_t two_w = w + w;
+
+  /*
+   * Shift mantissa and exponent into bits 23-28 and bits 13-22 so they become
+   * mantissa and exponent of a single-precision floating-point number:
+   *
+   *       S|Exponent |          Mantissa
+   *      +-+---+-----+------------+----------------+
+   *      |0|000|EEEEE|MM MMMM MMMM|0 0000 0000 0000|
+   *      +-+---+-----+------------+----------------+
+   * Bits   | 23-31   |           0-22
+   *
+   * Next, there are some adjustments to the exponent:
+   * - The exponent needs to be corrected by the difference in exponent bias
+   * between single-precision and half-precision formats (0x7F - 0xF = 0x70)
+   * - Inf and NaN values in the inputs should become Inf and NaN values after
+   * conversion to the single-precision number. Therefore, if the biased
+   * exponent of the half-precision input was 0x1F (max possible value), the
+   * biased exponent of the single-precision output must be 0xFF (max possible
+   * value). We do this correction in two steps:
+   *   - First, we adjust the exponent by (0xFF - 0x1F) = 0xE0 (see exp_offset
+   * below) rather than by 0x70 suggested by the difference in the exponent bias
+   * (see above).
+   *   - Then we multiply the single-precision result of exponent adjustment by
+   * 2**(-112) to reverse the effect of exponent adjustment by 0xE0 less the
+   * necessary exponent adjustment by 0x70 due to difference in exponent bias.
+   *     The floating-point multiplication hardware would ensure than Inf and
+   * NaN would retain their value on at least partially IEEE754-compliant
+   * implementations.
+   *
+   * Note that the above operations do not handle denormal inputs (where biased
+   * exponent == 0). However, they also do not operate on denormal inputs, and
+   * do not produce denormal results.
+   */
+  constexpr uint32_t exp_offset = UINT32_C(0xE0) << 23;
+  // const float exp_scale = 0x1.0p-112f;
+  constexpr uint32_t scale_bits = (uint32_t)15 << 23;
+  float exp_scale_val = 0;
+#if defined(_MSC_VER) && defined(__clang__)
+  __builtin_memcpy(&exp_scale_val, &scale_bits, sizeof(exp_scale_val));
+#else
+  std::memcpy(&exp_scale_val, &scale_bits, sizeof(exp_scale_val));
+#endif
+
+  const float exp_scale = exp_scale_val;
+  const float normalized_value =
+      fp32_from_bits((two_w >> 4) + exp_offset) * exp_scale;
+
+  /*
+   * Convert denormalized half-precision inputs into single-precision results
+   * (always normalized). Zero inputs are also handled here.
+   *
+   * In a denormalized number the biased exponent is zero, and mantissa has
+   * on-zero bits. First, we shift mantissa into bits 0-9 of the 32-bit word.
+   *
+   *                  zeros           |  mantissa
+   *      +---------------------------+------------+
+   *      |0000 0000 0000 0000 0000 00|MM MMMM MMMM|
+   *      +---------------------------+------------+
+   * Bits             10-31                0-9
+   *
+   * Now, remember that denormalized half-precision numbers are represented as:
+   *    FP16 = mantissa * 2**(-24).
+   * The trick is to construct a normalized single-precision number with the
+   * same mantissa and thehalf-precision input and with an exponent which would
+   * scale the corresponding mantissa bits to 2**(-24). A normalized
+   * single-precision floating-point number is represented as: FP32 = (1 +
+   * mantissa * 2**(-23)) * 2**(exponent - 127) Therefore, when the biased
+   * exponent is 126, a unit change in the mantissa of the input denormalized
+   * half-precision number causes a change of the constructed single-precision
+   * number by 2**(-24), i.e. the same amount.
+   *
+   * The last step is to adjust the bias of the constructed single-precision
+   * number. When the input half-precision number is zero, the constructed
+   * single-precision number has the value of FP32 = 1 * 2**(126 - 127) =
+   * 2**(-1) = 0.5 Therefore, we need to subtract 0.5 from the constructed
+   * single-precision number to get the numerical equivalent of the input
+   * half-precision number.
+   */
+  constexpr uint32_t magic_mask = UINT32_C(126) << 23;
+  constexpr float magic_bias = 0.5f;
+  const float denormalized_value =
+      fp32_from_bits((two_w >> 17) | magic_mask) - magic_bias;
+
+  /*
+   * - Choose either results of conversion of input as a normalized number, or
+   * as a denormalized number, depending on the input exponent. The variable
+   * two_w contains input exponent in bits 27-31, therefore if its smaller than
+   * 2**27, the input is either a denormal number, or zero.
+   * - Combine the result of conversion of exponent and mantissa with the sign
+   * of the input number.
+   */
+  constexpr uint32_t denormalized_cutoff = UINT32_C(1) << 27;
+  const uint32_t result = sign |
+      (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value)
+                                   : fp32_to_bits(normalized_value));
+  return fp32_from_bits(result);
+#endif // C10_X86_F16
+}
+
+/*
+ * Convert a 32-bit floating-point number in IEEE single-precision format to a
+ * 16-bit floating-point number in IEEE half-precision format, in bit
+ * representation.
+ *
+ * @note The implementation relies on IEEE-like (no assumption about rounding
+ * mode and no operations on denormals) floating-point operations and bitcasts
+ * between integer and floating-point variables.
+ */
+inline uint16_t fp16_ieee_from_fp32_value(float f) {
+#ifdef C10_X86_F16
+  return _cvtss_sh(f, _MM_FROUND_TO_NEAREST_INT);
+#else
+  // const float scale_to_inf = 0x1.0p+112f;
+  // const float scale_to_zero = 0x1.0p-110f;
+  constexpr uint32_t scale_to_inf_bits = (uint32_t)239 << 23;
+  constexpr uint32_t scale_to_zero_bits = (uint32_t)17 << 23;
+  float scale_to_inf_val = 0, scale_to_zero_val = 0;
+  std::memcpy(&scale_to_inf_val, &scale_to_inf_bits, sizeof(scale_to_inf_val));
+  std::memcpy(
+      &scale_to_zero_val, &scale_to_zero_bits, sizeof(scale_to_zero_val));
+  const float scale_to_inf = scale_to_inf_val;
+  const float scale_to_zero = scale_to_zero_val;
+
+#if defined(_MSC_VER) && _MSC_VER == 1916
+  float base = ((signbit(f) != 0 ? -f : f) * scale_to_inf) * scale_to_zero;
+#else
+  float base = (fabsf(f) * scale_to_inf) * scale_to_zero;
+#endif
+
+  const uint32_t w = fp32_to_bits(f);
+  const uint32_t shl1_w = w + w;
+  const uint32_t sign = w & UINT32_C(0x80000000);
+  uint32_t bias = shl1_w & UINT32_C(0xFF000000);
+  if (bias < UINT32_C(0x71000000)) {
+    bias = UINT32_C(0x71000000);
+  }
+
+  base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
+  const uint32_t bits = fp32_to_bits(base);
+  const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
+  const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
+  const uint32_t nonsign = exp_bits + mantissa_bits;
+  return static_cast<uint16_t>(
+      (sign >> 16) |
+      (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign));
+#endif // C10_X86_F16
+}
+
+/*
+ * Convert a 16-bit floating-point number in IEEE half-precision format, in bit
+ * representation, to a 32-bit floating-point number in IEEE single-precision
+ * format, in bit representation.
+ *
+ * @note The implementation doesn't use any floating-point operations.
+ */
+inline uint32_t fp16_ieee_to_fp32_bits(uint16_t h) {
+  /*
+   * Extend the half-precision floating-point number to 32 bits and shift to the
+   * upper part of the 32-bit word:
+   *      +---+-----+------------+-------------------+
+   *      | S |EEEEE|MM MMMM MMMM|0000 0000 0000 0000|
+   *      +---+-----+------------+-------------------+
+   * Bits  31  26-30    16-25            0-15
+   *
+   * S - sign bit, E - bits of the biased exponent, M - bits of the mantissa, 0
+   * - zero bits.
+   */
+  const uint32_t w = (uint32_t)h << 16;
+  /*
+   * Extract the sign of the input number into the high bit of the 32-bit word:
+   *
+   *      +---+----------------------------------+
+   *      | S |0000000 00000000 00000000 00000000|
+   *      +---+----------------------------------+
+   * Bits  31                 0-31
+   */
+  const uint32_t sign = w & UINT32_C(0x80000000);
+  /*
+   * Extract mantissa and biased exponent of the input number into the bits 0-30
+   * of the 32-bit word:
+   *
+   *      +---+-----+------------+-------------------+
+   *      | 0 |EEEEE|MM MMMM MMMM|0000 0000 0000 0000|
+   *      +---+-----+------------+-------------------+
+   * Bits  30  27-31     17-26            0-16
+   */
+  const uint32_t nonsign = w & UINT32_C(0x7FFFFFFF);
+  /*
+   * Renorm shift is the number of bits to shift mantissa left to make the
+   * half-precision number normalized. If the initial number is normalized, some
+   * of its high 6 bits (sign == 0 and 5-bit exponent) equals one. In this case
+   * renorm_shift == 0. If the number is denormalize, renorm_shift > 0. Note
+   * that if we shift denormalized nonsign by renorm_shift, the unit bit of
+   * mantissa will shift into exponent, turning the biased exponent into 1, and
+   * making mantissa normalized (i.e. without leading 1).
+   */
+#ifdef _MSC_VER
+  unsigned long nonsign_bsr;
+  _BitScanReverse(&nonsign_bsr, (unsigned long)nonsign);
+  uint32_t renorm_shift = (uint32_t)nonsign_bsr ^ 31;
+#else
+  uint32_t renorm_shift = __builtin_clz(nonsign);
+#endif
+  renorm_shift = renorm_shift > 5 ? renorm_shift - 5 : 0;
+  /*
+   * Iff half-precision number has exponent of 15, the addition overflows
+   * it into bit 31, and the subsequent shift turns the high 9 bits
+   * into 1. Thus inf_nan_mask == 0x7F800000 if the half-precision number
+   * had exponent of 15 (i.e. was NaN or infinity) 0x00000000 otherwise
+   */
+  const int32_t inf_nan_mask =
+      ((int32_t)(nonsign + 0x04000000) >> 8) & INT32_C(0x7F800000);
+  /*
+   * Iff nonsign is 0, it overflows into 0xFFFFFFFF, turning bit 31
+   * into 1. Otherwise, bit 31 remains 0. The signed shift right by 31
+   * broadcasts bit 31 into all bits of the zero_mask. Thus zero_mask ==
+   * 0xFFFFFFFF if the half-precision number was zero (+0.0h or -0.0h)
+   * 0x00000000 otherwise
+   */
+  const int32_t zero_mask = (int32_t)(nonsign - 1) >> 31;
+  /*
+   * 1. Shift nonsign left by renorm_shift to normalize it (if the input
+   * was denormal)
+   * 2. Shift nonsign right by 3 so the exponent (5 bits originally)
+   * becomes an 8-bit field and 10-bit mantissa shifts into the 10 high
+   * bits of the 23-bit mantissa of IEEE single-precision number.
+   * 3. Add 0x70 to the exponent (starting at bit 23) to compensate the
+   * different in exponent bias (0x7F for single-precision number less 0xF
+   * for half-precision number).
+   * 4. Subtract renorm_shift from the exponent (starting at bit 23) to
+   * account for renormalization. As renorm_shift is less than 0x70, this
+   * can be combined with step 3.
+   * 5. Binary OR with inf_nan_mask to turn the exponent into 0xFF if the
+   * input was NaN or infinity.
+   * 6. Binary ANDNOT with zero_mask to turn the mantissa and exponent
+   * into zero if the input was zero.
+   * 7. Combine with the sign of the input number.
+   */
+  return sign |
+      ((((nonsign << renorm_shift >> 3) + ((0x70 - renorm_shift) << 23)) |
+        inf_nan_mask) &
+       ~zero_mask);
+}
+
+#ifdef C10_X86_F16
+#undef C10_X86_F16
+#endif // C10_X86_F16
+
+#if defined(__aarch64__) && !defined(__CUDACC__)
+inline float16_t fp16_from_bits(uint16_t h) {
+  return c10::bit_cast<float16_t>(h);
+}
+
+inline uint16_t fp16_to_bits(float16_t f) {
+  return c10::bit_cast<uint16_t>(f);
+}
+
+// According to https://godbolt.org/z/frExdbsWG it would translate to single
+// fcvt s0, h0
+inline float native_fp16_to_fp32_value(uint16_t h) {
+  return static_cast<float>(fp16_from_bits(h));
+}
+
+inline uint16_t native_fp16_from_fp32_value(float f) {
+  return fp16_to_bits(static_cast<float16_t>(f));
+}
+#endif
+
+} // namespace detail
+
+//---------- below is copied from c10/util/Half-inl.h ----------------//
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+#if defined(__aarch64__) && !defined(__CUDACC__)
+/// Constructors
+inline Half::Half(float16_t value) : x(detail::fp16_to_bits(value)) {}
+inline Half::operator float16_t() const {
+  return detail::fp16_from_bits(x);
+}
+#else
+
+inline C10_HOST_DEVICE Half::Half(float value)
+    :
+#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
+      x(__half_as_short(__float2half(value)))
+#elif defined(__SYCL_DEVICE_ONLY__)
+      x(c10::bit_cast<uint16_t>(sycl::half(value)))
+#elif (defined(CPU_CAPABILITY_AVX2) || defined(CPU_CAPABILITY_AVX512)) && \
+    !defined(__APPLE__)
+      x(at::vec::float2half_scalar(value))
+#else
+      x(detail::fp16_ieee_from_fp32_value(value))
+#endif
+{
+}
+
+/// Implicit conversions
+
+inline C10_HOST_DEVICE Half::operator float() const {
+#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
+  return __half2float(*reinterpret_cast<const __half*>(&x));
+#elif defined(__SYCL_DEVICE_ONLY__)
+  return float(c10::bit_cast<sycl::half>(x));
+#elif (defined(CPU_CAPABILITY_AVX2) || defined(CPU_CAPABILITY_AVX512)) && \
+    !defined(__APPLE__)
+  return at::vec::half2float_scalar(x);
+#elif defined(__aarch64__) && !defined(__CUDACC__)
+  return detail::native_fp16_to_fp32_value(x);
+#else
+  return detail::fp16_ieee_to_fp32_value(x);
+#endif
+}
+
+#endif /* !defined(__aarch64__) || defined(__CUDACC__) \
+        */
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+inline C10_HOST_DEVICE Half::Half(const __half& value) {
+  x = *reinterpret_cast<const unsigned short*>(&value);
+}
+inline C10_HOST_DEVICE Half::operator __half() const {
+  return *reinterpret_cast<const __half*>(&x);
+}
+#endif
+
+#ifdef SYCL_LANGUAGE_VERSION
+inline C10_HOST_DEVICE Half::Half(const sycl::half& value) {
+  x = *reinterpret_cast<const unsigned short*>(&value);
+}
+inline C10_HOST_DEVICE Half::operator sycl::half() const {
+  return *reinterpret_cast<const sycl::half*>(&x);
+}
+#endif
+
+// CUDA intrinsics
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 350)) || \
+    (defined(__clang__) && defined(__CUDA__))
+inline __device__ Half __ldg(const Half* ptr) {
+  return __ldg(reinterpret_cast<const __half*>(ptr));
+}
+#endif
+
+/// Arithmetic
+
+inline C10_HOST_DEVICE Half operator+(const Half& a, const Half& b) {
+  return static_cast<float>(a) + static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Half operator-(const Half& a, const Half& b) {
+  return static_cast<float>(a) - static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Half operator*(const Half& a, const Half& b) {
+  return static_cast<float>(a) * static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Half operator/(const Half& a, const Half& b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE Half operator-(const Half& a) {
+#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530) || \
+    defined(__HIP_DEVICE_COMPILE__)
+  return __hneg(a);
+#elif defined(__SYCL_DEVICE_ONLY__)
+  return -c10::bit_cast<sycl::half>(a);
+#else
+  return -static_cast<float>(a);
+#endif
+}
+
+inline C10_HOST_DEVICE Half& operator+=(Half& a, const Half& b) {
+  a = a + b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Half& operator-=(Half& a, const Half& b) {
+  a = a - b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Half& operator*=(Half& a, const Half& b) {
+  a = a * b;
+  return a;
+}
+
+inline C10_HOST_DEVICE Half& operator/=(Half& a, const Half& b) {
+  a = a / b;
+  return a;
+}
+
+/// Arithmetic with floats
+
+inline C10_HOST_DEVICE float operator+(Half a, float b) {
+  return static_cast<float>(a) + b;
+}
+inline C10_HOST_DEVICE float operator-(Half a, float b) {
+  return static_cast<float>(a) - b;
+}
+inline C10_HOST_DEVICE float operator*(Half a, float b) {
+  return static_cast<float>(a) * b;
+}
+inline C10_HOST_DEVICE float operator/(Half a, float b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<float>(a) / b;
+}
+
+inline C10_HOST_DEVICE float operator+(float a, Half b) {
+  return a + static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator-(float a, Half b) {
+  return a - static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator*(float a, Half b) {
+  return a * static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float operator/(float a, Half b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<float>(b);
+}
+
+inline C10_HOST_DEVICE float& operator+=(float& a, const Half& b) {
+  return a += static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator-=(float& a, const Half& b) {
+  return a -= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator*=(float& a, const Half& b) {
+  return a *= static_cast<float>(b);
+}
+inline C10_HOST_DEVICE float& operator/=(float& a, const Half& b) {
+  return a /= static_cast<float>(b);
+}
+
+/// Arithmetic with doubles
+
+inline C10_HOST_DEVICE double operator+(Half a, double b) {
+  return static_cast<double>(a) + b;
+}
+inline C10_HOST_DEVICE double operator-(Half a, double b) {
+  return static_cast<double>(a) - b;
+}
+inline C10_HOST_DEVICE double operator*(Half a, double b) {
+  return static_cast<double>(a) * b;
+}
+inline C10_HOST_DEVICE double operator/(Half a, double b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return static_cast<double>(a) / b;
+}
+
+inline C10_HOST_DEVICE double operator+(double a, Half b) {
+  return a + static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator-(double a, Half b) {
+  return a - static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator*(double a, Half b) {
+  return a * static_cast<double>(b);
+}
+inline C10_HOST_DEVICE double operator/(double a, Half b)
+    __ubsan_ignore_float_divide_by_zero__ {
+  return a / static_cast<double>(b);
+}
+
+/// Arithmetic with ints
+
+inline C10_HOST_DEVICE Half operator+(Half a, int b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a + static_cast<Half>(b);
+}
+inline C10_HOST_DEVICE Half operator-(Half a, int b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a - static_cast<Half>(b);
+}
+inline C10_HOST_DEVICE Half operator*(Half a, int b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a * static_cast<Half>(b);
+}
+inline C10_HOST_DEVICE Half operator/(Half a, int b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a / static_cast<Half>(b);
+}
+
+inline C10_HOST_DEVICE Half operator+(int a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) + b;
+}
+inline C10_HOST_DEVICE Half operator-(int a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) - b;
+}
+inline C10_HOST_DEVICE Half operator*(int a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) * b;
+}
+inline C10_HOST_DEVICE Half operator/(int a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) / b;
+}
+
+//// Arithmetic with int64_t
+
+inline C10_HOST_DEVICE Half operator+(Half a, int64_t b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a + static_cast<Half>(b);
+}
+inline C10_HOST_DEVICE Half operator-(Half a, int64_t b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a - static_cast<Half>(b);
+}
+inline C10_HOST_DEVICE Half operator*(Half a, int64_t b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a * static_cast<Half>(b);
+}
+inline C10_HOST_DEVICE Half operator/(Half a, int64_t b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return a / static_cast<Half>(b);
+}
+
+inline C10_HOST_DEVICE Half operator+(int64_t a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) + b;
+}
+inline C10_HOST_DEVICE Half operator-(int64_t a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) - b;
+}
+inline C10_HOST_DEVICE Half operator*(int64_t a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) * b;
+}
+inline C10_HOST_DEVICE Half operator/(int64_t a, Half b) {
+  // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
+  return static_cast<Half>(a) / b;
+}
+
+/// NOTE: we do not define comparisons directly and instead rely on the implicit
+/// conversion from c10::Half to float.
+
+C10_CLANG_DIAGNOSTIC_POP()
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+
+using c10::Half;
+using c10::operator+;
+using c10::operator-;
+using c10::operator*;
+using c10::operator/;
+using c10::operator+=;
+using c10::operator-=;
+using c10::operator*=;
+using c10::operator/=;
+using c10::operator<<;
+
+namespace detail {
+#if defined(__aarch64__) && !defined(__CUDACC__)
+using c10::detail::fp16_from_bits;
+using c10::detail::fp16_to_bits;
+using c10::detail::native_fp16_from_fp32_value;
+using c10::detail::native_fp16_to_fp32_value;
+#endif
+
+using c10::detail::fp16_ieee_from_fp32_value;
+using c10::detail::fp16_ieee_to_fp32_bits;
+using c10::detail::fp16_ieee_to_fp32_value;
+} // namespace detail
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace std {
+
+template <>
+class numeric_limits<c10::Half> {
+ public:
+  static constexpr bool is_specialized = true;
+  static constexpr bool is_signed = true;
+  static constexpr bool is_integer = false;
+  static constexpr bool is_exact = false;
+  static constexpr bool has_infinity = true;
+  static constexpr bool has_quiet_NaN = true;
+  static constexpr bool has_signaling_NaN = true;
+  static constexpr auto has_denorm = numeric_limits<float>::has_denorm;
+  static constexpr auto has_denorm_loss =
+      numeric_limits<float>::has_denorm_loss;
+  static constexpr auto round_style = numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = true;
+  static constexpr bool is_bounded = true;
+  static constexpr bool is_modulo = false;
+  static constexpr int digits = 11;
+  static constexpr int digits10 = 3;
+  static constexpr int max_digits10 = 5;
+  static constexpr int radix = 2;
+  static constexpr int min_exponent = -13;
+  static constexpr int min_exponent10 = -4;
+  static constexpr int max_exponent = 16;
+  static constexpr int max_exponent10 = 4;
+  static constexpr auto traps = numeric_limits<float>::traps;
+  static constexpr auto tinyness_before =
+      numeric_limits<float>::tinyness_before;
+  static constexpr c10::Half min() {
+    return c10::Half(0x0400, c10::Half::from_bits());
+  }
+  static constexpr c10::Half lowest() {
+    return c10::Half(0xFBFF, c10::Half::from_bits());
+  }
+  static constexpr c10::Half max() {
+    return c10::Half(0x7BFF, c10::Half::from_bits());
+  }
+  static constexpr c10::Half epsilon() {
+    return c10::Half(0x1400, c10::Half::from_bits());
+  }
+  static constexpr c10::Half round_error() {
+    return c10::Half(0x3800, c10::Half::from_bits());
+  }
+  static constexpr c10::Half infinity() {
+    return c10::Half(0x7C00, c10::Half::from_bits());
+  }
+  static constexpr c10::Half quiet_NaN() {
+    return c10::Half(0x7E00, c10::Half::from_bits());
+  }
+  static constexpr c10::Half signaling_NaN() {
+    return c10::Half(0x7D00, c10::Half::from_bits());
+  }
+  static constexpr c10::Half denorm_min() {
+    return c10::Half(0x0001, c10::Half::from_bits());
+  }
+};
+
+} // namespace std
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/HeaderOnlyArrayRef.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/HeaderOnlyArrayRef.h
new file mode 100644
index 0000000000000000000000000000000000000000..751ffef32bb1da3a00f3735f9f7200d805f3f9d2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/HeaderOnlyArrayRef.h
@@ -0,0 +1,248 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include <algorithm>
+#include <array>
+#include <cstddef>
+#include <functional>
+#include <initializer_list>
+#include <iterator>
+#include <type_traits>
+#include <vector>
+
+namespace c10 {
+
+/// HeaderOnlyArrayRef - A subset of ArrayRef that is implemented only
+/// in headers. This will be a base class from which ArrayRef inherits, so that
+/// we can keep much of the implementation shared.
+///
+/// [HeaderOnlyArrayRef vs ArrayRef note]
+/// As HeaderOnlyArrayRef is a subset of ArrayRef, it has slightly less
+/// functionality than ArrayRef. We document the minor differences below:
+/// 1. ArrayRef has an extra convenience constructor for SmallVector.
+/// 2. ArrayRef uses TORCH_CHECK. HeaderOnlyArrayRef uses header-only
+///    STD_TORCH_CHECK, which will output a std::runtime_error vs a
+///    c10::Error. Consequently, you should use ArrayRef when possible
+///    and HeaderOnlyArrayRef only when necessary to support headeronly code.
+/// In all other aspects, HeaderOnlyArrayRef is identical to ArrayRef, with the
+/// positive benefit of being header-only and thus independent of libtorch.so.
+template <typename T>
+class HeaderOnlyArrayRef {
+ public:
+  using iterator = const T*;
+  using const_iterator = const T*;
+  using size_type = size_t;
+  using value_type = T;
+
+  using reverse_iterator = std::reverse_iterator<iterator>;
+
+ protected:
+  /// The start of the array, in an external buffer.
+  const T* Data;
+
+  /// The number of elements.
+  size_type Length;
+
+ public:
+  /// @name Constructors
+  /// @{
+
+  /// Construct an empty HeaderOnlyArrayRef.
+  /* implicit */ constexpr HeaderOnlyArrayRef() : Data(nullptr), Length(0) {}
+
+  /// Construct a HeaderOnlyArrayRef from a single element.
+  // TODO Make this explicit
+  constexpr HeaderOnlyArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {}
+
+  /// Construct a HeaderOnlyArrayRef from a pointer and length.
+  constexpr HeaderOnlyArrayRef(const T* data, size_t length)
+      : Data(data), Length(length) {}
+
+  /// Construct a HeaderOnlyArrayRef from a range.
+  constexpr HeaderOnlyArrayRef(const T* begin, const T* end)
+      : Data(begin), Length(end - begin) {}
+
+  template <
+      typename Container,
+      typename U = decltype(std::declval<Container>().data()),
+      typename = std::enable_if_t<
+          (std::is_same_v<U, T*> || std::is_same_v<U, T const*>)>>
+  /* implicit */ HeaderOnlyArrayRef(const Container& container)
+      : Data(container.data()), Length(container.size()) {}
+
+  /// Construct a HeaderOnlyArrayRef from a std::vector.
+  // The enable_if stuff here makes sure that this isn't used for
+  // std::vector<bool>, because ArrayRef can't work on a std::vector<bool>
+  // bitfield.
+  template <typename A>
+  /* implicit */ HeaderOnlyArrayRef(const std::vector<T, A>& Vec)
+      : Data(Vec.data()), Length(Vec.size()) {
+    static_assert(
+        !std::is_same_v<T, bool>,
+        "HeaderOnlyArrayRef<bool> cannot be constructed from a std::vector<bool> bitfield.");
+  }
+
+  /// Construct a HeaderOnlyArrayRef from a std::array
+  template <size_t N>
+  /* implicit */ constexpr HeaderOnlyArrayRef(const std::array<T, N>& Arr)
+      : Data(Arr.data()), Length(N) {}
+
+  /// Construct a HeaderOnlyArrayRef from a C array.
+  template <size_t N>
+  // NOLINTNEXTLINE(*c-arrays*)
+  /* implicit */ constexpr HeaderOnlyArrayRef(const T (&Arr)[N])
+      : Data(Arr), Length(N) {}
+
+  /// Construct a HeaderOnlyArrayRef from a std::initializer_list.
+  /* implicit */ constexpr HeaderOnlyArrayRef(
+      const std::initializer_list<T>& Vec)
+      : Data(
+            std::begin(Vec) == std::end(Vec) ? static_cast<T*>(nullptr)
+                                             : std::begin(Vec)),
+        Length(Vec.size()) {}
+
+  /// @}
+  /// @name Simple Operations
+  /// @{
+
+  constexpr iterator begin() const {
+    return this->Data;
+  }
+  constexpr iterator end() const {
+    return this->Data + this->Length;
+  }
+
+  // These are actually the same as iterator, since ArrayRef only
+  // gives you const iterators.
+  constexpr const_iterator cbegin() const {
+    return this->Data;
+  }
+  constexpr const_iterator cend() const {
+    return this->Data + this->Length;
+  }
+
+  constexpr reverse_iterator rbegin() const {
+    return reverse_iterator(end());
+  }
+  constexpr reverse_iterator rend() const {
+    return reverse_iterator(begin());
+  }
+
+  /// Check if all elements in the array satisfy the given expression
+  constexpr bool allMatch(const std::function<bool(const T&)>& pred) const {
+    return std::all_of(cbegin(), cend(), pred);
+  }
+
+  /// empty - Check if the array is empty.
+  constexpr bool empty() const {
+    return this->Length == 0;
+  }
+
+  constexpr const T* data() const {
+    return this->Data;
+  }
+
+  /// size - Get the array size.
+  constexpr size_t size() const {
+    return this->Length;
+  }
+
+  /// front - Get the first element.
+  constexpr const T& front() const {
+    STD_TORCH_CHECK(
+        !this->empty(),
+        "HeaderOnlyArrayRef: attempted to access front() of empty list");
+    return this->Data[0];
+  }
+
+  /// back - Get the last element.
+  constexpr const T& back() const {
+    STD_TORCH_CHECK(
+        !this->empty(),
+        "HeaderOnlyArrayRef: attempted to access back() of empty list");
+    return this->Data[this->Length - 1];
+  }
+
+  /// equals - Check for element-wise equality.
+  constexpr bool equals(HeaderOnlyArrayRef RHS) const {
+    return this->Length == RHS.Length &&
+        std::equal(begin(), end(), RHS.begin());
+  }
+
+  /// slice(n, m) - Take M elements of the array starting at element N
+  constexpr HeaderOnlyArrayRef<T> slice(size_t N, size_t M) const {
+    STD_TORCH_CHECK(
+        N + M <= this->size(),
+        "HeaderOnlyArrayRef: invalid slice, N = ",
+        N,
+        "; M = ",
+        M,
+        "; size = ",
+        this->size());
+    return HeaderOnlyArrayRef<T>(this->data() + N, M);
+  }
+
+  /// slice(n) - Chop off the first N elements of the array.
+  constexpr HeaderOnlyArrayRef<T> slice(size_t N) const {
+    STD_TORCH_CHECK(
+        N <= this->size(),
+        "HeaderOnlyArrayRef: invalid slice, N = ",
+        N,
+        "; size = ",
+        this->size());
+    return slice(N, this->size() - N);
+  }
+
+  /// @}
+  /// @name Operator Overloads
+  /// @{
+  constexpr const T& operator[](size_t Index) const {
+    return this->Data[Index];
+  }
+
+  /// Vector compatibility
+  constexpr const T& at(size_t Index) const {
+    STD_TORCH_CHECK(
+        Index < this->Length,
+        "HeaderOnlyArrayRef: invalid index Index = ",
+        Index,
+        "; Length = ",
+        this->Length);
+    return this->Data[Index];
+  }
+
+  /// Disallow accidental assignment from a temporary.
+  ///
+  /// The declaration here is extra complicated so that "arrayRef = {}"
+  /// continues to select the move assignment operator.
+  template <typename U>
+  std::enable_if_t<std::is_same_v<U, T>, HeaderOnlyArrayRef<T>>& operator=(
+      // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
+      U&& Temporary) = delete;
+
+  /// Disallow accidental assignment from a temporary.
+  ///
+  /// The declaration here is extra complicated so that "arrayRef = {}"
+  /// continues to select the move assignment operator.
+  template <typename U>
+  std::enable_if_t<std::is_same_v<U, T>, HeaderOnlyArrayRef<T>>& operator=(
+      std::initializer_list<U>) = delete;
+
+  /// @}
+  /// @name Expensive Operations
+  /// @{
+  std::vector<T> vec() const {
+    return std::vector<T>(this->Data, this->Data + this->Length);
+  }
+
+  /// @}
+};
+
+} // namespace c10
+
+namespace torch::headeronly {
+using c10::HeaderOnlyArrayRef;
+using IntHeaderOnlyArrayRef = HeaderOnlyArrayRef<int64_t>;
+} // namespace torch::headeronly
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Metaprogramming.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Metaprogramming.h
new file mode 100644
index 0000000000000000000000000000000000000000..2589f338d35dbcf569e15966a788dd16130c0681
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/Metaprogramming.h
@@ -0,0 +1,237 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/TypeList.h>
+#include <type_traits>
+
+namespace c10::guts {
+
+/**
+ * Access information about result type or arguments from a function type.
+ * Example:
+ * using A = function_traits<int (float, double)>::return_type // A == int
+ * using A = function_traits<int (float, double)>::parameter_types::tuple_type
+ * // A == tuple<float, double>
+ */
+template <class Func>
+struct function_traits {
+  static_assert(
+      !std::is_same_v<Func, Func>,
+      "In function_traits<Func>, Func must be a plain function type.");
+};
+template <class Result, class... Args>
+struct function_traits<Result(Args...)> {
+  using func_type = Result(Args...);
+  using return_type = Result;
+  using parameter_types = typelist::typelist<Args...>;
+  static constexpr auto number_of_parameters = sizeof...(Args);
+};
+
+/**
+ * infer_function_traits: creates a `function_traits` type for a simple
+ * function (pointer) or functor (lambda/struct). Currently does not support
+ * class methods.
+ */
+
+template <typename Functor>
+struct infer_function_traits {
+  using type = function_traits<
+      c10::guts::detail::strip_class_t<decltype(&Functor::operator())>>;
+};
+
+template <typename Result, typename... Args>
+struct infer_function_traits<Result (*)(Args...)> {
+  using type = function_traits<Result(Args...)>;
+};
+
+template <typename Result, typename... Args>
+struct infer_function_traits<Result(Args...)> {
+  using type = function_traits<Result(Args...)>;
+};
+
+template <typename T>
+using infer_function_traits_t = typename infer_function_traits<T>::type;
+
+/**
+ * make_function_traits: creates a `function_traits` type given a Return type
+ * and a typelist of Argument types
+ *
+ * Example:
+ * bool f(int, int);
+ *
+ * infer_function_traits_t<f> == make_function_traits_t<bool,
+ * typelist::typelist<int, int>>
+ */
+template <typename Result, typename ArgList>
+struct make_function_traits {
+  static_assert(
+      false_t<ArgList>::value,
+      "In guts::make_function_traits<Result, TypeList>, the ArgList argument must be typelist<...>.");
+};
+
+template <typename Result, typename... Args>
+struct make_function_traits<Result, typelist::typelist<Args...>> {
+  using type = function_traits<Result(Args...)>;
+};
+
+template <typename Result, typename ArgList>
+using make_function_traits_t =
+    typename make_function_traits<Result, ArgList>::type;
+
+/**
+ * make_offset_index_sequence<Start, N>
+ * Like make_index_sequence<N>, but starting from Start instead of 0.
+ *
+ * Example:
+ *  make_offset_index_sequence<10, 3> == std::index_sequence<10, 11, 12>
+ */
+template <size_t Start, size_t N, size_t... Is>
+struct make_offset_index_sequence_impl
+    : make_offset_index_sequence_impl<Start, N - 1, Start + N - 1, Is...> {
+  static_assert(
+      static_cast<int>(Start) >= 0,
+      "make_offset_index_sequence: Start < 0");
+  static_assert(static_cast<int>(N) >= 0, "make_offset_index_sequence: N < 0");
+};
+
+template <size_t Start, size_t... Is>
+struct make_offset_index_sequence_impl<Start, 0, Is...> {
+  typedef std::index_sequence<Is...> type;
+};
+
+template <size_t Start, size_t N>
+using make_offset_index_sequence =
+    typename make_offset_index_sequence_impl<Start, N>::type;
+
+/**
+ * Use tuple_elements to extract a position-indexed subset of elements
+ * from the argument tuple into a result tuple.
+ *
+ * Example:
+ *  std::tuple<int, const char*, double> t = std::make_tuple(0, "HEY", 2.0);
+ *  std::tuple<int, double> result = tuple_elements(t, std::index_sequence<0,
+ * 2>());
+ */
+template <class Tuple, size_t... Is>
+constexpr auto tuple_elements(Tuple t, std::index_sequence<Is...> /*unused*/) {
+  return std::tuple<std::tuple_element_t<Is, Tuple>...>(std::get<Is>(t)...);
+}
+
+/**
+ * Use tuple_take to extract the first or last n elements from the argument
+ * tuple into a result tuple.
+ *
+ * Example:
+ *  std::tuple<int, const char*, double> t = std::make_tuple(0, "HEY", 2.0);
+ *  std::tuple<int, const char*> first_two = tuple_take<decltype(t), 2>(t);
+ *  std::tuple<const char*, double> last_two = tuple_take<decltype(t), -2>(t);
+ */
+template <class Tuple, int N, class Enable = void>
+struct TupleTake {};
+
+template <class Tuple, int N>
+struct TupleTake<Tuple, N, std::enable_if_t<N >= 0, void>> {
+  static auto call(Tuple t) {
+    constexpr size_t size = std::tuple_size<Tuple>();
+    static_assert(N <= size, "tuple_take: N > size");
+    return tuple_elements(t, std::make_index_sequence<N>{});
+  }
+};
+
+template <class Tuple, int N>
+    struct TupleTake < Tuple,
+    N, std::enable_if_t<N<0, void>> {
+  static auto call(Tuple t) {
+    constexpr size_t size = std::tuple_size<Tuple>();
+    static_assert(-N <= size, "tuple_take: -N > size");
+    return tuple_elements(t, make_offset_index_sequence<size + N, -N>{});
+  }
+};
+
+template <class Tuple, int N>
+auto tuple_take(Tuple t) {
+  return TupleTake<Tuple, N>::call(t);
+}
+
+/**
+ * Use tuple_slice to extract a contiguous subtuple from the argument.
+ *
+ * Example:
+ *  std::tuple<int, const char*, double, bool> t = std::make_tuple(0,
+ * "HEY", 2.0, false); std::tuple<int, const char*> middle_two =
+ * tuple_slice<decltype(t), 1, 2>(t);
+ */
+template <class Tuple, size_t Start, size_t N>
+constexpr auto tuple_slice(Tuple t) {
+  constexpr size_t size = std::tuple_size<Tuple>();
+  static_assert(Start + N <= size, "tuple_slice: Start + N > size");
+  return tuple_elements(t, make_offset_index_sequence<Start, N>{});
+}
+
+/**
+ * Use tuple_map to run a mapping function over a tuple to get a new tuple.
+ *
+ * Example 1:
+ *   auto result = tuple_map(std::tuple<int32_t, int32_t, int32_t>(3, 4, 5), []
+ * (int32_t a) -> int16_t {return a+1;});
+ *   // result == std::tuple<int16_t, int16_t, int16_t>(4, 5, 6)
+ *
+ * Example 2:
+ *   struct Mapper {
+ *     std::string operator()(int32_t a) const {
+ *       return std::to_string(a);
+ *     }
+ *     int64_t operator()(const std::string& a) const {
+ *        return atoi(a.c_str());
+ *     }
+ *   };
+ *   auto result = tuple_map(std::tuple<int32_t, std::string>(3, "4"),
+ * Mapper());
+ *   // result == std::tuple<std::string, int64_t>("3", 4)
+ *
+ * Example 3:
+ *   struct A final {
+ *    int32_t func() {
+ *      return 5;
+ *    }
+ *  };
+ *  struct B final {
+ *    std::string func() {
+ *      return "5";
+ *    }
+ *  };
+ *  auto result = tuple_map(std::make_tuple(A(), B()), [] (auto a) { return
+ * a.func(); });
+ *  // result == std::tuple<int32_t, std::string>(5, "5");
+ */
+namespace detail {
+template <class Mapper, class... Args, size_t... Indices>
+auto tuple_map(
+    // NOLINTNEXTLINE(cppcoreguidelines-rvalue-reference-param-not-moved)
+    std::tuple<Args...>&& tuple,
+    const Mapper& mapper,
+    std::index_sequence<Indices...> /*unused*/) {
+  return std::tuple<decltype(mapper(std::forward<Args>(std::get<Indices>(
+      tuple))))...>(mapper(std::forward<Args>(std::get<Indices>(tuple)))...);
+}
+} // namespace detail
+
+template <class Mapper, class... Args>
+auto tuple_map(std::tuple<Args...>&& tuple, const Mapper& mapper) {
+  return detail::tuple_map(
+      std::move(tuple), mapper, std::index_sequence_for<Args...>());
+}
+
+} // namespace c10::guts
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly, guts)
+
+using c10::guts::function_traits;
+using c10::guts::infer_function_traits_t;
+using c10::guts::make_function_traits_t;
+using c10::guts::tuple_elements;
+using c10::guts::tuple_map;
+using c10::guts::tuple_slice;
+using c10::guts::tuple_take;
+
+HIDDEN_NAMESPACE_END(torch, headeronly, guts)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeList.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeList.h
new file mode 100644
index 0000000000000000000000000000000000000000..cd81f0cc1dcf97e0444ace259be265693c2935ec
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeList.h
@@ -0,0 +1,548 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/TypeTraits.h>
+#include <algorithm>
+#include <cstddef>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+
+namespace c10::guts {
+
+template <class... T>
+struct false_t : std::false_type {};
+template <template <class> class... T>
+struct false_higher_t : std::false_type {};
+
+namespace typelist {
+
+/**
+ * Type holding a list of types for compile time type computations
+ */
+template <class... Items>
+struct typelist final {
+ public:
+  typelist() = delete; // not for instantiation
+};
+
+/**
+ * Returns the number of types in a typelist
+ * Example:
+ *   3  ==  size<typelist<int, int, double>>::value
+ */
+template <class TypeList>
+struct size final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::size<T>, T must be typelist<...>.");
+};
+template <class... Types>
+struct size<typelist<Types...>> final {
+  static constexpr size_t value = sizeof...(Types);
+};
+
+/**
+ * Transforms a list of types into a tuple holding these types.
+ * Example:
+ *   std::tuple<int, string>  ==  to_tuple_t<typelist<int, string>>
+ */
+template <class TypeList>
+struct to_tuple final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::to_tuple<T>, T must be typelist<...>.");
+};
+template <class... Types>
+struct to_tuple<typelist<Types...>> final {
+  using type = std::tuple<Types...>;
+};
+template <class TypeList>
+using to_tuple_t = typename to_tuple<TypeList>::type;
+
+/**
+ * Creates a typelist containing the types of a given tuple.
+ * Example:
+ *   typelist<int, string>  ==  from_tuple_t<std::tuple<int, string>>
+ */
+template <class Tuple>
+struct from_tuple final {
+  static_assert(
+      false_t<Tuple>::value,
+      "In typelist::from_tuple<T>, T must be std::tuple<...>.");
+};
+template <class... Types>
+struct from_tuple<std::tuple<Types...>> final {
+  using type = typelist<Types...>;
+};
+template <class Tuple>
+using from_tuple_t = typename from_tuple<Tuple>::type;
+
+/**
+ * Concatenates multiple type lists.
+ * Example:
+ *   typelist<int, string, int>  ==  concat_t<typelist<int, string>,
+ * typelist<int>>
+ */
+template <class... TypeLists>
+struct concat final {
+  static_assert(
+      false_t<TypeLists...>::value,
+      "In typelist::concat<T1, ...>, the T arguments each must be typelist<...>.");
+};
+template <class... Head1Types, class... Head2Types, class... TailLists>
+struct concat<typelist<Head1Types...>, typelist<Head2Types...>, TailLists...>
+    final {
+  using type =
+      typename concat<typelist<Head1Types..., Head2Types...>, TailLists...>::
+          type;
+};
+template <class... HeadTypes>
+struct concat<typelist<HeadTypes...>> final {
+  using type = typelist<HeadTypes...>;
+};
+template <>
+struct concat<> final {
+  using type = typelist<>;
+};
+template <class... TypeLists>
+using concat_t = typename concat<TypeLists...>::type;
+
+/**
+ * Filters the types in a type list by a type trait.
+ * Examples:
+ *   typelist<int&, const string&&>  ==  filter_t<std::is_reference,
+ * typelist<void, string, int&, bool, const string&&, int>>
+ */
+template <template <class> class Condition, class TypeList>
+struct filter final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::filter<Condition, TypeList>, the TypeList argument must be typelist<...>.");
+};
+template <template <class> class Condition, class Head, class... Tail>
+struct filter<Condition, typelist<Head, Tail...>> final {
+  static_assert(
+      is_type_condition<Condition>::value,
+      "In typelist::filter<Condition, TypeList>, the Condition argument must be a condition type trait, i.e. have a static constexpr bool ::value member.");
+  using type = std::conditional_t<
+      Condition<Head>::value,
+      concat_t<
+          typelist<Head>,
+          typename filter<Condition, typelist<Tail...>>::type>,
+      typename filter<Condition, typelist<Tail...>>::type>;
+};
+template <template <class> class Condition>
+struct filter<Condition, typelist<>> final {
+  static_assert(
+      is_type_condition<Condition>::value,
+      "In typelist::filter<Condition, TypeList>, the Condition argument must be a condition type trait, i.e. have a static constexpr bool ::value member.");
+  using type = typelist<>;
+};
+template <template <class> class Condition, class TypeList>
+using filter_t = typename filter<Condition, TypeList>::type;
+
+/**
+ * Counts how many types in the list fulfill a type trait
+ * Examples:
+ *   2  ==  count_if<std::is_reference, typelist<void, string, int&, bool, const
+ * string&&, int>>
+ */
+template <template <class> class Condition, class TypeList>
+struct count_if final {
+  static_assert(
+      is_type_condition<Condition>::value,
+      "In typelist::count_if<Condition, TypeList>, the Condition argument must be a condition type trait, i.e. have a static constexpr bool ::value member.");
+  static_assert(
+      is_instantiation_of<typelist, TypeList>::value,
+      "In typelist::count_if<Condition, TypeList>, the TypeList argument must be typelist<...>.");
+  // TODO Direct implementation might be faster
+  static constexpr size_t value = size<filter_t<Condition, TypeList>>::value;
+};
+
+/**
+ * Checks if a typelist contains a certain type.
+ * Examples:
+ *  contains<typelist<int, string>, string> == true_type
+ *  contains<typelist<int, string>, double> == false_type
+ */
+namespace detail {
+template <class TypeList, class Type, class Enable = void>
+struct contains {};
+template <class Type>
+struct contains<typelist<>, Type, void> : std::false_type {};
+template <class Type, class Head, class... Tail>
+struct contains<
+    typelist<Head, Tail...>,
+    Type,
+    std::enable_if_t<std::is_same_v<Head, Type>>> : std::true_type {};
+template <class Type, class Head, class... Tail>
+struct contains<
+    typelist<Head, Tail...>,
+    Type,
+    std::enable_if_t<!std::is_same_v<Head, Type>>>
+    : contains<typelist<Tail...>, Type> {};
+} // namespace detail
+template <class TypeList, class Type>
+using contains = typename detail::contains<TypeList, Type>::type;
+
+/**
+ * Returns true iff the type trait is true for all types in the type list
+ * Examples:
+ *   true   ==  all<std::is_reference, typelist<int&, const float&&, const
+ * MyClass&>>::value false  ==  all<std::is_reference, typelist<int&, const
+ * float&&, MyClass>>::value
+ */
+template <template <class> class Condition, class TypeList>
+struct all {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::all<Condition, TypeList>, the TypeList argument must be typelist<...>.");
+};
+template <template <class> class Condition, class... Types>
+struct all<Condition, typelist<Types...>>
+    : std::conjunction<Condition<Types>...> {
+  static_assert(
+      is_type_condition<Condition>::value,
+      "In typelist::all<Condition, TypeList>, the Condition argument must be a condition type trait, i.e. have a static constexpr bool ::value member.");
+};
+
+/**
+ * Returns true iff the type trait is true for any type in the type list
+ * Examples:
+ *   true   ==  true_for_any_type<std::is_reference, typelist<int, const
+ * float&&, const MyClass>>::value false  ==
+ * true_for_any_type<std::is_reference, typelist<int, const float,
+ * MyClass>>::value
+ */
+template <template <class> class Condition, class TypeList>
+struct true_for_any_type final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::true_for_any_type<Condition, TypeList>, the TypeList argument must be typelist<...>.");
+};
+template <template <class> class Condition, class... Types>
+struct true_for_any_type<Condition, typelist<Types...>> final
+    : std::disjunction<Condition<Types>...> {
+  static_assert(
+      is_type_condition<Condition>::value,
+      "In typelist::true_for_any_type<Condition, TypeList>, the Condition argument must be a condition type trait, i.e. have a static constexpr bool ::value member.");
+};
+
+/**
+ * Maps types of a type list using a type trait
+ * Example:
+ *  typelist<int&, double&, string&>  ==  map_t<std::add_lvalue_reference_t,
+ * typelist<int, double, string>>
+ */
+template <template <class> class Mapper, class TypeList>
+struct map final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::map<Mapper, TypeList>, the TypeList argument must be typelist<...>.");
+};
+template <template <class> class Mapper, class... Types>
+struct map<Mapper, typelist<Types...>> final {
+  using type = typelist<Mapper<Types>...>;
+};
+template <template <class> class Mapper, class TypeList>
+using map_t = typename map<Mapper, TypeList>::type;
+
+/**
+ * Returns the first element of a type list.
+ * Example:
+ *   int  ==  head_t<typelist<int, string>>
+ */
+template <class TypeList>
+struct head final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::head<T>, the T argument must be typelist<...>.");
+};
+template <class Head, class... Tail>
+struct head<typelist<Head, Tail...>> final {
+  using type = Head;
+};
+template <class TypeList>
+using head_t = typename head<TypeList>::type;
+
+/**
+ * Returns the first element of a type list, or the specified default if the
+ * type list is empty. Example: int  ==  head_t<bool, typelist<int, string>>
+ *   bool  ==  head_t<bool, typelist<>>
+ */
+template <class Default, class TypeList>
+struct head_with_default final {
+  using type = Default;
+};
+template <class Default, class Head, class... Tail>
+struct head_with_default<Default, typelist<Head, Tail...>> final {
+  using type = Head;
+};
+template <class Default, class TypeList>
+using head_with_default_t = typename head_with_default<Default, TypeList>::type;
+
+/**
+ * Returns the N-th element of a type list.
+ * Example:
+ * int == element_t<1, typelist<float, int, char>>
+ */
+
+/// Base template.
+template <size_t Index, class TypeList>
+struct element final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::element<T>, the T argument must be typelist<...>.");
+};
+
+/// Successful case, we have reached the zero index and can "return" the head
+/// type.
+template <class Head, class... Tail>
+struct element<0, typelist<Head, Tail...>> {
+  using type = Head;
+};
+
+/// Error case, we have an index but ran out of types! It will only be selected
+/// if `Ts...` is actually empty!
+template <size_t Index, class... Ts>
+struct element<Index, typelist<Ts...>> {
+  static_assert(
+      Index < sizeof...(Ts),
+      "Index is out of bounds in typelist::element");
+};
+
+/// Shave off types until we hit the <0, Head, Tail...> or <Index> case.
+template <size_t Index, class Head, class... Tail>
+struct element<Index, typelist<Head, Tail...>>
+    : element<Index - 1, typelist<Tail...>> {};
+
+/// Convenience alias.
+template <size_t Index, class TypeList>
+using element_t = typename element<Index, TypeList>::type;
+
+/**
+ * Returns the last element of a type list.
+ * Example:
+ *   int  ==  last_t<typelist<int, string>>
+ */
+template <class TypeList>
+struct last final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::last<T>, the T argument must be typelist<...>.");
+};
+template <class Head, class... Tail>
+struct last<typelist<Head, Tail...>> final {
+  using type = typename last<typelist<Tail...>>::type;
+};
+template <class Head>
+struct last<typelist<Head>> final {
+  using type = Head;
+};
+template <class TypeList>
+using last_t = typename last<TypeList>::type;
+static_assert(std::is_same_v<int, last_t<typelist<double, float, int>>>);
+
+/**
+ * Take/drop a number of arguments from a typelist.
+ * Example:
+ *   typelist<int, string> == take_t<typelist<int, string, bool>, 2>
+ *   typelist<bool> == drop_t<typelist<int, string, bool>, 2>
+ */
+namespace detail {
+template <class TypeList, size_t offset, class IndexSequence>
+struct take_elements final {};
+
+template <class TypeList, size_t offset, size_t... Indices>
+struct take_elements<TypeList, offset, std::index_sequence<Indices...>> final {
+  using type = typelist<typename element<offset + Indices, TypeList>::type...>;
+};
+} // namespace detail
+
+template <class TypeList, size_t num>
+struct take final {
+  static_assert(
+      is_instantiation_of<typelist, TypeList>::value,
+      "In typelist::take<T, num>, the T argument must be typelist<...>.");
+  static_assert(
+      num <= size<TypeList>::value,
+      "Tried to typelist::take more elements than there are in the list");
+  using type = typename detail::
+      take_elements<TypeList, 0, std::make_index_sequence<num>>::type;
+};
+template <class TypeList, size_t num>
+using take_t = typename take<TypeList, num>::type;
+
+template <class TypeList, size_t num>
+struct drop final {
+  static_assert(
+      is_instantiation_of<typelist, TypeList>::value,
+      "In typelist::drop<T, num>, the T argument must be typelist<...>.");
+  static_assert(
+      num <= size<TypeList>::value,
+      "Tried to typelist::drop more elements than there are in the list");
+  using type = typename detail::take_elements<
+      TypeList,
+      num,
+      std::make_index_sequence<size<TypeList>::value - num>>::type;
+};
+template <class TypeList, size_t num>
+using drop_t = typename drop<TypeList, num>::type;
+
+/**
+ * Like drop, but returns an empty list rather than an assertion error if `num`
+ * is larger than the size of the TypeList.
+ * Example:
+ *   typelist<> == drop_if_nonempty_t<typelist<string, bool>, 2>
+ *   typelist<> == drop_if_nonempty_t<typelist<int, string, bool>, 3>
+ */
+template <class TypeList, size_t num>
+struct drop_if_nonempty final {
+  static_assert(
+      is_instantiation_of<typelist, TypeList>::value,
+      "In typelist::drop<T, num>, the T argument must be typelist<...>.");
+  using type = typename detail::take_elements<
+      TypeList,
+      std::min(num, size<TypeList>::value),
+      std::make_index_sequence<
+          size<TypeList>::value - std::min(num, size<TypeList>::value)>>::type;
+};
+template <class TypeList, size_t num>
+using drop_if_nonempty_t = typename drop_if_nonempty<TypeList, num>::type;
+
+/**
+ * Reverses a typelist.
+ * Example:
+ *   typelist<int, string>  == reverse_t<typelist<string, int>>
+ */
+template <class TypeList>
+struct reverse final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::reverse<T>, the T argument must be typelist<...>.");
+};
+template <class Head, class... Tail>
+struct reverse<typelist<Head, Tail...>> final {
+  using type =
+      concat_t<typename reverse<typelist<Tail...>>::type, typelist<Head>>;
+};
+template <>
+struct reverse<typelist<>> final {
+  using type = typelist<>;
+};
+template <class TypeList>
+using reverse_t = typename reverse<TypeList>::type;
+
+/**
+ * Find the index of the first type in a typelist fulfilling a type trait
+ * condition. Example:
+ *
+ * 2 == find_if<typelist<char, int, char&, int&>, std::is_reference>::value
+ */
+template <class TypeList, template <class> class Condition, class Enable = void>
+struct find_if final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::find_if<TypeList, Condition>, the TypeList argument must be typelist<...>.");
+};
+template <template <class> class Condition>
+struct find_if<typelist<>, Condition, void> final {
+  static_assert(
+      false_higher_t<Condition>::value,
+      "In typelist::find_if<Type/List, Condition>, didn't find any type fulfilling the Condition.");
+};
+template <class Head, class... Tail, template <class> class Condition>
+struct find_if<
+    typelist<Head, Tail...>,
+    Condition,
+    std::enable_if_t<Condition<Head>::value>>
+    final {
+  static constexpr size_t value = 0;
+};
+template <class Head, class... Tail, template <class> class Condition>
+struct find_if<
+    typelist<Head, Tail...>,
+    Condition,
+    std::enable_if_t<!Condition<Head>::value>>
+    final {
+  static constexpr size_t value =
+      1 + find_if<typelist<Tail...>, Condition>::value;
+};
+
+/**
+ * Maps a list of types into a list of values.
+ * Examples:
+ *   // Example 1
+ *   auto sizes =
+ *     map_types_to_values<typelist<int64_t, bool, uint32_t>>(
+ *       [] (auto t) { return sizeof(decltype(t)::type); }
+ *     );
+ *   //  sizes  ==  std::tuple<size_t, size_t, size_t>{8, 1, 4}
+ *
+ *   // Example 2
+ *   auto shared_ptrs =
+ *     map_types_to_values<typelist<int, double>>(
+ *       [] (auto t) { return make_shared<typename decltype(t)::type>(); }
+ *     );
+ *   // shared_ptrs == std::tuple<shared_ptr<int>, shared_ptr<double>>()
+ */
+namespace detail {
+template <class T>
+struct type_ final {
+  using type = T;
+};
+template <class TypeList>
+struct map_types_to_values final {
+  static_assert(
+      false_t<TypeList>::value,
+      "In typelist::map_types_to_values<T>, the T argument must be typelist<...>.");
+};
+template <class... Types>
+struct map_types_to_values<typelist<Types...>> final {
+  template <class Func>
+  static auto call(Func&& func) {
+    return std::tuple{std::forward<Func>(func)(type_<Types>())...};
+  }
+};
+} // namespace detail
+
+template <class TypeList, class Func>
+auto map_types_to_values(Func&& func) {
+  return detail::map_types_to_values<TypeList>::call(std::forward<Func>(func));
+}
+
+} // namespace typelist
+} // namespace c10::guts
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+
+namespace guts {
+using c10::guts::false_t;
+
+namespace typelist {
+using c10::guts::typelist::all;
+using c10::guts::typelist::concat_t;
+using c10::guts::typelist::contains;
+using c10::guts::typelist::count_if;
+using c10::guts::typelist::drop_if_nonempty_t;
+using c10::guts::typelist::drop_t;
+using c10::guts::typelist::filter_t;
+using c10::guts::typelist::find_if;
+using c10::guts::typelist::from_tuple_t;
+using c10::guts::typelist::head_t;
+using c10::guts::typelist::head_with_default_t;
+using c10::guts::typelist::map_t;
+using c10::guts::typelist::map_types_to_values;
+using c10::guts::typelist::reverse_t;
+using c10::guts::typelist::size;
+using c10::guts::typelist::take_t;
+using c10::guts::typelist::to_tuple_t;
+using c10::guts::typelist::true_for_any_type;
+using c10::guts::typelist::typelist;
+
+} // namespace typelist
+
+} // namespace guts
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeSafeSignMath.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeSafeSignMath.h
new file mode 100644
index 0000000000000000000000000000000000000000..c33a286bc5b5577ee9ecd138170688c2906c9dfe
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeSafeSignMath.h
@@ -0,0 +1,148 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <limits>
+#include <type_traits>
+
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wstring-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wstring-conversion")
+#endif
+#if C10_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
+#endif
+
+namespace c10 {
+
+/// Returns false since we cannot have x < 0 if x is unsigned.
+template <typename T>
+inline constexpr bool is_negative(
+    const T& /*x*/,
+    std::true_type /*is_unsigned*/) {
+  return false;
+}
+
+/// Returns true if a signed variable x < 0
+template <typename T>
+inline constexpr bool is_negative(const T& x, std::false_type /*is_unsigned*/) {
+  return x < T(0);
+}
+
+/// Returns true if x < 0
+/// NOTE: Will fail on an unsigned custom type
+///       For the most part it's possible to fix this if
+///       the custom type has a constexpr constructor.
+///       However, notably, c10::Half does not :-(
+template <typename T>
+inline constexpr bool is_negative(const T& x) {
+  return is_negative(x, std::is_unsigned<T>());
+}
+
+/// Returns the sign of an unsigned variable x as 0, 1
+template <typename T>
+inline constexpr int signum(const T& x, std::true_type /*is_unsigned*/) {
+  return T(0) < x;
+}
+
+/// Returns the sign of a signed variable x as -1, 0, 1
+template <typename T>
+inline constexpr int signum(const T& x, std::false_type /*is_unsigned*/) {
+  return (T(0) < x) - (x < T(0));
+}
+
+/// Returns the sign of x as -1, 0, 1
+/// NOTE: Will fail on an unsigned custom type
+///       For the most part it's possible to fix this if
+///       the custom type has a constexpr constructor.
+///       However, notably, c10::Half does not :-(
+template <typename T>
+inline constexpr int signum(const T& x) {
+  return signum(x, std::is_unsigned<T>());
+}
+
+/// Returns true if a and b are not both negative
+template <typename T, typename U>
+inline constexpr bool signs_differ(const T& a, const U& b) {
+  return is_negative(a) != is_negative(b);
+}
+
+// Suppress sign compare warning when compiling with GCC
+// as later does not account for short-circuit rule before
+// raising the warning, see https://godbolt.org/z/Tr3Msnz99
+#ifdef __GNUC__
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wsign-compare"
+#endif
+
+/// Returns true if x is greater than the greatest value of the type Limit
+template <typename Limit, typename T>
+inline constexpr bool greater_than_max(const T& x) {
+  constexpr bool can_overflow =
+      std::numeric_limits<T>::digits > std::numeric_limits<Limit>::digits;
+  return can_overflow && x > std::numeric_limits<Limit>::max();
+}
+
+#ifdef __GNUC__
+#pragma GCC diagnostic pop
+#endif
+
+/// Returns true if x < lowest(Limit). Standard comparison
+template <typename Limit, typename T>
+inline constexpr bool less_than_lowest(
+    const T& x,
+    std::false_type /*limit_is_unsigned*/,
+    std::false_type /*x_is_unsigned*/) {
+  return x < std::numeric_limits<Limit>::lowest();
+}
+
+/// Returns false since all the limit is signed and therefore includes
+/// negative values but x cannot be negative because it is unsigned
+template <typename Limit, typename T>
+inline constexpr bool less_than_lowest(
+    const T& /*x*/,
+    std::false_type /*limit_is_unsigned*/,
+    std::true_type /*x_is_unsigned*/) {
+  return false;
+}
+
+/// Returns true if x < 0, where 0 is constructed from T.
+/// Limit is not signed, so its lower value is zero
+template <typename Limit, typename T>
+inline constexpr bool less_than_lowest(
+    const T& x,
+    std::true_type /*limit_is_unsigned*/,
+    std::false_type /*x_is_unsigned*/) {
+  return x < T(0);
+}
+
+/// Returns false sign both types are unsigned
+template <typename Limit, typename T>
+inline constexpr bool less_than_lowest(
+    const T& /*x*/,
+    std::true_type /*limit_is_unsigned*/,
+    std::true_type /*x_is_unsigned*/) {
+  return false;
+}
+
+/// Returns true if x is less than the lowest value of type T
+/// NOTE: Will fail on an unsigned custom type
+///       For the most part it's possible to fix this if
+///       the custom type has a constexpr constructor.
+///       However, notably, c10::Half does not :
+template <typename Limit, typename T>
+inline constexpr bool less_than_lowest(const T& x) {
+  return less_than_lowest<Limit>(
+      x, std::is_unsigned<Limit>(), std::is_unsigned<T>());
+}
+
+} // namespace c10
+
+C10_CLANG_DIAGNOSTIC_POP()
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::greater_than_max;
+using c10::is_negative;
+using c10::less_than_lowest;
+using c10::signs_differ;
+using c10::signum;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeTraits.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeTraits.h
new file mode 100644
index 0000000000000000000000000000000000000000..152b3f76eb0944a1ddadbee002b4f0787c36d15f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/TypeTraits.h
@@ -0,0 +1,164 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+
+#include <functional>
+#include <type_traits>
+
+namespace c10::guts {
+
+/**
+ * is_equality_comparable<T> is true_type iff the equality operator is defined
+ * for T.
+ */
+template <class T, class Enable = void>
+struct is_equality_comparable : std::false_type {};
+template <class T>
+struct is_equality_comparable<
+    T,
+    std::void_t<decltype(std::declval<T&>() == std::declval<T&>())>>
+    : std::true_type {};
+template <class T>
+using is_equality_comparable_t = typename is_equality_comparable<T>::type;
+
+/**
+ * is_hashable<T> is true_type iff std::hash is defined for T
+ */
+template <class T, class Enable = void>
+struct is_hashable : std::false_type {};
+template <class T>
+struct is_hashable<T, std::void_t<decltype(std::hash<T>()(std::declval<T&>()))>>
+    : std::true_type {};
+template <class T>
+using is_hashable_t = typename is_hashable<T>::type;
+
+/**
+ * is_function_type<T> is true_type iff T is a plain function type (i.e.
+ * "Result(Args...)")
+ */
+template <class T>
+struct is_function_type : std::false_type {};
+template <class Result, class... Args>
+struct is_function_type<Result(Args...)> : std::true_type {};
+template <class T>
+using is_function_type_t = typename is_function_type<T>::type;
+
+/**
+ * is_instantiation_of<T, I> is true_type iff I is a template instantiation of T
+ * (e.g. vector<int> is an instantiation of vector) Example:
+ *    is_instantiation_of_t<vector, vector<int>> // true
+ *    is_instantiation_of_t<pair, pair<int, string>> // true
+ *    is_instantiation_of_t<vector, pair<int, string>> // false
+ */
+template <template <class...> class Template, class T>
+struct is_instantiation_of : std::false_type {};
+template <template <class...> class Template, class... Args>
+struct is_instantiation_of<Template, Template<Args...>> : std::true_type {};
+template <template <class...> class Template, class T>
+using is_instantiation_of_t = typename is_instantiation_of<Template, T>::type;
+
+namespace detail {
+/**
+ * strip_class: helper to remove the class type from pointers to `operator()`.
+ */
+
+template <typename T>
+struct strip_class {};
+template <typename Class, typename Result, typename... Args>
+struct strip_class<Result (Class::*)(Args...)> {
+  using type = Result(Args...);
+};
+template <typename Class, typename Result, typename... Args>
+struct strip_class<Result (Class::*)(Args...) const> {
+  using type = Result(Args...);
+};
+template <typename T>
+using strip_class_t = typename strip_class<T>::type;
+} // namespace detail
+
+/**
+ * Evaluates to true_type, iff the given class is a Functor
+ * (i.e. has a call operator with some set of arguments)
+ */
+
+template <class Functor, class Enable = void>
+struct is_functor : std::false_type {};
+template <class Functor>
+struct is_functor<
+    Functor,
+    std::enable_if_t<is_function_type<
+        detail::strip_class_t<decltype(&Functor::operator())>>::value>>
+    : std::true_type {};
+
+/**
+ * lambda_is_stateless<T> is true iff the lambda type T is stateless
+ * (i.e. does not have a closure).
+ * Example:
+ *  auto stateless_lambda = [] (int a) {return a;};
+ *  lambda_is_stateless<decltype(stateless_lambda)> // true
+ *  auto stateful_lambda = [&] (int a) {return a;};
+ *  lambda_is_stateless<decltype(stateful_lambda)> // false
+ */
+namespace detail {
+template <class LambdaType, class FuncType>
+struct is_stateless_lambda__ final {
+  static_assert(
+      !std::is_same_v<LambdaType, LambdaType>,
+      "Base case shouldn't be hit");
+};
+// implementation idea: According to the C++ standard, stateless lambdas are
+// convertible to function pointers
+template <class LambdaType, class C, class Result, class... Args>
+struct is_stateless_lambda__<LambdaType, Result (C::*)(Args...) const>
+    : std::is_convertible<LambdaType, Result (*)(Args...)> {};
+template <class LambdaType, class C, class Result, class... Args>
+struct is_stateless_lambda__<LambdaType, Result (C::*)(Args...)>
+    : std::is_convertible<LambdaType, Result (*)(Args...)> {};
+
+// case where LambdaType is not even a functor
+template <class LambdaType, class Enable = void>
+struct is_stateless_lambda_ final : std::false_type {};
+// case where LambdaType is a functor
+template <class LambdaType>
+struct is_stateless_lambda_<
+    LambdaType,
+    std::enable_if_t<is_functor<LambdaType>::value>>
+    : is_stateless_lambda__<LambdaType, decltype(&LambdaType::operator())> {};
+} // namespace detail
+template <class T>
+using is_stateless_lambda = detail::is_stateless_lambda_<std::decay_t<T>>;
+
+/**
+ * is_type_condition<C> is true_type iff C<...> is a type trait representing a
+ * condition (i.e. has a constexpr static bool ::value member) Example:
+ *   is_type_condition<std::is_reference>  // true
+ */
+template <template <class> class C, class Enable = void>
+struct is_type_condition : std::false_type {};
+template <template <class> class C>
+struct is_type_condition<
+    C,
+    std::enable_if_t<
+        std::is_same_v<bool, std::remove_cv_t<decltype(C<int>::value)>>>>
+    : std::true_type {};
+
+/**
+ * is_fundamental<T> is true_type iff the lambda type T is a fundamental type
+ * (that is, arithmetic type, void, or nullptr_t). Example: is_fundamental<int>
+ * // true We define it here to resolve a MSVC bug. See
+ * https://github.com/pytorch/pytorch/issues/30932 for details.
+ */
+template <class T>
+struct is_fundamental : std::is_fundamental<T> {};
+} // namespace c10::guts
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly, guts)
+
+using c10::guts::is_equality_comparable;
+using c10::guts::is_function_type;
+using c10::guts::is_hashable;
+using c10::guts::is_instantiation_of;
+using c10::guts::is_stateless_lambda;
+using c10::guts::is_type_condition;
+
+HIDDEN_NAMESPACE_END(torch, headeronly, guts)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/bit_cast.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/bit_cast.h
new file mode 100644
index 0000000000000000000000000000000000000000..3f357f8a06a985332cfa4c8850ceeac3a38dd585
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/bit_cast.h
@@ -0,0 +1,50 @@
+#pragma once
+
+#include <cstring>
+#include <type_traits>
+
+#include <torch/headeronly/macros/Macros.h>
+
+#if __has_include(<bit>) && (defined(__cpp_lib_bit_cast) && __cpp_lib_bit_cast >= 201806L)
+#include <bit>
+#define C10_HAVE_STD_BIT_CAST 1
+#else
+#define C10_HAVE_STD_BIT_CAST 0
+#endif // __has_include(<bit>) && (__cplusplus >= 202002L ||
+       // (defined(__cpp_lib_bit_cast) && __cpp_lib_bit_cast >= 201806L))
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+
+#if C10_HAVE_STD_BIT_CAST
+using std::bit_cast;
+#else
+// Implementations of std::bit_cast() from C++ 20.
+//
+// This is a less sketchy version of reinterpret_cast.
+//
+// See https://en.cppreference.com/w/cpp/numeric/bit_cast for more
+// information as well as the source of our implementations.
+template <class To, class From>
+C10_HOST_DEVICE std::enable_if_t<
+    sizeof(To) == sizeof(From) && std::is_trivially_copyable_v<From> &&
+        std::is_trivially_copyable_v<To>,
+    To>
+// constexpr support needs compiler magic
+bit_cast(const From& src) noexcept {
+  static_assert(
+      std::is_trivially_constructible_v<To>,
+      "This implementation additionally requires "
+      "destination type to be trivially constructible");
+
+  To dst;
+  std::memcpy(&dst, &src, sizeof(To));
+  return dst;
+}
+#endif // C10_HAVE_STD_BIT_CAST
+#undef C10_HAVE_STD_BIT_CAST
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+namespace c10 {
+using torch::headeronly::bit_cast;
+} // namespace c10
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/bits.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/bits.h
new file mode 100644
index 0000000000000000000000000000000000000000..ac3deb8a6ee500c795d86a73cb6362f63dd41a51
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/bits.h
@@ -0,0 +1,71 @@
+#pragma once
+#include <cstdint>
+
+#include <torch/headeronly/macros/Macros.h>
+
+namespace c10 {
+
+/**
+ * bits1x8 is an uninterpreted dtype of a tensor with 1 bit (packed to byte
+ * boundary), without any semantics defined.
+ */
+struct alignas(1) bits1x8 {
+  using underlying = uint8_t;
+  uint8_t val_;
+  bits1x8() = default;
+  C10_HOST_DEVICE explicit bits1x8(uint8_t val) : val_(val) {}
+};
+
+/**
+ * bits2x4 is an uninterpreted dtype of a tensor with 2 bits (packed to byte
+ * boundary), without any semantics defined.
+ */
+struct alignas(1) bits2x4 {
+  using underlying = uint8_t;
+  uint8_t val_;
+  bits2x4() = default;
+  C10_HOST_DEVICE explicit bits2x4(uint8_t val) : val_(val) {}
+};
+
+/**
+ * bits4x2 is an uninterpreted dtype of a tensor with 4 bits (packed to byte
+ * boundary), without any semantics defined.
+ */
+struct alignas(1) bits4x2 {
+  using underlying = uint8_t;
+  uint8_t val_;
+  bits4x2() = default;
+  C10_HOST_DEVICE explicit bits4x2(uint8_t val) : val_(val) {}
+};
+
+/**
+ * bits8 is an uninterpreted dtype of a tensor with 8 bits, without any
+ * semantics defined.
+ */
+struct alignas(1) bits8 {
+  uint8_t val_;
+  bits8() = default;
+  C10_HOST_DEVICE explicit bits8(uint8_t val) : val_(val) {}
+};
+
+/**
+ * bits16 is an uninterpreted dtype of a tensor with 16 bits, without any
+ * semantics defined.
+ */
+struct alignas(2) bits16 {
+  uint16_t val_;
+  bits16() = default;
+  C10_HOST_DEVICE explicit bits16(uint16_t val) : val_(val) {}
+};
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+
+using c10::bits16;
+using c10::bits1x8;
+using c10::bits2x4;
+using c10::bits4x2;
+using c10::bits8;
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/complex.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/complex.h
new file mode 100644
index 0000000000000000000000000000000000000000..733a22d5dbb7a52c9217dc21969d84135c91e4be
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/complex.h
@@ -0,0 +1,616 @@
+#pragma once
+
+#include <complex>
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/Half.h>
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+#include <thrust/complex.h>
+#endif
+
+C10_CLANG_DIAGNOSTIC_PUSH()
+#if C10_CLANG_HAS_WARNING("-Wimplicit-float-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-float-conversion")
+#endif
+#if C10_CLANG_HAS_WARNING("-Wfloat-conversion")
+C10_CLANG_DIAGNOSTIC_IGNORE("-Wfloat-conversion")
+#endif
+
+namespace c10 {
+
+// c10::complex is an implementation of complex numbers that aims
+// to work on all devices supported by PyTorch
+//
+// Most of the APIs duplicates std::complex
+// Reference: https://en.cppreference.com/w/cpp/numeric/complex
+//
+// [NOTE: Complex Operator Unification]
+// Operators currently use a mix of std::complex, thrust::complex, and
+// c10::complex internally. The end state is that all operators will use
+// c10::complex internally.  Until then, there may be some hacks to support all
+// variants.
+//
+//
+// [Note on Constructors]
+//
+// The APIs of constructors are mostly copied from C++ standard:
+//   https://en.cppreference.com/w/cpp/numeric/complex/complex
+//
+// Since C++14, all constructors are constexpr in std::complex
+//
+// There are three types of constructors:
+// - initializing from real and imag:
+//     `constexpr complex( const T& re = T(), const T& im = T() );`
+// - implicitly-declared copy constructor
+// - converting constructors
+//
+// Converting constructors:
+// - std::complex defines converting constructor between float/double/long
+// double,
+//   while we define converting constructor between float/double.
+// - For these converting constructors, upcasting is implicit, downcasting is
+//   explicit.
+// - We also define explicit casting from std::complex/thrust::complex
+//   - Note that the conversion from thrust is not constexpr, because
+//     thrust does not define them as constexpr ????
+//
+//
+// [Operator =]
+//
+// The APIs of operator = are mostly copied from C++ standard:
+//   https://en.cppreference.com/w/cpp/numeric/complex/operator%3D
+//
+// Since C++20, all operator= are constexpr. Although we are not building with
+// C++20, we also obey this behavior.
+//
+// There are three types of assign operator:
+// - Assign a real value from the same scalar type
+//   - In std, this is templated as complex& operator=(const T& x)
+//     with specialization `complex& operator=(T x)` for float/double/long
+//     double Since we only support float and double, on will use `complex&
+//     operator=(T x)`
+// - Copy assignment operator and converting assignment operator
+//   - There is no specialization of converting assignment operators, which type
+//   is
+//     convertible is solely dependent on whether the scalar type is convertible
+//
+// In addition to the standard assignment, we also provide assignment operators
+// with std and thrust
+//
+//
+// [Casting operators]
+//
+// std::complex does not have casting operators. We define casting operators
+// casting to std::complex and thrust::complex
+//
+//
+// [Operator ""]
+//
+// std::complex has custom literals `i`, `if` and `il` defined in namespace
+// `std::literals::complex_literals`. We define our own custom literals in the
+// namespace `c10::complex_literals`. Our custom literals does not follow the
+// same behavior as in std::complex, instead, we define _if, _id to construct
+// float/double complex literals.
+//
+//
+// [real() and imag()]
+//
+// In C++20, there are two overload of these functions, one it to return the
+// real/imag, another is to set real/imag, they are both constexpr. We follow
+// this design.
+//
+//
+// [Operator +=,-=,*=,/=]
+//
+// Since C++20, these operators become constexpr. In our implementation, they
+// are also constexpr.
+//
+// There are two types of such operators: operating with a real number, or
+// operating with another complex number. For the operating with a real number,
+// the generic template form has argument type `const T &`, while the overload
+// for float/double/long double has `T`. We will follow the same type as
+// float/double/long double in std.
+//
+// [Unary operator +-]
+//
+// Since C++20, they are constexpr. We also make them expr
+//
+// [Binary operators +-*/]
+//
+// Each operator has three versions (taking + as example):
+// - complex + complex
+// - complex + real
+// - real + complex
+//
+// [Operator ==, !=]
+//
+// Each operator has three versions (taking == as example):
+// - complex == complex
+// - complex == real
+// - real == complex
+//
+// Some of them are removed on C++20, but we decide to keep them
+//
+// [Operator <<, >>]
+//
+// These are implemented by casting to std::complex
+//
+//
+//
+// TODO(@zasdfgbnm): c10::complex<c10::Half> is not currently supported,
+// because:
+//  - lots of members and functions of c10::Half are not constexpr
+//  - thrust::complex only support float and double
+
+template <typename T>
+struct alignas(sizeof(T) * 2) complex {
+  using value_type = T;
+
+  T real_ = T(0);
+  T imag_ = T(0);
+
+  constexpr complex() = default;
+  C10_HOST_DEVICE constexpr complex(const T& re, const T& im = T())
+      : real_(re), imag_(im) {}
+  template <typename U>
+  explicit constexpr complex(const std::complex<U>& other)
+      : complex(other.real(), other.imag()) {}
+#if defined(__CUDACC__) || defined(__HIPCC__)
+  template <typename U>
+  explicit C10_HOST_DEVICE complex(const thrust::complex<U>& other)
+      : real_(other.real()), imag_(other.imag()) {}
+// NOTE can not be implemented as follow due to ROCm bug:
+//   explicit C10_HOST_DEVICE complex(const thrust::complex<U> &other):
+//   complex(other.real(), other.imag()) {}
+#endif
+
+  // Use SFINAE to specialize casting constructor for c10::complex<float> and
+  // c10::complex<double>
+  template <typename U = T>
+  C10_HOST_DEVICE explicit constexpr complex(
+      const std::enable_if_t<std::is_same_v<U, float>, complex<double>>& other)
+      : real_(other.real_), imag_(other.imag_) {}
+  template <typename U = T>
+  C10_HOST_DEVICE constexpr complex(
+      const std::enable_if_t<std::is_same_v<U, double>, complex<float>>& other)
+      : real_(other.real_), imag_(other.imag_) {}
+
+  constexpr complex<T>& operator=(T re) {
+    real_ = re;
+    imag_ = 0;
+    return *this;
+  }
+
+  constexpr complex<T>& operator+=(T re) {
+    real_ += re;
+    return *this;
+  }
+
+  constexpr complex<T>& operator-=(T re) {
+    real_ -= re;
+    return *this;
+  }
+
+  constexpr complex<T>& operator*=(T re) {
+    real_ *= re;
+    imag_ *= re;
+    return *this;
+  }
+
+  constexpr complex<T>& operator/=(T re) {
+    real_ /= re;
+    imag_ /= re;
+    return *this;
+  }
+
+  template <typename U>
+  constexpr complex<T>& operator=(const complex<U>& rhs) {
+    real_ = rhs.real();
+    imag_ = rhs.imag();
+    return *this;
+  }
+
+  template <typename U>
+  constexpr complex<T>& operator+=(const complex<U>& rhs) {
+    real_ += rhs.real();
+    imag_ += rhs.imag();
+    return *this;
+  }
+
+  template <typename U>
+  constexpr complex<T>& operator-=(const complex<U>& rhs) {
+    real_ -= rhs.real();
+    imag_ -= rhs.imag();
+    return *this;
+  }
+
+  template <typename U>
+  constexpr complex<T>& operator*=(const complex<U>& rhs) {
+    // (a + bi) * (c + di) = (a*c - b*d) + (a * d + b * c) i
+    T a = real_;
+    T b = imag_;
+    U c = rhs.real();
+    U d = rhs.imag();
+    real_ = a * c - b * d;
+    imag_ = a * d + b * c;
+    return *this;
+  }
+
+#ifdef __APPLE__
+#define FORCE_INLINE_APPLE __attribute__((always_inline))
+#else
+#define FORCE_INLINE_APPLE
+#endif
+  template <typename U>
+  constexpr FORCE_INLINE_APPLE complex<T>& operator/=(const complex<U>& rhs)
+      __ubsan_ignore_float_divide_by_zero__ {
+    // (a + bi) / (c + di) = (ac + bd)/(c^2 + d^2) + (bc - ad)/(c^2 + d^2) i
+    // the calculation below follows numpy's complex division
+    T a = real_;
+    T b = imag_;
+    U c = rhs.real();
+    U d = rhs.imag();
+
+#if defined(__GNUC__) && !defined(__clang__)
+    // std::abs is already constexpr by gcc
+    auto abs_c = std::abs(c);
+    auto abs_d = std::abs(d);
+#else
+    auto abs_c = c < 0 ? -c : c;
+    auto abs_d = d < 0 ? -d : d;
+#endif
+
+    if (abs_c >= abs_d) {
+      if (abs_c == U(0) && abs_d == U(0)) {
+        /* divide by zeros should yield a complex inf or nan */
+        real_ = a / abs_c;
+        imag_ = b / abs_d;
+      } else {
+        auto rat = d / c;
+        auto scl = U(1.0) / (c + d * rat);
+        real_ = (a + b * rat) * scl;
+        imag_ = (b - a * rat) * scl;
+      }
+    } else {
+      auto rat = c / d;
+      auto scl = U(1.0) / (d + c * rat);
+      real_ = (a * rat + b) * scl;
+      imag_ = (b * rat - a) * scl;
+    }
+    return *this;
+  }
+#undef FORCE_INLINE_APPLE
+
+  template <typename U>
+  constexpr complex<T>& operator=(const std::complex<U>& rhs) {
+    real_ = rhs.real();
+    imag_ = rhs.imag();
+    return *this;
+  }
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+  template <typename U>
+  C10_HOST_DEVICE complex<T>& operator=(const thrust::complex<U>& rhs) {
+    real_ = rhs.real();
+    imag_ = rhs.imag();
+    return *this;
+  }
+#endif
+
+  template <typename U>
+  explicit constexpr operator std::complex<U>() const {
+    return std::complex<U>(std::complex<T>(real(), imag()));
+  }
+
+#if defined(__CUDACC__) || defined(__HIPCC__)
+  template <typename U>
+  C10_HOST_DEVICE explicit operator thrust::complex<U>() const {
+    return static_cast<thrust::complex<U>>(thrust::complex<T>(real(), imag()));
+  }
+#endif
+
+  // consistent with NumPy behavior
+  explicit constexpr operator bool() const {
+    return real() || imag();
+  }
+
+  C10_HOST_DEVICE constexpr T real() const {
+    return real_;
+  }
+  constexpr void real(T value) {
+    real_ = value;
+  }
+  C10_HOST_DEVICE constexpr T imag() const {
+    return imag_;
+  }
+  constexpr void imag(T value) {
+    imag_ = value;
+  }
+};
+
+namespace complex_literals {
+
+constexpr complex<float> operator""_if(long double imag) {
+  return complex<float>(0.0f, static_cast<float>(imag));
+}
+
+constexpr complex<double> operator""_id(long double imag) {
+  return complex<double>(0.0, static_cast<double>(imag));
+}
+
+constexpr complex<float> operator""_if(unsigned long long imag) {
+  return complex<float>(0.0f, static_cast<float>(imag));
+}
+
+constexpr complex<double> operator""_id(unsigned long long imag) {
+  return complex<double>(0.0, static_cast<double>(imag));
+}
+
+} // namespace complex_literals
+
+template <typename T>
+constexpr complex<T> operator+(const complex<T>& val) {
+  return val;
+}
+
+template <typename T>
+constexpr complex<T> operator-(const complex<T>& val) {
+  return complex<T>(-val.real(), -val.imag());
+}
+
+template <typename T>
+constexpr complex<T> operator+(const complex<T>& lhs, const complex<T>& rhs) {
+  complex<T> result = lhs;
+  return result += rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator+(const complex<T>& lhs, const T& rhs) {
+  complex<T> result = lhs;
+  return result += rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator+(const T& lhs, const complex<T>& rhs) {
+  return complex<T>(lhs + rhs.real(), rhs.imag());
+}
+
+template <typename T>
+constexpr complex<T> operator-(const complex<T>& lhs, const complex<T>& rhs) {
+  complex<T> result = lhs;
+  return result -= rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator-(const complex<T>& lhs, const T& rhs) {
+  complex<T> result = lhs;
+  return result -= rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator-(const T& lhs, const complex<T>& rhs) {
+  complex<T> result = -rhs;
+  return result += lhs;
+}
+
+template <typename T>
+constexpr complex<T> operator*(const complex<T>& lhs, const complex<T>& rhs) {
+  complex<T> result = lhs;
+  return result *= rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator*(const complex<T>& lhs, const T& rhs) {
+  complex<T> result = lhs;
+  return result *= rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator*(const T& lhs, const complex<T>& rhs) {
+  complex<T> result = rhs;
+  return result *= lhs;
+}
+
+template <typename T>
+constexpr complex<T> operator/(const complex<T>& lhs, const complex<T>& rhs) {
+  complex<T> result = lhs;
+  return result /= rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator/(const complex<T>& lhs, const T& rhs) {
+  complex<T> result = lhs;
+  return result /= rhs;
+}
+
+template <typename T>
+constexpr complex<T> operator/(const T& lhs, const complex<T>& rhs) {
+  complex<T> result(lhs, T());
+  return result /= rhs;
+}
+
+// Define operators between integral scalars and c10::complex. std::complex does
+// not support this when T is a floating-point number. This is useful because it
+// saves a lot of "static_cast" when operate a complex and an integer. This
+// makes the code both less verbose and potentially more efficient.
+#define COMPLEX_INTEGER_OP_TEMPLATE_CONDITION                 \
+  typename std::enable_if_t<                                  \
+      std::is_floating_point_v<fT> && std::is_integral_v<iT>, \
+      int> = 0
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator+(const c10::complex<fT>& a, const iT& b) {
+  return a + static_cast<fT>(b);
+}
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator+(const iT& a, const c10::complex<fT>& b) {
+  return static_cast<fT>(a) + b;
+}
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator-(const c10::complex<fT>& a, const iT& b) {
+  return a - static_cast<fT>(b);
+}
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator-(const iT& a, const c10::complex<fT>& b) {
+  return static_cast<fT>(a) - b;
+}
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator*(const c10::complex<fT>& a, const iT& b) {
+  return a * static_cast<fT>(b);
+}
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator*(const iT& a, const c10::complex<fT>& b) {
+  return static_cast<fT>(a) * b;
+}
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator/(const c10::complex<fT>& a, const iT& b) {
+  return a / static_cast<fT>(b);
+}
+
+template <typename fT, typename iT, COMPLEX_INTEGER_OP_TEMPLATE_CONDITION>
+constexpr c10::complex<fT> operator/(const iT& a, const c10::complex<fT>& b) {
+  return static_cast<fT>(a) / b;
+}
+
+#undef COMPLEX_INTEGER_OP_TEMPLATE_CONDITION
+
+template <typename T>
+constexpr bool operator==(const complex<T>& lhs, const complex<T>& rhs) {
+  return (lhs.real() == rhs.real()) && (lhs.imag() == rhs.imag());
+}
+
+template <typename T>
+constexpr bool operator==(const complex<T>& lhs, const T& rhs) {
+  return (lhs.real() == rhs) && (lhs.imag() == T());
+}
+
+template <typename T>
+constexpr bool operator==(const T& lhs, const complex<T>& rhs) {
+  return (lhs == rhs.real()) && (T() == rhs.imag());
+}
+
+template <typename T>
+constexpr bool operator!=(const complex<T>& lhs, const complex<T>& rhs) {
+  return !(lhs == rhs);
+}
+
+template <typename T>
+constexpr bool operator!=(const complex<T>& lhs, const T& rhs) {
+  return !(lhs == rhs);
+}
+
+template <typename T>
+constexpr bool operator!=(const T& lhs, const complex<T>& rhs) {
+  return !(lhs == rhs);
+}
+
+template <typename T, typename CharT, typename Traits>
+std::basic_ostream<CharT, Traits>& operator<<(
+    std::basic_ostream<CharT, Traits>& os,
+    const complex<T>& x) {
+  return (os << static_cast<std::complex<T>>(x));
+}
+
+template <typename T, typename CharT, typename Traits>
+std::basic_istream<CharT, Traits>& operator>>(
+    std::basic_istream<CharT, Traits>& is,
+    complex<T>& x) {
+  std::complex<T> tmp;
+  is >> tmp;
+  x = tmp;
+  return is;
+}
+
+template <typename T>
+C10_HOST_DEVICE complex<T> polar(const T& r, const T& theta = T()) {
+#if defined(__CUDACC__) || defined(__HIPCC__)
+  return static_cast<complex<T>>(thrust::polar(r, theta));
+#else
+  // std::polar() requires r >= 0, so spell out the explicit implementation to
+  // avoid a branch.
+  return complex<T>(r * std::cos(theta), r * std::sin(theta));
+#endif
+}
+
+template <>
+struct alignas(4) complex<Half> {
+  Half real_;
+  Half imag_;
+
+  // Constructors
+  complex() = default;
+  // Half constructor is not constexpr so the following constructor can't
+  // be constexpr
+  C10_HOST_DEVICE explicit inline complex(const Half& real, const Half& imag)
+      : real_(real), imag_(imag) {}
+  C10_HOST_DEVICE inline complex(const c10::complex<float>& value)
+      : real_(value.real()), imag_(value.imag()) {}
+
+  // Conversion operator
+  inline C10_HOST_DEVICE operator c10::complex<float>() const {
+    return {real_, imag_};
+  }
+
+  constexpr C10_HOST_DEVICE Half real() const {
+    return real_;
+  }
+  constexpr C10_HOST_DEVICE Half imag() const {
+    return imag_;
+  }
+
+  C10_HOST_DEVICE complex<Half>& operator+=(const complex<Half>& other) {
+    real_ = static_cast<float>(real_) + static_cast<float>(other.real_);
+    imag_ = static_cast<float>(imag_) + static_cast<float>(other.imag_);
+    return *this;
+  }
+
+  C10_HOST_DEVICE complex<Half>& operator-=(const complex<Half>& other) {
+    real_ = static_cast<float>(real_) - static_cast<float>(other.real_);
+    imag_ = static_cast<float>(imag_) - static_cast<float>(other.imag_);
+    return *this;
+  }
+
+  C10_HOST_DEVICE complex<Half>& operator*=(const complex<Half>& other) {
+    auto a = static_cast<float>(real_);
+    auto b = static_cast<float>(imag_);
+    auto c = static_cast<float>(other.real());
+    auto d = static_cast<float>(other.imag());
+    real_ = a * c - b * d;
+    imag_ = a * d + b * c;
+    return *this;
+  }
+};
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::complex;
+using c10::operator+;
+using c10::operator-;
+using c10::operator*;
+using c10::operator/;
+using c10::operator+=;
+using c10::operator-=;
+using c10::operator*=;
+using c10::operator/=;
+using c10::operator==;
+using c10::operator!=;
+using c10::operator<<;
+using c10::operator>>;
+using c10::polar;
+
+namespace complex_literals {
+using c10::complex_literals::operator""_if;
+using c10::complex_literals::operator""_id;
+} // namespace complex_literals
+
+HIDDEN_NAMESPACE_END(torch, headeronly)
+
+C10_CLANG_DIAGNOSTIC_POP()
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/floating_point_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/floating_point_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..1e60bd85c10fc529243db43aa30522f368a85816
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/floating_point_utils.h
@@ -0,0 +1,38 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+#include <torch/headeronly/util/bit_cast.h>
+#include <cstdint>
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly, detail)
+
+C10_HOST_DEVICE inline float fp32_from_bits(uint32_t w) {
+#if defined(__OPENCL_VERSION__)
+  return as_float(w);
+#elif defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
+  return __uint_as_float((unsigned int)w);
+#elif defined(__INTEL_COMPILER)
+  return _castu32_f32(w);
+#else
+  return torch::headeronly::bit_cast<float>(w);
+#endif
+}
+
+C10_HOST_DEVICE inline uint32_t fp32_to_bits(float f) {
+#if defined(__OPENCL_VERSION__)
+  return as_uint(f);
+#elif defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
+  return (uint32_t)__float_as_uint(f);
+#elif defined(__INTEL_COMPILER)
+  return _castf32_u32(f);
+#else
+  return torch::headeronly::bit_cast<uint32_t>(f);
+#endif
+}
+
+HIDDEN_NAMESPACE_END(torch, headeronly, detail)
+
+namespace c10::detail {
+using torch::headeronly::detail::fp32_from_bits;
+using torch::headeronly::detail::fp32_to_bits;
+} // namespace c10::detail
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/qint32.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/qint32.h
new file mode 100644
index 0000000000000000000000000000000000000000..b1d89d3d9e1f3dc6ca0d095b9fc66e8e69137365
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/qint32.h
@@ -0,0 +1,22 @@
+#pragma once
+#include <cstdint>
+
+#include <torch/headeronly/macros/Macros.h>
+
+namespace c10 {
+
+/**
+ * qint32 is for signed 32 bit quantized Tensors
+ */
+struct alignas(4) qint32 {
+  using underlying = int32_t;
+  int32_t val_;
+  qint32() = default;
+  C10_HOST_DEVICE explicit qint32(int32_t val) : val_(val) {}
+};
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::qint32;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/qint8.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/qint8.h
new file mode 100644
index 0000000000000000000000000000000000000000..36178b6b78c69461407e6782754d30dd4538c1c4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/qint8.h
@@ -0,0 +1,24 @@
+#pragma once
+#include <cstdint>
+
+#include <torch/headeronly/macros/Macros.h>
+
+namespace c10 {
+
+/**
+ * This is the data type for quantized Tensors. Right now we only have
+ * qint8 which is for 8 bit Tensors, and qint32 for 32 bit int Tensors,
+ * we might have 4 bit, 2 bit or 1 bit data types in the future.
+ */
+struct alignas(1) qint8 {
+  using underlying = int8_t;
+  int8_t val_;
+  qint8() = default;
+  C10_HOST_DEVICE explicit qint8(int8_t val) : val_(val) {}
+};
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::qint8;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint2x4.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint2x4.h
new file mode 100644
index 0000000000000000000000000000000000000000..26517084985f362cb3446a855172fc180589ab31
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint2x4.h
@@ -0,0 +1,23 @@
+#pragma once
+#include <cstdint>
+
+#include <torch/headeronly/macros/Macros.h>
+
+namespace c10 {
+
+/**
+ * quint2x4 is for un-signed 2 bit quantized Tensors that are packed to byte
+ * boundary.
+ */
+struct alignas(1) quint2x4 {
+  using underlying = uint8_t;
+  uint8_t val_;
+  quint2x4() = default;
+  C10_HOST_DEVICE explicit quint2x4(uint8_t val) : val_(val) {}
+};
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::quint2x4;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint4x2.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint4x2.h
new file mode 100644
index 0000000000000000000000000000000000000000..cd83fe009a8a3c6d5d9725cd2c860e3f4529f58b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint4x2.h
@@ -0,0 +1,23 @@
+#pragma once
+#include <cstdint>
+
+#include <torch/headeronly/macros/Macros.h>
+
+namespace c10 {
+
+/**
+ * quint4x2 is for un-signed 4 bit quantized Tensors that are packed to byte
+ * boundary.
+ */
+struct alignas(1) quint4x2 {
+  using underlying = uint8_t;
+  uint8_t val_;
+  quint4x2() = default;
+  C10_HOST_DEVICE explicit quint4x2(uint8_t val) : val_(val) {}
+};
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::quint4x2;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint8.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint8.h
new file mode 100644
index 0000000000000000000000000000000000000000..2c253d6686c9ef312244ef9af8a14d96c00d5c59
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/quint8.h
@@ -0,0 +1,22 @@
+#pragma once
+#include <cstdint>
+
+#include <torch/headeronly/macros/Macros.h>
+
+namespace c10 {
+
+/**
+ * quint8 is for unsigned 8 bit quantized Tensors
+ */
+struct alignas(1) quint8 {
+  using underlying = uint8_t;
+  uint8_t val_;
+  quint8() = default;
+  C10_HOST_DEVICE explicit quint8(uint8_t val) : val_(val) {}
+};
+
+} // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::quint8;
+HIDDEN_NAMESPACE_END(torch, headeronly)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/shim_utils.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/shim_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..e0bf304329fe3c3a8d852a6ffbf289cab0c13491
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/util/shim_utils.h
@@ -0,0 +1,29 @@
+#pragma once
+
+#include <torch/headeronly/macros/Macros.h>
+
+#include <sstream>
+#include <stdexcept>
+
+#define TORCH_SUCCESS 0
+#define TORCH_FAILURE 1
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly, detail)
+[[maybe_unused]] C10_NOINLINE static void throw_exception(
+    const char* call,
+    const char* file,
+    int64_t line) {
+  std::stringstream ss;
+  ss << call << " API call failed at " << file << ", line " << line;
+  throw std::runtime_error(ss.str());
+}
+HIDDEN_NAMESPACE_END(torch, headeronly, detail)
+
+// This API is 100% inspired by AOTI_TORCH_ERROR_CODE_CHECK defined in
+// pytorch/torch/csrc/inductor/aoti_runtime/utils.h to handle the returns
+// of the APIs in the shim. We are genericizing this for more global use
+// of the shim beyond AOTI, for examples, see torch/csrc/stable/ops.h.
+#define TORCH_ERROR_CODE_CHECK(call)                                       \
+  if ((call) != TORCH_SUCCESS) {                                           \
+    torch::headeronly::detail::throw_exception(#call, __FILE__, __LINE__); \
+  }
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/version.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/version.h
new file mode 100644
index 0000000000000000000000000000000000000000..90b28cd39d82bc7d0bd2532a8f0ff3b26f769d18
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/headeronly/version.h
@@ -0,0 +1,26 @@
+#pragma once
+
+/// Indicates the major version of LibTorch.
+#define TORCH_VERSION_MAJOR 2
+
+/// Indicates the minor version of LibTorch.
+#define TORCH_VERSION_MINOR 12
+
+/// Indicates the patch version of LibTorch.
+#define TORCH_VERSION_PATCH 1
+
+/// Indicates the ABI version tag of LibTorch.
+#define TORCH_VERSION_ABI_TAG 0
+
+/// Indicates the version of LibTorch as a string literal.
+#define TORCH_VERSION \
+  "2.12.1"
+
+/// Indicates the ABI version of LibTorch as a single uint64.
+/// [ byte ][ byte ][ byte ][ byte ][ byte ][ byte ][ byte ][ byte ]
+/// [ MAJ  ][ MIN  ][ PATCH][                              ABI TAG ]
+#define TORCH_ABI_VERSION ( \
+  ((0ULL + TORCH_VERSION_MAJOR) << 56) | \
+  ((0ULL + TORCH_VERSION_MINOR) << 48) | \
+  ((0ULL + TORCH_VERSION_PATCH) << 40) | \
+  ((0ULL + TORCH_VERSION_ABI_TAG) << 0))
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/library.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/library.h
new file mode 100644
index 0000000000000000000000000000000000000000..d9446a2749ea5186b23c9fcbe5848c575ec94b48
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/library.h
@@ -0,0 +1,1113 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+/// \file
+///
+/// This header provides an API for extending PyTorch's core library
+/// of operators with user defined operators and data types.  This
+/// API can be used in a few ways:
+///
+/// * You can define new custom operators and classes with TORCH_LIBRARY(),
+///   making them available for use in both eager Python as well as in
+///   TorchScript. This API is modeled off of pybind11's `PYBIND11_MODULE`
+///   macro, as the provided functionality is similar (pybind11 lets you bind
+///   C++ to Python only; `torch/library.h` lets you bind C++ simultaneously to
+///   Python and TorchScript).
+///
+/// * You can override existing operators with TORCH_LIBRARY_IMPL(),
+///   providing a new implementation for these operators for a custom
+///   backend (e.g., XLA).  When you pass operators with tensors of your custom
+///   backend, your overridden implementations will be called instead
+///   of the standard implementations.
+///
+/// * You can use both capabilities at the same time, allowing you
+///   to write custom operators that register CPU/CUDA/Autograd
+///   implementations without having to write the boilerplate
+///   conditionals yourself.
+///
+/// For a tutorial style introduction to the library API, check
+/// out the [Extending TorchScript with Custom C++
+/// Operators](https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html)
+/// tutorial.
+///
+/// ```
+/// // Define a library whose operators live in the namespace 'myops'.
+/// // You must define all of the operators for this library in
+/// // this namespace.
+/// TORCH_LIBRARY(myops, m) {
+///   // Define a operator with exactly one implementation for all backends.
+///   m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl);
+///
+///   // Define a schema for an operator, but provide no implementation
+///   // (use this syntax if you want to use the dispatcher)
+///   m.def("mul(Tensor self, Tensor other) -> Tensor");
+///
+///   // Provide an implementation for a defined operator (you can
+///   // provide multiple; one per backend).  The dispatcher takes care of
+///   // calling the correct implementation depending on if we get a CPU
+///   // tensor or a CUDA tensor
+///   m.impl("mul", torch::kCPU, &mul_cpu_impl);
+///   m.impl("mul", torch::kCUDA, &mul_cuda_impl);
+/// }
+///
+/// // Define implementations for operators for a non-standard backend,
+/// // e.g., XLA (valid values are entries of DispatchKey).  This can
+/// // be used to define operators in a different file than the initial
+/// // TORCH_LIBRARY definition (e.g., if it is in an external library)
+/// TORCH_LIBRARY_IMPL(myops, XLA, m) {
+///   m.impl("mul", &mul_xla_impl);
+/// }
+/// ```
+
+#include <ATen/core/op_registration/infer_schema.h>
+#include <ATen/core/op_registration/op_allowlist.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <c10/core/DispatchKey.h>
+#include <torch/csrc/jit/frontend/function_schema_parser.h>
+
+// Just for inferFunctionSchemaFromFunctor
+#include <ATen/core/enum_tag.h>
+#include <ATen/core/op_registration/op_registration.h>
+
+namespace torch {
+
+#if defined C10_MOBILE
+/**
+ * The NoInferSchemaTag is a type name used to indicate that this call to the
+ * CppFunction constructor should not trigger schema inference from functor.
+ * Schema inference from functor utilizes template meta-programming, and is
+ * costly from a size perspective. Ideally, one would expect that the schema
+ * inference would require very little binary size since most of the
+ * computation can be done by the compiler at build time, but that isn't
+ * necessarily the case.
+ *
+ * Schema inference is elided only for mobile use-cases where we don't need
+ * the additional runtime cost or size overhead on client devices.
+ *
+ */
+struct NoInferSchemaTag {};
+#endif
+
+#define HAS_PT2_COMPLIANT_TAG
+
+// For multipy/torchdeploy use case  // codespell:ignore multipy
+enum class _RegisterOrVerify { REGISTER, VERIFY };
+
+template <class CurClass>
+class class_;
+
+#define HAS_IMPL_ABSTRACT_PYSTUB
+
+/// Represents a C++ function that implements an operator.  Most users won't
+/// interact directly with this class, except via error messages: the
+/// constructors this function define the set of permissible "function"-like
+/// things you can bind via the interface.
+///
+/// This class erases the type of the passed in function, but durably records
+/// the type via an inferred schema for the function.
+class TORCH_API CppFunction final {
+  // TODO: This is morally the same thing as KernelRegistrationConfig, but it's
+  // opaque to the user.
+
+ public:
+  /// This overload accepts function pointers, e.g., `CppFunction(&add_impl)`
+  template <typename Func>
+  explicit CppFunction(
+      Func* f,
+      std::enable_if_t<
+          c10::guts::is_function_type<Func>::value,
+          std::nullptr_t>  /*unused*/= nullptr)
+      : func_(c10::KernelFunction::makeFromUnboxedRuntimeFunction(f)),
+        cpp_signature_(c10::impl::CppSignature::make<Func>()),
+        schema_(
+            c10::detail::inferFunctionSchemaFromFunctor<std::decay_t<Func>>())
+        {}
+
+  /// This overload accepts compile time function pointers, e.g.,
+  /// `CppFunction(TORCH_FN(add_impl))`
+  template <typename FuncPtr>
+  explicit CppFunction(
+      FuncPtr f,
+      std::enable_if_t<
+          c10::is_compile_time_function_pointer<FuncPtr>::value,
+          std::nullptr_t>  /*unused*/= nullptr)
+      : func_(c10::KernelFunction::makeFromUnboxedFunction(f)),
+        cpp_signature_(
+            c10::impl::CppSignature::make<typename FuncPtr::FuncType>()),
+        schema_(c10::detail::inferFunctionSchemaFromFunctor<
+                typename FuncPtr::FuncType>())
+        {}
+
+  /// This overload accepts lambdas, e.g., `CppFunction([](const Tensor& self) {
+  /// ... })`
+  template <typename Lambda>
+  explicit CppFunction(
+      Lambda&& f,
+      std::enable_if_t<
+          c10::guts::is_functor<std::decay_t<Lambda>>::value,
+          std::nullptr_t>  /*unused*/= nullptr)
+      : func_(c10::KernelFunction::makeFromUnboxedLambda(
+            std::forward<Lambda>(f))),
+        cpp_signature_(c10::impl::CppSignature::make<Lambda>()),
+        schema_(c10::detail::inferFunctionSchemaFromFunctor<
+                std::decay_t<Lambda>>())
+        {}
+
+#if defined C10_MOBILE
+  /// This overload accepts function pointers, e.g., `CppFunction(&add_impl,
+  /// NoInferSchemaTag())`
+  template <typename Func>
+  explicit CppFunction(
+      Func* f,
+      NoInferSchemaTag,
+      std::enable_if_t<
+          c10::guts::is_function_type<Func>::value,
+          std::nullptr_t> = nullptr)
+      : func_(c10::KernelFunction::makeFromUnboxedRuntimeFunction(f)),
+        cpp_signature_(c10::impl::CppSignature::make<Func>())
+        // TODO: Don't go through WrapRuntimeKernelFunctor
+        ,
+        schema_(nullptr),
+        debug_() {}
+
+  /// This overload accepts compile time function pointers, e.g.,
+  /// `CppFunction(TORCH_FN(add_impl), NoInferSchemaTag())`
+  template <typename FuncPtr>
+  explicit CppFunction(
+      FuncPtr f,
+      NoInferSchemaTag,
+      std::enable_if_t<
+          c10::is_compile_time_function_pointer<FuncPtr>::value,
+          std::nullptr_t> = nullptr)
+      : func_(c10::KernelFunction::makeFromUnboxedFunction(f)),
+        cpp_signature_(
+            c10::impl::CppSignature::make<typename FuncPtr::FuncType>())
+        // TODO: Don't go through WrapRuntimeKernelFunctor
+        ,
+        schema_(nullptr),
+        debug_() {}
+
+  /// This overload accepts lambdas, e.g., `CppFunction([](const Tensor& self) {
+  /// ... }. NoInferSchemaTag())`
+  template <typename Lambda>
+  explicit CppFunction(
+      Lambda&& f,
+      NoInferSchemaTag,
+      std::enable_if_t<
+          c10::guts::is_functor<std::decay_t<Lambda>>::value,
+          std::nullptr_t> = nullptr)
+      : func_(c10::KernelFunction::makeFromUnboxedLambda(
+            std::forward<Lambda>(f))),
+        cpp_signature_(c10::impl::CppSignature::make<Lambda>())
+        // TODO: Don't go through WrapRuntimeKernelFunctor
+        ,
+        schema_(nullptr),
+        debug_() {}
+#endif
+
+  ~CppFunction();
+
+  CppFunction(const CppFunction&) = delete;
+  CppFunction& operator=(const CppFunction&) = delete;
+
+  CppFunction(CppFunction&&) noexcept = default;
+
+  CppFunction& operator=(CppFunction&&) = default;
+
+  /// \private
+  /// Creates a function from a type-erased boxed kernel.
+  static CppFunction makeFromBoxedKernel(c10::BoxedKernel kernel) {
+    return CppFunction(
+        c10::KernelFunction::makeFromBoxedKernel(std::move(kernel)),
+        /* cpp_signature */ std::nullopt, // not known for boxed functions
+        /* schema */ nullptr);
+  }
+
+  /// This creates a fallthrough function.  Fallthrough functions
+  /// immediately redispatch to the next available dispatch key,
+  /// but are implemented more efficiently than a hand written
+  /// function done in the same way.
+  static CppFunction makeFallthrough() {
+    return makeFromBoxedKernel(c10::BoxedKernel::makeFallthrough());
+  }
+
+  /// \private
+  ///
+  /// Creates a function that raises an error saying that named tensors
+  /// are not supported when called.
+  static CppFunction makeNamedNotSupported() {
+    return makeFromBoxedKernel(c10::BoxedKernel::makeNamedNotSupported());
+  }
+
+  /// Create a function from a boxed kernel function with signature
+  /// `void(const OperatorHandle&, Stack*)`; i.e., they receive a
+  /// stack of arguments in a boxed calling convention, rather than
+  /// in the native C++ calling convention.  Boxed functions are
+  /// typically only used to register backend fallbacks via
+  /// torch::Library::fallback().
+  template <c10::BoxedKernel::BoxedKernelFunction* func>
+  static CppFunction makeFromBoxedFunction() {
+    return makeFromBoxedKernel(c10::BoxedKernel::makeFromFunction<func>());
+  }
+
+  // Variant that takes in a boxed kernel function with a plumbed
+  // DispatchKeySet. See Note [Plumbing Keys Through The Dispatcher] for
+  // details.
+  template <c10::BoxedKernel::BoxedKernelFunction_withDispatchKeys* func>
+  static CppFunction makeFromBoxedFunction() {
+    return makeFromBoxedKernel(c10::BoxedKernel::makeFromFunction<func>());
+  }
+
+  /// Create a function from a boxed kernel functor which defines
+  /// `operator()(const OperatorHandle&, DispatchKeySet, Stack*)`
+  /// (receiving arguments from boxed calling convention) and inherits
+  /// from `c10::OperatorKernel`.  Unlike makeFromBoxedFunction, functions
+  /// registered in this way can also carry additional state which
+  /// is managed by the functor; this is useful if you're writing an
+  /// adapter to some other implementation, e.g., a Python callable, which
+  /// is dynamically associated with the registered kernel.
+  template <class KernelFunctor>
+  static CppFunction makeFromBoxedFunctor(
+      std::unique_ptr<KernelFunctor> kernelFunctor) {
+    return makeFromBoxedKernel(
+        c10::BoxedKernel::makeFromFunctor(std::move(kernelFunctor)));
+  }
+
+  /// Create a function from an unboxed kernel function.
+  /// This is typically used to register common operators.
+  template <
+      typename FuncPtr,
+      std::enable_if_t<
+          c10::guts::is_function_type<FuncPtr>::value,
+          std::nullptr_t> = nullptr>
+  static CppFunction makeFromUnboxedFunction(FuncPtr* f) {
+    return CppFunction(f);
+  }
+
+  /// Create a function from a compile time unboxed kernel function pointer.
+  /// This is typically used to register common operators.
+  /// Compile time function pointers can be used to allow the compiler
+  /// to optimize (e.g. inline) calls to it.
+  template <
+      typename FuncPtr,
+      std::enable_if_t<
+          c10::is_compile_time_function_pointer<FuncPtr>::value,
+          std::nullptr_t> = nullptr>
+  static CppFunction makeFromUnboxedFunction(FuncPtr f) {
+    return CppFunction(f);
+  }
+
+  CppFunction&& debug(std::string d) && {
+    debug_ = std::move(d);
+    return std::move(*this);
+  }
+
+ private:
+  std::optional<c10::DispatchKey> dispatch_key_;
+  c10::KernelFunction func_;
+  std::optional<c10::impl::CppSignature> cpp_signature_;
+  std::unique_ptr<c10::FunctionSchema> schema_;
+  std::string debug_;
+
+  // The "setter" for dispatch_key_
+  template <typename Func>
+  friend CppFunction dispatch(c10::DispatchKey /*k*/, Func&& /*raw_f*/);
+
+  // The only class which actually pulls out values from CppFunction (does so
+  // destructively, felt too lazy to write accessors that I don't even
+  // want users to use)
+  friend class Library;
+
+  CppFunction(
+      c10::KernelFunction func,
+      std::optional<c10::impl::CppSignature> cpp_signature,
+      std::unique_ptr<c10::FunctionSchema> schema);
+};
+
+/// \defgroup torch-dispatch-overloads torch::dispatch overloads
+
+/// Create a torch::CppFunction which is associated with a specific
+/// dispatch key.  torch::CppFunctions that are tagged with a
+/// c10::DispatchKey don't get invoked unless the dispatcher determines
+/// that this particular c10::DispatchKey is the one that should be
+/// dispatched to.
+///
+/// This function is generally not used directly, instead, prefer using
+/// TORCH_LIBRARY_IMPL(), which will implicitly set the c10::DispatchKey
+/// for all registration calls inside of its body.
+///
+/// \ingroup torch-dispatch-overloads
+template <typename Func>
+inline CppFunction dispatch(c10::DispatchKey k, Func&& raw_f) {
+  CppFunction f(std::forward<Func>(raw_f));
+  if (k == c10::DispatchKey::CatchAll) {
+    f.dispatch_key_ = std::nullopt;
+  } else {
+    f.dispatch_key_ = k;
+  }
+  return f;
+}
+
+/// Convenience overload of dispatch() which accepts c10::DeviceType
+///
+/// \ingroup torch-dispatch-overloads
+template <typename Func>
+inline CppFunction dispatch(c10::DeviceType type, Func&& raw_f) {
+  auto deviceTypeToDispatchKey = [](c10::DeviceType t) {
+    C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")
+    switch (t) {
+      // This list is synchronized with the k-constants in c10/core/DeviceType.h
+      case c10::DeviceType::CPU:
+        return c10::DispatchKey::CPU;
+      case c10::DeviceType::CUDA:
+        return c10::DispatchKey::CUDA;
+      case c10::DeviceType::IPU:
+        return c10::DispatchKey::IPU;
+      case c10::DeviceType::XLA:
+        return c10::DispatchKey::XLA;
+      case c10::DeviceType::Lazy:
+        return c10::DispatchKey::Lazy;
+      case c10::DeviceType::XPU:
+        return c10::DispatchKey::XPU;
+      case c10::DeviceType::MPS:
+        return c10::DispatchKey::MPS;
+      case c10::DeviceType::Meta:
+        return c10::DispatchKey::Meta;
+      case c10::DeviceType::HIP:
+        return c10::DispatchKey::HIP;
+      case c10::DeviceType::MAIA:
+        return c10::DispatchKey::MAIA;
+      case c10::DeviceType::HPU:
+        return c10::DispatchKey::HPU;
+      case c10::DeviceType::MTIA:
+        return c10::DispatchKey::MTIA;
+      case c10::DeviceType::PrivateUse1:
+        return c10::DispatchKey::PrivateUse1;
+      default:
+        TORCH_CHECK(
+            false,
+            "Device type ",
+            t,
+            " cannot be overloaded at dispatch time, "
+            "please file a bug report explaining what you were trying to do.");
+    }
+    C10_DIAGNOSTIC_POP()
+  };
+  return dispatch(deviceTypeToDispatchKey(type), std::forward<Func>(raw_f));
+}
+
+/// \defgroup torch-schema-overloads torch::schema overloads
+
+/// Construct a c10::FunctionSchema from a string, with an explicitly
+/// specified c10::AliasAnalysisKind.  Ordinarily, schemas are simply
+/// passed in as strings, but if you need to specify a custom alias
+/// analysis, you can replace the string with a call to this function.
+///
+/// ```
+/// // Default alias analysis (FROM_SCHEMA)
+/// m.def("def3(Tensor self) -> Tensor");
+/// // Pure function alias analysis
+/// m.def(torch::schema("def3(Tensor self) -> Tensor",
+/// c10::AliasAnalysisKind::PURE_FUNCTION));
+/// ```
+///
+/// \ingroup torch-schema-overloads
+inline c10::FunctionSchema schema(const char* str, c10::AliasAnalysisKind k, bool allow_typevars=false) {
+  c10::FunctionSchema s = torch::jit::parseSchema(str, /*allow_typevars*/allow_typevars);
+  s.setAliasAnalysis(k);
+  return s;
+}
+
+/// Function schemas can be directly constructed from string literals.
+///
+/// \ingroup torch-schema-overloads
+inline c10::FunctionSchema schema(const char* s, bool allow_typevars=false) {
+  return schema(s, c10::AliasAnalysisKind::FROM_SCHEMA, allow_typevars);
+}
+
+/// \private
+///
+/// Already constructed function schemas are accepted if they are
+/// rvalues.
+///
+/// \ingroup torch-schema-overloads
+inline c10::FunctionSchema&& schema(c10::FunctionSchema&& s) {
+  return std::move(s);
+}
+
+namespace detail {
+
+inline std::variant<c10::OperatorName, c10::FunctionSchema> constructSchemaOrName(
+    c10::FunctionSchema&& s) {
+  return std::move(s);
+}
+inline std::variant<c10::OperatorName, c10::FunctionSchema> constructSchemaOrName(
+    c10::OperatorName&& n) {
+  return std::move(n);
+}
+inline std::variant<c10::OperatorName, c10::FunctionSchema>
+constructSchemaOrName(const char* str) {
+  auto s = torch::jit::parseSchemaOrName(str);
+  if (std::holds_alternative<c10::FunctionSchema>(s)) {
+    std::get<c10::FunctionSchema>(s).setAliasAnalysis(
+        c10::AliasAnalysisKind::FROM_SCHEMA);
+  }
+  return s;
+}
+
+class TorchLibraryInit;
+
+} // namespace detail
+
+// Note [Selective build]
+// ~~~~~~~~~~~~~~~~~~~~~~
+// In some settings, especially mobile, it is important to avoid compiling any
+// references to functions that you aren't actually going to use, so that they
+// can be eliminated by the linker.  We call this capability "selective build".
+//
+// A very easy way to implement selective build which results in a lot of
+// boilerplate is to just add ifdef's around every registration call, but this
+// means you have to write a lot of extra lines of code at every registration
+// site, and it also means you have to define some munging scheme to map
+// operators to macros.
+//
+// Instead of doing this, we have a different mechanism centered around the
+// concept of a SelectiveStr.  A selective name is like a const char* string,
+// except it also carries at compile time a boolean saying whether or not a
+// registration should actually happen or not.  We then have extra overloads
+// which bypass registration entirely if a selective name is disabled.  We do a
+// constexpr test to see if a operator should be enabled or not; this is
+// currently implemented in ATen/core/op_registration/op_allowlist.h
+
+namespace detail {
+
+// dummy class for non selected custom torchbind classes
+class ClassNotSelected {
+ public:
+  ClassNotSelected& def_pickle(...) {
+    return *this;
+  }
+  ClassNotSelected& def(...) {
+    return *this;
+  }
+};
+
+// A SelectiveStr is like a const char*, except that it also comes
+// with a type brand that says whether or not the name is enabled or
+// not.  If the string is disabled, then (at compile time) we DON'T generate
+// a registration call for it.  This class is not intended to be called
+// directly; use TORCH_SELECTIVE_NAME or TORCH_SELECTIVE_SCHEMA macros below
+// to create it.
+template <bool enabled>
+class SelectiveStr {
+ public:
+  constexpr explicit SelectiveStr(const char* name) : name_(name) {}
+  constexpr operator const char*() {
+    return name_;
+  }
+
+ private:
+  const char* name_;
+};
+
+#define TORCH_SELECTIVE_CLASS(n) \
+  torch::detail::SelectiveStr<c10::impl::custom_class_allowlist_check(n)>(n)
+#define TORCH_SELECTIVE_NAME(n) \
+  torch::detail::SelectiveStr<c10::impl::op_allowlist_check(n)>(n)
+#define TORCH_SELECTIVE_SCHEMA(n) \
+  torch::detail::SelectiveStr<c10::impl::schema_allowlist_check(n)>(n)
+
+} // namespace detail
+
+/// This object provides the API for defining operators and providing
+/// implementations at dispatch keys.  Typically, a torch::Library
+/// is not allocated directly; instead it is created by the
+/// TORCH_LIBRARY() or TORCH_LIBRARY_IMPL() macros.
+///
+/// Most methods on torch::Library return a reference to itself,
+/// supporting method chaining.
+///
+/// ```
+/// // Examples:
+///
+/// TORCH_LIBRARY(torchvision, m) {
+///    // m is a torch::Library
+///    m.def("roi_align", ...);
+///    ...
+/// }
+///
+/// TORCH_LIBRARY_IMPL(aten, XLA, m) {
+///    // m is a torch::Library
+///    m.impl("add", ...);
+///    ...
+/// }
+/// ```
+///
+class TORCH_API Library final {
+ public:
+  /// \private
+  ///
+  /// Which type of macro produced this Library
+  enum Kind {
+    DEF, // from TORCH_LIBRARY (no qualifier)
+    IMPL,
+    FRAGMENT,
+  };
+
+  /// \private
+  ///
+  /// Use TORCH_LIBRARY() or TORCH_LIBRARY_IMPL() instead of using these
+  /// constructors directly
+  Library(
+      Kind kind,
+      std::string ns,
+      std::optional<c10::DispatchKey> k,
+      const char* file,
+      uint32_t line);
+
+  Library(const Library&) = delete;
+  Library& operator=(const Library&) = delete;
+  Library(Library&&) = default;
+  Library& operator=(Library&&) = default;
+  ~Library() = default;
+
+  // Some notes about the API design here.  We had the following constraints:
+  //
+  //  - We need to support multiple "types" of arguments for schema and
+  //    functions (e.g., unnamed lambda types, regular functions, const char*,
+  //    fully instantiated schemas)
+  //  - We don't want to write exponentially many overloads
+  //  - We don't want to rely on implicit conversion to a common type,
+  //    because the C++ compiler will only be willing to do a single
+  //    implicit conversion (reducing the set of valid types which you
+  //    can invoke with); also error messages are worse when an implicit
+  //    conversion is not selected (as the compiler will not explain
+  //    why it didn't select an implicit conversion; this is different
+  //    from overloads where it will explain each candidate overload and
+  //    why it didn't apply)
+  //
+  // To solve all of these constraints at the same time, we use a trick taken
+  // from the pybind11 library: template over the argument in the user visible
+  // API, and inside of the templated function explicitly call an overloaded
+  // function to resolve the argument to a real type.  You get the good error
+  // messages from overloads, but at the same time you only need to write the
+  // overload for any given argument type once.
+
+  /// Declare an operator with a schema, but don't provide any implementations
+  /// for it.  You're expected to then provide implementations using the
+  /// impl() method.  All template arguments are inferred.
+  ///
+  /// \param raw_schema The schema of the operator to be defined.
+  ///     Typically, this is a `const char*` string literal, but any type
+  ///     accepted by torch::schema() is accepted here.
+  ///
+  /// ```
+  /// // Example:
+  /// TORCH_LIBRARY(myops, m) {
+  ///   m.def("add(Tensor self, Tensor other) -> Tensor");
+  /// }
+  /// ```
+
+  Library& def(
+      c10::FunctionSchema&& s,
+      const std::vector<at::Tag>& tags = {},
+      _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) & {
+    return _def(std::move(s), nullptr, tags, rv);
+  }
+
+  Library& def(
+      const char* raw_schema,
+      const std::vector<at::Tag>& tags = {},
+      _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) & {
+    return _def(schema(raw_schema), nullptr, tags, rv);
+  }
+
+  /// Declares that for all operators that are subsequently def'ed, their
+  /// fake impls may be found in the given Python module (pymodule).
+  /// This registers some help text that is used if the fake impl
+  /// cannot be found.
+  ///
+  /// Args:
+  /// - pymodule: the python module
+  /// - context: We may include this in the error message.
+  Library& set_python_module(const char* pymodule, const char* context = "") {
+    python_module_ = {pymodule, context};
+    return *this;
+  }
+
+  /// Deprecated; use set_python_module instead
+  Library& impl_abstract_pystub(const char* pymodule, const char* context = "") {
+    return set_python_module(pymodule, context);
+  }
+
+  /// Define an operator for a schema and then register an implementation for
+  /// it.  This is typically what you would use if you aren't planning
+  /// on making use of the dispatcher to structure your operator
+  /// implementation.  It's roughly equivalent to calling def() and
+  /// then impl(), but if you omit the schema of the operator, we will
+  /// infer it from the type of your C++ function.  All template
+  /// arguments are inferred.
+  ///
+  /// \param raw_name_or_schema The schema of the operator to be
+  ///   defined, or just the name of the operator if the schema is to be
+  ///   inferred from `raw_f`.  Typically a `const char*` literal.
+  /// \param raw_f The C++ function that implements this operator.
+  ///   Any valid constructor of torch::CppFunction is accepted here;
+  ///   typically you provide a function pointer or lambda.
+  ///
+  /// ```
+  /// // Example:
+  /// TORCH_LIBRARY(myops, m) {
+  ///   m.def("add", add_fn);
+  /// }
+  /// ```
+  template <typename NameOrSchema, typename Func>
+  Library& def(NameOrSchema&& raw_name_or_schema, Func&& raw_f,
+      const std::vector<at::Tag>& tags = {}) & {
+    CppFunction f(std::forward<Func>(raw_f));
+    return _def(
+        detail::constructSchemaOrName(
+            ::std::forward<NameOrSchema>(raw_name_or_schema)),
+        ::std::move(f), tags);
+  }
+
+  /// Register an implementation for an operator.  You may register multiple
+  /// implementations for a single operator at different dispatch keys
+  /// (see torch::dispatch()).  Implementations must have a corresponding
+  /// declaration (from def()), otherwise they are invalid.  If you plan
+  /// to register multiple implementations, DO NOT provide a function
+  /// implementation when you def() the operator.
+  ///
+  /// \param name The name of the operator to implement.  Do NOT provide
+  ///   schema here.
+  /// \param raw_f The C++ function that implements this operator.  Any
+  ///   valid constructor of torch::CppFunction is accepted here;
+  ///   typically you provide a function pointer or lambda.
+  ///
+  /// ```
+  /// // Example:
+  /// TORCH_LIBRARY_IMPL(myops, CUDA, m) {
+  ///   m.impl("add", add_cuda);
+  /// }
+  /// ```
+  template <typename Name, typename Func>
+  Library& impl(
+      Name name,
+      Func&& raw_f,
+      _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) & {
+    // TODO: need to raise an error when you impl a function that has a
+    // catch all def
+#if defined C10_MOBILE
+    CppFunction f(std::forward<Func>(raw_f), NoInferSchemaTag());
+#else
+    CppFunction f(std::forward<Func>(raw_f));
+#endif
+    return _impl(name, std::move(f), rv);
+  }
+
+#if defined C10_MOBILE
+  // Note: This overload is needed only for C10_MOBILE, since the automatically
+  // defined copy constructor for the CppFunction doesn't have the additional
+  // NoInferSchemaTag argument. We define the overload for the impl() function
+  // to accept a CppFunction&& argument. The already constructed CppFunction
+  // object may or may not have the inferred schema, but it doesn't matter
+  // for our purposes since if it already has the inferred schema, then we
+  // might as well just pass it through directly.
+  //
+  template <typename Name>
+  Library& impl(Name name, CppFunction&& raw_f) & {
+    // TODO: need to raise an error when you impl a function that has a
+    // catch all def
+    CppFunction f(std::forward<CppFunction>(raw_f));
+    return _impl(name, std::move(f));
+  }
+#endif
+
+  // Helper for getting an OperatorName for a const char*.  You probably
+  // don't need this.
+  c10::OperatorName _resolve(const char* name) const;
+
+  /// \private
+  ///
+  /// Convenience overload for directly specifying the dispatch key when
+  /// impl().  You probably don't need this; instead, prefer specifying
+  /// the dispatch key for the entire block in TORCH_LIBRARY_IMPL()
+  template <typename Name, typename Dispatch, typename Func>
+  Library& impl(Name name, Dispatch&& key, Func&& raw_f) & {
+    return impl(
+        name, dispatch(std::forward<Dispatch>(key), std::forward<Func>(raw_f)));
+  }
+
+  template <typename Name, typename Func>
+  Library& impl_UNBOXED(Name /*name*/, Func* /*raw_f*/) & {
+    static_assert(
+        c10::guts::false_t<Func>(),
+        ".impl_UNBOXED(...) was removed. Please use .impl(...) instead.");
+    return *this;
+  }
+
+  // These overloads cover cases when a SelectiveStr (see Note [Selective
+  // build]) has been disabled at compile time.  In that case, don't generate
+  // any code referencing the passed in functions at all.
+  Library& def(detail::SelectiveStr<false> /*unused*/, const std::vector<at::Tag>& tags [[maybe_unused]] = {}) & {
+    return *this;
+  }
+  Library& def(detail::SelectiveStr<true> raw_schema, const std::vector<at::Tag>& tags = {}) & {
+    return def(raw_schema.operator const char*(), tags);
+  }
+  template <typename Func>
+  Library& def(detail::SelectiveStr<false> /*unused*/, Func&& /*raw_f*/, const std::vector<at::Tag>& tags [[maybe_unused]] = {}) & {
+    return *this;
+  }
+  template <typename Func>
+  Library& def(detail::SelectiveStr<true> raw_name_or_schema, Func&& raw_f, const std::vector<at::Tag>& tags = {}) & {
+    return def(
+        raw_name_or_schema.operator const char*(), std::forward<Func>(raw_f), tags);
+  }
+
+  template <typename Func>
+  // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
+  Library& impl(detail::SelectiveStr<false> /*unused*/, Func&& /*raw_f*/) & {
+    return *this;
+  }
+  template <typename Dispatch, typename Func>
+  Library& impl(
+      detail::SelectiveStr<false> /*unused*/,
+      // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
+      Dispatch&& /*key*/,
+      // NOLINTNEXTLINE(cppcoreguidelines-missing-std-forward)
+      Func&& /*raw_f*/) & {
+    return *this;
+  }
+  template <typename Func>
+  Library& impl_UNBOXED(
+      detail::SelectiveStr<false> /*name*/,
+      Func* /*raw_f*/) & {
+    static_assert(
+        c10::guts::false_t<Func>(),
+        ".impl_UNBOXED(...) was removed. Please use .impl(...) instead.");
+    return *this;
+  }
+
+  template <typename Func>
+  Library& impl(detail::SelectiveStr<true> name, Func&& raw_f) & {
+    return impl(name.operator const char*(), std::forward<Func>(raw_f));
+  }
+  template <typename Dispatch, typename Func>
+  Library& impl(
+      detail::SelectiveStr<true> name,
+      Dispatch&& key,
+      Func&& raw_f) & {
+    return impl(
+        name.operator const char*(),
+        std::forward<Dispatch>(key),
+        std::forward<Func>(raw_f));
+  }
+  template <typename Func>
+  Library& impl_UNBOXED(
+      detail::SelectiveStr<true> /*name*/,
+      Func* /*raw_f*/) & {
+    static_assert(
+        c10::guts::false_t<Func>(),
+        ".impl_UNBOXED(...) was removed. Please use .impl(...) instead.");
+    return *this;
+  }
+
+  /// Register a fallback implementation for all operators which will be used
+  /// if there is not a specific implementation for an operator available.
+  /// There MUST be a DispatchKey associated with a fallback; e.g.,
+  /// only call this from TORCH_LIBRARY_IMPL() with namespace `_`.
+  ///
+  /// \param raw_f The function that implements the fallback.  Unboxed
+  ///   functions typically do not work as fallback functions, as
+  ///   fallback functions must work for every operator (even though
+  ///   they have varying type signatures).  Typical arguments are
+  ///   CppFunction::makeFallthrough() or
+  ///   CppFunction::makeFromBoxedFunction()
+  ///
+  /// ```
+  /// // Example:
+  ///
+  /// TORCH_LIBRARY_IMPL(_, AutogradXLA, m) {
+  ///   // If there is not a kernel explicitly registered
+  ///   // for AutogradXLA, fallthrough to the next
+  ///   // available kernel
+  ///   m.fallback(torch::CppFunction::makeFallthrough());
+  /// }
+  ///
+  /// // See aten/src/ATen/core/dispatch/backend_fallback_test.cpp
+  /// // for a full example of boxed fallback
+  /// ```
+  template <typename Func>
+  Library& fallback(Func&& raw_f) & {
+    CppFunction f((std::forward<Func>(raw_f)));
+    return _fallback(std::move(f));
+  }
+
+  template <class CurClass>
+  inline torch::class_<CurClass> class_(const std::string& className);
+
+  // These overloads enable the use of selective build on classes registered
+  // within a library. The API is the same as before with 1 minor change.
+  // Instead of m.class_<foo>("foo") you instead do
+  // m.class_<foo>(TORCH_SELECTIVE_CLASS("foo"))
+  template <class CurClass>
+  inline torch::class_<CurClass> class_(detail::SelectiveStr<true> className);
+
+  template <class CurClass>
+  inline detail::ClassNotSelected class_(detail::SelectiveStr<false> className);
+
+  // De-registers all registrations created with this Library
+  void reset();
+
+ private:
+  Kind kind_;
+  std::optional<std::string> ns_;
+  std::optional<c10::DispatchKey> dispatch_key_;
+  std::optional<std::pair<const char*, const char*>> python_module_;
+  const char* file_;
+  uint32_t line_;
+
+  std::vector<c10::RegistrationHandleRAII> registrars_;
+
+  friend class detail::TorchLibraryInit;
+
+  // Non-user visible actual implementations of functions.  These aren't
+  // public because we only implement & qualifier and not && qualifier
+  Library& _def(
+      c10::FunctionSchema&& schema,
+      c10::OperatorName* out_name = nullptr,
+      const std::vector<at::Tag>& tags = {},
+      _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &;
+  Library& _def(
+      std::variant<c10::OperatorName, c10::FunctionSchema>&& /*name_or_schema*/,
+      CppFunction&& f,
+      const std::vector<at::Tag>& tags = {}) &;
+  Library& _impl(
+      const char* name,
+      CppFunction&& f,
+      _RegisterOrVerify rv = _RegisterOrVerify::REGISTER) &;
+  Library& _fallback(CppFunction&& f) &;
+
+  at::OperatorName _parseNameForLib(const char* name_str) const;
+};
+
+#if defined(TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT) && defined(C10_MOBILE)
+void initialize_torch_libraries();
+#endif
+
+namespace detail {
+
+#if defined(TORCH_LIBRARY_THREAD_UNSAFE_LAZY_INIT) && defined(C10_MOBILE)
+// This is an experimental feature to defer TorchLibraryInit cost to run either
+// at model load time, or when a client application explicitly calls
+// torch::initialize_torch_libraries().
+//
+// This is not thread safe, the client is required to ensure that libraries
+// containing TORCH_LIBRARY initializers are loaded in a thread safe manner.
+extern std::vector<TorchLibraryInit*> torch_library_initializers;
+class TorchLibraryInit final {
+    private:
+      using InitFn = void(Library&);
+      Library::Kind kind;
+      InitFn* init_function;
+      const char* ns;
+      std::optional<c10::DispatchKey> key;
+      const char* file;
+      uint32_t line;
+      std::unique_ptr<Library> lib = nullptr;
+
+    public:
+      TorchLibraryInit(
+            Library::Kind kind,
+            InitFn* fn,
+            const char* ns,
+            std::optional<c10::DispatchKey> k,
+            const char* file,
+            uint32_t line) : kind(kind), init_function(fn), ns(ns), key(k), file(file), line(line) {
+              torch_library_initializers.push_back(this);
+            }
+
+      void initialize() {
+        lib = std::make_unique<Library>(kind, ns, key, file, line);
+        init_function(*lib);
+      }
+};
+#else
+class TorchLibraryInit final {
+ private:
+  using InitFn = void(Library&);
+  Library lib_;
+
+ public:
+  TorchLibraryInit(
+      Library::Kind kind,
+      InitFn* fn,
+      const char* ns,
+      std::optional<c10::DispatchKey> k,
+      const char* file,
+      uint32_t line)
+      : lib_(kind, ns, k, file, line) {
+    fn(lib_);
+  }
+};
+#endif
+
+} // namespace detail
+
+} // namespace torch
+
+// NB: The EXACT NAMING of the initializer functions (e.g.,
+// TORCH_LIBRARY_init_aten) matters for the code analyzer;
+// see the regexes at tools/code_analyzer/run_analyzer.sh
+
+/// Macro for defining a function that will be run at static
+/// initialization time to define a library of operators in the
+/// namespace `ns` (must be a valid C++ identifier, no quotes).
+/// Use this macro when you want to define a new set of custom operators
+/// that do not already exist in PyTorch.
+///
+/// Example usage:
+///
+/// ```
+/// TORCH_LIBRARY(myops, m) {
+///   // m is a torch::Library; methods on it will define
+///   // operators in the myops namespace
+///   m.def("add", add_impl);
+/// }
+/// ```
+///
+/// The `m` argument is bound to a torch::Library that is used to
+/// register operators.  There may only be one TORCH_LIBRARY()
+/// for any given namespace.
+#define TORCH_LIBRARY(ns, m)                                                   \
+  static void TORCH_LIBRARY_init_##ns(torch::Library&);                        \
+  static const torch::detail::TorchLibraryInit TORCH_LIBRARY_static_init_##ns( \
+      torch::Library::DEF,                                                     \
+      &TORCH_LIBRARY_init_##ns,                                                \
+      C10_STRINGIZE(ns),                                                                     \
+      std::nullopt,                                                            \
+      __FILE__,                                                                \
+      __LINE__);                                                               \
+  void TORCH_LIBRARY_init_##ns(torch::Library& m)
+
+/// \private
+///
+/// This macro is a version of TORCH_LIBRARY() that doesn't enforce that there
+/// is only one library (it is a "fragment").  This is used inside the
+/// PerOpRegistration.cpp file, as well as in places where all op registrations
+/// within the same namespace cannot be easily put into one macro block
+/// (this is mostly the case for custom ops in fbcode that were ported from
+/// the old API)
+#define TORCH_LIBRARY_FRAGMENT(ns, m) _TORCH_LIBRARY_FRAGMENT(ns, m, C10_UID)
+
+/// \private
+///
+/// The above macro requires an extra unique identifier (uid) to prevent
+/// variable name collisions This can happen if TORCH_LIBRARY_FRAGMENT is called
+/// multiple times with the same namespace in the same translation unit. Note
+/// that the TORCH_LIBRARY variant doesn't run into this problem, because it
+/// enforces that it can only be called once for a given namespace.
+#define _TORCH_LIBRARY_FRAGMENT(ns, m, uid)                       \
+  static void C10_CONCATENATE(                                    \
+      TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid)(torch::Library&); \
+  static const torch::detail::TorchLibraryInit C10_CONCATENATE(   \
+      TORCH_LIBRARY_FRAGMENT_static_init_##ns##_, uid)(           \
+      torch::Library::FRAGMENT,                                   \
+      &C10_CONCATENATE(TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid), \
+      C10_STRINGIZE(ns),                                                        \
+      std::nullopt,                                               \
+      __FILE__,                                                   \
+      __LINE__);                                                  \
+  void C10_CONCATENATE(                                           \
+      TORCH_LIBRARY_FRAGMENT_init_##ns##_, uid)(torch::Library & m)
+
+/// Macro for defining a function that will be run at static
+/// initialization time to define operator overrides for dispatch key
+/// `k` (must be an unqualified enum member of c10::DispatchKey) in
+/// namespace `ns` (must be a valid C++ identifier, no quotes).  Use this
+/// macro when you want to implement a preexisting set of custom
+/// operators on a new dispatch key (e.g., you want to provide CUDA
+/// implementations of already existing operators).  One common usage
+/// pattern is to use TORCH_LIBRARY() to define schema for all new
+/// operators you want to define, and then use several
+/// TORCH_LIBRARY_IMPL() blocks to provide implementations of the
+/// operator for CPU, CUDA and Autograd.
+///
+/// In some cases, you need to define something that applies to all namespaces,
+/// not just one namespace (usually a fallback).  In that case, use the reserved
+/// namespace _, e.g.,
+///
+/// ```
+/// TORCH_LIBRARY_IMPL(_, XLA, m) {
+///    m.fallback(xla_fallback);
+/// }
+/// ```
+///
+/// Example usage:
+///
+/// ```
+/// TORCH_LIBRARY_IMPL(myops, CPU, m) {
+///   // m is a torch::Library; methods on it will define
+///   // CPU implementations of operators in the myops namespace.
+///   // It is NOT valid to call torch::Library::def()
+///   // in this context.
+///   m.impl("add", add_cpu_impl);
+/// }
+/// ```
+///
+/// If ``add_cpu_impl`` is an overloaded function, use a
+/// ``static_cast`` to specify which overload you want
+/// (by providing the full type).
+///
+// NB: if the dispatch key is not whitelisted, we simply omit the Library
+// call entirely
+#define TORCH_LIBRARY_IMPL(ns, k, m) _TORCH_LIBRARY_IMPL(ns, k, m, C10_UID)
+
+/// \private
+///
+/// The above macro requires an extra unique identifier (uid) to prevent
+/// variable name collisions. This can happen if TORCH_LIBRARY_IMPL is called
+/// multiple times with the same namespace and dispatch key in the same
+/// translation unit.
+#define _TORCH_LIBRARY_IMPL(ns, k, m, uid)                                \
+  static void C10_CONCATENATE(                                            \
+      TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid)(torch::Library&);       \
+  static const torch::detail::TorchLibraryInit C10_CONCATENATE(           \
+      TORCH_LIBRARY_IMPL_static_init_##ns##_##k##_, uid)(                 \
+      torch::Library::IMPL,                                               \
+      &C10_CONCATENATE(TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid),       \
+      C10_STRINGIZE(ns),                                                                \
+      std::make_optional(c10::DispatchKey::k),                            \
+      __FILE__,                                                           \
+      __LINE__);                                                          \
+  void C10_CONCATENATE(                                                   \
+      TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid)(torch::Library & m)
+
+// These are variants of the macros above which are to be used for testing (they
+// don't setup the static initializer, so you can control the visibility of
+// the allocated library yourself).
+//
+// DO NOT use these in production code, they are NOT understood by the
+// code analyzer and will be incorrectly analyzed in those situations.
+
+/// \private
+#define MAKE_TORCH_LIBRARY(ns) \
+  torch::Library(torch::Library::DEF, #ns, std::nullopt, __FILE__, __LINE__)
+/// \private
+#define MAKE_TORCH_LIBRARY_IMPL(ns, k)         \
+  torch::Library(                              \
+      torch::Library::IMPL,                    \
+      #ns,                                     \
+      std::make_optional(c10::DispatchKey::k), \
+      __FILE__,                                \
+      __LINE__)
+
+// Make the custom class API visible, so it is available from
+// torch::Library.
+
+#include <torch/custom_class.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/script.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/script.h
new file mode 100644
index 0000000000000000000000000000000000000000..1f0916234cf3f38187fbec23fc18111051ec430b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/include/torch/script.h
@@ -0,0 +1,18 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <torch/csrc/api/include/torch/types.h>
+#include <torch/csrc/autograd/InferenceMode.h>
+#include <torch/csrc/autograd/custom_function.h>
+#include <torch/csrc/autograd/generated/variable_factories.h>
+#include <torch/csrc/autograd/grad_mode.h>
+#include <torch/csrc/jit/runtime/custom_operator.h>
+#include <torch/csrc/jit/serialization/import.h>
+#include <torch/csrc/jit/serialization/pickle.h>
+#include <torch/custom_class.h>
+
+#include <ATen/ATen.h>
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..00257c6faf52ef51a196470b61cd8e00233354c9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/__init__.py
@@ -0,0 +1,315 @@
+# mypy: allow-untyped-defs
+import warnings
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any
+
+import torch._C
+
+# These are imported so users can access them from the `torch.jit` module
+from torch._jit_internal import (
+    _Await,
+    _drop,
+    _IgnoreContextManager,
+    _isinstance,
+    _overload,
+    _overload_method,
+    export,
+    Final,
+    Future,
+    ignore,
+    is_scripting,
+    unused,
+)
+from torch.jit._async import fork, wait
+from torch.jit._await import _awaitable, _awaitable_nowait, _awaitable_wait
+from torch.jit._decomposition_utils import _register_decomposition
+from torch.jit._freeze import freeze, optimize_for_inference, run_frozen_optimizations
+from torch.jit._fuser import (
+    fuser,
+    last_executed_optimized_graph,
+    optimized_execution,
+    set_fusion_strategy,
+)
+from torch.jit._ir_utils import _InsertPoint
+from torch.jit._script import (
+    _ScriptProfile,
+    _unwrap_optional,
+    Attribute,
+    CompilationUnit,
+    interface,
+    RecursiveScriptClass,
+    RecursiveScriptModule,
+    script,
+    script_method,
+    ScriptFunction,
+    ScriptModule,
+    ScriptWarning,
+)
+from torch.jit._serialization import (
+    jit_module_from_flatbuffer,
+    load,
+    save,
+    save_jit_module_to_flatbuffer,
+)
+from torch.jit._trace import (
+    _flatten,
+    _get_trace_graph,
+    _script_if_tracing,
+    _unique_state_dict,
+    is_tracing,
+    ONNXTracedModule,
+    TopLevelTracedModule,
+    trace,
+    trace_module,
+    TracedModule,
+    TracerWarning,
+    TracingCheckError,
+)
+from torch.utils import set_module
+
+
+__all__ = [
+    "Attribute",
+    "CompilationUnit",
+    "Error",
+    "Future",
+    "ScriptFunction",
+    "ScriptModule",
+    "annotate",
+    "enable_onednn_fusion",
+    "export",
+    "export_opnames",
+    "fork",
+    "freeze",
+    "interface",
+    "ignore",
+    "isinstance",
+    "load",
+    "onednn_fusion_enabled",
+    "optimize_for_inference",
+    "save",
+    "script",
+    "script_if_tracing",
+    "set_fusion_strategy",
+    "strict_fusion",
+    "trace",
+    "trace_module",
+    "unused",
+    "wait",
+]
+
+# For backwards compatibility
+_fork = fork
+_wait = wait
+_set_fusion_strategy = set_fusion_strategy
+
+
+def export_opnames(m):
+    r"""
+    Generate new bytecode for a Script module.
+
+    Returns what the op list would be for a Script Module based off the current code base.
+
+    If you have a LiteScriptModule and want to get the currently present
+    list of ops call _export_operator_list instead.
+    """
+    return torch._C._export_opnames(m._c)
+
+
+# torch.jit.Error
+Error = torch._C.JITException
+set_module(Error, "torch.jit")
+# This is not perfect but works in common cases
+Error.__name__ = "Error"
+Error.__qualname__ = "Error"
+
+
+# for use in python if using annotate
+def annotate(the_type, the_value):
+    """Use to give type of `the_value` in TorchScript compiler.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    This method is a pass-through function that returns `the_value`, used to hint TorchScript
+    compiler the type of `the_value`. It is a no-op when running outside of TorchScript.
+
+    Though TorchScript can infer correct type for most Python expressions, there are some cases where
+    type inference can be wrong, including:
+
+    - Empty containers like `[]` and `{}`, which TorchScript assumes to be container of `Tensor`
+    - Optional types like `Optional[T]` but assigned a valid value of type `T`, TorchScript would assume
+      it is type `T` rather than `Optional[T]`
+
+    Note that `annotate()` does not help in `__init__` method of `torch.nn.Module` subclasses because it
+    is executed in eager mode. To annotate types of `torch.nn.Module` attributes,
+    use :meth:`~torch.jit.Attribute` instead.
+
+    Example:
+
+    .. testcode::
+
+        import torch
+        from typing import Dict
+
+        @torch.jit.script
+        def fn():
+            # Telling TorchScript that this empty dictionary is a (str -> int) dictionary
+            # instead of default dictionary type of (str -> Tensor).
+            d = torch.jit.annotate(Dict[str, int], {})
+
+            # Without `torch.jit.annotate` above, following statement would fail because of
+            # type mismatch.
+            d["name"] = 20
+
+    .. testcleanup::
+
+        del fn
+
+    Args:
+        the_type: Python type that should be passed to TorchScript compiler as type hint for `the_value`
+        the_value: Value or expression to hint type for.
+
+    Returns:
+        `the_value` is passed back as return value.
+    """
+    return the_value
+
+
+def script_if_tracing(fn):
+    """
+    Compiles ``fn`` when it is first called during tracing.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    ``torch.jit.script`` has a non-negligible start up time when it is first called due to
+    lazy-initializations of many compiler builtins. Therefore you should not use
+    it in library code. However, you may want to have parts of your library work
+    in tracing even if they use control flow. In these cases, you should use
+    ``@torch.jit.script_if_tracing`` to substitute for
+    ``torch.jit.script``.
+
+    Args:
+        fn: A function to compile.
+
+    Returns:
+        If called during tracing, a :class:`ScriptFunction` created by `torch.jit.script` is returned.
+        Otherwise, the original function `fn` is returned.
+    """
+    return _script_if_tracing(fn)
+
+
+# for torch.jit.isinstance
+def isinstance(obj, target_type):
+    """
+    Provide container type refinement in TorchScript.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    It can refine parameterized containers of the List, Dict, Tuple, and Optional types. E.g. ``List[str]``,
+    ``Dict[str, List[torch.Tensor]]``, ``Optional[Tuple[int,str,int]]``. It can also
+    refine basic types such as bools and ints that are available in TorchScript.
+
+    Args:
+        obj: object to refine the type of
+        target_type: type to try to refine obj to
+    Returns:
+        ``bool``: True if obj was successfully refined to the type of target_type,
+            False otherwise with no new type refinement
+
+
+    Example (using ``torch.jit.isinstance`` for type refinement):
+    .. testcode::
+
+        import torch
+        from typing import Any, Dict, List
+
+        class MyModule(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def forward(self, input: Any): # note the Any type
+                if torch.jit.isinstance(input, List[torch.Tensor]):
+                    for t in input:
+                        y = t.clamp(0, 0.5)
+                elif torch.jit.isinstance(input, Dict[str, str]):
+                    for val in input.values():
+                        print(val)
+
+        m = torch.jit.script(MyModule())
+        x = [torch.rand(3,3), torch.rand(4,3)]
+        m(x)
+        y = {"key1":"val1","key2":"val2"}
+        m(y)
+    """
+    return _isinstance(obj, target_type)
+
+
+class strict_fusion:
+    """
+    Give errors if not all nodes have been fused in inference, or symbolically differentiated in training.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    Example:
+    Forcing fusion of additions.
+
+    .. code-block:: python
+
+        @torch.jit.script
+        def foo(x):
+            with torch.jit.strict_fusion():
+                return x + x + x
+
+    """
+
+    def __init__(self) -> None:
+        if not torch._jit_internal.is_scripting():
+            warnings.warn("Only works in script mode", stacklevel=2)
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, type: Any, value: Any, tb: Any) -> None:
+        pass
+
+
+# Context manager for globally hiding source ranges when printing graphs.
+# Note that these functions are exposed to Python as static members of the
+# Graph class, so mypy checks need to be skipped.
+@contextmanager
+def _hide_source_ranges() -> Iterator[None]:
+    old_enable_source_ranges = torch._C.Graph.global_print_source_ranges  # type: ignore[attr-defined]
+    try:
+        torch._C.Graph.set_global_print_source_ranges(False)  # type: ignore[attr-defined]
+        yield
+    finally:
+        torch._C.Graph.set_global_print_source_ranges(old_enable_source_ranges)  # type: ignore[attr-defined]
+
+
+def enable_onednn_fusion(enabled: bool) -> None:
+    """Enable or disables onednn JIT fusion based on the parameter `enabled`.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+    """
+    torch._C._jit_set_llga_enabled(enabled)
+
+
+def onednn_fusion_enabled():
+    """Return whether onednn JIT fusion is enabled.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+    """
+    return torch._C._jit_llga_enabled()
+
+
+del Any
+
+if not torch._C._jit_init():
+    raise RuntimeError("JIT initialization failed")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_async.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_async.py
new file mode 100644
index 0000000000000000000000000000000000000000..cedcd0fe3a27e7ef6805d3a32d990c8b04e0f9c2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_async.py
@@ -0,0 +1,131 @@
+# mypy: allow-untyped-defs
+"""Async API.
+
+This module contains the API for parallelism in TorchScript, notably:
+    * torch.jit.fork
+    * torch.jit.wait
+
+This is not intended to be imported directly; please use the exposed
+functionalities in `torch.jit`.
+"""
+
+import warnings
+
+import torch
+from torch._jit_internal import Future
+from torch.jit._builtins import _register_builtin
+from torch.utils import set_module
+
+
+set_module(Future, "torch.jit")
+
+
+def fork(func, *args, **kwargs):
+    r"""
+    Create an asynchronous task executing `func` and a reference to the value of the result of this execution.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    `fork` will return immediately, so the return value of `func` may not have been computed yet. To force completion
+    of the task and access the return value invoke `torch.jit.wait` on the Future. `fork` invoked
+    with a `func` which returns `T` is typed as `torch.jit.Future[T]`. `fork` calls can be arbitrarily
+    nested, and may be invoked with positional and keyword arguments.
+    Asynchronous execution will only occur when run in TorchScript. If run in pure python,
+    `fork` will not execute in parallel. `fork` will also not execute in parallel when invoked
+    while tracing, however the `fork` and `wait` calls will be captured in the exported IR Graph.
+
+    .. warning::
+        `fork` tasks will execute non-deterministically. We recommend only spawning
+        parallel fork tasks for pure functions that do not modify their inputs,
+        module attributes, or global state.
+
+    Args:
+        func (callable or torch.nn.Module):  A Python function or `torch.nn.Module`
+            that will be invoked. If executed in TorchScript, it will execute asynchronously,
+            otherwise it will not. Traced invocations of fork will be captured in the IR.
+        ``*args``, ``**kwargs``: arguments to invoke `func` with.
+    Returns:
+        `torch.jit.Future[T]`: a reference to the execution of `func`. The value `T`
+        can only be accessed by forcing completion of `func` through `torch.jit.wait`.
+
+    Example (fork a free function):
+
+    .. code-block:: python
+
+        import torch
+        from torch import Tensor
+
+
+        def foo(a: Tensor, b: int) -> Tensor:
+            return a + b
+
+
+        def bar(a):
+            fut: torch.jit.Future[Tensor] = torch.jit.fork(foo, a, b=2)
+            return torch.jit.wait(fut)
+
+
+        script_bar = torch.jit.script(bar)
+        input = torch.tensor(2)
+        # only the scripted version executes asynchronously
+        assert script_bar(input) == bar(input)
+        # trace is not run asynchronously, but fork is captured in IR
+        graph = torch.jit.trace(bar, (input,)).graph
+        assert "fork" in str(graph)
+
+    Example (fork a module method):
+
+    .. code-block:: python
+
+        import torch
+        from torch import Tensor
+
+
+        class AddMod(torch.nn.Module):
+            def forward(self, a: Tensor, b: int):
+                return a + b
+
+
+        class Mod(torch.nn.Module):
+            def __init__(self) -> None:
+                super(self).__init__()
+                self.mod = AddMod()
+
+            def forward(self, input):
+                fut = torch.jit.fork(self.mod, a, b=2)
+                return torch.jit.wait(fut)
+
+
+        input = torch.tensor(2)
+        mod = Mod()
+        assert mod(input) == torch.jit.script(mod).forward(input)
+    """
+    warnings.warn(
+        "`torch.jit.fork` is deprecated. Please use `torch.compile` instead.",
+        DeprecationWarning,
+    )
+    return torch._C.fork(func, *args, **kwargs)
+
+
+def wait(future):
+    r"""
+    Force completion of a `torch.jit.Future[T]` asynchronous task, returning the result of the task.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    See :func:`~fork` for docs and examples.
+    Args:
+        future (torch.jit.Future[T]): an asynchronous task reference, created through `torch.jit.fork`
+    Returns:
+        `T`: the return value of the completed task
+    """
+    warnings.warn(
+        "`torch.jit.wait` is deprecated. Please use `torch.compile` instead.",
+        DeprecationWarning,
+    )
+    return torch._C.wait(future)
+
+
+_register_builtin(wait, "aten::wait")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_await.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_await.py
new file mode 100644
index 0000000000000000000000000000000000000000..b00da7b3384b7a263aeab613b3e613c4bd3c3667
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_await.py
@@ -0,0 +1,27 @@
+# mypy: allow-untyped-defs
+import torch
+from torch._jit_internal import _Await
+from torch.jit._builtins import _register_builtin
+from torch.utils import set_module
+
+
+set_module(_Await, "torch.jit")
+
+
+def _awaitable(func, *args, **kwargs):
+    r"""Create Await object that will call specified functioni with specified args, when it is requested for the result."""
+    return torch._C._awaitable(func, *args, **kwargs)
+
+
+def _awaitable_wait(aw):
+    r"""Request await the result of execution, if Await is not completed yet, the func will be called immediately."""
+    return torch._C._awaitable_wait(aw)
+
+
+def _awaitable_nowait(o):
+    r"""Create completed Await with specified result."""
+    return torch._C._awaitable_nowait(o)
+
+
+_register_builtin(_awaitable_wait, "prim::awaitable_wait")
+_register_builtin(_awaitable_nowait, "prim::awaitable_nowait")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_builtins.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_builtins.py
new file mode 100644
index 0000000000000000000000000000000000000000..34caf60599a94637b1478f0d2c0277febceadc52
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_builtins.py
@@ -0,0 +1,203 @@
+# mypy: allow-untyped-defs
+import cmath
+import math
+import warnings
+from collections import OrderedDict
+
+import torch
+import torch.backends.cudnn as cudnn
+from torch.nn.modules.utils import (
+    _list_with_default,
+    _pair,
+    _quadruple,
+    _single,
+    _triple,
+)
+
+
+_builtin_table: dict[int, str] | None = None
+
+_modules_containing_builtins = (
+    torch,
+    torch._C._nn,
+    torch._C._fft,  # type: ignore[attr-defined]
+    torch._C._linalg,  # type: ignore[attr-defined]
+    torch._C._nested,  # type: ignore[attr-defined]
+    torch._C._sparse,  # type: ignore[attr-defined]
+    torch._C._special,  # type: ignore[attr-defined]
+)
+
+_builtin_ops = [
+    # Pairs of (function, op_name)
+    (_pair, "aten::_pair"),
+    (_quadruple, "aten::_quadruple"),
+    (_single, "aten::_single"),
+    (_triple, "aten::_triple"),
+    (_list_with_default, "aten::list_with_default"),
+    (OrderedDict, "aten::dict"),
+    (dict, "aten::dict"),
+    (cudnn.is_acceptable, "aten::cudnn_is_acceptable"),
+    (math.ceil, "aten::ceil"),
+    (math.copysign, "aten::copysign"),
+    (math.erf, "aten::erf"),
+    (math.erfc, "aten::erfc"),
+    (math.exp, "aten::exp"),
+    (math.expm1, "aten::expm1"),
+    (math.fabs, "aten::fabs"),
+    (math.floor, "aten::floor"),
+    (math.gamma, "aten::gamma"),
+    (math.lgamma, "aten::lgamma"),
+    (math.log, "aten::log"),
+    (math.log10, "aten::log10"),
+    (math.log1p, "aten::log1p"),
+    (math.pow, "aten::pow"),
+    (math.sqrt, "aten::sqrt"),
+    (math.isnan, "aten::isnan"),
+    (math.asinh, "aten::asinh"),
+    (math.atanh, "aten::atanh"),
+    (math.cosh, "aten::cosh"),
+    (math.sinh, "aten::sinh"),
+    (math.tanh, "aten::tanh"),
+    (math.acos, "aten::acos"),
+    (math.asin, "aten::asin"),
+    (math.atan, "aten::atan"),
+    (math.atan2, "aten::atan2"),
+    (math.cos, "aten::cos"),
+    (math.sin, "aten::sin"),
+    (math.tan, "aten::tan"),
+    (math.asinh, "aten::asinh"),
+    (math.atanh, "aten::atanh"),
+    (math.acosh, "aten::acosh"),
+    (math.fmod, "aten::fmod"),
+    (math.modf, "aten::modf"),
+    (math.factorial, "aten::factorial"),
+    (math.frexp, "aten::frexp"),
+    (math.isinf, "aten::isinf"),
+    (math.degrees, "aten::degrees"),
+    (math.radians, "aten::radians"),
+    (cmath.isnan, "aten::isnan"),
+    (cmath.isfinite, "aten::isfinite"),
+    (cmath.isinf, "aten::isinf"),
+    (cmath.phase, "aten::angle"),
+    (cmath.rect, "aten::polar"),
+    (cmath.log, "aten::log"),
+    (cmath.log10, "aten::log10"),
+    (cmath.sqrt, "aten::sqrt"),
+    (cmath.exp, "aten::exp"),
+    (cmath.sin, "aten::sin"),
+    (cmath.tan, "aten::tan"),
+    (cmath.cos, "aten::cos"),
+    (cmath.asin, "aten::asin"),
+    (cmath.acos, "aten::acos"),
+    (cmath.atan, "aten::atan"),
+    (cmath.sinh, "aten::sinh"),
+    (cmath.cosh, "aten::cosh"),
+    (cmath.tanh, "aten::tanh"),
+    (cmath.asinh, "aten::asinh"),
+    (cmath.acosh, "aten::acosh"),
+    (cmath.atanh, "aten::atanh"),
+    (math.ldexp, "aten::ldexp"),
+    (torch._assert, "aten::_assert"),
+    (torch.autograd.grad, "aten::grad"),
+    (torch.autograd.backward, "aten::backward"),
+    (torch._C._infer_size, "aten::_infer_size"),
+    (
+        torch.nn.functional._no_grad_embedding_renorm_,  # type: ignore[attr-defined]
+        "aten::_no_grad_embedding_renorm_",
+    ),
+    (torch.nn.functional.assert_int_or_pair, "aten::_assert_int_or_pair"),
+    (torch.nn.init._no_grad_fill_, "aten::_no_grad_fill_"),
+    (torch.nn.init._no_grad_normal_, "aten::_no_grad_normal_"),
+    (torch.nn.init._no_grad_uniform_, "aten::_no_grad_uniform_"),
+    (torch.nn.init._no_grad_zero_, "aten::_no_grad_zero_"),
+    (torch._C._get_tracing_state, "aten::_get_tracing_state"),
+    (torch._C._get_cpu_capability, "aten::_get_cpu_capability"),
+    (warnings.warn, "aten::warn"),
+    (torch._VF.stft, "aten::stft"),  # type: ignore[attr-defined]
+    (torch._VF.istft, "aten::istft"),  # type: ignore[attr-defined]
+    (torch._VF.cdist, "aten::cdist"),  # type: ignore[attr-defined]
+    (torch._VF.norm, "aten::norm"),  # type: ignore[attr-defined]
+    (torch._VF.unique_dim, "aten::unique_dim"),
+    (torch._VF.unique_consecutive, "aten::unique_consecutive"),  # type: ignore[attr-defined]
+    (torch._VF.nuclear_norm, "aten::nuclear_norm"),
+    (torch._VF.frobenius_norm, "aten::frobenius_norm"),
+    (torch._VF.tensordot, "aten::tensordot"),  # type: ignore[attr-defined]
+]
+
+# ops in torch.functional are bound to torch
+# in these cases, we want to resolve the function to their python implementation
+# instead looking up a builtin "aten::" schema
+
+
+def _gen_torch_functional_registered_ops():
+    # eventually ops should encompass all of torch/functional.py, (torch.functional.__all__)
+    # but we are currently only able to compile some of the functions. additionally,
+    # some functions directly map to their aten:: implementations.
+    # TODO: add support for more ops
+    ops = [
+        "stft",
+        "istft",
+        "lu",
+        "cdist",
+        "norm",
+        "unique",
+        "unique_consecutive",
+        "tensordot",
+    ]
+    return {getattr(torch.functional, name) for name in ops}
+
+
+_functional_registered_ops = _gen_torch_functional_registered_ops()
+
+
+def _is_special_functional_bound_op(fn):
+    return fn in _functional_registered_ops
+
+
+# lazily built to ensure the correct initialization order
+def _get_builtin_table():
+    global _builtin_table
+    if _builtin_table is not None:
+        return _builtin_table
+    _builtin_table = {}
+
+    def register_all(mod) -> None:
+        for name in dir(mod):
+            v = getattr(mod, name)
+            if (
+                callable(v)
+                and not _is_special_functional_bound_op(v)
+                and v is not torch.no_grad
+                and v is not torch.autocast
+            ):
+                # Fixup inconsistency in segment_reduce
+                if name == "_segment_reduce":
+                    name = name[1:]
+                _builtin_ops.append((v, "aten::" + name))
+
+    for mod in _modules_containing_builtins:
+        register_all(mod)
+
+    _builtin_ops.append((math.gcd, "aten::gcd"))
+    _builtin_ops.append((math.isfinite, "aten::isfinite"))
+    _builtin_ops.append((math.remainder, "aten::mathremainder"))  # type: ignore[attr-defined]
+
+    import torch.distributed.autograd as dist_autograd
+
+    if dist_autograd.is_available():
+        _builtin_ops.append((dist_autograd.get_gradients, "aten::get_gradients"))
+        _builtin_ops.append((dist_autograd.backward, "aten::dist_backward"))
+
+    # populate the _builtin_table from _builtin_ops
+    for builtin, aten_op in _builtin_ops:
+        _builtin_table[id(builtin)] = aten_op
+
+    return _builtin_table
+
+
+def _register_builtin(fn, op) -> None:
+    _get_builtin_table()[id(fn)] = op
+
+
+def _find_builtin(fn):
+    return _get_builtin_table().get(id(fn))
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_check.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_check.py
new file mode 100644
index 0000000000000000000000000000000000000000..36440769f063ffe391cfe745d2d625276f4eac81
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_check.py
@@ -0,0 +1,261 @@
+# mypy: allow-untyped-defs
+import ast
+import inspect
+import sys
+import textwrap
+import warnings
+
+import torch
+
+
+class AttributeTypeIsSupportedChecker(ast.NodeVisitor):
+    """Check the ``__init__`` method of a given ``nn.Module``.
+
+    It ensures that all instance-level attributes can be properly initialized.
+
+    Specifically, we do type inference based on attribute values...even
+    if the attribute in question has already been typed using
+    Python3-style annotations or ``torch.jit.annotate``. This means that
+    setting an instance-level attribute to ``[]`` (for ``List``),
+    ``{}`` for ``Dict``), or ``None`` (for ``Optional``) isn't enough
+    information for us to properly initialize that attribute.
+
+    An object of this class can walk a given ``nn.Module``'s AST and
+    determine if it meets our requirements or not.
+
+    Known limitations
+    1. We can only check the AST nodes for certain constructs; we can't
+    ``eval`` arbitrary expressions. This means that function calls,
+    class instantiations, and complex expressions that resolve to one of
+    the "empty" values specified above will NOT be flagged as
+    problematic.
+    2. We match on string literals, so if the user decides to use a
+    non-standard import (e.g. `from typing import List as foo`), we
+    won't catch it.
+
+    Example:
+        .. code-block:: python
+
+            class M(torch.nn.Module):
+                def fn(self):
+                    return []
+
+                def __init__(self) -> None:
+                    super().__init__()
+                    self.x: List[int] = []
+
+                def forward(self, x: List[int]):
+                    self.x = x
+                    return 1
+
+        The above code will pass the ``AttributeTypeIsSupportedChecker``
+        check since we have a function call in ``__init__``. However,
+        it will still fail later with the ``RuntimeError`` "Tried to set
+        nonexistent attribute: x. Did you forget to initialize it in
+        __init__()?".
+
+    Args:
+        nn_module - The instance of ``torch.nn.Module`` whose
+            ``__init__`` method we wish to check
+    """
+
+    def check(self, nn_module: torch.nn.Module) -> None:
+        source_lines = inspect.getsource(nn_module.__class__.__init__)
+
+        # Ignore comments no matter the indentation
+        def is_useless_comment(line):
+            line = line.strip()
+            return line.startswith("#") and not line.startswith("# type:")
+
+        source_lines = "\n".join(
+            [l for l in source_lines.split("\n") if not is_useless_comment(l)]
+        )
+
+        # This AST only contains the `__init__` method of the nn.Module
+        init_ast = ast.parse(textwrap.dedent(source_lines))
+
+        # Get items annotated in the class body
+        if sys.version_info >= (3, 14):
+            import annotationlib
+
+            self.class_level_annotations = list(
+                annotationlib.get_annotations(
+                    nn_module, format=annotationlib.Format.FORWARDREF
+                ).keys()
+            )
+        else:
+            self.class_level_annotations = list(nn_module.__annotations__.keys())
+
+        # Flag for later
+        self.visiting_class_level_ann = False
+
+        self.visit(init_ast)
+
+    def _is_empty_container(self, node: ast.AST, ann_type: str) -> bool:
+        if ann_type == "List":
+            # Assigning `[]` to a `List` type gives you a Node where
+            # value=List(elts=[], ctx=Load())
+            if not isinstance(node, ast.List):
+                return False
+            if node.elts:
+                return False
+        elif ann_type == "Dict":
+            # Assigning `{}` to a `Dict` type gives you a Node where
+            # value=Dict(keys=[], values=[])
+            if not isinstance(node, ast.Dict):
+                return False
+            if node.keys:
+                return False
+        elif ann_type == "Optional":
+            # Assigning `None` to an `Optional` type gives you a
+            # Node where value=Constant(value=None, kind=None)
+            if not isinstance(node, ast.Constant):
+                return False
+            if node.value:  # type: ignore[attr-defined]
+                return False
+
+        return True
+
+    def visit_Assign(self, node) -> None:
+        """Store assignment state when assigning to a Call Node.
+
+        If we're visiting a Call Node (the right-hand side of an
+        assignment statement), we won't be able to check the variable
+        that we're assigning to (the left-hand side of an assignment).
+        Because of this, we need to store this state in visitAssign.
+        (Luckily, we only have to do this if we're assigning to a Call
+        Node, i.e. ``torch.jit.annotate``. If we're using normal Python
+        annotations, we'll be visiting an AnnAssign Node, which has its
+        target built in.)
+        """
+        try:
+            if (
+                isinstance(node.value, ast.Call)
+                and node.targets[0].attr in self.class_level_annotations
+            ):
+                self.visiting_class_level_ann = True
+        except AttributeError:
+            return
+        self.generic_visit(node)
+        self.visiting_class_level_ann = False
+
+    def visit_AnnAssign(self, node) -> None:
+        """Visit an AnnAssign node in an ``nn.Module``'s ``__init__`` method.
+
+        It checks if it conforms to our attribute annotation rules."""
+        # If we have a local variable
+        try:
+            if node.target.value.id != "self":
+                return
+        except AttributeError:
+            return
+
+        # If we have an attribute that's already been annotated at the
+        # class level
+        if node.target.attr in self.class_level_annotations:
+            return
+
+        # TODO @ansley: add `Union` once landed
+
+        # NB: Even though `Tuple` is a "container", we don't want to
+        # check for it here. `Tuple` functions as an type with an
+        # "infinite" number of subtypes, in the sense that you can have
+        # `Tuple[())]`, `Tuple[T1]`, `Tuple[T2]`, `Tuple[T1, T2]`,
+        # `Tuple[T2, T1]` and so on, and none of these subtypes can be
+        # used in place of the other. Therefore, assigning an empty
+        # tuple in `__init__` CORRECTLY means that that variable
+        # cannot be reassigned later to a non-empty tuple. Same
+        # deal with `NamedTuple`
+
+        containers = {"List", "list", "Dict", "dict", "Optional"}
+
+        # If we're not evaluating one of the specified problem types
+        try:
+            if node.annotation.value.id not in containers:
+                return
+        except AttributeError:
+            # To evaluate a base type (`str`, `int`, etc.), we would
+            # have needed to get the name through `node.annotation.id`
+            # instead of `node.annotation.value.id`. Seems that we're
+            # not evaluating one of our "containers"
+            return
+
+        # Check if the assigned variable is empty
+        ann_type = node.annotation.value.id
+        if not self._is_empty_container(node.value, ann_type):
+            return
+
+        warnings.warn(
+            "The TorchScript type system doesn't support "
+            "instance-level annotations on empty non-base "
+            "types in `__init__`. Instead, either 1) use a "
+            "type annotation in the class body, or 2) wrap "
+            "the type in `torch.jit.Attribute`.",
+            stacklevel=2,
+        )
+
+    def visit_Call(self, node) -> None:
+        """Determine if a Call node is 'torch.jit.annotate' in __init__.
+
+        Visit a Call node in an ``nn.Module``'s ``__init__``
+        method and determine if it's ``torch.jit.annotate``. If so,
+        see if it conforms to our attribute annotation rules.
+        """
+        # If we have an attribute that's already been annotated at the
+        # class level
+        if self.visiting_class_level_ann:
+            return
+
+        # If this isn't a call to `torch.jit.annotate`
+        try:
+            if (
+                node.func.value.value.id != "torch"
+                or node.func.value.attr != "jit"
+                or node.func.attr != "annotate"
+            ):
+                self.generic_visit(node)
+            elif (
+                node.func.value.value.id != "jit" or node.func.value.attr != "annotate"
+            ):
+                self.generic_visit(node)
+        except AttributeError:
+            # Looks like we didn't even have the right node structure
+            # to check for `torch.jit.annotate` in the first place
+            self.generic_visit(node)
+
+        # Invariant: we have a `torch.jit.annotate` or a
+        # `torch.annotate` call
+
+        # A Call Node for `torch.jit.annotate` should have an `args`
+        # list of length 2 where args[0] represents the annotation and
+        # args[1] represents the actual value
+        if len(node.args) != 2:
+            return
+
+        if not isinstance(node.args[0], ast.Subscript):
+            return
+
+        # See notes in `visit_AnnAssign` r.e. containers
+
+        containers = {"List", "Dict", "Optional"}
+
+        try:
+            ann_type = node.args[0].value.id  # type: ignore[attr-defined]
+        except AttributeError:
+            return
+
+        if ann_type not in containers:
+            return
+
+        # Check if the assigned variable is empty
+        if not self._is_empty_container(node.args[1], ann_type):
+            return
+
+        warnings.warn(
+            "The TorchScript type system doesn't support "
+            "instance-level annotations on empty non-base "
+            "types in `__init__`. Instead, either 1) use a "
+            "type annotation in the class body, or 2) wrap "
+            "the type in `torch.jit.Attribute`.",
+            stacklevel=2,
+        )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_dataclass_impls.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_dataclass_impls.py
new file mode 100644
index 0000000000000000000000000000000000000000..67da5e20206327f23ea3c92ae48a0452c344fead
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_dataclass_impls.py
@@ -0,0 +1,190 @@
+# mypy: allow-untyped-defs
+# Functions for synthesizing magic methods for JIT-compiled dataclasses
+import ast
+import dataclasses
+import inspect
+import os
+from collections.abc import Callable
+from functools import partial
+
+from torch._jit_internal import FAKE_FILENAME_PREFIX, is_optional
+from torch._sources import ParsedDef, SourceContext
+
+
+def _get_fake_filename(cls, method_name):
+    return os.path.join(FAKE_FILENAME_PREFIX, cls.__name__, method_name)
+
+
+def compose_fn(cls, name: str, body_lines: list[str], signature: str) -> ParsedDef:
+    body = "\n".join(f"  {b}" for b in body_lines)
+    decl = f"def {name}{signature}:\n{body}"
+
+    # Parse the function declaration
+    try:
+        py_ast = ast.parse(decl)
+    except SyntaxError as e:
+        # This should only happen if there's some unforeseeable change
+        # in the dataclasses module that makes our synthesized code fail
+        raise RuntimeError(
+            f"TorchScript failed to synthesize dataclass method '{name}' for class '{cls.__name__}'. "
+            "Please file a bug report at <https://github.com/pytorch/pytorch/issues>"
+        ) from e
+    fake_filename = _get_fake_filename(cls, name)
+    # Parse the function
+    return ParsedDef(
+        py_ast,
+        ctx=SourceContext(
+            source=decl, filename=fake_filename, file_lineno=0, leading_whitespace_len=0
+        ),
+        source=decl,
+        filename=fake_filename,
+        file_lineno=0,
+    )
+
+
+def synthesize__init__(cls) -> ParsedDef:
+    # Supporting default factories in the way that people expect would sort of require us to
+    # allow compiling lambda functions, which is not currently supported.
+    if any(
+        field.default_factory is not dataclasses.MISSING
+        for field in dataclasses.fields(cls)
+    ):
+        raise NotImplementedError(
+            "Default factory initializers are not supported in TorchScript dataclasses"
+        )
+
+    # Simply read off the generated __init__ signature from CPython's implementation. It'll be
+    # almost correct except for InitVar annotations, which we need to handle specially.
+    signature = inspect.signature(cls.__init__)
+
+    # Handle InitVars if needed (only works on Python 3.8+, when a `type` attribute was added to InitVar);
+    # see CPython commit here https://github.com/python/cpython/commit/01ee12ba35a333e8a6a25c4153c4a21838e9585c
+    init_vars: list[str] = []
+    params = []
+    for name, param in signature.parameters.items():
+        ann = param.annotation
+
+        if isinstance(ann, dataclasses.InitVar):
+            # The TorchScript interpreter can't handle InitVar annotations, so we unwrap the underlying type here
+            init_vars.append(name)
+            params.append(param.replace(annotation=ann.type))  # type: ignore[attr-defined]
+        else:
+            params.append(param)
+
+    signature = signature.replace(parameters=params)
+
+    body = [
+        # Assign all attributes to self
+        f"self.{field.name} = {field.name}"
+        for field in dataclasses.fields(cls)
+        if field.init and field.name not in init_vars
+    ]
+    # Call user's impl of __post_init__ if it exists
+    if hasattr(cls, "__post_init__"):
+        body.append("self.__post_init__(" + ", ".join(init_vars) + ")")
+
+    return compose_fn(cls, "__init__", body or ["pass"], signature=str(signature))
+
+
+# This is a placeholder at the moment since the TorchScript interpreter doesn't call __repr__
+def synthesize__repr__(cls) -> ParsedDef:
+    return compose_fn(
+        cls,
+        "__repr__",
+        [
+            f"return '{cls.__name__}("
+            + ", ".join(
+                [
+                    f"{field.name}=self.{field.name}"
+                    for field in dataclasses.fields(cls)
+                    if field.repr
+                ]
+            )
+            + ")'"
+        ],
+        signature="(self) -> str",
+    )
+
+
+def synthesize__hash__(cls) -> ParsedDef:
+    return compose_fn(
+        cls,
+        "__hash__",
+        [
+            # This is just a placeholder to prevent compilation from failing; this won't even get called at
+            # all right now because the TorchScript interpreter doesn't call custom __hash__ implementations
+            "raise NotImplementedError('__hash__ is not supported for dataclasses in TorchScript')"
+        ],
+        signature="(self) -> int",
+    )
+
+
+# Implementation for __eq__ and __ne__
+def synthesize_equality(cls, name: str, converse: str) -> ParsedDef:
+    return synthesize_comparison(
+        cls,
+        name,
+        allow_eq=True,
+        raise_on_none=False,
+        inner=[f"if val1 {converse} val2: return False"],
+    )
+
+
+def synthesize_inequality(cls, name: str, op: str, allow_eq: bool) -> ParsedDef:
+    return synthesize_comparison(
+        cls,
+        name,
+        allow_eq,
+        raise_on_none=True,
+        inner=[
+            f"if val1 {op} val2: return True",
+            f"elif val2 {op} val1: return False",
+        ],
+    )
+
+
+def synthesize_comparison(
+    cls, name: str, allow_eq: bool, raise_on_none: bool, inner: list[str]
+) -> ParsedDef:
+    body = []
+    for field in dataclasses.fields(cls):
+        if not field.compare:
+            continue
+
+        body.extend(
+            [
+                f"val1 = self.{field.name}",
+                f"val2 = other.{field.name}",
+            ]
+        )
+        body.extend(
+            inner
+            if not is_optional(field.type)
+            else [
+                # Type refinement for optional fields; we need this to avoid type errors from the interpreter
+                "if val1 is not None and val2 is not None:",
+                *["  " + line for line in inner],
+                "elif (val1 is None) != (val2 is None):",
+                f"  raise TypeError('Cannot compare {cls.__name__} with None')"
+                if raise_on_none
+                else "  return False",
+            ]
+        )
+
+    body.append(f"return {allow_eq}")
+    return compose_fn(
+        cls, name, body, signature=f"(self, other: {cls.__name__}) -> bool"
+    )
+
+
+DATACLASS_MAGIC_METHODS: dict[str, Callable] = {
+    "__init__": synthesize__init__,
+    "__repr__": synthesize__repr__,
+    "__hash__": synthesize__hash__,
+    "__eq__": partial(synthesize_equality, name="__eq__", converse="!="),
+    "__ne__": partial(synthesize_equality, name="__ne__", converse="=="),
+    "__lt__": partial(synthesize_inequality, name="__lt__", op="<", allow_eq=False),
+    "__le__": partial(synthesize_inequality, name="__le__", op="<", allow_eq=True),
+    "__gt__": partial(synthesize_inequality, name="__gt__", op=">", allow_eq=False),
+    "__ge__": partial(synthesize_inequality, name="__ge__", op=">", allow_eq=True),
+}
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_decomposition_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_decomposition_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f83698300441a0539a2634c6351856ee0aca4e7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_decomposition_utils.py
@@ -0,0 +1,14 @@
+# mypy: allow-untyped-defs
+import torch
+from torch._ops import OpOverload, OpOverloadPacket
+
+
+def _register_decomposition(op: OpOverload, graph: torch._C.Graph) -> None:
+    if isinstance(op, OpOverloadPacket):
+        raise AssertionError(
+            f"Must pass specific op overload, not overload packet, found {op}"
+        )
+    if not isinstance(op, OpOverload):
+        raise AssertionError(f"Expected OpOverload, got {type(op)}")
+
+    torch._C._jit_register_decomposition_for_schema(op._schema, graph)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_decompositions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_decompositions.py
new file mode 100644
index 0000000000000000000000000000000000000000..0e3d5a7f9824c357b1858921f7400e82166b7a88
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_decompositions.py
@@ -0,0 +1,143 @@
+# mypy: allow-untyped-defs
+import torch
+from torch import Tensor
+
+
+aten = torch.ops.aten
+import inspect
+import warnings
+from collections.abc import Callable
+from typing import TypeVar
+from typing_extensions import ParamSpec
+
+from torch.types import Number
+
+
+decomposition_table: dict[str, torch.jit.ScriptFunction] = {}
+function_name_set: set[str] = set()
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+
+def check_decomposition_has_type_annotations(f) -> None:
+    inspect_empty = inspect._empty  # type: ignore[attr-defined]
+    sig = inspect.signature(f)
+    for param in sig.parameters.values():
+        if param.annotation == inspect_empty:
+            raise AssertionError(
+                f"No signature on param {param.name} for function {f.name}"
+            )
+
+    if sig.return_annotation == inspect_empty:
+        raise AssertionError(f"No return annotation for function {f.name}")
+
+
+def signatures_match(decomposition_sig, torch_op_sig):
+    decomp_params = decomposition_sig.parameters
+    op_params = torch_op_sig.parameters
+
+    if len(decomp_params) != len(op_params):
+        return False
+
+    for decomp_param, op_param in zip(decomp_params.values(), op_params.values()):
+        # can't check full equality yet because not all fields are correctly deduced
+        # in the torch_op_sig - like default value
+        # can't check 'kind' bc
+        # kwarg-only values with defaults not yet supported in TS
+        inspect_empty = inspect._empty  # type: ignore[attr-defined]
+        for field in ["name", "annotation"]:
+            if field == "name" and decomp_param.name == "self":
+                warnings.warn(
+                    "PyTorch uses 'input' instead of 'self' on public api", stacklevel=2
+                )
+
+            if getattr(decomp_param, field) != getattr(op_param, field):
+                return False
+
+        decomp_default = decomp_param.default
+        op_default = op_param.default
+        # default value not always correctly inferred as being present on torch schema,
+        # but if specified on both they should be equal
+        if decomp_default != inspect_empty and op_default != inspect_empty:
+            if decomp_default != op_default:
+                return False
+
+    return decomposition_sig.return_annotation == torch_op_sig.return_annotation
+
+
+def register_decomposition(
+    aten_op: torch._ops.OpOverload,
+    registry: dict[str, torch.jit.ScriptFunction] | None = None,
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
+    def decomposition_decorator(f: Callable[_P, _T]) -> Callable[_P, _T]:
+        nonlocal registry
+        if registry is None:
+            registry = decomposition_table
+
+        if not isinstance(aten_op, torch._ops.OpOverload):
+            raise AssertionError(
+                f"Expected aten_op to be OpOverload, got {type(aten_op)}"
+            )
+
+        # Need unique name for jit function serialization
+        if f.__name__ in function_name_set:
+            raise AssertionError(f"Duplicated function name {f.__name__}")
+        function_name_set.add(f.__name__)
+
+        scripted_func = torch.jit.script(f)
+        torch._C._jit_pass_inline(scripted_func.graph)
+
+        for _ in range(2):
+            torch._C._jit_pass_peephole(scripted_func.graph)
+            torch._C._jit_pass_constant_propagation(scripted_func.graph)
+
+        registry[str(aten_op._schema)] = scripted_func
+        return f
+
+    return decomposition_decorator
+
+
+# TODO: replace torch.sigmoid -> aten.sigmoid
+
+
+@register_decomposition(aten.var.correction)
+def var_decomposition(
+    input: Tensor,
+    dim: list[int] | None = None,
+    correction: Number | None = None,
+    keepdim: bool = False,
+) -> Tensor:
+    if dim is None:
+        dim_i: list[int] = []
+        dim = dim_i
+
+    if isinstance(dim, (tuple, list)) and len(dim) == 0:
+        n = input.numel()
+    else:
+        n = 1
+        for dim_i in dim:  # type: ignore[assignment]
+            n *= input.shape[dim_i]  # type: ignore[call-overload]
+
+    mean = aten.mean(input, dim, True)
+    sub = input - mean
+    sq = sub * sub
+    sum = aten.sum(sq, dim, keepdim)
+
+    if correction is None:
+        denom = float(n - 1)
+    else:
+        if isinstance(correction, int):
+            denom = float(n - correction)
+        elif isinstance(correction, float):
+            denom = float(n) - correction
+        else:
+            raise RuntimeError("correction must be int or float")
+
+    # pyrefly: ignore [no-matching-overload]
+    return sum / max(0, denom)
+
+
+@register_decomposition(aten.var.default)
+def var(input: Tensor, unbiased: bool = True) -> Tensor:
+    return var_decomposition(input, correction=(1 if unbiased else 0))
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_freeze.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_freeze.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2344dc8a7beca7bdb8c526a2f57b04fb5c52a68
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_freeze.py
@@ -0,0 +1,248 @@
+# mypy: allow-untyped-defs
+"""Freezing.
+
+This is not intended to be imported directly; please use the exposed
+functionalities in `torch.jit`.
+"""
+
+import warnings
+
+import torch
+from torch.jit._script import RecursiveScriptModule, ScriptModule
+
+
+def freeze(
+    mod, preserved_attrs: list[str] | None = None, optimize_numerics: bool = True
+):
+    r"""Freeze ScriptModule, inline submodules, and attributes as constants.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    Freezing a :class:`ScriptModule` will clone it and attempt to inline the cloned
+    module's submodules, parameters, and attributes as constants in the TorchScript IR Graph.
+    By default, `forward` will be preserved, as well as attributes & methods specified in
+    `preserved_attrs`. Additionally, any attribute that is modified within a preserved
+    method will be preserved.
+
+    Freezing currently only accepts ScriptModules that are in eval mode.
+
+    Freezing applies generic optimization that will speed up your model regardless of machine.
+    To further optimize using server-specific settings, run `optimize_for_inference` after
+    freezing.
+
+    Args:
+        mod (:class:`ScriptModule`): a module to be frozen
+        preserved_attrs (Optional[List[str]]): a list of attributes to preserve in addition to the forward method.
+            Attributes modified in preserved methods will also be preserved.
+        optimize_numerics (bool): If ``True``, a set of optimization passes will be run that does not strictly
+            preserve numerics. Full details of optimization can be found at `torch.jit.run_frozen_optimizations`.
+
+    Returns:
+        Frozen :class:`ScriptModule`.
+
+    Example (Freezing a simple module with a Parameter):
+
+    .. testcode::
+        import torch
+        class MyModule(torch.nn.Module):
+            def __init__(self, N, M):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.rand(N, M))
+                self.linear = torch.nn.Linear(N, M)
+
+            def forward(self, input):
+                output = self.weight.mm(input)
+                output = self.linear(output)
+                return output
+
+        scripted_module = torch.jit.script(MyModule(2, 3).eval())
+        frozen_module = torch.jit.freeze(scripted_module)
+        # parameters have been removed and inlined into the Graph as constants
+        assert len(list(frozen_module.named_parameters())) == 0
+        # See the compiled graph as Python code
+        print(frozen_module.code)
+
+    Example (Freezing a module with preserved attributes)
+
+    .. testcode::
+        import torch
+        class MyModule2(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.modified_tensor = torch.tensor(10.)
+                self.version = 1
+
+            def forward(self, input):
+                self.modified_tensor += 1
+                return input + self.modified_tensor
+
+        scripted_module = torch.jit.script(MyModule2().eval())
+        frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=["version"])
+        # we've manually preserved `version`, so it still exists on the frozen module and can be modified
+        assert frozen_module.version == 1
+        frozen_module.version = 2
+        # `modified_tensor` is detected as being mutated in the forward, so freezing preserves
+        # it to retain model semantics
+        assert frozen_module(torch.tensor(1)) == torch.tensor(12)
+        # now that we've run it once, the next result will be incremented by one
+        assert frozen_module(torch.tensor(1)) == torch.tensor(13)
+
+    Note:
+        Freezing submodule attributes is also supported:
+        frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=["submodule.version"])
+
+    Note:
+        If you're not sure why an attribute is not being inlined as a constant, you can run
+        `dump_alias_db` on frozen_module.forward.graph to see if freezing has detected the
+        attribute is being modified.
+
+    Note:
+        Because freezing makes weights constants and removes module hierarchy, `to` and other
+        nn.Module methods to manipulate device or dtype no longer work. As a workaround,
+        You can remap devices by specifying `map_location` in `torch.jit.load`, however
+        device-specific logic may have been baked into the model.
+    """
+    warnings.warn(
+        "`torch.jit.freeze` is deprecated. Please use `torch.compile` instead.",
+        DeprecationWarning,
+    )
+    if not isinstance(mod, ScriptModule):
+        raise RuntimeError(
+            "Freezing expects a ScriptModule as input. "
+            "Please use torch.jit.script or torch.jit.trace to script your 'nn.Module'."
+        )
+
+    if mod.training:
+        raise RuntimeError(
+            "Freezing is currently only implemented for modules in eval mode. "
+            "Please call .eval() on your module before freezing."
+        )
+
+    preserved_attrs = preserved_attrs if preserved_attrs is not None else []
+
+    out = RecursiveScriptModule(torch._C._freeze_module(mod._c, preserved_attrs))
+    RecursiveScriptModule._finalize_scriptmodule(out)
+
+    preserved_methods = [x for x in preserved_attrs if mod._c._has_method(x)]
+    run_frozen_optimizations(out, optimize_numerics, preserved_methods)
+
+    return out
+
+
+def run_frozen_optimizations(
+    mod, optimize_numerics: bool = True, preserved_methods: list[str] | None = None
+) -> None:
+    r"""
+    Run a series of optimizations looking for patterns that occur in frozen graphs.
+
+    The current set of optimizations includes:
+        - Dropout Removal
+        - Pretranspose Linear Layers
+        - Concat Linear Layers with same input Tensor
+        - Conv -> Batchnorm folding
+        - Conv -> Add/Sub folding
+        - Conv -> Mul/Div folding
+
+    Args:
+        mod (:class:`ScriptModule`): a frozen module to be optimized
+
+        optimize_numerics (bool): If ``True``, a set of optimization passes will be run that does not strictly
+        preserve numerics. These optimizations preserve default rtol and atol of `torch.testing.assert_close`
+        when applied on a single transformation, however in a module where many transformations are applied
+        the rtol or atol may no longer fall within the default `assert_close` tolerance. Conv -> Batchnorm folding,
+        Conv-Add/Sub, and Conv -> Mul/Div folding all may alter numerics.
+
+    Returns:
+        None
+
+    Note:
+        In rare occasions, this can result in slower execution.
+
+    Example (Freezing a module with Conv->Batchnorm)
+    .. code-block:: python
+        import torch
+
+        in_channels, out_channels = 3, 32
+        conv = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=2, bias=True
+        )
+        bn = torch.nn.BatchNorm2d(out_channels, eps=0.001)
+        mod = torch.nn.Sequential(conv, bn)
+        # set optimize to False here, by default freezing runs run_frozen_optimizations
+        frozen_mod = torch.jit.freeze(torch.jit.script(mod.eval()), optimize=False)
+        # inspect frozen mod
+        assert "batch_norm" in str(frozen_mod.graph)
+        torch.jit.run_frozen_optimizations(frozen_mod)
+        assert "batch_norm" not in str(frozen_mod.graph)
+
+    """
+    if mod._c._has_method("forward"):
+        torch._C._jit_pass_optimize_frozen_graph(mod.graph, optimize_numerics)
+
+    if preserved_methods is None:
+        preserved_methods = []
+
+    for method in preserved_methods:
+        torch._C._jit_pass_optimize_frozen_graph(
+            mod.__getattr__(method).graph, optimize_numerics
+        )
+
+
+def optimize_for_inference(
+    mod: ScriptModule, other_methods: list[str] | None = None
+) -> ScriptModule:
+    """
+    Perform a set of optimization passes to optimize a model for the purposes of inference.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    If the model is not already frozen, optimize_for_inference
+    will invoke `torch.jit.freeze` automatically.
+
+    In addition to generic optimizations that should speed up your model regardless
+    of environment, prepare for inference will also bake in build specific settings
+    such as the presence of CUDNN or MKLDNN, and may in the future make transformations
+    which speed things up on one machine but slow things down on another. Accordingly,
+    serialization is not implemented following invoking `optimize_for_inference` and
+    is not guaranteed.
+
+    This is still in prototype, and may have the potential to slow down your model.
+    Primary use cases that have been targeted so far have been vision models on cpu
+    and gpu to a lesser extent.
+
+    Example (optimizing a module with Conv->Batchnorm)::
+
+        import torch
+
+        in_channels, out_channels = 3, 32
+        conv = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=2, bias=True
+        )
+        bn = torch.nn.BatchNorm2d(out_channels, eps=0.001)
+        mod = torch.nn.Sequential(conv, bn)
+        frozen_mod = torch.jit.optimize_for_inference(torch.jit.script(mod.eval()))
+        assert "batch_norm" not in str(frozen_mod.graph)
+        # if built with MKLDNN, convolution will be run with MKLDNN weights
+        assert "MKLDNN" in frozen_mod.graph
+    """
+    warnings.warn(
+        "`torch.jit.optimize_for_inference` is deprecated. Please use `torch.compile` instead.",
+        DeprecationWarning,
+    )
+    if not isinstance(mod, ScriptModule):
+        raise RuntimeError(
+            "optimize_for_inference expects a ScriptModule as input. "
+            "Please use torch.jit.script or torch.jit.trace to script your 'nn.Module'."
+        )
+
+    if other_methods is None:
+        other_methods = []
+
+    if hasattr(mod, "training"):
+        mod = freeze(mod.eval(), preserved_attrs=other_methods)
+
+    torch._C._jit_pass_optimize_for_inference(mod._c, other_methods)
+
+    return mod
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_fuser.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_fuser.py
new file mode 100644
index 0000000000000000000000000000000000000000..fdbe41859aced3bf30fc56a9a0231e3a2a15f846
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_fuser.py
@@ -0,0 +1,173 @@
+# mypy: allow-untyped-defs
+import contextlib
+import warnings
+
+import torch
+
+
+@contextlib.contextmanager
+def optimized_execution(should_optimize):
+    """Context manager that controls whether the JIT's executor will run optimizations before executing a function."""
+    stored_flag = torch._C._get_graph_executor_optimize()
+    torch._C._set_graph_executor_optimize(should_optimize)
+    try:
+        yield
+    finally:
+        torch._C._set_graph_executor_optimize(stored_flag)
+
+
+@contextlib.contextmanager
+def fuser(name):
+    """Context manager that facilitates switching between backend fusers.
+
+    Valid names:
+    * ``fuser0`` - enables only legacy fuser
+    * ``fuser1`` - enables only NNC
+    * ``fuser2`` - enables only nvFuser
+    * ``fuser3`` - enables oneDNN Graph
+    """
+    old_cpu_fuse = torch._C._jit_can_fuse_on_cpu()
+    old_gpu_fuse = torch._C._jit_can_fuse_on_gpu()
+    old_texpr_fuser_state = torch._C._jit_texpr_fuser_enabled()
+    old_nvfuser_state = torch._C._jit_nvfuser_enabled()
+    old_llga_state = torch._C._jit_llga_enabled()
+    if name == "fuser0":  # legacy fuser
+        torch._C._jit_override_can_fuse_on_cpu(True)
+        torch._C._jit_override_can_fuse_on_gpu(True)
+        torch._C._jit_set_texpr_fuser_enabled(False)
+        torch._C._jit_set_nvfuser_enabled(False)
+        torch._C._jit_set_llga_enabled(False)
+    elif name == "fuser1":  # NNC
+        old_profiling_executor = torch._C._jit_set_profiling_executor(True)
+        old_profiling_mode = torch._C._get_graph_executor_optimize(True)
+        torch._C._jit_override_can_fuse_on_cpu(True)
+        torch._C._jit_override_can_fuse_on_gpu(True)
+        torch._C._jit_set_texpr_fuser_enabled(True)
+        torch._C._jit_set_nvfuser_enabled(False)
+        torch._C._jit_set_llga_enabled(False)
+    elif name == "fuser2":  # nvFuser
+        torch._C._jit_override_can_fuse_on_cpu(False)
+        torch._C._jit_override_can_fuse_on_gpu(False)
+        torch._C._jit_set_texpr_fuser_enabled(False)
+        torch._C._jit_set_nvfuser_enabled(True)
+        torch._C._jit_set_llga_enabled(False)
+    elif name == "fuser3":  # oneDNN Graph
+        old_profiling_executor = torch._C._jit_set_profiling_executor(True)
+        old_profiling_mode = torch._C._get_graph_executor_optimize(True)
+        torch._C._jit_override_can_fuse_on_cpu(True)
+        torch._C._jit_override_can_fuse_on_gpu(False)
+        torch._C._jit_set_texpr_fuser_enabled(True)
+        torch._C._jit_set_nvfuser_enabled(False)
+        torch._C._jit_set_llga_enabled(True)
+    elif name == "none":  # Turn Pytorch fuser off
+        torch._C._jit_override_can_fuse_on_cpu(False)
+        torch._C._jit_override_can_fuse_on_gpu(False)
+        torch._C._jit_set_texpr_fuser_enabled(False)
+        torch._C._jit_set_nvfuser_enabled(False)
+        torch._C._jit_set_llga_enabled(False)
+    else:
+        raise Exception(f"unrecognized fuser option (name: {name})")  # noqa: TRY002
+    try:
+        yield
+    finally:
+        if name in ["fuser1", "fuser3"]:  # NNC or oneDNN Graph
+            torch._C._jit_set_profiling_executor(old_profiling_executor)  # type: ignore[possibly-undefined]
+            torch._C._get_graph_executor_optimize(old_profiling_mode)  # type: ignore[possibly-undefined]
+        # recover the previous values
+        torch._C._jit_override_can_fuse_on_cpu(old_cpu_fuse)
+        torch._C._jit_override_can_fuse_on_gpu(old_gpu_fuse)
+        torch._C._jit_set_texpr_fuser_enabled(old_texpr_fuser_state)
+        torch._C._jit_set_nvfuser_enabled(old_nvfuser_state)
+        torch._C._jit_set_llga_enabled(old_llga_state)
+
+
+last_executed_optimized_graph = torch._C._last_executed_optimized_graph
+
+
+def _get_differentiable_graph_node(node, diff_node) -> None:
+    if node.kind() == "prim::DifferentiableGraph":
+        diff_node.append(node)
+    else:
+        for block in node.blocks():
+            for n in block.nodes():
+                _get_differentiable_graph_node(n, diff_node)
+
+
+def _graph_for(self, *args, **kwargs):
+    return _script_method_graph_for(self, self, *args, **kwargs)
+
+
+def _script_method_graph_for(self, parent, *args, **kwargs):
+    try:
+        dbs = parent.get_debug_state()
+        eps = list(dbs.execution_plans.values())
+        if len(eps) != 1:
+            raise AssertionError(f"Expected exactly 1 execution plan, got {len(eps)}")
+        graph = eps[0].graph.copy()
+
+        # graph_executor_states for differentiable node
+        fw_states = eps[0].code.differentiable_op_executor_states()
+        diff_nodes: list[torch._C.Node] = []
+        for n in graph.nodes():
+            _get_differentiable_graph_node(n, diff_nodes)
+
+        if len(fw_states) != len(diff_nodes):
+            raise AssertionError(
+                f"Expected fw_states ({len(fw_states)}) and diff_nodes "
+                f"({len(diff_nodes)}) to have the same length"
+            )
+        # swap each differentiable graph with optimized graph in their execution plan
+        for n, state in zip(diff_nodes, fw_states):
+            fw_execution_plans = list(state.execution_plans.values())
+            # we can only update the subgraph when there's a unique execution
+            # plan. Avoid assert here so we would skip the ones that can't be
+            # updated while try the best effort to update other nodes.
+            if len(fw_execution_plans) == 1:
+                n.g_("Subgraph", fw_execution_plans[0].graph)
+
+        return graph
+    except Exception:
+        # fallback approach, we just ran the graph and return the recorded optimized
+        # graph
+        self(*args, **kwargs)
+        return last_executed_optimized_graph()
+
+
+def set_fusion_strategy(strategy: list[tuple[str, int]]):
+    """Set the type and number of specializations that can occur during fusion.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    Usage: provide a list of pairs (type, depth) where type is one of "STATIC" or "DYNAMIC"
+    and depth is an integer.
+
+    Behavior - static vs dynamic:
+        In STATIC fusion, fused ops are compiled to have fixed input shapes. The shape is determined
+        based on some initial profiling runs.
+        In DYNAMIC fusion, fused ops are compiled to have variable input shapes, so that multiple
+        shapes are possible.
+
+    In both cases, we also recompile on new striding behavior, device, or dtype.
+
+    Behavior - fallback functions & depth:
+        When an input doesn't match the format required by the specialized compiled op, it will run
+        a fallback function. Fallback functions are recursively be compiled and specialized based
+        on the observed tensor shapes. Since compilation can be slow, the "depth" parameter is provided to
+        limit the number of specializations that can be compiled, before giving up on recompiling and
+        falling back to a completely un-fused, un-specialized implementation.
+
+    The list of (type, depth) pairs controls the type of specializations and the number of
+    specializations. For example: [("STATIC", 2), ("DYNAMIC", 2)] indicates that the first
+    two specializations will use static fusions, the following two specializations will use
+    dynamic fusion, and any inputs that satisfy none of the 4 options will run an
+    unfused implementation.
+
+    NB: in the future, if more as more fusion backends are added there may be more granular
+    apis for specific fusers.
+    """
+    warnings.warn(
+        "`torch.jit.set_fusion_strategy` is deprecated. Please use `torch.compile` instead.",
+        DeprecationWarning,
+    )
+    return torch._C._jit_set_fusion_strategy(strategy)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_ir_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_ir_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b2c55288ebde2b3597ede507ff9a8e2384d03bc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_ir_utils.py
@@ -0,0 +1,32 @@
+from types import TracebackType
+
+import torch
+
+
+class _InsertPoint:
+    def __init__(
+        self,
+        insert_point_graph: torch._C.Graph,
+        insert_point: torch._C.Node | torch._C.Block,
+    ) -> None:
+        self.insert_point = insert_point
+        self.g = insert_point_graph
+        self.guard = None
+
+    def __enter__(self) -> None:
+        self.prev_insert_point = self.g.insertPoint()
+        self.g.setInsertPoint(self.insert_point)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self.g.setInsertPoint(self.prev_insert_point)
+
+
+def insert_point_guard(
+    self: torch._C.Graph, insert_point: torch._C.Node | torch._C.Block
+) -> _InsertPoint:
+    return _InsertPoint(self, insert_point)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_logging.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0687023386ee7745af4e2cb0731f76528125477
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_logging.py
@@ -0,0 +1,11 @@
+import torch
+
+
+add_stat_value = torch.ops.prim.AddStatValue
+
+set_logger = torch._C._logging_set_logger
+LockingLogger = torch._C.LockingLogger
+AggregationType = torch._C.AggregationType
+NoopLogger = torch._C.NoopLogger
+
+time_point = torch.ops.prim.TimePoint
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_monkeytype_config.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_monkeytype_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4cf1089494cea6141ccaf32bc7fd3a61548bc1e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_monkeytype_config.py
@@ -0,0 +1,195 @@
+# mypy: allow-untyped-defs
+import inspect
+import sys
+import typing
+from collections import defaultdict
+from collections.abc import Iterable
+from pathlib import Path
+from types import CodeType
+
+import torch
+
+
+_IS_MONKEYTYPE_INSTALLED = True
+try:
+    import monkeytype  # type: ignore[import]
+
+    # pyrefly: ignore [import-error, missing-import]
+    from monkeytype import trace as monkeytype_trace
+    from monkeytype.config import _startswith, LIB_PATHS  # type: ignore[import]
+    from monkeytype.db.base import (  # type: ignore[import]
+        CallTraceStore,
+        CallTraceStoreLogger,
+        CallTraceThunk,
+    )
+    from monkeytype.tracing import CallTrace, CodeFilter  # type: ignore[import]
+except ImportError:
+    _IS_MONKEYTYPE_INSTALLED = False
+
+
+# Checks whether a class is defined in `torch.*` modules
+def is_torch_native_class(cls):
+    if not hasattr(cls, "__module__"):
+        return False
+
+    parent_modules = cls.__module__.split(".")
+    if not parent_modules:
+        return False
+
+    root_module = sys.modules.get(parent_modules[0])
+    return root_module is torch
+
+
+def get_type(type):
+    """Convert the given type to a torchScript acceptable format."""
+    if isinstance(type, str):
+        return type
+    elif inspect.getmodule(type) == typing:
+        # If the type is a type imported from typing
+        # like Tuple, List, Dict then replace `typing.`
+        # with a null string. This needs to be done since
+        # typing.List is not accepted by TorchScript.
+        type_to_string = str(type)
+        return type_to_string.replace(type.__module__ + ".", "")
+    elif is_torch_native_class(type):
+        # If the type is a subtype of torch module, then TorchScript expects a fully qualified name
+        # for the type which is obtained by combining the module name and type name.
+        return type.__module__ + "." + type.__name__
+    else:
+        # For all other types use the name for the type.
+        return type.__name__
+
+
+def get_optional_of_element_type(types):
+    """Extract element type, return as `Optional[element type]` from consolidated types.
+
+    Helper function to extracts the type of the element to be annotated to Optional
+    from the list of consolidated types and returns `Optional[element type]`.
+    TODO: To remove this check once Union support lands.
+    """
+    elem_type = types[1] if type(None) is types[0] else types[0]
+    elem_type = get_type(elem_type)
+
+    # Optional type is internally converted to Union[type, NoneType], which
+    # is not supported yet in TorchScript. Hence, representing the optional type as string.
+    return "Optional[" + elem_type + "]"
+
+
+def get_qualified_name(func):
+    return func.__qualname__
+
+
+if _IS_MONKEYTYPE_INSTALLED:
+
+    class JitTypeTraceStoreLogger(CallTraceStoreLogger):
+        """A JitTypeCallTraceLogger that stores logged traces in a CallTraceStore."""
+
+        def log(self, trace: CallTrace) -> None:
+            # pyrefly: ignore [missing-attribute]
+            self.traces.append(trace)
+
+    class JitTypeTraceStore(CallTraceStore):
+        def __init__(self) -> None:
+            super().__init__()
+            # A dictionary keeping all collected CallTrace
+            # key is fully qualified name of called function
+            # value is list of all CallTrace
+            self.trace_records: dict[str, list] = defaultdict(list)
+
+        def add(self, traces: Iterable[CallTrace]) -> None:
+            for t in traces:
+                qualified_name = get_qualified_name(t.func)
+                self.trace_records[qualified_name].append(t)
+
+        def filter(
+            self,
+            qualified_name: str,
+            qualname_prefix: str | None = None,
+            limit: int = 2000,
+        ) -> list[CallTraceThunk]:
+            return self.trace_records[qualified_name]
+
+        def analyze(self, qualified_name: str) -> dict:
+            # Analyze the types for the given module
+            # and create a dictionary of all the types
+            # for arguments.
+            records = self.trace_records[qualified_name]
+            all_args = defaultdict(set)
+            for record in records:
+                for arg, arg_type in record.arg_types.items():
+                    all_args[arg].add(arg_type)
+            return all_args
+
+        def consolidate_types(self, qualified_name: str) -> dict:
+            all_args = self.analyze(qualified_name)
+            # If there are more types for an argument,
+            # then consolidate the type to `Any` and replace the entry
+            # by type `Any`.
+            for arg, types in all_args.items():
+                types = list(types)
+                type_length = len(types)
+                if type_length == 2 and type(None) in types:
+                    # TODO: To remove this check once Union support in TorchScript lands.
+                    all_args[arg] = get_optional_of_element_type(types)
+                elif type_length > 1:
+                    all_args[arg] = "Any"
+                elif type_length == 1:
+                    all_args[arg] = get_type(types[0])
+            return all_args
+
+        def get_args_types(self, qualified_name: str) -> dict:
+            return self.consolidate_types(qualified_name)
+
+    class JitTypeTraceConfig(monkeytype.config.Config):
+        def __init__(self, s: JitTypeTraceStore) -> None:
+            super().__init__()
+            self.s = s
+
+        def trace_logger(self) -> JitTypeTraceStoreLogger:
+            """Return a JitCallTraceStoreLogger that logs to the configured trace store."""
+            # pyrefly: ignore [bad-argument-count]
+            return JitTypeTraceStoreLogger(self.trace_store())
+
+        def trace_store(self) -> CallTraceStore:
+            return self.s
+
+        def code_filter(self) -> CodeFilter | None:
+            return jit_code_filter
+
+else:
+    # When MonkeyType is not installed, we provide dummy class definitions
+    # for the below classes.
+    class JitTypeTraceStoreLogger:  # type:  ignore[no-redef]
+        def __init__(self) -> None:
+            pass
+
+    class JitTypeTraceStore:  # type:  ignore[no-redef]
+        def __init__(self) -> None:
+            self.trace_records = None
+
+    class JitTypeTraceConfig:  # type:  ignore[no-redef]
+        def __init__(self) -> None:
+            pass
+
+    monkeytype_trace = None  # type: ignore[assignment]  # noqa: F811
+
+
+def jit_code_filter(code: CodeType) -> bool:
+    """Codefilter for Torchscript to trace forward calls.
+
+    The custom CodeFilter is required while scripting a FX Traced forward calls.
+    FX Traced forward calls have `code.co_filename` start with '<' which is used
+    to exclude tracing of stdlib and site-packages in the default code filter.
+    Since we need all forward calls to be traced, this custom code filter
+    checks for code.co_name to be 'forward' and enables tracing for all such calls.
+    The code filter is similar to default code filter for monkeytype and
+    excludes tracing of stdlib and site-packages.
+    """
+    # Filter code without a source file and exclude this check for 'forward' calls.
+    if code.co_name != "forward" and (
+        not code.co_filename or code.co_filename[0] == "<"
+    ):
+        return False
+
+    filename = Path(code.co_filename).resolve()
+    return not any(_startswith(filename, lib_path) for lib_path in LIB_PATHS)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_passes/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_passes/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_passes/_property_propagation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_passes/_property_propagation.py
new file mode 100644
index 0000000000000000000000000000000000000000..c410b8fbb7fd329442aa867c0e39c03cd4f15199
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_passes/_property_propagation.py
@@ -0,0 +1,46 @@
+"""
+Tools to help with tensor property propagation.
+
+This is not intended to be imported directly; please use the exposed
+functionalities in `torch.jit`.
+"""
+
+from typing import Any
+
+import torch
+from torch import TensorType
+from torch._C import Graph
+
+
+def apply_input_props_using_example(graph: Graph, example_input: list[Any]) -> None:
+    """
+    Applies properties for each tensor in the graph inputs
+    using the example supplied.
+    """
+    graph_inputs = list(graph.inputs())
+    if len(graph_inputs) == 0:
+        return
+
+    # Strip self args off for methods
+    in_0 = graph_inputs[0]
+    if isinstance(in_0.type(), torch._C.ClassType) and in_0.debugName() == "self":
+        graph_inputs = graph_inputs[1:]
+
+    if not len(graph_inputs) == len(example_input):
+        raise RuntimeError(
+            "Number of inputs in graph does not match number of inputs in the example"
+        )
+
+    for i, (graph_i, example_i) in enumerate(zip(graph_inputs, example_input)):
+        if example_i is None:
+            continue  # Skip the type check
+
+        if isinstance(example_i, torch.Tensor) != isinstance(
+            graph_i.type(), TensorType
+        ):
+            raise RuntimeError(
+                f"Input {i} does not match type of example", graph_i, example_i
+            )
+
+        if isinstance(example_i, torch.Tensor):
+            graph_i.setType(TensorType.create_from_tensor(example_i))  # type: ignore[arg-type]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_pickle.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_pickle.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d73254f74994aaad6e6da8c09f3783baac09025
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_pickle.py
@@ -0,0 +1,38 @@
+# These functions are referenced from the pickle archives produced by
+# ScriptModule.save()
+
+
+# These (`build_*`) functions used to be used by `pickler.cpp` to specify
+# the type of the list for certain special types, but now all lists get
+# a type attached and restored via `restore_type_tag` below. The legacy
+# functions should stick around for backwards-compatibility.
+
+
+def build_intlist(data: list[int]) -> list[int]:
+    return data
+
+
+def build_tensorlist(data: list[object]) -> list[object]:
+    return data
+
+
+def build_doublelist(data: list[float]) -> list[float]:
+    return data
+
+
+def build_boollist(data: list[bool]) -> list[bool]:
+    return data
+
+
+def build_tensor_from_id(data: int | object) -> int | None:
+    if isinstance(data, int):
+        # just the id, can't really do anything
+        return data
+    return None
+
+
+def restore_type_tag(value: object, type_str: str) -> object:
+    # The type_ptr is used by the jit unpickler to restore the full static type
+    # to container types like list when they are re-loaded, but this doesn't
+    # matter for Python, so just return the plain value
+    return value
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_recursive.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_recursive.py
new file mode 100644
index 0000000000000000000000000000000000000000..a3a93c278f633bac3b10477a6ef57bbd6ac9fa2d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_recursive.py
@@ -0,0 +1,1073 @@
+# mypy: allow-untyped-defs
+import collections
+import functools
+import inspect
+import textwrap
+import types
+import warnings
+
+import torch
+import torch._jit_internal as _jit_internal
+from torch._sources import fake_range
+from torch.jit._builtins import _find_builtin
+from torch.jit._check import AttributeTypeIsSupportedChecker
+from torch.jit._state import _add_script_class, _get_script_class, _python_cu
+from torch.jit.frontend import (
+    get_class_properties,
+    get_default_args,
+    get_jit_class_def,
+    get_jit_def,
+)
+from torch.nn import Module
+
+
+ScriptMethodStub = collections.namedtuple(
+    "ScriptMethodStub", ("resolution_callback", "def_", "original_method")
+)
+PropertyStub = collections.namedtuple("PropertyStub", ("resolution_callback", "def_"))
+
+
+# TODO: there should be a more principled way of doing this.
+ignored_attributes = [
+    "_version",
+    "_parameters",
+    "_buffers",
+    "_non_persistent_buffers_set",
+    "_backward_hooks",
+    "_backward_pre_hooks",
+    "_forward_hooks",
+    "_forward_hooks_with_kwargs",
+    "_forward_pre_hooks",
+    "_forward_pre_hooks_with_kwargs",
+    "_forward_hooks_always_called",
+    "_state_dict_hooks",
+    "_state_dict_pre_hooks",
+    "_load_state_dict_pre_hooks",
+    "_load_state_dict_post_hooks",
+    "_modules",
+    "_initializing",
+    "dump_patches",
+]
+
+
+def _compile_and_register_class(obj, rcb, qualified_name):
+    script_class = _get_script_class(obj)
+
+    if not script_class:
+        ast = get_jit_class_def(obj, obj.__name__)
+        defaults = torch.jit.frontend.get_default_args_for_class(obj)
+        script_class = torch._C._jit_script_class_compile(
+            qualified_name, ast, defaults, rcb
+        )
+        _add_script_class(obj, script_class)
+
+    return script_class
+
+
+def make_stub(func, name):
+    rcb = _jit_internal.createResolutionCallbackFromClosure(func)
+    ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
+    return ScriptMethodStub(rcb, ast, func)
+
+
+def make_stub_from_method(nn_module, method_name):
+    func = getattr(nn_module, method_name)
+    if isinstance(func, ScriptMethodStub):
+        return func
+    # Make sure the name present in the resulting AST will match the name
+    # requested here. The only time they don't match is if you do something
+    # like:
+    #   def _forward(self):
+    #       pass
+    #   forward = _forward
+    # In this case, the actual function object will have the name `_forward`,
+    # even though we requested a stub for `forward`.
+    return make_stub(func, method_name)
+
+
+def make_stubs_from_exported_methods(mod):
+    stubs = []
+    for name in dir(mod):
+        item = getattr(mod, name, None)
+        if (
+            _jit_internal.get_torchscript_modifier(item)
+            is _jit_internal.FunctionModifiers.EXPORT
+        ):
+            stubs.append(make_stub_from_method(mod, name))
+
+    return stubs
+
+
+def jit_ignored_properties(module):
+    user_annotated_ignored_attributes = getattr(
+        module, "__jit_ignored_attributes__", []
+    )
+
+    def get_properties_names(module):
+        return {k for k, v in vars(module).items() if isinstance(v, property)}
+
+    properties = get_properties_names(type(module))
+    user_annoted_ignored_properties = set()
+
+    for ignored_attr in user_annotated_ignored_attributes:
+        if ignored_attr in properties:
+            user_annoted_ignored_properties.add(ignored_attr)
+    return user_annoted_ignored_properties
+
+
+# base types that can be constants
+# in addition, tuples and lists of these base types are also considered constants
+# If you edit this list, then you also need to edit the handlers in
+# ConstantValue in jit/script/init.cpp
+_constant_types = (
+    bool,
+    float,
+    int,
+    str,
+    type(None),
+    torch.device,
+    torch.layout,
+    torch.dtype,
+    torch.qscheme,
+)
+
+
+def _get_valid_constant(attr, v, owner_type):
+    if isinstance(v, _constant_types):
+        return v
+    elif isinstance(v, (tuple, list)):
+        return tuple(_get_valid_constant(attr, x, owner_type) for x in v)
+    constants = ", ".join(torch.typename(typ) for typ in _constant_types)
+    raise TypeError(
+        textwrap.dedent(
+            f"""
+        '{torch.typename(type(v))}' object in attribute '{owner_type}.{attr}' is not a valid constant.
+        Valid constants are:
+        1. a nn.ModuleList
+        2. a value of type {{{constants}}}
+        3. a list or tuple of (2)
+        """
+        )
+    )
+
+
+class SourceContext(torch._C._jit_tree_views.SourceRangeFactory):
+    pass
+
+
+def get_annotations(obj):
+    # In Python-3.10+ it is recommended to use inspect.get_annotations
+    # See https://docs.python.org/3.10/howto/annotations.html
+    # But also, in 3.10 annotations from base class are not inherited
+    # by unannotated derived one, so they must be manually extracted
+    annotations = inspect.get_annotations(obj)
+    if annotations:
+        return annotations
+
+    def get_cls_annotations(cls):
+        cls_annotations = inspect.get_annotations(cls)
+        if cls_annotations:
+            return cls_annotations
+        for base in cls.__bases__:
+            cls_annotations = get_cls_annotations(base)
+            if cls_annotations:
+                return cls_annotations
+        return {}
+
+    cls = obj if isinstance(obj, type) else type(obj)
+    return get_cls_annotations(cls)
+
+
+def infer_concrete_type_builder(nn_module, share_types=True):
+    """
+    Build a ConcreteModuleTypeBuilder from an nn.Module.
+
+    This ConcreteModuleType doesn't have a JIT type associated with it yet, it
+    must be filled in by the caller.
+    """
+    concrete_type_builder = torch._C.ConcreteModuleTypeBuilder(type(nn_module))
+    if isinstance(nn_module, (torch.nn.ModuleDict)):
+        concrete_type_builder.set_module_dict()
+    if isinstance(nn_module, (torch.nn.ModuleList, torch.nn.Sequential)):
+        concrete_type_builder.set_module_list()
+    if isinstance(nn_module, (torch.nn.ParameterList)):
+        concrete_type_builder.set_parameter_list()
+    if isinstance(nn_module, (torch.nn.ParameterDict)):
+        concrete_type_builder.set_parameter_dict()
+
+    class_annotations = get_annotations(nn_module)
+    if isinstance(nn_module, (torch.ao.quantization.QuantWrapper)):
+        class_annotations = {}
+
+    # Get user-annotated ignored attributes.
+    user_annotated_ignored_attributes = getattr(
+        nn_module, "__jit_ignored_attributes__", []
+    )
+    concrete_type_builder.add_ignored_attributes(user_annotated_ignored_attributes)
+    ignored_properties = jit_ignored_properties(nn_module)
+
+    # try to infer the type from type annotation or from the object itself
+    def infer_type(name, item):
+        # The forward function from Module is special; never use this annotations; we
+        # need to infer type directly using JIT.  I originally wanted to write
+        # this test as isinstance(class_annotations[name], Callable) but
+        # isinstance on typing things doesn't seem to work: isinstance(list, Callable)
+        # is also true!
+        inferred = False
+        try:
+            if (
+                name in class_annotations
+                and class_annotations[name]
+                != torch.nn.Module.__annotations__["forward"]
+            ):
+                ann_to_type = torch.jit.annotations.ann_to_type(
+                    class_annotations[name], fake_range()
+                )
+                attr_type = torch._C.InferredType(ann_to_type)
+            elif isinstance(item, torch.jit.Attribute):
+                ann_to_type = torch.jit.annotations.ann_to_type(item.type, fake_range())
+                attr_type = torch._C.InferredType(ann_to_type)
+            else:
+                attr_type = torch._C._jit_try_infer_type(item)
+                inferred = True
+        except RuntimeError as re:
+            raise RuntimeError(f"Error inferring type for {name}: {item}: {re}") from re
+
+        return attr_type, inferred
+
+    added_names = set()
+
+    for name, item in nn_module._parameters.items():
+        if name in user_annotated_ignored_attributes:
+            continue
+
+        if not (item is None or isinstance(item, torch.Tensor)):
+            raise AssertionError(
+                f"Expected parameter '{name}' to be None or Tensor, got {type(item)}"
+            )
+        attr_type, _ = infer_type(name, item)
+        # We currently have the invariant in various places in our code
+        # that parameters must be Tensors. However, the nn.Module API also
+        # allows NoneType parameters. These parameters are not returned as
+        # part of `parameters()` and its variants, but are available
+        # through direct attribute access.
+        concrete_type_builder.add_attribute(name, attr_type.type(), True, False)
+        added_names.add(name)
+
+    for name, item in nn_module._buffers.items():
+        if name in user_annotated_ignored_attributes:
+            continue
+
+        if not (item is None or isinstance(item, torch.Tensor)):
+            raise AssertionError(
+                f"Expected buffer '{name}' to be None or Tensor, got {type(item)}"
+            )
+        attr_type, _ = infer_type(name, item)
+        concrete_type_builder.add_attribute(name, attr_type.type(), False, True)
+        added_names.add(name)
+
+    for name, item in nn_module._modules.items():
+        if name in user_annotated_ignored_attributes:
+            continue
+
+        attr_type, _ = infer_type(name, item)
+        if item is None:
+            # Modules can be None. We don't have direct support for optional
+            # Modules, so the register it as an NoneType attribute instead.
+            concrete_type_builder.add_attribute(name, attr_type.type(), False, False)
+            continue
+        if attr_type.success():
+            if not attr_type.type().is_interface_type():
+                raise AssertionError(
+                    f"Expected inferred type to be interface type for '{name}'"
+                )
+            # if the type can be inferred, it should be a module interface type
+            sub_concrete_type = torch._C.ConcreteModuleType.from_jit_type(
+                attr_type.type()
+            )
+        else:
+            # otherwise we get the concrete module type for item and add it to concrete_type
+            sub_concrete_type = get_module_concrete_type(item, share_types)
+        concrete_type_builder.add_module(name, sub_concrete_type)
+
+        added_names.add(name)
+
+    # populate constants_set
+    constants_set = set(getattr(nn_module, "__constants__", ()))
+
+    # Constants annotated via `Final[T]` rather than being added to `__constants__`
+    for name, ann in class_annotations.items():
+        if torch._jit_internal.is_final(ann):
+            constants_set.add(name)
+
+    for name in constants_set:
+        if name in added_names:
+            # TODO: We should really error in this case, but its bc-breaking so
+            # we need to warn for at least one release
+            if name in nn_module._modules:
+                hint = "submodule"
+            elif name in nn_module._buffers:
+                hint = "buffer"
+            elif name in nn_module._parameters:
+                hint = "parameter"
+            else:
+                raise AssertionError(
+                    "added_names must be submodule, parameter, or buffer"
+                )
+
+            warnings.warn(
+                f"'{name}' was found in ScriptModule constants, "
+                f" but it is a non-constant {hint}. Consider removing it.",
+                stacklevel=2,
+            )
+            continue
+        if not hasattr(nn_module, name):
+            # TODO: We should really error in this case, but its bc-breaking so
+            # we need to warn for at least one release
+            warnings.warn(
+                f"'{name}' was found in ScriptModule constants, "
+                "but was not actually set in __init__. "
+                "Consider removing it.",
+                stacklevel=2,
+            )
+            continue
+        value = getattr(nn_module, name)
+        concrete_type_builder.add_constant(
+            name, _get_valid_constant(name, value, type(nn_module).__name__)
+        )
+        added_names.add(name)
+
+    # populate overloads
+    overloads = getattr(nn_module, "__overloads__", {})
+    # update with any annotated overloads
+    overloads.update(
+        get_overload_name_mapping(
+            get_overload_annotations(nn_module, ignored_properties)
+        )
+    )
+    for name, overloaded_names in overloads.items():
+        concrete_type_builder.add_overload(name, overloaded_names)
+
+    for name, value in nn_module.__dict__.items():
+        if name in ignored_attributes or name.startswith("__"):
+            # Python objects have lots of random attributes attached to them;
+            # PyTorch adds a few more. Prevent these from getting compiled.
+            continue
+
+        if name in user_annotated_ignored_attributes:
+            continue
+
+        if name in added_names:
+            # Don't re-add anything we already added
+            continue
+
+        isoverloadpacket = isinstance(value, torch._ops.OpOverloadPacket)
+        if isoverloadpacket:
+            value = value.op
+        # Handle Python function attributes
+        if inspect.isfunction(value):
+            try:
+                scripted_fn = torch.jit.script(value)
+                concrete_type_builder.add_function_attribute(
+                    name, torch._C._jit_try_infer_type(scripted_fn).type(), value
+                )
+            except Exception as e:
+                # If we fail to script the function, it isn't a hard error.
+                # Instead, we will add it to the list of attributes we failed
+                # to convert, with the compilation error.
+                hint = (
+                    "(This function exists as an attribute on the Python module, "
+                    "but we failed to compile it to a TorchScript function. "
+                    f"\nThe error stack is reproduced here:\n{e})"
+                )
+                concrete_type_builder.add_failed_attribute(name, hint)
+
+            continue
+
+        # Handle calls to builtin functions (either bespoke builtins from torch.jit._builtins or
+        # a call to an aten function like torch.add)
+        builtin_symbol_name = _find_builtin(value)
+        if builtin_symbol_name:
+            concrete_type_builder.add_builtin_function(name, builtin_symbol_name)
+            continue
+
+        # Handle Script function attributes
+        if isinstance(value, torch.jit.ScriptFunction):
+            concrete_type_builder.add_function_attribute(
+                name, torch._C._jit_try_infer_type(value).type(), value
+            )
+            continue
+
+        # If we got here, this is a regular "data" attribute, add it to the concrete type
+        attr_type, inferred = infer_type(name, value)
+        if attr_type.success():
+            concrete_type_builder.add_attribute(name, attr_type.type(), False, False)
+        else:
+            # TODO: could add more detail here. For example, what the user should do
+            # when the pytype is `list` or `NoneType`
+            inferred_msg = (
+                "Its type was inferred; try adding a type annotation for the attribute."
+                if inferred
+                else ""
+            )
+            additional_info = f"{attr_type.reason()}. {inferred_msg}"
+            hint = (
+                "(This attribute exists on the Python module, "
+                f"but we failed to convert Python type: '{torch.typename(type(value))}' "
+                f"to a TorchScript type. {additional_info})"
+            )
+            concrete_type_builder.add_failed_attribute(name, hint)
+
+    # add hooks to concrete type
+    for hook in nn_module._forward_hooks.values():
+        concrete_type_builder.add_forward_hook(hook)
+    for pre_hook in nn_module._forward_pre_hooks.values():
+        concrete_type_builder.add_forward_pre_hook(pre_hook)
+
+    return concrete_type_builder
+
+
+class ConcreteTypeStore:
+    type_store: dict[type[Module], list[torch._C.ConcreteModuleType]]
+    methods_compiled: set[torch._C.ConcreteModuleType]
+
+    def __init__(self) -> None:
+        # Python module type => List[ConcreteModuleType)]
+        self.type_store = {}
+        # ConcreteTypes that have had their methods already compiled
+        self.methods_compiled = set()
+
+    def get_or_create_concrete_type(self, nn_module):
+        """Infer a ConcreteType from this `nn.Module` instance. Underlying JIT types are reused if possible."""
+        concrete_type_builder = infer_concrete_type_builder(nn_module)
+
+        nn_module_type = type(nn_module)
+        if nn_module_type not in self.type_store:
+            self.type_store[nn_module_type] = []
+
+        # Search the type store for an already-available JIT type
+        known_types = self.type_store[nn_module_type]
+        for known_type in known_types:
+            if known_type.equals(concrete_type_builder):
+                return known_type
+
+        # We didn't find anything; generate a new JIT type from this concrete type
+        concrete_type = concrete_type_builder.build()
+        self.type_store[nn_module_type].append(concrete_type)
+        return concrete_type
+
+
+concrete_type_store = ConcreteTypeStore()
+
+
+def create_methods_and_properties_from_stubs(
+    concrete_type, method_stubs, property_stubs
+) -> None:
+    method_defs = [m.def_ for m in method_stubs]
+    method_rcbs = [m.resolution_callback for m in method_stubs]
+    method_defaults = [get_default_args(m.original_method) for m in method_stubs]
+
+    property_defs = [p.def_ for p in property_stubs]
+    property_rcbs = [p.resolution_callback for p in property_stubs]
+
+    concrete_type._create_methods_and_properties(
+        property_defs, property_rcbs, method_defs, method_rcbs, method_defaults
+    )
+
+
+def create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs) -> None:
+    hook_defs = [h.def_ for h in hook_stubs]
+    hook_rcbs = [h.resolution_callback for h in hook_stubs]
+
+    pre_hook_defs = [h.def_ for h in pre_hook_stubs]
+    pre_hook_rcbs = [h.resolution_callback for h in pre_hook_stubs]
+
+    concrete_type._create_hooks(hook_defs, hook_rcbs, pre_hook_defs, pre_hook_rcbs)
+
+
+def get_module_concrete_type(nn_module, share_types=True):
+    """
+    Get a concrete type for nn_modules.
+
+    If share_types is True, the concrete type is fetched from concrete_type_store.
+    If it is False, a new concrete type is created without first searching concrete_type_store.
+
+    Args:
+        nn_module:  The original Python nn.Module that we are creating a ScriptModule for.
+        share_types = Whether to share underlying JIT types between modules (if possible).
+
+    Returns:
+        A concrete type for nn_module.
+    """
+    if not isinstance(nn_module, Module):
+        raise AssertionError(f"Expected Module, got {type(nn_module)}")
+    if isinstance(nn_module, torch.jit.ScriptModule) and hasattr(
+        nn_module, "_concrete_type"
+    ):
+        return nn_module._concrete_type
+
+    if share_types:
+        # Look into the store of cached JIT types
+        concrete_type = concrete_type_store.get_or_create_concrete_type(nn_module)
+    else:
+        # Get a concrete type directly, without trying to reuse an existing JIT
+        # type from the type store.
+        concrete_type_builder = infer_concrete_type_builder(nn_module, share_types)
+        concrete_type_builder.set_poisoned()
+        concrete_type = concrete_type_builder.build()
+
+    return concrete_type
+
+
+def create_script_class(obj):
+    """
+    Create and return a RecursiveScriptClass instance from a Python object.
+
+    Arguments:
+        obj: A Python object.
+    """
+    qualified_class_name = _jit_internal._qualified_name(type(obj))
+    rcb = _jit_internal.createResolutionCallbackForClassMethods(type(obj))
+    # Script the type of obj if it hasn't already been scripted.
+    _compile_and_register_class(type(obj), rcb, qualified_class_name)
+    class_ty = _python_cu.get_class(qualified_class_name)
+    # Create an empty torch._C.ScriptObject with the scripted type.
+    cpp_object = torch._C._create_object_with_type(class_ty)
+    # Copy all of the attributes over to the torch._C.ScriptObject.
+    for name, value in obj.__dict__.items():
+        cpp_object.setattr(name, value)
+
+    # Wrap the torch._C.ScriptObject in a RecursiveScriptClass instance.
+    return wrap_cpp_class(cpp_object)
+
+
+def create_script_module(nn_module, stubs_fn, share_types=True, is_tracing=False):
+    """
+    Create a new ScriptModule from an nn.Module.
+
+    Args:
+        nn_module:  The original Python nn.Module that we are creating a ScriptModule for.
+        stubs_fn:  Lambda that takes an nn.Module and generates a list of ScriptMethodStubs to compile.
+        share_types:  Whether to share underlying JIT types between modules (if possible).
+            NOTE: Only set to False this when we cannot guarantee type sharing will work
+                correctly. This only happens today for traced modules, where the same
+                module can produce different traced methods depending on the inputs.
+        is_tracing: Whether this function is called during tracing or scripting. If tracing,
+                we don't need to do AttributeTypeIsSupportedChecker because all the unsupported
+                attributes will be baked as constant in the tracing graph. In addition,
+                this check significantly slows down the traced modules when the module size is big.
+    """
+    if isinstance(nn_module, torch.jit.RecursiveScriptModule):
+        raise AssertionError("Cannot script a RecursiveScriptModule (already compiled)")
+    check_module_initialized(nn_module)
+    concrete_type = get_module_concrete_type(nn_module, share_types)
+    if not is_tracing:
+        AttributeTypeIsSupportedChecker().check(nn_module)
+    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
+
+
+def create_script_module_impl(nn_module, concrete_type, stubs_fn):
+    """
+    Convert an nn.Module to a RecursiveScriptModule.
+
+    Args:
+        nn_module:  The original Python nn.Module that we are creating a ScriptModule for.
+        concrete_type:  The fully initialized ConcreteType of the module.
+        stubs_fn:  Lambda that takes an nn.Module and generates a list of ScriptMethodStubs to compile.
+    """
+    cpp_module = torch._C._create_module_with_type(concrete_type.jit_type)
+    method_stubs = stubs_fn(nn_module)
+    property_stubs = get_property_stubs(nn_module)
+    hook_stubs, pre_hook_stubs = get_hook_stubs(nn_module)
+    ignored_properties = jit_ignored_properties(nn_module)
+
+    def init_fn(script_module) -> None:
+        # Initialize the ScriptModule:
+        # 1. Copy the attributes/parameters/buffers from the original `nn_module` to the new ScriptModule.
+        for name in concrete_type.get_attributes():
+            orig_value = getattr(nn_module, name)
+            orig_value = (
+                orig_value.value
+                if isinstance(orig_value, torch.jit.Attribute)
+                else orig_value
+            )
+            cpp_module.setattr(name, orig_value)
+
+        # 2. Copy the submodules from the original `nn_module` to the new ScriptModule,
+        #    recursively scripting them.
+        for name, sub_concrete_type in concrete_type.get_modules():
+            orig_value = getattr(nn_module, name)
+            if not isinstance(orig_value, Module):
+                raise AssertionError(f"Expected Module but got {type(orig_value)}")
+            module_type = sub_concrete_type.jit_type
+            if isinstance(module_type, torch._C.InterfaceType):
+                # use the interface inference rule to compile the module
+                scripted = interface_script(module_type, orig_value)
+            elif isinstance(orig_value, torch.jit.ScriptModule):
+                scripted = orig_value
+            else:
+                # always reuse the provided stubs_fn to infer the methods to compile
+                scripted = create_script_module_impl(
+                    orig_value, sub_concrete_type, stubs_fn
+                )
+
+            cpp_module.setattr(name, scripted)
+            script_module._modules[name] = scripted
+
+        # 3. Copy @ignored/@unused methods and attrs from the original `nn_module` to the new ScriptModule.
+        #    This ensures we can access these Python methods on the ScriptModule.
+        for name in dir(nn_module):
+            if name in ignored_properties:
+                continue
+            item = getattr(nn_module, name, None)
+            if inspect.ismethod(item) and _jit_internal.is_ignored_fn(item):
+                unbound_function = getattr(nn_module, name).__func__
+                bound_method = unbound_function.__get__(script_module)
+                setattr(script_module, name, bound_method)
+            elif concrete_type.is_ignored_attribute(name):
+                setattr(script_module, name, item)
+
+        # For convenience, attach the concrete type to the new ScriptModule
+        script_module._concrete_type = concrete_type
+
+    # Actually create the ScriptModule, initializing it with the function we just defined
+    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
+
+    # Compile methods if necessary
+    if concrete_type not in concrete_type_store.methods_compiled:
+        create_methods_and_properties_from_stubs(
+            concrete_type, method_stubs, property_stubs
+        )
+        # Create hooks after methods to ensure no name collisions between hooks and methods.
+        # If done before, hooks can overshadow methods that aren't exported.
+        create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs)
+        torch._C._run_emit_module_hook(cpp_module)
+        concrete_type_store.methods_compiled.add(concrete_type)
+
+    # Copy the forward hooks and pre-hooks to the new ScriptModule
+    # to allow the hooks to be run from eager as ScriptFunctions
+    for idx, fn in enumerate(script_module._c._get_forward_pre_hooks()):
+        script_module._forward_pre_hooks[idx] = fn
+    for idx, fn in enumerate(script_module._c._get_forward_hooks()):
+        script_module._forward_hooks[idx] = fn
+
+    # Special handling so methods like __len__ work in script methods on classes derived from containers
+    if (
+        isinstance(
+            nn_module, (torch.nn.ModuleList, torch.nn.Sequential, torch.nn.ModuleDict)
+        )
+        and "__len__" not in cpp_module._method_names()
+    ):
+        script_module.define(f"def __len__(self):\n   return {len(nn_module)}\n")
+    if (
+        isinstance(nn_module, torch.nn.ModuleDict)
+        and "__contains__" not in cpp_module._method_names()
+    ):
+        if len(nn_module.keys()):
+            keys = repr(list(nn_module.keys()))
+            script_module.define(
+                f"def __contains__(self, key: str):\n   return key in {keys}\n"
+            )
+        else:
+            script_module.define("def __contains__(self, key: str):\n   return False\n")
+
+    # Make the compiled methods available to the Python ScriptModule class.
+    for method_stub in method_stubs:
+        if method_stub.original_method is None:
+            # define()'d methods don't have an Python original_method, so we
+            # don't need to do any Python re-wrapping stuff
+            continue
+
+        name = method_stub.original_method.__name__
+        if name != method_stub.def_.name().name:
+            # TODO: Why skip this? Because @torch.jit._overload_method will
+            # mangle the name of the function.
+            continue
+        script_method = cpp_module._get_method(name)
+
+        # Wrap the original to propagate docstrings and such.
+        # TODO: we don't currently do this functions that are recursively
+        # compiled, we should.
+        wrapped_script_method = functools.wraps(method_stub.original_method)(
+            script_method
+        )
+
+        # Add the methods to the script_module directly. This ensures they will
+        # be found first when `name` is looked up (as opposed to the stubs or
+        # nn.Module.forward)
+        script_module.__dict__[name] = wrapped_script_method
+
+    # Make module properties available on the Python ScriptModule class.
+    for property_stub in property_stubs:
+        property_name = property_stub.def_.name().name
+        fget = cpp_module._get_method(property_stub.def_.getter_name().name)
+        # Setter is optional, so it may not exist.
+        setter_name = property_stub.def_.setter_name()
+        fset = cpp_module._get_method(setter_name.name) if setter_name else None
+        script_module.__dict__[property_name] = property(property_name, fget, fset)  # type: ignore[arg-type]
+
+    # copy over python methods to script module if they aren't defined on the script module
+    # this is currently an internal api used only on module containers
+    for name in dir(nn_module):
+        if name in ignored_properties:
+            continue
+        item = getattr(nn_module, name, None)
+        if (
+            _jit_internal.get_torchscript_modifier(item)
+            is _jit_internal.FunctionModifiers.COPY_TO_SCRIPT_WRAPPER
+        ):
+            add_python_attr_to_scripted_model(script_module, nn_module, name)
+
+    return script_module
+
+
+# We define shims of certain attributes on the RecursiveScriptModule to support
+# magic methods. To check if a script model defines an attribute we need
+# to also check that the attribute is not the shim
+def script_model_defines_attr(script_model, attr):
+    script_attr = getattr(script_model, attr, None)
+    if script_attr is None:
+        return False
+    default_attr = getattr(torch.jit.RecursiveScriptModule, attr, None)
+    if default_attr is None:
+        return False
+    return script_attr != default_attr
+
+
+def add_python_attr_to_scripted_model(script_model, orig, attr) -> None:
+    if hasattr(orig, attr) and script_model_defines_attr(script_model, attr):
+        setattr(script_model, attr, getattr(orig, attr))
+
+
+def get_overload_annotations(mod, jit_ignored_properties):
+    # original function => [(mangled overload name, overload function)]
+    overloads = {}
+
+    for name in dir(type(mod)):
+        if name in jit_ignored_properties:
+            continue
+        item = getattr(mod, name, None)
+        if not callable(item):
+            continue
+
+        # builtin functions like repr() in python 2 do not have __module__ defined
+        if hasattr(item, "__module__") and item.__module__ is not None:
+            method_overloads = _jit_internal._get_overloaded_methods(
+                item, mod.__class__
+            )
+            if method_overloads is None:
+                continue
+
+            # pyrefly: ignore [missing-attribute]
+            if item.__func__ in method_overloads:
+                raise RuntimeError(
+                    _jit_internal.get_overload_no_implementation_error_message(
+                        "method", item.__func__
+                    )
+                )
+
+            names = [name + "__" + str(i) for i in range(len(method_overloads))]
+            overloads[item] = list(zip(names, method_overloads))
+
+    return overloads
+
+
+def get_overload_name_mapping(overload_info):
+    # Same format as __overloads__
+    # original function => [overload names]
+    overload_name_mappings: dict[str, list[str]] = {}
+    for orig_fn, overloads in overload_info.items():
+        original_name = orig_fn.__name__
+        if original_name not in overload_name_mappings:
+            overload_name_mappings[original_name] = []
+
+        for overload_name, _ in overloads:
+            overload_name_mappings[original_name].append(overload_name)
+    return overload_name_mappings
+
+
+def _check_no_signature(func) -> None:
+    signature = torch.jit.annotations.get_signature(
+        func, None, fake_range(), inspect.ismethod(func)
+    )
+    if signature is None:
+        qual_name = _jit_internal._qualified_name(func)
+        raise RuntimeError(
+            f"Must explicitly add type annotations to overloaded functions: {qual_name}"
+        )
+
+
+def make_stubs_for_overloads(overload_info):
+    overload_stubs = []
+    for orig_fn, overloads in overload_info.items():
+        orig_ast = get_jit_def(
+            orig_fn, orig_fn.__name__, self_name="RecursiveScriptModule"
+        )
+        for overload_name, overload_fn in overloads:
+            _check_no_signature(overload_fn)
+            over_ast = get_jit_def(
+                overload_fn, overload_fn.__name__, self_name="RecursiveScriptModule"
+            )
+            new_ast = torch._C._replace_overloaded_method_decl(
+                over_ast.decl(), orig_ast, overload_name
+            )
+            _rcb = _jit_internal.createResolutionCallbackFromClosure(orig_fn)
+            overload_stubs.append(ScriptMethodStub(_rcb, new_ast, overload_fn))
+    return overload_stubs
+
+
+def check_module_initialized(mod) -> None:
+    if not isinstance(mod, torch.nn.Module):
+        raise AssertionError(f"Expected torch.nn.Module, got {type(mod)}")
+    if not hasattr(mod, "_parameters"):
+        raise RuntimeError(
+            f"'{torch.typename(type(mod))}' has not been initialized, did you forget to call 'super()'?"
+        )
+
+    # This is to avoid importing torch.distributed.nn
+    if not hasattr(mod, "remote_parameters"):
+        for name, param in mod._parameters.items():
+            if param is not None and torch.nn.parameter.is_lazy(param):
+                raise RuntimeError(
+                    f"'{torch.typename(type(mod))}' has uninitialized parameters {name}. Did you forget to run a forward pass?"
+                )
+        for name, buf in mod._buffers.items():
+            if buf is not None and torch.nn.parameter.is_lazy(buf):
+                raise RuntimeError(
+                    f"'{torch.typename(type(mod))}' has uninitialized buffers {name}. Did you forget to run a forward pass?"
+                )
+
+
+def infer_methods_to_compile(nn_module):
+    """Implement the default rules for which methods should act as starting points for compilation.
+
+    (TODO add a link when the rules are published).
+    """
+    check_module_initialized(nn_module)
+    ignored_properties = jit_ignored_properties(nn_module)
+
+    methods: list[str] = []
+    if hasattr(nn_module, "forward") and not _jit_internal.is_ignored_fn(
+        nn_module.forward
+    ):
+        forward_func = getattr(nn_module.forward, "__func__", None)
+        module_forward = getattr(torch.nn.Module, "forward", None)
+        if forward_func != module_forward:
+            methods = ["forward"]
+
+    exported = []
+    for name in dir(nn_module):
+        if name in ignored_properties:
+            continue
+        item = getattr(nn_module, name, None)
+        if (
+            _jit_internal.get_torchscript_modifier(item)
+            is _jit_internal.FunctionModifiers.EXPORT
+        ):
+            exported.append(name)
+
+    methods = methods + exported
+
+    overload_name_mappings = dict(getattr(nn_module, "__overloads__", {}))
+    overload_info = get_overload_annotations(nn_module, ignored_properties)
+    overload_name_mappings.update(get_overload_name_mapping(overload_info))
+    overload_stubs = make_stubs_for_overloads(overload_info)
+
+    nn_module.__overloads__ = overload_name_mappings
+
+    # we shouldn't directly compile overloaded methods, just its overloads
+    def ignore_overloaded(method_name):
+        return method_name not in overload_name_mappings
+
+    filtered_methods = filter(ignore_overloaded, methods)
+
+    # Unique the methods. We don't want to use a set to store the methods because it
+    # introduces non-determinism to compile order.
+    uniquer: set[str] = set()
+    uniqued_methods = []
+    for name in filtered_methods:
+        if name in uniquer:
+            continue
+        uniqued_methods.append(name)
+        uniquer.add(name)
+
+    stubs = [make_stub_from_method(nn_module, method) for method in uniqued_methods]
+    return overload_stubs + stubs
+
+
+def get_hook_stubs(nn_module):
+    """Return forward hook and pre_hook ScriptModuleStubs."""
+    check_module_initialized(nn_module)
+    hook_map: dict = {}
+
+    hook_stubs = []
+    for hook in nn_module._forward_hooks.values():
+        if hook.__name__ in hook_map:
+            if id(hook) != id(hook_map[hook.__name__]):
+                raise RuntimeError(
+                    f"Hook '{hook.__name__}' on {type(nn_module).__name__} "
+                    "has at least two different python definitions."
+                    " Please use unique names for all hooks."
+                )
+        else:
+            hook_map[hook.__name__] = hook
+        hook_stubs.append(make_stub(hook, hook.__name__))
+
+    pre_hook_stubs = []
+    for pre_hook in nn_module._forward_pre_hooks.values():
+        if pre_hook.__name__ in hook_map:
+            if id(pre_hook) != id(hook_map[pre_hook.__name__]):
+                raise RuntimeError(
+                    f"Pre-hook '{pre_hook.__name__}' on {type(nn_module).__name__} "
+                    "has at least two different python definitions."
+                    " Please use unique names for all hooks."
+                )
+        else:
+            hook_map[pre_hook.__name__] = pre_hook
+        pre_hook_stubs.append(make_stub(pre_hook, pre_hook.__name__))
+
+    return hook_stubs, pre_hook_stubs
+
+
+def get_property_stubs(nn_module):
+    """Create property stubs for the properties of the module by creating method stubs for the getter and setter."""
+    module_ty = type(nn_module)
+    properties_asts = get_class_properties(module_ty, self_name="RecursiveScriptModule")
+    rcbs = {}
+
+    for name in dir(module_ty):
+        item = getattr(module_ty, name, None)
+        if isinstance(item, property):
+            if not item.fget:
+                raise RuntimeError(
+                    f"Property {name} of {nn_module.__name__} must have a getter"
+                )
+
+            rcbs[name] = _jit_internal.createResolutionCallbackFromClosure(item.fget)
+
+    stubs = [PropertyStub(rcbs[ast.name().name], ast) for ast in properties_asts]
+    return stubs
+
+
+def interface_script(mod_interface, nn_module):
+    """
+    Make a ScriptModule from an nn.Module, using the interface methods rule for determining which methods to compile.
+
+    Args:
+        mod_interface: the interface type that the module have
+        nn_module:  The original Python nn.Module that we are creating a ScriptModule for.
+    """
+    if isinstance(nn_module, torch.jit.ScriptModule):
+        return nn_module
+
+    check_module_initialized(nn_module)
+
+    def infer_interface_methods_to_compile(nn_module):
+        """Rule to infer the methods from the interface type.
+
+        It is used to know which methods need to act as starting points for compilation.
+        """
+        stubs = [
+            make_stub_from_method(nn_module, method)
+            for method in mod_interface.getMethodNames()
+        ]
+        return stubs
+
+    return create_script_module(nn_module, infer_interface_methods_to_compile)
+
+
+def try_compile_fn(fn, loc):
+    if _jit_internal.is_ignored_fn(fn):
+        # Don't do anything for @ignore'd functions
+        return None
+
+    if isinstance(fn, torch.nn.Module):
+        # Since modules are callable pybind recognizes them as functions, but
+        # don't do anything for them
+        return None
+
+    if not inspect.isfunction(fn) and not inspect.ismethod(fn):
+        raise RuntimeError(
+            f"`{fn}` is not a function. Recursive scripting only supports "
+            "Python functions or methods currently.\n"
+            f"Consider manually annotating `{fn}` with @torch.jit.script."
+        )
+
+    # The object returned by __prepare_scriptable__ might have a different closure.
+    # Resolve it here to get the right resolution callback.
+    fn = fn.__prepare_scriptable__() if hasattr(fn, "__prepare_scriptable__") else fn  # type: ignore[operator]
+
+    # We don't have the actual scope where the function was defined, but we can
+    # extract the necessary info from the closed over variables on the function
+    # object
+    rcb = _jit_internal.createResolutionCallbackFromClosure(fn)
+    return torch.jit.script(fn, _rcb=rcb)
+
+
+def wrap_cpp_class(cpp_class):
+    """Wrap this torch._C.Object in a Python RecursiveScriptClass."""
+    return torch.jit.RecursiveScriptClass(cpp_class)
+
+
+def wrap_cpp_module(cpp_module):
+    """Wrap this torch._C.ScriptModule in a Python ScriptModule, recursively for all submodules."""
+
+    def init_fn(script_module) -> None:
+        for name, cpp_module in torch._C.ModuleDict(script_module._c).items():
+            setattr(script_module, name, wrap_cpp_module(cpp_module))
+        script_module._concrete_type = torch._C.ConcreteModuleType.from_jit_type(
+            script_module._c._type()
+        )
+
+        for idx, fn in enumerate(script_module._c._get_forward_pre_hooks()):
+            script_module._forward_pre_hooks[idx] = fn
+        for idx, fn in enumerate(script_module._c._get_forward_hooks()):
+            script_module._forward_hooks[idx] = fn
+
+    return torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
+
+
+def compile_unbound_method(concrete_type, fn):
+    if _jit_internal.is_ignored_fn(fn):
+        return None
+    stub = make_stub(fn, fn.__name__)
+    with torch._jit_internal._disable_emit_hooks():
+        # We don't want to call the hooks here since the graph that is calling
+        # this function is not yet complete
+        create_methods_and_properties_from_stubs(concrete_type, (stub,), ())
+    return stub
+
+
+def lazy_bind(concrete_type, unbound_method):
+    """
+    Return a function that lazily binds `unbound_method` to a provided Module IValue, then invokes the method.
+
+    We do this so that any Python shenanigans that
+    will poison type sharing are impossible at compile time.
+    """
+
+    def lazy_binding_method(cpp_module, *args):
+        def init_fn(script_module) -> None:
+            orig_class = concrete_type.py_class
+
+            # Copy @ignored/@unused methods from the original module to the new one.
+            # This ensures they are available during execution.
+            for name in dir(orig_class):
+                item = getattr(orig_class, name, None)
+                if _jit_internal.is_ignored_fn(item):
+                    setattr(script_module, name, item)
+
+            # Copy constants over so they are available during execution.
+            for name, value in concrete_type.get_constants().items():
+                setattr(script_module, name, value)
+
+        script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
+        method = types.MethodType(unbound_method, script_module)
+        return method(*args)
+
+    # make the lazy binding method "look like" the original method
+    lazy_binding_method.original_fn = unbound_method  # type: ignore[attr-defined]
+    lazy_binding_method.__name__ = unbound_method.__name__
+    torch._jit_internal.copy_torchscript_modifier(unbound_method, lazy_binding_method)
+
+    return lazy_binding_method
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_script.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_script.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb723c62ff6582a7218550fef189c052b8182344
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_script.py
@@ -0,0 +1,1803 @@
+"""TorchScript.
+
+This module contains functionality to support the JIT's scripting frontend, notably:
+    - torch.jit.script
+
+This is not intended to be imported directly; please use the exposed
+functionalities in `torch.jit`.
+"""
+
+import collections
+import copy
+import enum
+import functools
+import inspect
+import pickle
+import sys
+import warnings
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from typing import Any, TypeVar
+from typing_extensions import deprecated, Self
+
+import torch
+import torch._jit_internal as _jit_internal
+from torch._classes import classes
+from torch._jit_internal import _get_model_id, _qualified_name
+from torch._utils_internal import log_torchscript_usage
+from torch.jit._builtins import _register_builtin
+from torch.jit._fuser import _graph_for, _script_method_graph_for
+from torch.jit._monkeytype_config import (
+    JitTypeTraceConfig,
+    JitTypeTraceStore,
+    monkeytype_trace,
+)
+from torch.jit._recursive import (
+    _compile_and_register_class,
+    infer_methods_to_compile,
+    ScriptMethodStub,
+    wrap_cpp_module,
+)
+from torch.jit._state import (
+    _enabled,
+    _set_jit_function_cache,
+    _set_jit_overload_cache,
+    _try_get_jit_cached_function,
+    _try_get_jit_cached_overloads,
+)
+from torch.jit.frontend import get_default_args, get_jit_class_def, get_jit_def
+from torch.nn import Module
+from torch.overrides import (
+    has_torch_function,
+    has_torch_function_unary,
+    has_torch_function_variadic,
+)
+from torch.package import PackageExporter, PackageImporter
+from torch.utils import set_module
+
+from ._serialization import validate_map_location
+
+
+_T = TypeVar("_T")
+
+
+type_trace_db = JitTypeTraceStore()  # DB to hold all call traces from MonkeyType
+
+torch._C.ScriptMethod.graph_for = _script_method_graph_for  # type: ignore[attr-defined]
+torch._C.ScriptFunction.graph_for = _graph_for  # type: ignore[attr-defined]
+ScriptFunction = torch._C.ScriptFunction
+ScriptFunction.__doc__ = """
+Functionally equivalent to a :class:`ScriptModule`, but represents a single
+function and does not have any attributes or Parameters.
+"""
+ScriptFunction.__name__ = "ScriptFunction"
+ScriptFunction.__qualname__ = "torch.jit.ScriptFunction"
+set_module(ScriptFunction, "torch.jit")
+
+
+# Throws an error if a jit function is pickled.
+# Helps to avoid Python crashes for Python versions 3.9.5 + when protocol 0 or 1 is given as an argument.
+def _reduce(cls):
+    raise pickle.PickleError("ScriptFunction cannot be pickled")
+
+
+ScriptFunction.__reduce__ = _reduce  # type: ignore[assignment]
+
+
+if _enabled:
+    Attribute = collections.namedtuple("Attribute", ["value", "type"])
+else:
+
+    def Attribute(value, type):  # type: ignore[no-redef]
+        return value
+
+
+Attribute.__doc__ = """
+    This method is a pass-through function that returns `value`, mostly
+    used to indicate to the TorchScript compiler that the left-hand side
+    expression is a class instance attribute with type of `type`. Note that
+    `torch.jit.Attribute` should only be used in `__init__` method of `jit.ScriptModule`
+    subclasses.
+
+    Though TorchScript can infer correct type for most Python expressions, there are some cases where
+    type inference can be wrong, including:
+
+    - Empty containers like `[]` and `{}`, which TorchScript assumes to be container of `Tensor`
+    - Optional types like `Optional[T]` but assigned a valid value of type `T`, TorchScript would assume
+      it is type `T` rather than `Optional[T]`
+
+    In eager mode, it is simply a pass-through function that returns `value`
+    without other implications.
+
+    Example:
+
+    .. testcode::
+
+        import torch
+        from typing import Dict
+
+        class AttributeModule(torch.jit.ScriptModule):
+            def __init__(self) -> None:
+                super().__init__()
+                self.foo = torch.jit.Attribute(0.1, float)
+
+                # we should be able to use self.foo as a float here
+                assert 0.0 < self.foo
+
+                self.names_ages = torch.jit.Attribute({}, Dict[str, int])
+                self.names_ages["someone"] = 20
+                assert isinstance(self.names_ages["someone"], int)
+
+        m = AttributeModule()
+        # m will contain two attributes
+        # 1. foo of type float
+        # 2. names_ages of type Dict[str, int]
+
+    .. testcleanup::
+
+        del AttributeModule
+        del m
+
+    Note: it's now preferred to instead use type annotations instead of `torch.jit.Attribute`:
+
+    .. testcode::
+
+        import torch
+        from typing import Dict
+
+        class AttributeModule(torch.nn.Module):
+            names: Dict[str, int]
+
+            def __init__(self) -> None:
+                super().__init__()
+                self.names = {}
+
+        m = AttributeModule()
+
+    .. testcleanup::
+
+        del AttributeModule
+        del m
+
+    Args:
+        value: An initial value to be assigned to attribute.
+        type: A Python type
+
+    Returns:
+        Returns `value`
+"""
+
+
+def _get_type_trace_db():
+    # This is a private API. Use of this for external purposes is discouraged.
+    return type_trace_db
+
+
+# Gets a function from the name of a method on a type
+def _get_function_from_type(cls, name):
+    return getattr(cls, name, None)
+
+
+# ScriptClasses must be new-style classes because we construct them using their
+# __new__ method.
+def _is_new_style_class(cls):
+    if hasattr(cls, "__class__"):
+        return "__dict__" in dir(cls) or hasattr(cls, "__slots__")
+
+
+# These OrderedDictWrapper classes replace the actual OrderedDicts in
+# module with versions that get/set properties inside of Module.
+# This allows us to reuse most of nn.Module while still storing the
+# data in C++.
+# Each OrderedDict needs to support:
+#  x not in view
+#  x in view
+#  view[name] = ...
+#  view.values()
+#  del view[name]
+#  view.items()
+#  view.keys()
+#  len(view)
+
+
+class OrderedDictWrapper:
+    def __init__(self, _c):
+        self._c = _c
+
+    def keys(self):
+        return [k for k, v in self.items()]
+
+    def values(self):
+        return [v for k, v in self.items()]
+
+    def __len__(self):
+        return len(self.values())
+
+    def __delitem__(self, k):
+        raise RuntimeError("cannot delete methods or parameters of a script module")
+
+    def items(self):
+        return self._c.items()
+
+    def __setitem__(self, k, v):
+        if k not in self:
+            raise RuntimeError(
+                f"Can't add a new parameter after ScriptModule construction. Tried to add '{k}"
+            )
+        self._c.setattr(k, v)
+
+    def __contains__(self, k):
+        return self._c.contains(k)
+
+    def __getitem__(self, k):
+        if k not in self:
+            raise KeyError(k)
+        return self._c.getattr(k)
+
+
+class OrderedModuleDict(OrderedDictWrapper):
+    def __init__(self, module, python_dict):
+        super().__init__(torch._C.ModuleDict(module))
+        # contains _both_ script modules and non-script python-only modules
+
+        # because script modules are subclassed in python and the
+        # C++ Module class will not hold references to them,
+        # to ensure that you always get the same python value here
+        # we store it in the python dict as well
+        self._python_modules = python_dict
+
+    def items(self):
+        r = self._python_modules.items()
+        return r
+
+    def __contains__(self, k):
+        return k in self._python_modules
+
+    def __setitem__(self, k, v):
+        # Cases where sub-module can be re-assigned after ScriptModule construction
+        # 1. If the attr is an module interface type, it's guaranteed that the module is
+        #    not inlined in the graph, so it's safe to swap a new ScriptModule in.
+        # 2. if the new value if a ScriptModule with the same JIT type, IR won't change
+        #    and it's legit to swap a new module in.
+        # In these two cases we allow swapping a new scripted module and update the
+        # corresponding python module dict to keep sync.
+        # Note: the value to be swapped in has to be ScriptModule instead of nn.Module,
+        # otherwise it's illegal and we throw error.
+        if isinstance(v, ScriptModule):
+            self._c.setattr(k, v)
+            self._python_modules[k] = v
+        else:
+            raise RuntimeError(
+                "Cannot re-assign modules in a ScriptModule with non-scripted "
+                f"module, tried to replace existing module '{k}': {v}"
+            )
+
+    def __getitem__(self, k):
+        return self._python_modules[k]
+
+
+# For each user-defined class that subclasses ScriptModule, this meta-class:
+# (1) finds all the methods annotated with @script_method in a ScriptModule and
+#     removes them from the class attributes
+# (2) puts a wrapper around the class's __init__ method to recursively compile
+#     all of the script_methods with the module after the original __init__ has
+#     run. This has to occur after the user-defined __init__ so that submodules and
+#     parameters are initialized _before_ the script compiler resolve references to
+#     `self.param` or `self.module`.
+class ScriptMeta(type):
+    def __init__(cls, name, bases, attrs):  # noqa: B902
+        # Aggregate all the ScriptMethods and constants from superclasses
+        cls._methods: dict[str, Any] = {}
+        cls._constants_set = set(getattr(cls, "__constants__", ()))
+        for base in reversed(bases):
+            for k, v in getattr(base, "_methods", {}).items():
+                cls._methods[k] = v
+            base_constants: set = getattr(base, "_constants_set", set())
+            cls._constants_set = cls._constants_set.union(base_constants)
+
+        # find all the script methods of the current class
+        for k, v in sorted(attrs.items()):
+            if isinstance(v, ScriptMethodStub):
+                delattr(cls, k)
+                cls._methods[v.original_method.__name__] = v
+
+        if getattr(cls, "_disable_script_meta", False):
+            # We leave built-in ScriptModule types alone, since this metaclass
+            # is only for compiling user classes that inherit from
+            # ScriptModule.
+            super().__init__(name, bases, attrs)
+            return
+
+        original_init = getattr(cls, "__init__", lambda self: None)
+
+        @functools.wraps(original_init)
+        def init_then_script(self, *args, **kwargs):
+            num_methods = len(cls._methods)
+            original_init(self, *args, **kwargs)
+            added_methods_in_init = len(cls._methods) > num_methods
+
+            if type(self) is cls:
+
+                def make_stubs(module):
+                    cls = type(module)
+                    if hasattr(cls, "_methods"):
+                        return [v for k, v in sorted(cls._methods.items())]
+                    else:
+                        return infer_methods_to_compile(module)
+
+                self.__dict__["_actual_script_module"] = (
+                    torch.jit._recursive.create_script_module(
+                        self, make_stubs, share_types=not added_methods_in_init
+                    )
+                )
+
+                # Delete the Python attributes that now shadow the ScriptModule
+                # ones, so that __getattr__ and __setattr__ will properly find
+                # the scripted versions.
+                concrete_type = self._actual_script_module._concrete_type
+                for name in concrete_type.get_attributes():
+                    delattr(self, name)
+                for name, _ in concrete_type.get_modules():
+                    delattr(self, name)
+                for name in ("_parameters", "_buffers", "_modules"):
+                    delattr(self, name)
+
+        cls.__init__ = init_then_script  # type: ignore[misc]
+        super().__init__(name, bases, attrs)
+
+
+class _CachedForward:
+    def __get__(self, obj, cls):
+        return self.__getattr__("forward")  # type: ignore[attr-defined]
+
+
+class ScriptWarning(Warning):
+    pass
+
+
+def script_method(fn):
+    if sys.version_info >= (3, 14):
+        warnings.warn(
+            "`torch.jit.script_method` is not supported in Python 3.14+ and may break. "
+            "Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    else:
+        warnings.warn(
+            "`torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    if not _enabled:
+        return fn
+    # NOTE: we need to traverse two frames here because the meta-class frame
+    # for ScriptModule will be present, as opposed to invoking @script on a
+    # a function or invoking define() on a CompilationUnit.
+    # The stack will look like:
+    #
+    # 0. createResolutionCallback()
+    # 1. script_method()
+    # 2. ScriptModule metaclass frame
+    # 3. Surrounding scope
+    #
+    # createResolutionCallback internally adds 1 to get us to the scope of this
+    # function (the calling function). Adding 2 gets us to the proper surrounding scope.
+    _rcb = _jit_internal.createResolutionCallbackFromFrame(frames_up=2)
+    ast = get_jit_def(fn, fn.__name__, self_name="ScriptModule")
+    return ScriptMethodStub(_rcb, ast, fn)
+
+
+class ConstMap:
+    def __init__(self, const_mapping: Mapping[str, Any]) -> None:
+        self.const_mapping = const_mapping
+
+    def __getattr__(self, attr: str) -> Any:
+        return self.const_mapping[attr]
+
+
+def unpackage_script_module(
+    importer: PackageImporter, script_module_id: str
+) -> torch.nn.Module:
+    """
+    Call by ``torch.package.PackageImporter``'s Pickler's ``persistent_load`` function.
+
+    Performs work of loading and returning a ScriptModule from a ``torch.package`` archive.
+    """
+    if not isinstance(importer.zip_reader, torch._C.PyTorchFileReader):
+        raise RuntimeError(
+            "Loading ScriptObjects from a PackageImporter created from a "
+            "directory is not supported. Use a package archive file instead."
+        )
+    cu = torch._C.CompilationUnit()
+    cpp_module = torch._C._import_ir_module_from_package(
+        cu,
+        importer.zip_reader,
+        importer.storage_context,
+        validate_map_location(importer.last_map_location),
+        script_module_id,
+    )
+    return wrap_cpp_module(cpp_module)
+
+
+if _enabled:
+    _magic_methods = [
+        "__iter__",
+        "__len__",
+        "__neg__",
+        "__mul__",
+        "__contains__",
+        "__add__",
+        "__sub__",
+        "__pow__",
+        "__truediv__",
+        "__mod__",
+        "__ne__",
+        "__eq__",
+        "__lt__",
+        "__gt__",
+        "__le__",
+        "__ge__",
+        "__and__",
+        "__or__",
+        "__xor__",
+        "__getitem__",
+        "__setitem__",
+        "__call__",
+        "__int__",
+        "__float__",
+        "__bool__",
+        "__str__",
+        "__enter__",
+        "__exit__",
+    ]
+
+    class RecursiveScriptClass:
+        """Wrapper for a TorchScript class instance for use in Python.
+
+        An analogue of RecursiveScriptModule for regular objects that are not modules.
+        This class is a wrapper around a torch._C.ScriptObject that represents an instance
+        of a TorchScript class and allows it to be used in Python.
+
+        Attributes:
+            _c [torch._C.ScriptObject]: The C++ object to which attribute lookups and method
+                calls are forwarded.
+            _props [Dict[str, property]]: A dictionary of properties fetched from self._c and
+                exposed on this wrppaer.
+        """
+
+        def __init__(self, cpp_class):
+            super().__init__()
+            self.__dict__["_initializing"] = True
+            self._c = cpp_class
+
+            # Add wrapped object's properties to this class instance.
+            self._props = {
+                prop.name: property(prop.getter, prop.setter)
+                for prop in self._c._properties()
+            }
+
+            self.__dict__["_initializing"] = False
+
+        def __getattr__(self, attr: str) -> Any:
+            if self.__dict__.get("_initializing"):
+                return super().__getattr__(attr)  # type: ignore[misc]
+
+            if attr in self._props:
+                return self._props[attr].fget()  # type: ignore[call-arg, misc]
+
+            return getattr(self._c, attr)
+
+        def __setattr__(self, attr: str, value: Any) -> None:
+            if self.__dict__.get("_initializing"):
+                return super().__setattr__(attr, value)
+
+            if attr in self._props:
+                return self._props[attr].fset(value)  # type: ignore[call-arg, misc]
+
+            setattr(self._c, attr, value)
+
+        # Delegate calls to magic methods like __len__ to the C++ module backing the
+        # RecursiveScriptClass.
+        def forward_magic_method(
+            self, method_name: str, *args: Any, **kwargs: Any
+        ) -> Any:
+            if not self._c._has_method(method_name):
+                raise TypeError
+
+            self_method = self.__getattr__(method_name)
+            return self_method(*args, **kwargs)
+
+        def __getstate__(self):
+            raise pickle.PickleError("ScriptClasses cannot be pickled")
+
+        def __iadd__(self, other: Self) -> Self:
+            if self._c._has_method("__iadd__"):
+                return self.forward_magic_method("__iadd__", other)
+            else:
+                return self.forward_magic_method("__add__", other)
+
+    for method_name in _magic_methods:
+
+        def method_template(self, *args, **kwargs):
+            return self.forward_magic_method(method_name, *args, **kwargs)
+
+        setattr(RecursiveScriptClass, method_name, method_template)
+
+    # this is a Python 'non-data descriptor' that causes the first access
+    # to ScriptModule's forward to look up the forward method and stash
+    # it in the objects dict. Due to the standard rules for attribute lookup,
+    # subsequent lookups will just directly return the previously looked up method.
+    # This is necessary because nn.Module defines forward as a method. If we
+    # did nothing, __getattr__ would not be called. Instead we'd get nn.Module.forward
+    # which always throws an exception.
+
+    class ScriptModule(Module, metaclass=ScriptMeta):
+        r"""Wrapper for C++ torch::jit::Module with methods, attributes, and parameters.
+
+        A wrapper around C++ ``torch::jit::Module``. ``ScriptModule``\s
+        contain methods, attributes, parameters, and
+        constants. These can be accessed the same way as on a normal ``nn.Module``.
+        """
+
+        __jit_unused_properties__ = [
+            "code",
+            "code_with_constants",
+            "graph",
+            "inlined_graph",
+            "original_name",
+        ]
+
+        def __init__(self) -> None:
+            super().__init__()
+
+        forward: Callable[..., Any] = _CachedForward()  # type: ignore[assignment]
+
+        def __getattr__(self, attr: str) -> Any:
+            if "_actual_script_module" not in self.__dict__:
+                return super().__getattr__(attr)
+            return getattr(self._actual_script_module, attr)
+
+        def __setattr__(self, attr: str, value: Any) -> None:
+            if "_actual_script_module" not in self.__dict__:
+                # Unwrap torch.jit.Attribute into a regular setattr + record
+                # the provided type in __annotations__.
+                #
+                # This ensures that if we use the attr again in `__init__`, it
+                # will look like the actual value, not an instance of Attribute.
+                # pyrefly: ignore [invalid-argument]
+                if isinstance(value, Attribute):
+                    # NB: Ensure that we set __annotations__ on the specific
+                    # class in question, and not on a superclass (which would
+                    # be wrong wrong wrong!).
+                    # See also https://github.com/pytorch/pytorch/issues/39463
+                    if "__annotations__" not in self.__class__.__dict__:
+                        self.__class__.__annotations__ = {}
+                    self.__annotations__[attr] = value.type
+                    value = value.value
+                return super().__setattr__(attr, value)
+
+            setattr(self._actual_script_module, attr, value)
+
+        def define(self, src):
+            if "_actual_script_module" in self.__dict__:
+                # If we have completed initialization, just defer to the
+                # backing RecursiveScriptModule to eagerly compile the provided
+                # source.
+                return self._actual_script_module.define(src)
+
+            # Otherwise, we are still in the object's __init__.
+            # In that case, add `src` as a stub to be compiled.
+            #
+            # We use frames_up=1 to get to the proper surrounding scope. The stack
+            # will look like:
+            # 0. createResolutionCallback
+            # 1. define()
+            # 2. surrounding scope.
+            #
+            # createResolutionCallback internally adds 1 to get us to our frame, then
+            # we add 1 to get to the proper surrounding scope.
+            rcb = _jit_internal.createResolutionCallbackFromFrame(frames_up=1)
+            ast = torch._C._parse_source_def(src)
+            self._methods[ast.name().name] = ScriptMethodStub(rcb, ast, None)
+
+        def _replicate_for_data_parallel(self):
+            return self._actual_script_module._replicate_for_data_parallel()
+
+        def __reduce_package__(self, exporter: PackageExporter):
+            """Save a ScriptModule inside of a ``torch.package`` archive.
+
+            Called by ``torch.package.PackageExporter``'s Pickler's ``persistent_id`` when
+            saving TorchScript objects. Performs act of saving a ScriptModule inside of
+            a ``torch.package`` archive.
+
+            Returns method to load the ScriptModule from a ``torch.package.PackageImporter``'s
+            Pickler's ``persistent_load`` function.
+            """
+            script_module_id = exporter.get_unique_id()
+            exporter.script_module_serializer.serialize(self._c, int(script_module_id))
+            return (unpackage_script_module, (script_module_id,))
+
+    class RecursiveScriptModule(ScriptModule):
+        # XXX: RecursiveScriptModule inherits from ScriptModule for the sole
+        # reason that it retains the existing isinstance(ScriptModule)
+        # behavior.
+        r"""Retain the existing isinstance(ScriptModule) behavior.
+
+        The core data structure in TorchScript is the ``ScriptModule``. It is an
+        analogue of torch's ``nn.Module`` and represents an entire model as a tree of
+        submodules. Like normal modules, each individual module in a ``ScriptModule`` can
+        have submodules, parameters, and methods. In ``nn.Module``\s methods are implemented
+        as Python functions, but in ``ScriptModule``\s methods are implemented as
+        TorchScript functions, a statically-typed subset of Python that contains all
+        of PyTorch's built-in Tensor operations. This difference allows your
+        ``ScriptModule``\s code to run without the need for a Python interpreter.
+
+        ``ScriptModule``\s should not be created manually, instead use
+        either :func:`tracing <torch.jit.trace>` or :func:`scripting <torch.jit.script>`.
+        Tracing and scripting can be applied incrementally and :ref:`composed as necessary <Types>`.
+
+        * Tracing records the tensor operations as executed with a set of example inputs and uses these
+          operations to construct a computation graph. You can use the full dynamic behavior of Python with tracing,
+          but values other than Tensors and control flow aren't captured in the graph.
+
+        * Scripting inspects the Python code of the model
+          and compiles it to TorchScript. Scripting allows the use of many `types`_ of values and supports dynamic control flow.
+          Many, but not all features of Python are supported by the compiler, so changes to the source code may be necessary.
+        """
+
+        _disable_script_meta = True
+
+        def __init__(self, cpp_module):
+            self.__dict__["_initializing"] = True
+            self._c = cpp_module
+            super().__init__()
+            # Delete the 'training' attribute set up by `Module.__init__`. It
+            # will get set on the underlying cpp module, so we delete it here
+            # to avoid this version shadowing the cpp module version.
+            delattr(self, "training")
+
+        @staticmethod
+        def _construct(cpp_module, init_fn):
+            """
+            Construct a RecursiveScriptModule that's ready for use.
+
+            PyTorch code should use this to construct a RecursiveScriptModule instead
+            of instead of calling `__init__` directly, as it makes sure the
+            object is properly finalized (and in the future, we may take
+            control of how the RecursiveScriptModule instance is created).
+
+            Args:
+                cpp_module:  The C++ Module that will hold the actual state of
+                             this RecursiveScriptModule instance.
+                init_fn:  Lambda that initializes the RecursiveScriptModule passed to it.
+            """
+            script_module = RecursiveScriptModule(cpp_module)
+            init_fn(script_module)
+
+            # Finalize the ScriptModule: replace the nn.Module state with our
+            # custom implementations and flip the _initializing bit.
+            # pyrefly: ignore [missing-attribute]
+            RecursiveScriptModule._finalize_scriptmodule(script_module)
+            return script_module
+
+        @staticmethod
+        def _finalize_scriptmodule(script_module):
+            script_module._parameters = OrderedDictWrapper(
+                torch._C.ParameterDict(script_module._c)
+            )
+            script_module._buffers = OrderedDictWrapper(
+                torch._C.BufferDict(script_module._c)
+            )
+            script_module._modules = OrderedModuleDict(
+                script_module._c, script_module._modules
+            )
+            script_module._initializing = False
+
+        def _reconstruct(self, cpp_module):
+            """
+            Re-construct an instance of RecursiveScriptModule using an instance of a C++ module.
+
+            Args:
+                cpp_module: The C++ module that this RecursiveScriptModule will be rebuilt around.
+            """
+            self.__init__(cpp_module)  # type: ignore[misc]
+
+            # Copy the concrete type from the C++ module to this ScriptModule.
+            self._concrete_type = torch._C.ConcreteModuleType.from_jit_type(
+                self._c._type()
+            )
+
+            # Copy submodules from the C++ module to this ScriptModule.
+            modules = {}
+            for name, cpp_module in torch._C.ModuleDict(self._c).items():
+                modules[name] = wrap_cpp_module(cpp_module)
+            self._modules = OrderedModuleDict(self._c, modules)  # type: ignore[assignment]
+
+            # Copy parameters and buffers.
+            self._parameters = OrderedDictWrapper(torch._C.ParameterDict(self._c))  # type: ignore[assignment]
+            self._buffers = OrderedDictWrapper(torch._C.BufferDict(self._c))  # type: ignore[assignment]
+
+            # Get rid of the functions from the old C++ module.
+            self.__dict__ = {
+                k: v
+                for k, v in self.__dict__.items()
+                if not isinstance(v, torch._C.ScriptMethod)
+            }
+            self.__dict__["_initializing"] = False
+
+        @property
+        def graph(self):
+            r"""Return a string representation of the internal graph for the ``forward`` method."""
+            return self._c._get_method("forward").graph
+
+        @property
+        def inlined_graph(self):
+            r"""
+            Return a string representation of the internal graph for the ``forward`` method.
+
+            This graph will be preprocessed to inline all function and method calls.
+            """
+            return self.forward.inlined_graph  # type: ignore[attr-defined]
+
+        @property
+        def code(self):
+            r"""
+            Return a pretty-printed representation (as valid Python syntax) of the internal graph for the ``forward`` method.
+
+            """
+            return self.forward.code  # type: ignore[attr-defined]
+
+        @property
+        def code_with_constants(self):
+            r"""Return a tuple.
+
+            Returns a tuple of:
+
+            [0] a pretty-printed representation (as valid Python syntax) of
+            the internal graph for the ``forward`` method. See `code`.
+            [1] a ConstMap following the CONSTANT.cN format of the output in [0].
+            The indices in the [0] output are keys to the underlying constant's values.
+
+            """
+            r = self.forward.code_with_constants  # type: ignore[attr-defined]
+            return (r[0], ConstMap(r[1]))
+
+        def save(self, f, **kwargs):
+            r"""Save with a file-like object.
+
+            save(f, _extra_files={})
+
+            See :func:`torch.jit.save <torch.jit.save>` which accepts a file-like object.
+            This function, torch.save(), converts the object to a string, treating it as a path.
+            DO NOT confuse these two functions when it comes to the 'f' parameter functionality.
+            """
+            return self._c.save(str(f), **kwargs)
+
+        @deprecated(
+            "Lite Interpreter is deprecated. Please consider switching to ExecuTorch. \
+            https://docs.pytorch.org/executorch/stable/getting-started.html"
+        )
+        def _save_for_lite_interpreter(self, *args, **kwargs):
+            r"""Add (or update) the bytecode session to the script model.
+
+            _save_for_lite_interpreter(f)
+
+            The updated model is used
+            in lite interpreter for mobile applications.
+
+            Args:
+                f: a string containing a file name.
+                _extra_files: Map from filename to contents which will be stored as part of 'f'.
+
+            """
+            warnings.warn(
+                "Lite Interpreter is deprecated. Please consider switching to ExecuTorch. \
+                https://docs.pytorch.org/executorch/stable/getting-started.html",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return self._c._save_for_mobile(*args, **kwargs)
+
+        @deprecated(
+            "Lite Interpreter is deprecated. Please consider switching to ExecuTorch. \
+            https://docs.pytorch.org/executorch/stable/getting-started.html"
+        )
+        def _save_to_buffer_for_lite_interpreter(self, *args, **kwargs):
+            warnings.warn(
+                "Lite Interpreter is deprecated. Please consider switching to ExecuTorch. \
+                https://docs.pytorch.org/executorch/stable/getting-started.html",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return self._c._save_to_buffer_for_mobile(*args, **kwargs)
+
+        def save_to_buffer(self, *args: Any, **kwargs: Any) -> Any:
+            return self._c.save_to_buffer(*args, **kwargs)
+
+        def get_debug_state(self, *args: Any, **kwargs: Any) -> Any:
+            return self._c.get_debug_state()
+
+        def extra_repr(self) -> str:
+            return f"original_name={self.original_name}"
+
+        def graph_for(self, *args: Any, **kwargs: Any) -> Any:
+            return self.forward.graph_for(self, *args, **kwargs)  # type: ignore[attr-defined]
+
+        @property
+        def original_name(self):
+            # pyrefly: ignore [unnecessary-comparison]
+            if type(self) is str(self._c._type().name()):
+                return ""
+            return str(self._c._type().name())
+
+        def define(self, src):
+            # We use frames_up=1 to get to the proper surrounding scope. The stack
+            # will look like:
+            # 0. createResolutionCallback
+            # 1. define()
+            # 2. surrounding scope.
+            #
+            # createResolutionCallback internally adds 1 to get us to our frame, then
+            # we add 1 to get to the proper surrounding scope.
+            rcb = _jit_internal.createResolutionCallbackFromFrame(frames_up=1)
+            self._c._define(self._concrete_type, src, rcb)
+
+        def __getattr__(self, attr: str) -> Any:
+            if "_initializing" not in self.__dict__:
+                raise RuntimeError(
+                    "ScriptModule has not been initialized, did you forget to call super's init?"
+                )
+
+            if self._initializing:
+                return super().__getattr__(attr)
+
+            # _modules check is before hasattr since modules are included as attributes in _c,
+            # but we want to get the python wrapper from _modules instead of the raw _c object.
+            if attr in self._modules:
+                return self._modules[attr]
+            elif self._c.hasattr(attr):
+                return self._c.getattr(attr)
+            elif self._c._has_method(attr):
+                script_method = self._c._get_method(attr)
+                # cache method so future calls do not go through __getattr__
+                # to improve invocation performance
+                self.__dict__[attr] = script_method
+                return script_method
+
+            return super().__getattr__(attr)
+
+        def __setattr__(self, attr: str, value: Any) -> None:
+            if self._initializing:
+                return super().__setattr__(attr, value)
+
+            if attr in self._modules:
+                self._modules[attr] = value
+            elif self._c.hasattr(attr):
+                self._c.setattr(attr, value)
+            elif (
+                hasattr(self, "_concrete_type")
+                and attr in self._concrete_type.get_constants()
+            ):
+                # TODO: we don't have _concrete_type set after load(), and in general we lose constant information.
+                # We should encode constants as class type attributes (or something) so it persists across save/load.
+                raise AttributeError(
+                    f"Cannot mutate TorchScript constant value: '{attr}'. Value: '{value}'"
+                )
+            else:
+                # We allow setting Python attributes on the ScriptModule, for
+                # when people want to stash some convenience info on it.
+                # TODO: it's possible that the following is confusing:
+                #   s = torch.jit.script(...)
+                #   s.python_attr = ...
+                #   s.save()   <--- this doesn't have `python_attr`
+                # It's fairly trivial to save enough info to warn in this case.
+                return super().__setattr__(attr, value)
+
+        def __copy__(self) -> Self:
+            return torch.jit._recursive.wrap_cpp_module(copy.copy(self._c))
+
+        def __deepcopy__(self, memo: dict[int, Any] | None) -> Self:
+            return torch.jit._recursive.wrap_cpp_module(copy.deepcopy(self._c, memo))
+
+        # Python magic methods do method lookups on an object's class type, instead of looking up
+        # the method defines on the class instance. In order to continue to expose the magic methods
+        # of builtin-containers (ModuleList, Sequential, ModuleDict) to Python, we
+        # define magic methods here as a shim to the correct attribute.
+        def forward_magic_method(
+            self, method_name: str, *args: Any, **kwargs: Any
+        ) -> Any:
+            self_method = getattr(self, method_name)
+            if getattr(self_method, "__func__", None) == getattr(
+                RecursiveScriptModule, method_name
+            ):
+                raise NotImplementedError
+            return self_method(*args, **kwargs)
+
+        def __iter__(self) -> Iterator[Any]:
+            return self.forward_magic_method("__iter__")
+
+        def __getitem__(self, idx: int) -> Any:
+            return self.forward_magic_method("__getitem__", idx)
+
+        def __len__(self):
+            return self.forward_magic_method("__len__")
+
+        def __contains__(self, key):
+            return self.forward_magic_method("__contains__", key)
+
+        # dir is defined by the base nn.Module, so instead of throwing if
+        # it is not overridden, we call into the nn.Module __dir__ method
+        def __dir__(self) -> Sequence[str]:
+            self_method = self.__dir__
+            if (
+                self_method.__func__  # type: ignore[attr-defined]
+                is _get_function_from_type(RecursiveScriptModule, "__dir__")
+            ):
+                return super().__dir__()
+            return self_method()
+
+        # to resolve bool(value), Python looks if __bool__ is defined then __iter__
+        # is defined then returns true for classes. Since __iter__() on this
+        # class throws if it isn't overridden, we define __bool__ to preserve default behavior
+        def __bool__(self):
+            self_method = self.__bool__
+            if (
+                self_method.__func__  # type: ignore[attr-defined]
+                is _get_function_from_type(RecursiveScriptModule, "__bool__")
+            ):
+                return True
+            return self_method()
+
+        def _replicate_for_data_parallel(self):
+            # we have to initialize ScriptModule properly so that
+            # it works with pybind11
+            def init_fn(script_module):
+                # Don't do anything here, we'll initialize the ScriptModule below
+                return
+
+            # pyrefly: ignore [missing-attribute]
+            return RecursiveScriptModule._construct(
+                self._c._replicate_for_data_parallel(), init_fn
+            )
+
+    # Need to copy all RecursiveScriptModule methods to ScriptModule.
+    #
+    # This is because `super().foo()` does not use
+    # `__getattr__` to look up `foo`. So we need to make each method available on
+    # the ScriptModule manually.
+    # pyrefly: ignore [missing-attribute]
+    for name, item in RecursiveScriptModule.__dict__.items():
+        if not callable(item) and not isinstance(item, property):
+            continue
+        if name.startswith("__") or hasattr(ScriptModule, name):
+            continue
+        # We can copy over the implementation wholesale because besides the
+        # `super()` thing above, ScriptModule behaves exactly like
+        # RecursiveScriptModule
+        setattr(ScriptModule, name, item)
+
+    def _get_methods(cls):
+        import inspect
+
+        # In Python 3 unbound methods are functions, but in Python 2 they are methods
+        return inspect.getmembers(
+            cls, predicate=lambda x: inspect.isfunction(x) or inspect.ismethod(x)
+        )
+
+    _compiled_methods_allowlist = {
+        "forward",
+        "register_buffer",
+        "register_parameter",
+        "register_module",
+        "add_module",
+        "_apply",
+        "apply",
+        "cuda",
+        "cpu",
+        "to",
+        "type",
+        "float",
+        "double",
+        "half",
+        "state_dict",
+        "_save_to_state_dict",
+        "load_state_dict",
+        "_load_from_state_dict",
+        "_named_members",
+        "parameters",
+        "named_parameters",
+        "buffers",
+        "named_buffers",
+        "children",
+        "named_children",
+        "modules",
+        "named_modules",
+        "zero_grad",
+        "share_memory",
+        "_get_name",
+        "extra_repr",
+        "_slow_forward",
+        "_tracing_name",
+        "eval",
+        "train",
+        "get_extra_state",
+        "set_extra_state",
+    }
+
+    def _make_fail(name):
+        def fail(self, *args, **kwargs):
+            raise RuntimeError(name + " is not supported on ScriptModules")
+
+        return fail
+
+    for name, method in _get_methods(torch.nn.Module):
+        if name.startswith("__") or name.endswith("_call_impl"):
+            continue
+        if (
+            # pyrefly: ignore [missing-attribute]
+            name not in RecursiveScriptModule.__dict__
+            and name not in _compiled_methods_allowlist
+        ):
+            setattr(RecursiveScriptModule, method.__name__, _make_fail(name))
+
+
+else:
+    # TODO MAKE SURE THAT DISABLING WORKS
+    class RecursiveScriptClass:  # type: ignore[no-redef]
+        pass
+
+    class ScriptModule(torch.nn.Module):  # type: ignore[no-redef]
+        def __init__(self, arg=None):
+            super().__init__()
+
+    class RecursiveScriptModule(ScriptModule):  # type: ignore[no-redef]
+        def __init__(self, arg=None):
+            super().__init__()
+
+
+def call_prepare_scriptable_func_impl(obj, memo):
+    if not isinstance(obj, torch.nn.Module):
+        return obj
+
+    obj_id = id(obj)
+
+    # If obj_id is in memo, obj has already been prepared or is being
+    # prepared in another call up the stack.
+    if obj_id in memo:
+        return memo[id(obj)]
+
+    obj = (
+        # pyrefly: ignore [not-callable]
+        obj.__prepare_scriptable__() if hasattr(obj, "__prepare_scriptable__") else obj
+    )  # type: ignore[operator]
+    # Record obj in memo to avoid infinite recursion in the case of cycles in the module
+    # hierarchy when recursing below.
+    memo[obj_id] = obj
+
+    new_obj_dict = {}
+
+    for name, sub_module in obj.__dict__.items():
+        if name == "_modules":
+            for k, v in sub_module.items():
+                sub_module[k] = call_prepare_scriptable_func_impl(v, memo)
+            new_obj_dict[name] = sub_module
+        elif isinstance(sub_module, torch.nn.Module) and not isinstance(
+            sub_module, ScriptModule
+        ):
+            new_obj_dict[name] = call_prepare_scriptable_func_impl(sub_module, memo)
+        else:
+            new_obj_dict[name] = sub_module
+
+    for v in new_obj_dict.values():
+        obj.__dict__[name] = v
+
+    return obj
+
+
+def call_prepare_scriptable_func(obj):
+    memo: dict[int, torch.nn.Module] = {}
+    return call_prepare_scriptable_func_impl(obj, memo)
+
+
+def create_script_dict(obj):
+    """
+    Create a ``torch._C.ScriptDict`` instance with the data from ``obj``.
+
+    Args:
+        obj (dict): The Python dictionary that is used to initialize the ``ScriptDict``
+                    returned by this function.
+
+    Returns:
+        An instance of ``torch._C.ScriptDict`` that has the same data as ``obj``
+        and can be passed between Python and TorchScript with reference semantics and
+        zero copy overhead.
+    """
+    return torch._C.ScriptDict(obj)  # type: ignore[attr-defined]
+
+
+def create_script_list(obj, type_hint=None):
+    """
+    Create a ``torch._C.ScriptList`` instance with the data from ``obj``.
+
+    Args:
+        obj (dict): The Python list that is used to initialize the ``ScriptList``
+                    returned by this function.
+    Returns:
+        An instance of ``torch._C.ScriptList`` that has the same data as ``obj``
+        and can be passed between Python and TorchScript with reference semantics and
+        zero copy overhead.
+    """
+    return torch._C.ScriptList(obj)  # type: ignore[attr-defined]
+
+
+_TOPLEVEL: bool = True
+
+
+def _script_impl(
+    obj,
+    optimize=None,
+    _frames_up=0,
+    _rcb=None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+):
+    global type_trace_db
+
+    if optimize is not None:
+        warnings.warn(
+            "`optimize` is deprecated and has no effect. "
+            "Use `with torch.jit.optimized_execution()` instead",
+            FutureWarning,
+            stacklevel=3,
+        )
+
+    # No-op for modules, functions, class instances that are already scripted
+    if isinstance(obj, RecursiveScriptClass):
+        return obj
+    if isinstance(obj, ScriptModule):
+        return obj
+    if isinstance(obj, ScriptFunction):
+        return obj
+
+    if example_inputs:
+        # If MonkeyType is installed, enable profile directed type annotation
+        # Check if example_inputs are defined and generate call traces
+        # for the method by running eager mode version of the method with
+        # the provide example inputs. This logs all the traces in type_trace_db
+        type_trace_db = JitTypeTraceStore()
+        if monkeytype_trace:
+            # pyrefly: ignore [bad-argument-count]
+            monkeytype_config = JitTypeTraceConfig(type_trace_db)
+            with monkeytype_trace(monkeytype_config):
+                if isinstance(example_inputs, dict):
+                    # If the obj is an nn.Module or a class, then each method is
+                    # executed with the arguments provided in the example inputs.
+                    # example inputs here will be of type Dict(class.method, (arguments))
+                    # This is used to infer type annotations for those methods
+                    # which are not called directly under the hood of monkeytype.
+                    for module, example_input in example_inputs.items():
+                        for example in example_input:
+                            module(*example)
+                elif isinstance(example_inputs, list):
+                    for examples in example_inputs:
+                        obj(*examples)
+                else:
+                    raise ValueError(
+                        "Error: Unable to infer types. Please format the inputs to type `List[Tuple]`"
+                        " or `Dict[Callable, List[Tuple]]` to be run with MonkeyType."
+                    )
+        else:
+            warnings.warn(
+                "Warning: monkeytype is not installed. Please install https://github.com/Instagram/MonkeyType "
+                "to enable Profile-Directed Typing in TorchScript. Refer to "
+                "https://github.com/Instagram/MonkeyType/blob/master/README.rst to install MonkeyType. ",
+                stacklevel=2,
+            )
+
+    if isinstance(obj, torch.nn.Module):
+        obj = call_prepare_scriptable_func(obj)
+        return torch.jit._recursive.create_script_module(
+            obj, torch.jit._recursive.infer_methods_to_compile
+        )
+    else:
+        obj = (
+            obj.__prepare_scriptable__()
+            if hasattr(obj, "__prepare_scriptable__")
+            else obj
+        )  # type: ignore[operator]
+
+    if isinstance(obj, dict):
+        return create_script_dict(obj)
+    if isinstance(obj, list):
+        return create_script_list(obj)
+
+    if inspect.isclass(obj):
+        qualified_name = _qualified_name(obj)
+        # If this type is a `nn.Module` subclass, they probably meant to pass
+        # an instance instead of a Module
+        if issubclass(obj, torch.nn.Module):
+            raise RuntimeError(
+                f"Type '{obj}' cannot be compiled since it inherits from nn.Module, pass an instance instead"
+            )
+
+        # Enums are automatically usable in TorchScript, explicitly scripting
+        # is not necessary, but not harmful either.
+        if issubclass(obj, enum.Enum):
+            return obj
+
+        if not _is_new_style_class(obj):
+            raise RuntimeError(
+                "TorchScript classes must be new-style classes. "
+                "Please inherit from 'object'."
+            )
+        if len(obj.mro()) > 2:
+            raise RuntimeError(
+                "TorchScript classes does not support inheritance yet. "
+                "Please directly inherit from 'object'."
+            )
+        if _rcb is None:
+            _rcb = _jit_internal.createResolutionCallbackFromFrame(_frames_up + 1)
+        _compile_and_register_class(obj, _rcb, qualified_name)
+        return obj
+    elif inspect.isfunction(obj) or inspect.ismethod(obj):
+        qualified_name = _qualified_name(obj)
+        # this is a decorated fn, and we need to the underlying fn and its rcb
+        if hasattr(obj, "__script_if_tracing_wrapper"):
+            obj = obj.__original_fn  # type: ignore[union-attr]
+            _rcb = _jit_internal.createResolutionCallbackFromClosure(obj)
+
+        # some functions are explicitly marked as not supported in script mode
+        if hasattr(obj, "__script_unsupported"):
+            raise RuntimeError("TorchScript error: " + obj.__script_unsupported)
+
+        _check_directly_compile_overloaded(obj)
+        maybe_already_compiled_fn = _try_get_jit_cached_function(obj)
+        if maybe_already_compiled_fn:
+            maybe_already_compiled_fn._torchdynamo_inline = obj  # type: ignore[attr-defined]
+            return maybe_already_compiled_fn
+        ast = get_jit_def(obj, obj.__name__)
+        if _rcb is None:
+            _rcb = _jit_internal.createResolutionCallbackFromClosure(obj)
+        fn = torch._C._jit_script_compile(
+            qualified_name, ast, _rcb, get_default_args(obj)
+        )
+        # Forward docstrings
+        fn.__doc__ = obj.__doc__
+        fn.__name__ = "ScriptFunction"
+        fn.__qualname__ = "torch.jit.ScriptFunction"
+        # Allow torch.compile() to inline
+        fn._torchdynamo_inline = obj  # type: ignore[attr-defined]
+        _set_jit_function_cache(obj, fn)
+        return fn
+    else:
+        return torch.jit._recursive.create_script_class(obj)
+
+
+def script(
+    obj: Any,
+    optimize: None = None,
+    _frames_up: int = 0,
+    _rcb: Callable[[str], Any] | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> Any:
+    r"""Script the function.
+
+    Scripting a function or ``nn.Module`` will inspect the source code, compile
+    it as TorchScript code using the TorchScript compiler, and return a :class:`ScriptModule` or
+    :class:`ScriptFunction`. TorchScript itself is a subset of the Python language, so not all
+    features in Python work, but we provide enough functionality to compute on
+    tensors and do control-dependent operations. For a complete guide, see the
+    :ref:`language-reference`.
+
+    Scripting a dictionary or list copies the data inside it into a TorchScript instance than can be
+    subsequently passed by reference between Python and TorchScript with zero copy overhead.
+
+    ``torch.jit.script`` can be used as a function for modules, functions, dictionaries and lists
+     and as a decorator ``@torch.jit.script`` for torchscript-classes and functions.
+
+    Args:
+        obj (Callable, class, or nn.Module):  The ``nn.Module``, function, class type,
+                                                  dictionary, or list to compile.
+        example_inputs (Union[List[Tuple], Dict[Callable, List[Tuple]], None]): Provide example inputs
+            to annotate the arguments for a function or ``nn.Module``.
+
+    Returns:
+        If ``obj`` is ``nn.Module``, ``script`` returns
+        a :class:`ScriptModule` object. The returned :class:`ScriptModule` will
+        have the same set of sub-modules and parameters as the
+        original ``nn.Module``. If ``obj`` is a standalone function,
+        a :class:`ScriptFunction` will be returned. If ``obj`` is a ``dict``, then
+        ``script`` returns an instance of `torch._C.ScriptDict`. If ``obj`` is a ``list``,
+        then ``script`` returns an instance of `torch._C.ScriptList`.
+
+    **Scripting a function**
+        The ``@torch.jit.script`` decorator will construct a :class:`ScriptFunction`
+        by compiling the body of the function.
+
+        Example (scripting a function):
+
+        .. testcode::
+
+            import torch
+
+            @torch.jit.script
+            def foo(x, y):
+                if x.max() > y.max():
+                    r = x
+                else:
+                    r = y
+                return r
+
+            print(type(foo))  # torch.jit.ScriptFunction
+
+            # See the compiled graph as Python code
+            print(foo.code)
+
+            # Call the function using the TorchScript interpreter
+            foo(torch.ones(2, 2), torch.ones(2, 2))
+
+        .. testoutput::
+            :hide:
+
+            ...
+
+    ****Scripting a function using example_inputs**
+        Example inputs can be used to annotate a function arguments.
+
+        Example (annotating a function before scripting):
+
+        .. testcode::
+
+            import torch
+
+            def test_sum(a, b):
+                return a + b
+
+            # Annotate the arguments to be int
+            scripted_fn = torch.jit.script(test_sum, example_inputs=[(3, 4)])
+
+            print(type(scripted_fn))  # torch.jit.ScriptFunction
+
+            # See the compiled graph as Python code
+            print(scripted_fn.code)
+
+            # Call the function using the TorchScript interpreter
+            scripted_fn(20, 100)
+
+        .. testoutput::
+            :hide:
+
+            ...
+
+    **Scripting an nn.Module**
+        Scripting an ``nn.Module`` by default will compile the ``forward`` method and recursively
+        compile any methods, submodules, and functions called by ``forward``. If a ``nn.Module`` only uses
+        features supported in TorchScript, no changes to the original module code should be necessary. ``script``
+        will construct :class:`ScriptModule` that has copies of the attributes, parameters, and methods of
+        the original module.
+
+        Example (scripting a simple module with a Parameter):
+
+        .. testcode::
+
+            import torch
+
+            class MyModule(torch.nn.Module):
+                def __init__(self, N, M):
+                    super().__init__()
+                    # This parameter will be copied to the new ScriptModule
+                    self.weight = torch.nn.Parameter(torch.rand(N, M))
+
+                    # When this submodule is used, it will be compiled
+                    self.linear = torch.nn.Linear(N, M)
+
+                def forward(self, input):
+                    output = self.weight.mv(input)
+
+                    # This calls the `forward` method of the `nn.Linear` module, which will
+                    # cause the `self.linear` submodule to be compiled to a `ScriptModule` here
+                    output = self.linear(output)
+                    return output
+
+            scripted_module = torch.jit.script(MyModule(2, 3))
+
+        Example (scripting a module with traced submodules):
+
+        .. testcode::
+
+            import torch
+            import torch.nn as nn
+            import torch.nn.functional as F
+
+            class MyModule(nn.Module):
+                def __init__(self) -> None:
+                    super().__init__()
+                    # torch.jit.trace produces a ScriptModule's conv1 and conv2
+                    self.conv1 = torch.jit.trace(nn.Conv2d(1, 20, 5), torch.rand(1, 1, 16, 16))
+                    self.conv2 = torch.jit.trace(nn.Conv2d(20, 20, 5), torch.rand(1, 20, 16, 16))
+
+                def forward(self, input):
+                    input = F.relu(self.conv1(input))
+                    input = F.relu(self.conv2(input))
+                    return input
+
+            scripted_module = torch.jit.script(MyModule())
+
+        To compile a method other than ``forward`` (and recursively compile anything it calls), add
+        the :func:`@torch.jit.export <torch.jit.export>` decorator to the method. To opt out of compilation
+        use :func:`@torch.jit.ignore <torch.jit.ignore>` or :func:`@torch.jit.unused <torch.jit.unused>`.
+
+        Example (an exported and ignored method in a module)::
+
+            import torch
+            import torch.nn as nn
+
+
+            class MyModule(nn.Module):
+                def __init__(self) -> None:
+                    super().__init__()
+
+                @torch.jit.export
+                def some_entry_point(self, input):
+                    return input + 10
+
+                @torch.jit.ignore
+                def python_only_fn(self, input):
+                    # This function won't be compiled, so any
+                    # Python APIs can be used
+                    import pdb
+
+                    pdb.set_trace()
+
+                def forward(self, input):
+                    if self.training:
+                        self.python_only_fn(input)
+                    return input * 99
+
+
+            scripted_module = torch.jit.script(MyModule())
+            print(scripted_module.some_entry_point(torch.randn(2, 2)))
+            print(scripted_module(torch.randn(2, 2)))
+
+        Example ( Annotating forward of nn.Module using example_inputs)::
+
+            import torch
+            import torch.nn as nn
+            from typing import NamedTuple
+
+            class MyModule(NamedTuple):
+            result: List[int]
+
+            class TestNNModule(torch.nn.Module):
+                def forward(self, a) -> MyModule:
+                    result = MyModule(result=a)
+                    return result
+
+            pdt_model = TestNNModule()
+
+            # Runs the pdt_model in eager model with the inputs provided and annotates the arguments of forward
+            scripted_model = torch.jit.script(pdt_model, example_inputs={pdt_model: [([10, 20, ], ), ], })
+
+            # Run the scripted_model with actual inputs
+            print(scripted_model([20]))
+    """
+    if sys.version_info >= (3, 14):
+        warnings.warn(
+            "`torch.jit.script` is not supported in Python 3.14+ and may break. "
+            "Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    else:
+        warnings.warn(
+            "`torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    if not _enabled:
+        return obj
+    try:
+        global _TOPLEVEL
+        prev = _TOPLEVEL
+        _TOPLEVEL = False
+        ret = _script_impl(
+            obj=obj,
+            optimize=optimize,
+            _frames_up=_frames_up + 1,
+            _rcb=_rcb,
+            example_inputs=example_inputs,
+        )
+
+        if prev:
+            log_torchscript_usage("script", model_id=_get_model_id(ret))
+
+        return ret
+    finally:
+        _TOPLEVEL = prev
+
+
+# overloads are registered in _jit_internal and compiled here so that _overload
+# can be used in nn/functional.py without an import cycle
+
+
+def _check_overload_defaults(impl_defaults, overload_defaults, loc):
+    for name, overload_value in overload_defaults.items():
+        if name not in impl_defaults or impl_defaults[name] != overload_value:
+            raise torch.jit.frontend.FrontendError(
+                loc,
+                "Default parameters on overloads do not affect the runtime so they "
+                "must equal to the default parameter on the implementation function. Found on "
+                f"parameter {name}",
+            )
+
+
+def _compile_function_with_overload(overload_fn, qual_name, impl_fn):
+    overload_decl = get_jit_def(overload_fn, overload_fn.__name__).decl()
+    overload_signature = torch.jit.annotations.get_signature(
+        overload_fn, None, None, inspect.ismethod(overload_fn)
+    )
+    impl_ast = get_jit_def(impl_fn, impl_fn.__name__)
+    overload_defaults = get_default_args(overload_fn)
+    implementation_defaults = get_default_args(impl_fn)
+    _rcb = _jit_internal.createResolutionCallbackFromClosure(impl_fn)
+    _check_overload_defaults(
+        implementation_defaults, overload_defaults, overload_decl.range()
+    )
+    fn = torch._C._jit_script_compile_overload(
+        qual_name,
+        overload_decl,
+        impl_ast,
+        _rcb,
+        implementation_defaults,
+        overload_signature,
+    )
+    return fn
+
+
+def _get_overloads(obj):
+    # check for cached compiled fns
+    existing_compiled_fns = _try_get_jit_cached_overloads(obj)
+    qual_name = _qualified_name(obj)
+    uncompiled_overloads = _jit_internal._get_fn_overloads(qual_name)
+    if uncompiled_overloads is None:
+        return existing_compiled_fns
+
+    if obj in uncompiled_overloads:
+        raise RuntimeError(
+            _jit_internal.get_overload_no_implementation_error_message("function", obj)
+        )
+
+    compiled_fns = [
+        _compile_function_with_overload(overload_fn, qual_name, obj)
+        for overload_fn in uncompiled_overloads
+    ]
+
+    if existing_compiled_fns:
+        compiled_fns = existing_compiled_fns + compiled_fns
+
+    # cache compilation, remove information stored to do compilation
+    _set_jit_overload_cache(obj, compiled_fns)
+    _jit_internal._clear_fn_overloads(qual_name)
+    return compiled_fns
+
+
+def _check_directly_compile_overloaded(obj):
+    qual_name = _qualified_name(obj)
+    if _jit_internal._get_fn_overloads(qual_name) or _try_get_jit_cached_overloads(obj):
+        raise RuntimeError(
+            f"Function {qual_name} cannot be directly compiled because it"
+            " is overloaded. It must be used in a context of a function"
+            " where its inputs can determine which overload to call."
+        )
+
+
+def interface(obj: _T) -> _T:
+    r"""Decorate to annotate classes or modules of different types.
+
+    .. deprecated:: 2.5
+        TorchScript is deprecated, please use ``torch.compile`` instead.
+
+    This decorator can be used to define an interface that can be used to annotate
+    classes or modules of different types. This can be used for to annotate a submodule
+    or attribute class that could have different types that implement the same
+    interface, or which could be swapped at runtime; or to store a list of modules or
+    classes of varying types.
+
+    It is sometimes used to implement "Callables" - functions or modules that implement
+    an interface but whose implementations differ and which can be swapped out.
+
+    Example:
+    .. testcode::
+
+        import torch
+        from typing import List
+
+        @torch.jit.interface
+        class InterfaceType:
+            def run(self, x: torch.Tensor) -> torch.Tensor:
+                pass
+
+        # implements InterfaceType
+        @torch.jit.script
+        class Impl1:
+            def run(self, x: torch.Tensor) -> torch.Tensor:
+                return x.relu()
+
+        class Impl2(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.val = torch.rand(())
+
+            @torch.jit.export
+            def run(self, x: torch.Tensor) -> torch.Tensor:
+                return x + self.val
+
+        def user_fn(impls: List[InterfaceType], idx: int, val: torch.Tensor) -> torch.Tensor:
+            return impls[idx].run(val)
+
+        user_fn_jit = torch.jit.script(user_fn)
+
+        impls = [Impl1(), torch.jit.script(Impl2())]
+        val = torch.rand(4, 4)
+        user_fn_jit(impls, 0, val)
+        user_fn_jit(impls, 1, val)
+    """
+    warnings.warn(
+        "`torch.jit.interface` is deprecated. Please use `torch.compile` instead.",
+        DeprecationWarning,
+    )
+    if not inspect.isclass(obj):
+        raise RuntimeError("interface must be applied to a class")
+    if not _is_new_style_class(obj):
+        raise RuntimeError("TorchScript interfaces must inherit from 'object'")
+
+    # Expected MRO is:
+    #   User module
+    #   torch.nn.modules.module.Module
+    #   object
+    is_module_interface = issubclass(obj, torch.nn.Module) and len(obj.mro()) == 3
+
+    if not is_module_interface and len(obj.mro()) > 2:
+        raise RuntimeError(
+            "TorchScript interface does not support inheritance yet. "
+            "Please directly inherit from 'object' or 'nn.Module'."
+        )
+
+    qualified_name = _qualified_name(obj)
+    rcb = _jit_internal.createResolutionCallbackFromFrame(1)
+    # if this type is a `nn.Module` subclass, generate a module interface type
+    # instead of a class interface type; a module interface type only compiles
+    # the user provided methods as part of the interface
+    ast = get_jit_class_def(obj, obj.__name__)
+    mangled_classname = torch._C._jit_script_interface_compile(
+        qualified_name, ast, rcb, is_module_interface
+    )
+    obj.__torch_script_interface__ = mangled_classname
+    return obj
+
+
+def _recursive_compile_class(obj, loc):
+    _qual_name = _qualified_name(obj)
+    # We're starting a new compilation, so update the error call stack in
+    # case it fails
+    error_stack = torch._C.CallStack(_qual_name, loc)  # noqa: F841
+    rcb = _jit_internal.createResolutionCallbackForClassMethods(obj)
+    return _compile_and_register_class(obj, rcb, _qual_name)
+
+
+CompilationUnit = torch._C.CompilationUnit
+set_module(CompilationUnit, "torch.jit")
+
+
+def pad(s: str, padding: int, offset: int = 0, char: str = " "):
+    if padding >= len(s):
+        padding -= len(s)
+    return "".join([char for _ in range(padding + offset)]) + s
+
+
+class _ScriptProfileColumn:
+    def __init__(self, header: str, alignment: int = 4, offset: int = 0):
+        self.header = header
+        self.alignment = alignment
+        self.offset = offset
+        self.rows: dict[int, Any] = {}
+
+    def add_row(self, lineno: int, value: Any):
+        self.rows[lineno] = value
+
+    def materialize(self):
+        max_length = len(self.header)
+        rows: list[tuple[int, str]] = []
+        for key, value in self.rows.items():
+            cell = str(value)
+            rows.append((key, cell))
+            max_length = max(len(cell), max_length)
+
+        if self.alignment > 0:
+            padding = max_length + self.alignment
+            padding -= padding % self.alignment
+        else:
+            padding = 0
+
+        rows = [(key, pad(cell, padding, self.offset)) for key, cell in rows]
+        return pad(self.header, padding, self.offset), rows
+
+
+class _ScriptProfileTable:
+    def __init__(self, cols: list[_ScriptProfileColumn], source_range: list[int]):
+        self.cols = cols
+        self.source_range = source_range
+
+    def dump_string(self):
+        outputs: list[str] = []
+        cells: list[tuple[str, dict[int, str]]] = []
+        header_buffer = ""
+        for col in self.cols:
+            header, rows = col.materialize()
+            header_buffer += header
+            cells.append((header, dict(rows)))
+
+        outputs.append(header_buffer)
+        outputs.append(pad("", len(header_buffer), 0, "="))
+        for line in self.source_range:
+            row_buffer = ""
+            for header, rows in cells:
+                cell = rows.get(line)
+                if cell is None:
+                    row_buffer += pad("", len(header))
+                else:
+                    row_buffer += cell
+            outputs.append(row_buffer)
+        return "\n".join(outputs)
+
+
+class _ScriptProfile:
+    def __init__(self) -> None:
+        self.profile = classes.profiling._ScriptProfile()
+
+    def enable(self):
+        self.profile.enable()
+
+    def disable(self):
+        self.profile.disable()
+
+    def dump_string(self) -> str:
+        outputs: list[str] = []
+        for source_stats in self.profile._dump_stats():
+            source_ref = source_stats.source()
+            source_lines = source_ref.text().splitlines()
+            dedent = min(len(line) - len(line.lstrip(" ")) for line in source_lines)
+            source_lines = [line[dedent:] for line in source_lines]
+
+            start_line = source_ref.starting_lineno()
+            end_line = start_line + len(source_lines)
+            source_range = range(start_line, end_line)
+            lineno = _ScriptProfileColumn("Line #")
+            hits = _ScriptProfileColumn("Hits")
+            time_ns = _ScriptProfileColumn("Time (ns)")
+            line_contents = _ScriptProfileColumn("Line Contents", 0, 1)
+            stats = source_stats.line_map()
+            for line in source_range:
+                lineno.add_row(line, line)
+                line_contents.add_row(line, source_lines[line - start_line])
+                stat = stats.get(line)
+                if stat is not None:
+                    hits.add_row(line, stat.count())
+                    time_ns.add_row(line, stat.duration_ns())
+
+            table = _ScriptProfileTable(
+                [lineno, hits, time_ns, line_contents], list(source_range)
+            )
+            outputs.append(table.dump_string())
+        return "\n\n".join(outputs)
+
+    def dump(self):
+        print(self.dump_string())
+
+
+def _unwrap_optional(x):
+    if x is None:
+        raise AssertionError("Unwrapping null optional")
+    return x
+
+
+_register_builtin(_unwrap_optional, "aten::_unwrap_optional")
+_register_builtin(_jit_internal.is_scripting, "aten::is_scripting")
+_register_builtin(has_torch_function, "aten::has_torch_function")
+_register_builtin(has_torch_function_unary, "aten::has_torch_function")
+_register_builtin(has_torch_function_variadic, "aten::has_torch_function")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_script.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_script.pyi
new file mode 100644
index 0000000000000000000000000000000000000000..ff3856accd4e719a6bb612b49703d71a8f0317ea
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_script.pyi
@@ -0,0 +1,304 @@
+# mypy: allow-untyped-defs
+# mypy: disable-error-code="type-arg"
+from collections.abc import Callable, Iterator, Mapping
+from typing import Any, NamedTuple, overload, TypeAlias, TypeVar
+from typing_extensions import Never, Self
+
+from _typeshed import Incomplete
+
+import torch
+from torch._classes import classes as classes
+from torch._jit_internal import _qualified_name as _qualified_name
+from torch.jit._builtins import _register_builtin as _register_builtin
+from torch.jit._fuser import (
+    _graph_for as _graph_for,
+    _script_method_graph_for as _script_method_graph_for,
+)
+from torch.jit._monkeytype_config import (
+    JitTypeTraceConfig as JitTypeTraceConfig,
+    JitTypeTraceStore as JitTypeTraceStore,
+    monkeytype_trace as monkeytype_trace,
+)
+from torch.jit._recursive import (
+    _compile_and_register_class as _compile_and_register_class,
+    infer_methods_to_compile as infer_methods_to_compile,
+    ScriptMethodStub as ScriptMethodStub,
+    wrap_cpp_module as wrap_cpp_module,
+)
+from torch.jit._serialization import validate_map_location as validate_map_location
+from torch.jit._state import (
+    _enabled as _enabled,
+    _set_jit_function_cache as _set_jit_function_cache,
+    _set_jit_overload_cache as _set_jit_overload_cache,
+    _try_get_jit_cached_function as _try_get_jit_cached_function,
+    _try_get_jit_cached_overloads as _try_get_jit_cached_overloads,
+)
+from torch.jit.frontend import (
+    get_default_args as get_default_args,
+    get_jit_class_def as get_jit_class_def,
+    get_jit_def as get_jit_def,
+)
+from torch.nn import Module as Module
+from torch.overrides import (
+    has_torch_function as has_torch_function,
+    has_torch_function_unary as has_torch_function_unary,
+    has_torch_function_variadic as has_torch_function_variadic,
+)
+from torch.package import (
+    PackageExporter as PackageExporter,
+    PackageImporter as PackageImporter,
+)
+from torch.utils import set_module as set_module
+
+ScriptFunction = torch._C.ScriptFunction
+
+type_trace_db: JitTypeTraceStore
+
+# Defined in torch/csrc/jit/python/script_init.cpp
+ResolutionCallback: TypeAlias = Callable[[str], Callable[..., Any]]
+_ClassVar = TypeVar("_ClassVar", bound=type)
+_T = TypeVar("_T")
+
+def _reduce(cls) -> None: ...
+
+class Attribute(NamedTuple):
+    value: Incomplete
+    type: Incomplete
+
+def _get_type_trace_db(): ...
+def _get_function_from_type(cls, name): ...
+def _is_new_style_class(cls): ...
+
+class OrderedDictWrapper:
+    _c: Incomplete
+    def __init__(self, _c) -> None: ...
+    def keys(self): ...
+    def values(self): ...
+    def __len__(self) -> int: ...
+    def __delitem__(self, k) -> None: ...
+    def items(self): ...
+    def __setitem__(self, k, v) -> None: ...
+    def __contains__(self, k) -> bool: ...
+    def __getitem__(self, k): ...
+
+class OrderedModuleDict(OrderedDictWrapper):
+    _python_modules: Incomplete
+    def __init__(self, module, python_dict) -> None: ...
+    def items(self): ...
+    def __contains__(self, k) -> bool: ...
+    def __setitem__(self, k, v) -> None: ...
+    def __getitem__(self, k): ...
+
+class ScriptMeta(type):
+    def __init__(cls, name, bases, attrs) -> None: ...
+
+class _CachedForward:
+    def __get__(self, obj, cls): ...
+
+class ScriptWarning(Warning): ...
+
+def script_method(fn): ...
+
+class ConstMap:
+    const_mapping: Mapping[str, Any]
+    def __init__(self, const_mapping: Mapping[str, Any]) -> None: ...
+    def __getattr__(self, attr: str) -> Any: ...
+
+def unpackage_script_module(
+    importer: PackageImporter,
+    script_module_id: str,
+) -> torch.nn.Module: ...
+
+_magic_methods: Incomplete
+
+class RecursiveScriptClass:
+    _c: Incomplete
+    _props: Incomplete
+    def __init__(self, cpp_class) -> None: ...
+    def __getattr__(self, attr: str) -> Any: ...
+    def __setattr__(self, attr: str, value: Any) -> None: ...
+    def forward_magic_method(
+        self, method_name: str, *args: Any, **kwargs: Any
+    ) -> Any: ...
+    def __getstate__(self) -> None: ...
+    def __iadd__(self, other: Self) -> Self: ...
+
+def method_template(self, *args, **kwargs): ...
+
+class ScriptModule(Module, metaclass=ScriptMeta):
+    __jit_unused_properties__: Incomplete
+    def __init__(self) -> None: ...
+    forward: Callable[..., Any]
+    def __getattr__(self, attr: str) -> Any: ...
+    def __setattr__(self, attr: str, value: Any) -> None: ...
+    def define(self, src) -> None: ...
+    def _replicate_for_data_parallel(self): ...
+    def __reduce_package__(
+        self, exporter: PackageExporter
+    ) -> tuple[Any, tuple[Any, ...]]: ...
+    # add __jit_unused_properties__
+    @property
+    def code(self) -> str: ...
+    @property
+    def code_with_constants(self) -> tuple[str, ConstMap]: ...
+    @property
+    def graph(self) -> torch.Graph: ...
+    @property
+    def inlined_graph(self) -> torch.Graph: ...
+    @property
+    def original_name(self) -> str: ...
+
+class RecursiveScriptModule(ScriptModule):
+    _disable_script_meta: bool
+    _c: Incomplete
+    def __init__(self, cpp_module) -> None: ...
+    @staticmethod
+    def _construct(cpp_module, init_fn): ...
+    @staticmethod
+    def _finalize_scriptmodule(script_module) -> None: ...
+    _concrete_type: Incomplete
+    _modules: Incomplete
+    _parameters: Incomplete
+    _buffers: Incomplete
+    __dict__: Incomplete
+    def _reconstruct(self, cpp_module) -> None: ...
+    def save(self, f, **kwargs): ...
+    def _save_for_lite_interpreter(self, *args, **kwargs): ...
+    def _save_to_buffer_for_lite_interpreter(self, *args, **kwargs): ...
+    def save_to_buffer(self, *args, **kwargs): ...
+    def get_debug_state(self, *args, **kwargs): ...
+    def extra_repr(self) -> str: ...
+    def graph_for(self, *args, **kwargs): ...
+    def define(self, src) -> None: ...
+    def __getattr__(self, attr: str) -> Any: ...
+    def __setattr__(self, attr: str, value: Any) -> None: ...
+    def __copy__(self) -> Self: ...
+    def __deepcopy__(self, memo: dict[int, Any] | None) -> Self: ...
+    def forward_magic_method(
+        self, method_name: str, *args: Any, **kwargs: Any
+    ) -> Any: ...
+    def __iter__(self) -> Iterator[Any]: ...
+    def __getitem__(self, idx: int) -> Any: ...
+    def __len__(self) -> int: ...
+    def __contains__(self, key: str) -> bool: ...
+    def __dir__(self) -> list[str]: ...
+    def __bool__(self) -> bool: ...
+    def _replicate_for_data_parallel(self): ...
+
+def _get_methods(cls): ...
+
+_compiled_methods_allowlist: Incomplete
+
+def _make_fail(name): ...
+def call_prepare_scriptable_func_impl(obj, memo): ...
+def call_prepare_scriptable_func(obj): ...
+def create_script_dict(obj): ...
+def create_script_list(obj, type_hint: Incomplete | None = ...): ...
+@overload
+def script(
+    obj: type[Module],
+    optimize: bool | None = None,
+    _frames_up: int = 0,
+    _rcb: ResolutionCallback | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> Never: ...
+@overload
+def script(
+    obj: dict,
+    optimize: bool | None = None,
+    _frames_up: int = 0,
+    _rcb: ResolutionCallback | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> torch.ScriptDict: ...
+@overload
+def script(
+    obj: list,
+    optimize: bool | None = None,
+    _frames_up: int = 0,
+    _rcb: ResolutionCallback | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> torch.ScriptList: ...
+@overload
+def script(  # type: ignore[overload-overlap]
+    obj: Module,
+    optimize: bool | None = None,
+    _frames_up: int = 0,
+    _rcb: ResolutionCallback | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> RecursiveScriptModule: ...
+@overload
+def script(  # type: ignore[overload-overlap]
+    obj: _ClassVar,
+    optimize: bool | None = None,
+    _frames_up: int = 0,
+    _rcb: ResolutionCallback | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> _ClassVar: ...
+@overload
+def script(
+    obj: Callable,
+    optimize: bool | None = None,
+    _frames_up: int = 0,
+    _rcb: ResolutionCallback | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> ScriptFunction: ...
+@overload
+def script(
+    obj: Any,
+    optimize: bool | None = None,
+    _frames_up: int = 0,
+    _rcb: ResolutionCallback | None = None,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = None,
+) -> RecursiveScriptClass: ...
+@overload
+def script(
+    obj,
+    optimize: Incomplete | None = ...,
+    _frames_up: int = ...,
+    _rcb: Incomplete | None = ...,
+    example_inputs: list[tuple] | dict[Callable, list[tuple]] | None = ...,
+): ...
+def _check_overload_defaults(impl_defaults, overload_defaults, loc) -> None: ...
+def _compile_function_with_overload(overload_fn, qual_name, impl_fn): ...
+def _get_overloads(obj): ...
+def _check_directly_compile_overloaded(obj) -> None: ...
+def interface(obj: _T) -> _T: ...
+def _recursive_compile_class(obj, loc): ...
+
+CompilationUnit: Incomplete
+
+def pad(s: str, padding: int, offset: int = ..., char: str = ...): ...
+
+class _ScriptProfileColumn:
+    header: Incomplete
+    alignment: Incomplete
+    offset: Incomplete
+    rows: Incomplete
+    def __init__(
+        self,
+        header: str,
+        alignment: int = ...,
+        offset: int = ...,
+    ) -> None: ...
+    def add_row(self, lineno: int, value: Any): ...
+    def materialize(self): ...
+
+class _ScriptProfileTable:
+    cols: Incomplete
+    source_range: Incomplete
+    def __init__(
+        self,
+        cols: list[_ScriptProfileColumn],
+        source_range: list[int],
+    ) -> None: ...
+    def dump_string(self): ...
+
+class _ScriptProfile:
+    profile: Incomplete
+    def __init__(self) -> None: ...
+    def enable(self) -> None: ...
+    def disable(self) -> None: ...
+    def dump_string(self) -> str: ...
+    def dump(self) -> None: ...
+
+def _unwrap_optional(x): ...
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_serialization.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_serialization.py
new file mode 100644
index 0000000000000000000000000000000000000000..e086a2a651294a2fc6c7f879f000f32383434b34
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_serialization.py
@@ -0,0 +1,314 @@
+# mypy: allow-untyped-defs
+"""Serialization.
+
+This module contains functionality for serializing TorchScript modules, notably:
+    * torch.jit.save
+    * torch.jit.load
+
+This is not intended to be imported directly; please use the exposed
+functionalities in `torch.jit`.
+"""
+
+import os
+import sys
+import warnings
+
+import torch
+from torch._jit_internal import _get_model_id
+from torch._utils_internal import log_torchscript_usage
+from torch.jit._recursive import wrap_cpp_module
+from torch.serialization import validate_cuda_device
+
+
+def save(m, f, _extra_files=None) -> None:
+    r"""
+    Save an offline version of this module for use in a separate process.
+
+    The saved module serializes all of the methods, submodules, parameters, and
+    attributes of this module. It can be loaded into the C++ API using
+    ``torch::jit::load(filename)`` or into the Python API with
+    :func:`torch.jit.load <torch.jit.load>`.
+
+    To be able to save a module, it must not make any calls to native Python
+    functions.  This means that all submodules must be subclasses of
+    :class:`ScriptModule` as well.
+
+    .. DANGER::
+        All modules, no matter their device, are always loaded onto the CPU
+        during loading.  This is different from :func:`torch.load`'s semantics
+        and may change in the future.
+
+    Args:
+        m: A :class:`ScriptModule` to save.
+        f: A file-like object (has to implement write and flush) or a string
+           containing a file name.
+        _extra_files: Map from filename to contents which will be stored as part of `f`.
+
+    .. note::
+        torch.jit.save attempts to preserve the behavior of some operators
+        across versions. For example, dividing two integer tensors in
+        PyTorch 1.5 performed floor division, and if the module
+        containing that code is saved in PyTorch 1.5 and loaded in PyTorch 1.6
+        its division behavior will be preserved. The same module saved in
+        PyTorch 1.6 will fail to load in PyTorch 1.5, however, since the
+        behavior of division changed in 1.6, and 1.5 does not know how to
+        replicate the 1.6 behavior.
+
+    Example:
+    .. testcode::
+
+        import torch
+        import io
+
+        class MyModule(torch.nn.Module):
+            def forward(self, x):
+                return x + 10
+
+        m = torch.jit.script(MyModule())
+
+        # Save to file
+        torch.jit.save(m, 'scriptmodule.pt')
+        # This line is equivalent to the previous
+        m.save("scriptmodule.pt")
+
+        # Save to io.BytesIO buffer
+        buffer = io.BytesIO()
+        torch.jit.save(m, buffer)
+
+        # Save with extra files
+        extra_files = {'foo.txt': b'bar'}
+        torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files)
+    """
+    if sys.version_info >= (3, 14):
+        warnings.warn(
+            "`torch.jit.save` is not supported in Python 3.14+ and may break. "
+            "Please switch to `torch.export`.",
+            DeprecationWarning,
+        )
+    else:
+        warnings.warn(
+            "`torch.jit.save` is deprecated. Please switch to `torch.export`.",
+            DeprecationWarning,
+        )
+    log_torchscript_usage("save", model_id=_get_model_id(m))
+    if _extra_files is None:
+        _extra_files = {}
+    if isinstance(f, (str, os.PathLike)):
+        m.save(f, _extra_files=_extra_files)
+    else:
+        ret = m.save_to_buffer(_extra_files=_extra_files)
+        f.write(ret)
+
+
+def load(f, map_location=None, _extra_files=None, _restore_shapes=False):
+    r"""
+    Load a :class:`ScriptModule` or :class:`ScriptFunction` previously saved with :func:`torch.jit.save <torch.jit.save>`.
+
+    All previously saved modules, no matter their device, are first loaded onto CPU,
+    and then are moved to the devices they were saved from. If this fails (e.g.
+    because the run time system doesn't have certain devices), an exception is
+    raised.
+
+    Args:
+        f: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+        map_location (string or torch.device): A simplified version of
+            ``map_location`` in `torch.jit.save` used to dynamically remap
+            storages to an alternative set of devices.
+        _extra_files (dictionary of filename to content): The extra
+            filenames given in the map would be loaded and their content
+            would be stored in the provided map.
+        _restore_shapes (bool): Whether or not to retrace the module on load using stored inputs
+
+    Returns:
+        A :class:`ScriptModule` object.
+
+    .. warning::
+        It is possible to construct malicious pickle data which will execute arbitrary code
+        during func:`torch.jit.load`. Never load data that could have come from an untrusted
+        source, or that could have been tampered with. **Only load data you trust**.
+
+    Example:
+    .. testcode::
+
+        import torch
+        import io
+
+        torch.jit.load('scriptmodule.pt')
+
+        # Load ScriptModule from io.BytesIO object
+        with open('scriptmodule.pt', 'rb') as f:
+            buffer = io.BytesIO(f.read())
+
+        # Load all tensors to the original device
+        torch.jit.load(buffer)
+
+        # Load all tensors onto CPU, using a device
+        buffer.seek(0)
+        torch.jit.load(buffer, map_location=torch.device('cpu'))
+
+        # Load all tensors onto CPU, using a string
+        buffer.seek(0)
+        torch.jit.load(buffer, map_location='cpu')
+
+        # Load with extra files.
+        extra_files = {'foo.txt': ''}  # values will be replaced with data
+        torch.jit.load('scriptmodule.pt', _extra_files=extra_files)
+        print(extra_files['foo.txt'])
+
+    .. testoutput::
+        :hide:
+
+        ...
+
+    .. testcleanup::
+
+        import os
+        os.remove("scriptmodule.pt")
+    """
+    if sys.version_info >= (3, 14):
+        warnings.warn(
+            "`torch.jit.load` is not supported in Python 3.14+ and may break. "
+            "Please switch to `torch.export`.",
+            DeprecationWarning,
+        )
+    else:
+        warnings.warn(
+            "`torch.jit.load` is deprecated. Please switch to `torch.export`.",
+            DeprecationWarning,
+        )
+    if isinstance(f, (str, os.PathLike)):
+        if not os.path.exists(f):
+            raise ValueError(f"The provided filename {f} does not exist")
+        if os.path.isdir(f):
+            raise ValueError(f"The provided filename {f} is a directory")
+
+    map_location = validate_map_location(map_location)
+    if _extra_files is None:
+        _extra_files = {}
+
+    cu = torch._C.CompilationUnit()
+    if isinstance(f, (str, os.PathLike)):
+        cpp_module = torch._C.import_ir_module(
+            cu,
+            os.fspath(f),
+            map_location,
+            _extra_files,
+            # pyrefly: ignore [bad-argument-count]
+            _restore_shapes,
+        )  # type: ignore[call-arg]
+    else:
+        cpp_module = torch._C.import_ir_module_from_buffer(
+            cu,
+            f.read(),
+            map_location,
+            _extra_files,
+            # pyrefly: ignore [bad-argument-count]
+            _restore_shapes,
+        )  # type: ignore[call-arg]
+
+    # TODO: Pretty sure this approach loses ConstSequential status and such
+    ret = wrap_cpp_module(cpp_module)
+    log_torchscript_usage("load", model_id=_get_model_id(ret))
+    return ret
+
+
+def validate_map_location(map_location=None):
+    if isinstance(map_location, str):
+        map_location = torch.device(map_location)
+    elif not (map_location is None or isinstance(map_location, torch.device)):
+        raise ValueError(
+            "map_location should be either None, string or torch.device, "
+            "but got type: " + str(type(map_location))
+        )
+
+    if str(map_location).startswith("cuda"):
+        validate_cuda_device(map_location)
+
+    return map_location
+
+
+def jit_module_from_flatbuffer(f):
+    if isinstance(f, (str, os.PathLike)):
+        f = os.fspath(f)
+        return wrap_cpp_module(torch._C._load_jit_module_from_file(f))
+    else:
+        return wrap_cpp_module(torch._C._load_jit_module_from_bytes(f.read()))
+
+
+def save_jit_module_to_flatbuffer(m, f, _extra_files=None) -> None:
+    r"""
+    Save an offline version of this module for use in a separate process.
+
+    The saved module serializes all of the methods, submodules, parameters, and
+    attributes of this module. It can be loaded into the C++ API using
+    ``torch::jit::load_jit_module_from_file(filename)`` or into the Python API with
+    :func:`torch.jit.jit_module_from_flatbuffer<torch.jit.jit_module_from_flatbuffer>`.
+
+    To be able to save a module, it must not make any calls to native Python
+    functions.  This means that all submodules must be subclasses of
+    :class:`ScriptModule` as well.
+
+    .. DANGER::
+        All modules, no matter their device, are always loaded onto the CPU
+        during loading.  This is different from :func:`torch.load`'s semantics
+        and may change in the future.
+
+    Args:
+        m: A :class:`ScriptModule` to save.
+        f: A string for file path
+
+
+    Example:
+    .. testcode::
+
+        import torch
+        import io
+
+        class MyModule(torch.nn.Module):
+            def forward(self, x):
+                return x + 10
+
+        m = torch.jit.script(MyModule())
+
+        # Save to file
+        torch.jit.save_jit_module_to_flatbuffer(m, 'scriptmodule.ff')
+    """
+    extra_files = _extra_files
+    if extra_files is None:
+        extra_files = {}
+
+    if isinstance(f, (str, os.PathLike)):
+        f = os.fspath(f)
+        torch._C._save_jit_module(m._c, f, extra_files)
+    else:
+        s = torch._C._save_jit_module_to_bytes(m._c, extra_files)
+        f.write(s)
+
+
+def get_flatbuffer_module_info(path_or_file):
+    r"""Get some information regarding a model file in flatbuffer format.
+
+    Args:
+        path_or_file: Either str, Path or file like object (BytesIO OK).
+            If it's str or Path, we will read the file referenced by that
+            path as Bytes.
+
+    Returns:
+        A dict with metadata on what that file contains, currently looks like
+        this:
+        {
+            'bytecode_version': 4,  # int
+            'operator_version': 4,  # int
+            'function_names': {
+                '__torch__.___torch_mangle_0.Foo.forward'}, # set
+            'type_names': set(),  # set
+            'opname_to_num_args': {'aten::linear': 3} # Dict[str, int]
+        }
+    """
+    if isinstance(path_or_file, (str, os.PathLike)):
+        with open(path_or_file, "rb") as f:
+            all_bytes = f.read()
+    else:
+        all_bytes = path_or_file.read()
+    return torch._C._get_module_info_from_flatbuffer(all_bytes)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_shape_functions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_shape_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..26c458526c9c91f2f63df490bea39d3b712abf27
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_shape_functions.py
@@ -0,0 +1,1633 @@
+# mypy: allow-untyped-defs
+import math
+from collections.abc import Callable
+from typing import Any, Optional, Union
+
+
+number = Union[int, float]
+# flake8: noqa
+
+###
+# There are generated files that depend on this file
+# To re-generate, please run from the root of the repo:
+# python torchgen/shape_functions/gen_jit_shape_functions.py
+
+# How to test:
+# After regenerating files, compile PyTorch.
+# Then run: ./build/bin/test_jit --gtest_filter=TestShapeGraphLinting.Basic
+# If you have enabled opinfo testing for the op, also run:
+# python test/test_ops_jit.py TestJitCPU.test_variant_consistency_jit_[FAILING_OP]_cpu_float32
+# to reproduce errors from opinfo tests.
+
+# Example PR: https://github.com/pytorch/pytorch/pull/80860/files
+####
+
+import torch
+
+
+def broadcast(a: list[int], b: list[int]):
+    dimsA = len(a)
+    dimsB = len(b)
+    ndim = max(dimsA, dimsB)
+    expandedSizes: list[int] = []
+
+    for i in range(ndim):
+        offset = ndim - 1 - i
+        dimA = dimsA - 1 - offset
+        dimB = dimsB - 1 - offset
+        sizeA = a[dimA] if (dimA >= 0) else 1
+        sizeB = b[dimB] if (dimB >= 0) else 1
+
+        if sizeA != sizeB and sizeA != 1 and sizeB != 1:
+            # TODO: only assertion error is bound in C++ compilation right now
+            raise AssertionError(
+                f"The size of tensor a {sizeA} must match the size of tensor b ({sizeB}) at non-singleton dimension {i}"
+            )
+
+        expandedSizes.append(sizeB if sizeA == 1 else sizeA)
+
+    return expandedSizes
+
+
+def broadcast_three(a: list[int], b: list[int], c: list[int]):
+    return broadcast(broadcast(a, b), c)
+
+
+def broadcast_one_three(a: list[int], b: Any, c: list[int]):
+    return broadcast(a, c)
+
+
+def adaptive_avg_pool2d(self: list[int], out: list[int]):
+    if len(out) != 2:
+        raise AssertionError(f"Expected out to have length 2, but got {len(out)}")
+    if not (len(self) == 3 or len(self) == 4):
+        raise AssertionError(
+            f"Expected self to have length 3 or 4, but got {len(self)}"
+        )
+    for i in range(1, len(self)):
+        if self[i] == 0:
+            raise AssertionError(f"Expected self[{i}] to be non-zero, but got 0")
+
+    shape: list[int] = []
+    for i in range(0, len(self) - 2):
+        shape.append(self[i])
+    for elem in out:
+        shape.append(elem)
+    return shape
+
+
+def _copy(self: list[int]):
+    out: list[int] = []
+    for elem in self:
+        out.append(elem)
+    return out
+
+
+def unary(self: list[int]):
+    return _copy(self)
+
+
+def broadcast_inplace(a: list[int], b: list[int]):
+    dimsA = len(a)
+    dimsB = len(b)
+    if dimsB > dimsA:
+        raise AssertionError(
+            f"The dims of tensor b ({dimsB}) must be less than or equal to the dims of tensor a ({dimsA}) "
+        )
+    for dimA in range(dimsA):
+        dimB = dimsB - dimsA + dimA
+        sizeA = a[dimA]
+        sizeB = b[dimB] if (dimB >= 0) else 1
+        if sizeA != sizeB and sizeB != 1:
+            # TODO: only assertion error is bound in C++ compilation right now
+            raise AssertionError(
+                "The size of tensor a {} must match the size of tensor b ("
+                "{}) at non-singleton dimension {}".format(sizeA, sizeB, dimA)
+            )
+    return _copy(a)
+
+
+def expand(self: list[int], sizes: list[int]):
+    if len(sizes) < len(self):
+        raise AssertionError(
+            f"Expected len(sizes) ({len(sizes)}) >= len(self) ({len(self)})"
+        )
+    ndim = len(sizes)
+    tensor_dim = len(self)
+    if ndim == 0:
+        return _copy(sizes)
+    out: list[int] = []
+    for i in range(ndim):
+        offset = ndim - 1 - i
+        dim = tensor_dim - 1 - offset
+        size = self[dim] if dim >= 0 else 1
+        targetSize = sizes[i]
+        if targetSize == -1:
+            if dim < 0:
+                raise AssertionError(f"Expected dim ({dim}) >= 0 when targetSize is -1")
+            targetSize = size
+        if size != targetSize:
+            if size != 1:
+                raise AssertionError(
+                    f"Expected size ({size}) == 1 when size != targetSize ({targetSize})"
+                )
+            size = targetSize
+        out.append(size)
+    return out
+
+
+def expand_one_unused(self: list[int], sizes: list[int], inp0: Any):
+    return expand(self, sizes)
+
+
+def infer_size_impl(shape: list[int], numel: int) -> list[int]:
+    newsize = 1
+    infer_dim: Optional[int] = None
+    for dim in range(len(shape)):
+        if shape[dim] == -1:
+            if infer_dim is not None:
+                raise AssertionError("only one dimension can be inferred")
+            infer_dim = dim
+        elif shape[dim] >= 0:
+            newsize *= shape[dim]
+        else:
+            raise AssertionError("invalid shape dimensions")
+    if not (
+        numel == newsize
+        or (infer_dim is not None and newsize > 0 and numel % newsize == 0)
+    ):
+        raise AssertionError("invalid shape")
+    out = _copy(shape)
+    if infer_dim is not None:
+        out[infer_dim] = numel // newsize
+    return out
+
+
+def numel(sizes: list[int]):
+    numel = 1
+    for elem in sizes:
+        numel *= elem
+    return numel
+
+
+def view(self: list[int], sizes: list[int]):
+    return infer_size_impl(sizes, numel(self))
+
+
+def view_one_unused(self: list[int], sizes: list[int], *, implicit: bool = False):
+    return view(self, sizes)
+
+
+def sum_mean_dim(
+    self: list[int], opt_dims: Optional[list[int]], keep_dim: bool, dt: Any
+):
+    out: list[int] = []
+    if opt_dims is None or len(opt_dims) == 0:
+        dims: list[int] = list(range(len(self)))
+    else:
+        dims = opt_dims
+
+    for idx in range(len(self)):
+        is_mean_dim: bool = False
+        for reduce_dim in dims:
+            if idx == maybe_wrap_dim(reduce_dim, len(self)):
+                is_mean_dim = True
+        if is_mean_dim:
+            if keep_dim:
+                out.append(1)
+        else:
+            out.append(self[idx])
+    return out
+
+
+def max_dim(self: list[int], dim: int, keep_dim: bool):
+    out = sum_mean_dim(self, [dim], keep_dim, None)
+    return out, out
+
+
+# note: python already rounds down towards negative infinity on integer division, special arithmetic not needed
+def div_rtn(x: int, y: int):
+    return x // y
+
+
+def pooling_output_shape_pad_lr(
+    inputSize: int,
+    kernelSize: int,
+    pad_l: int,
+    pad_r: int,
+    stride: int,
+    dilation: int,
+    ceil_mode: bool,
+):
+    outputSize = (
+        div_rtn(
+            inputSize
+            + pad_l
+            + pad_r
+            - dilation * (kernelSize - 1)
+            - 1
+            + (stride - 1 if ceil_mode else 0),
+            stride,
+        )
+        + 1
+    )
+    if ceil_mode:
+        if (outputSize - 1) * stride >= inputSize + pad_l:
+            outputSize = outputSize - 1
+    return outputSize
+
+
+def pooling_output_shape(
+    inputSize: int,
+    kernelSize: int,
+    pad_l: int,
+    stride: int,
+    dilation: int,
+    ceil_mode: bool,
+):
+    if stride == 0:
+        raise AssertionError("stride should not be zero")
+    return pooling_output_shape_pad_lr(
+        inputSize, kernelSize, pad_l, pad_l, stride, dilation, ceil_mode
+    )
+
+
+def pool2d_shape_check(
+    input: list[int],
+    kH: int,
+    kW: int,
+    dH: int,
+    dW: int,
+    padH: int,
+    padW: int,
+    dilationH: int,
+    dilationW: int,
+    nInputPlane: int,
+    inputHeight: int,
+    inputWidth: int,
+    outputHeight: int,
+    outputWidth: int,
+):
+    ndim = len(input)
+
+    if not (kW > 0 and kH > 0):
+        raise AssertionError(f"Expected kW ({kW}) > 0 and kH ({kH}) > 0")
+    if not (dW > 0 and dH > 0):
+        raise AssertionError(f"Expected dW ({dW}) > 0 and dH ({dH}) > 0")
+    if not (dilationH > 0 and dilationW > 0):
+        raise AssertionError(
+            f"Expected dilationH ({dilationH}) > 0 and dilationW ({dilationW}) > 0"
+        )
+
+    valid_dims = input[1] != 0 and input[2] != 0
+    if not (
+        ndim == 3
+        and input[0] != 0
+        and valid_dims
+        or (ndim == 4 and valid_dims and input[3] != 0)
+    ):
+        raise AssertionError(f"Invalid input dimensions: ndim={ndim}, input={input}")
+
+    if not (kW // 2 >= padW and kH // 2 >= padH):
+        raise AssertionError(
+            f"Expected kW//2 ({kW // 2}) >= padW ({padW}) and "
+            f"kH//2 ({kH // 2}) >= padH ({padH})"
+        )
+    if not (outputWidth >= 1 and outputHeight >= 1):
+        raise AssertionError(
+            f"Expected outputWidth ({outputWidth}) >= 1 and "
+            f"outputHeight ({outputHeight}) >= 1"
+        )
+
+
+def max_pool2d(
+    input: list[int],
+    kernel_size: list[int],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    ceil_mode: bool,
+):
+    if not (len(kernel_size) == 1 or len(kernel_size) == 2):
+        raise AssertionError(
+            "max_pool2d: kernel_size must either be a single int, or a tuple of two ints"
+        )
+    kH = kernel_size[0]
+    kW = kH if len(kernel_size) == 1 else kernel_size[1]
+
+    if not (len(stride) == 0 or len(stride) == 1 or len(stride) == 2):
+        raise AssertionError(
+            "max_pool2d: stride must either be omitted, a single int, or a tuple of two ints"
+        )
+    dH = kH if len(stride) == 0 else stride[0]
+    if len(stride) == 0:
+        dW = kW
+    elif len(stride) == 1:
+        dW = dH
+    else:
+        dW = stride[1]
+
+    if not (len(padding) == 1 or len(padding) == 2):
+        raise AssertionError(
+            "max_pool2d: padding must either be a single int, or a tuple of two ints"
+        )
+    padH = padding[0]
+    padW = padH if len(padding) == 1 else padding[1]
+
+    if not (len(dilation) == 1 or len(dilation) == 2):
+        raise AssertionError(
+            "max_pool2d: dilation must be either a single int, or a tuple of two ints"
+        )
+    dilationH = dilation[0]
+    dilationW = dilationH if len(dilation) == 1 else dilation[1]
+
+    if not (len(input) == 3 or len(input) == 4):
+        raise AssertionError(f"Expected input length 3 or 4, but got {len(input)}")
+
+    nbatch = input[-4] if len(input) == 4 else 1
+    nInputPlane = input[-3]
+    inputHeight = input[-2]
+    inputWidth = input[-1]
+
+    outputHeight = pooling_output_shape(inputHeight, kH, padH, dH, dilationH, ceil_mode)
+    outputWidth = pooling_output_shape(inputWidth, kW, padW, dW, dilationW, ceil_mode)
+
+    pool2d_shape_check(
+        input,
+        kH,
+        kW,
+        dH,
+        dW,
+        padH,
+        padW,
+        dilationH,
+        dilationW,
+        nInputPlane,
+        inputHeight,
+        inputWidth,
+        outputHeight,
+        outputWidth,
+    )
+
+    if len(input) == 3:
+        return [nInputPlane, outputHeight, outputWidth]
+    else:
+        return [nbatch, nInputPlane, outputHeight, outputWidth]
+
+
+def max_pool2d_with_indices(
+    input: list[int],
+    kernel_size: list[int],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    ceil_mode: bool,
+):
+    out = max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
+    return (out, out)
+
+
+def upsample_nearest2d(
+    input: list[int],
+    output_size: Optional[list[int]],
+    scale_factors: Optional[list[float]],
+):
+    out: list[int] = []
+    out.append(input[0])
+    out.append(input[1])
+
+    if scale_factors is None and output_size is None:
+        raise AssertionError("Either output_size or scale_factors must be presented")
+
+    if output_size is not None:
+        if scale_factors is not None:
+            raise AssertionError(
+                "Must specify exactly one of output_size and scale_factors"
+            )
+        if len(output_size) != 2:
+            raise AssertionError(
+                f"Expected output_size to have length 2, but got {len(output_size)}"
+            )
+        out.append(output_size[0])
+        out.append(output_size[1])
+
+    if scale_factors is not None:
+        if output_size is not None:
+            raise AssertionError(
+                "Must specify exactly one of output_size and scale_factors"
+            )
+        if len(scale_factors) != 2:
+            raise AssertionError(
+                f"Expected scale_factors to have length 2, but got {len(scale_factors)}"
+            )
+        out.append(int(input[2] * scale_factors[0]))
+        out.append(int(input[3] * scale_factors[1]))
+
+    return out
+
+
+def mm(self: list[int], mat2: list[int]):
+    if len(self) != 2:
+        raise AssertionError(f"self must be a matrix (got {len(self)} dimensions)")
+    if len(mat2) != 2:
+        raise AssertionError(f"mat2 must be a matrix (got {len(mat2)} dimensions)")
+
+    if self[1] != mat2[0]:
+        raise AssertionError(
+            f"Matrix dimensions don't match for mm: self[1]={self[1]}, mat2[0]={mat2[0]}"
+        )
+    return [self[0], mat2[1]]
+
+
+def dot(self: list[int], tensor: list[int]):
+    if not (len(self) == 1 and len(tensor) == 1):
+        raise AssertionError(
+            f"Expected 1D tensors for dot, got len(self)={len(self)}, "
+            f"len(tensor)={len(tensor)}"
+        )
+    if self[0] != tensor[0]:
+        raise AssertionError(
+            f"Dot product dimension mismatch: self[0]={self[0]}, tensor[0]={tensor[0]}"
+        )
+    out: list[int] = []
+    return out
+
+
+def mv(self: list[int], vec: list[int]):
+    if not (len(self) == 2 and len(vec) == 1):
+        raise AssertionError(
+            f"Expected 2D matrix and 1D vector, got len(self)={len(self)}, "
+            f"len(vec)={len(vec)}"
+        )
+    if self[1] != vec[0]:
+        raise AssertionError(
+            f"Matrix-vector dimension mismatch: self[1]={self[1]}, vec[0]={vec[0]}"
+        )
+    # TODO: return self
+    return [self[0]]
+
+
+def unsqueeze(li: list[int], dim: int):
+    dim = maybe_wrap_dim(dim, len(li) + 1)
+    out = _copy(li)
+    out.insert(dim, 1)
+    return out
+
+
+def squeeze_nodim(li: list[int]):
+    out: list[int] = []
+    for i in range(len(li)):
+        if li[i] != 1:
+            out.append(li[i])
+    return out
+
+
+def squeeze(li: list[int], dim: int):
+    out: list[int] = []
+    wrapped_dim = maybe_wrap_dim(dim, len(li))
+    for i in range(len(li)):
+        if i == wrapped_dim:
+            if li[i] != 1:
+                out.append(li[i])
+        else:
+            out.append(li[i])
+    return out
+
+
+def squeeze_dims(li: list[int], dims: list[int]):
+    if len(dims) == 0:
+        return li
+    wrapped_dims = _copy(dims)
+    for i in range(len(dims)):
+        wrapped_dims[i] = maybe_wrap_dim(wrapped_dims[i], len(li))
+    result: list[int] = []
+    for i in range(len(li)):
+        if li[i] == 1:
+            if i not in wrapped_dims:
+                result.append(li[i])
+        else:
+            result.append(li[i])
+    return result
+
+
+def index_select(self: list[int], dim: int, index: list[int]):
+    dim = maybe_wrap_dim(dim, len(self))
+    numel = multiply_integers(index)
+    if len(index) > 1:
+        raise AssertionError(f"Expected len(index) <= 1, but got {len(index)}")
+    if not (dim == 0 or dim < len(self)):
+        raise AssertionError(
+            f"Expected dim ({dim}) == 0 or dim < len(self) ({len(self)})"
+        )
+    result_size: list[int] = []
+    for i in range(len(self)):
+        if dim == i:
+            result_size.append(numel)
+        else:
+            result_size.append(self[i])
+    return result_size
+
+
+def embedding(
+    weight: list[int],
+    indices: list[int],
+    padding_idx: int = -1,
+    scale_grad_by_freq: bool = False,
+    sparse: bool = False,
+):
+    if len(weight) != 2:
+        raise AssertionError(f"Expected weight to be 2D, but got {len(weight)}D")
+    if len(indices) == 1:
+        return index_select(weight, 0, indices)
+    size = _copy(indices)
+    size.append(weight[1])
+    return size
+
+
+def max_int():
+    return 9223372036854775807
+
+
+def slice(
+    self: list[int], dim: int, start: Optional[int], end: Optional[int], step: int
+):
+    ndim = len(self)
+    if ndim == 0:
+        raise AssertionError("Cannot slice a 0-dimensional tensor")
+    dim = maybe_wrap_dim(dim, ndim)
+    start_val = start if start is not None else 0
+    end_val = end if end is not None else max_int()
+    if step <= 0:
+        raise AssertionError(f"Expected step > 0, but got {step}")
+    if start_val == max_int():
+        start_val = 0
+    if start_val < 0:
+        start_val += self[dim]
+    if end_val < 0:
+        end_val += self[dim]
+    if start_val < 0:
+        start_val = 0
+    elif start_val > self[dim]:
+        start_val = self[dim]
+    if end_val < start_val:
+        end_val = start_val
+    elif end_val >= self[dim]:
+        end_val = self[dim]
+    slice_len = end_val - start_val
+    out = _copy(self)
+    out[dim] = (slice_len + step - 1) // step
+    return out
+
+
+def check_cat_no_zero_dim(tensors: list[list[int]]):
+    for tensor in tensors:
+        if len(tensor) <= 0:
+            raise AssertionError("Cannot concatenate tensor with 0 dimensions")
+
+
+def legacy_cat_wrap_dim(dim: int, tensor_sizes: list[list[int]]):
+    out_dim: Optional[int] = None
+    for size in tensor_sizes:
+        if not (len(size) == 1 and size[0] == 0):
+            if out_dim is None:
+                out_dim = maybe_wrap_dim(dim, len(size))
+    if out_dim is None:
+        out_dim = dim
+    return out_dim
+
+
+def should_skip(tensor: list[int]):
+    return numel(tensor) == 0 and len(tensor) == 1
+
+
+def check_cat_shape_except_dim(
+    first: list[int], second: list[int], dimension: int, index: int
+):
+    first_dims = len(first)
+    second_dims = len(second)
+    if first_dims != second_dims:
+        raise AssertionError(
+            f"Tensors must have same number of dimensions, got {first_dims} and "
+            f"{second_dims}"
+        )
+    for dim in range(0, first_dims):
+        if dim != dimension:
+            if first[dim] != second[dim]:
+                raise AssertionError(
+                    f"Sizes of tensors must match except in dimension {dimension}, "
+                    f"got {first[dim]} and {second[dim]} at dimension {dim}"
+                )
+
+
+def cat(tensors: list[list[int]], dim: int):
+    check_cat_no_zero_dim(tensors)
+    dim = legacy_cat_wrap_dim(dim, tensors)
+    if len(tensors) <= 0:
+        raise AssertionError("Cannot concatenate empty list of tensors")
+    not_skipped_tensor: Optional[list[int]] = None
+    for tensor in tensors:
+        if not should_skip(tensor):
+            not_skipped_tensor = tensor
+    if not_skipped_tensor is None:
+        return [0]
+
+    cat_dim_size = 0
+
+    for i in range(len(tensors)):
+        tensor = tensors[i]
+        if not should_skip(tensor):
+            check_cat_shape_except_dim(not_skipped_tensor, tensor, dim, i)
+            cat_dim_size = cat_dim_size + tensor[dim]
+
+    result_size = _copy(not_skipped_tensor)
+    result_size[dim] = cat_dim_size
+    return result_size
+
+
+def stack(tensors: list[list[int]], dim: int):
+    unsqueezed_tensors: list[list[int]] = []
+    for tensor in tensors:
+        unsqueezed = unsqueeze(tensor, dim)
+        unsqueezed_tensors.append(unsqueezed)
+    return cat(unsqueezed_tensors, dim)
+
+
+def select(self: list[int], dim: int, index: int):
+    ndim = len(self)
+    if ndim == 0:
+        raise AssertionError("Cannot select from a 0-dimensional tensor")
+    dim = maybe_wrap_dim(dim, ndim)
+    size = self[dim]
+    if index < -size or index >= size:
+        raise AssertionError(
+            f"Index {index} is out of bounds for dimension {dim} with size {size}"
+        )
+    if index < 0:
+        index += size
+    out: list[int] = []
+    for i in range(ndim):
+        if i != dim:
+            out.append(self[i])
+    return out
+
+
+def matmul(tensor1: list[int], tensor2: list[int]):
+    dim_tensor1 = len(tensor1)
+    dim_tensor2 = len(tensor2)
+    if dim_tensor1 == 1 and dim_tensor2 == 1:
+        return dot(tensor1, tensor2)
+    elif dim_tensor1 == 2 and dim_tensor2 == 1:
+        return mv(tensor1, tensor2)
+    elif dim_tensor1 == 1 and dim_tensor2 == 2:
+        return squeeze(mm(unsqueeze(tensor1, 0), tensor2), 0)
+    elif dim_tensor1 == 2 and dim_tensor2 == 2:
+        return mm(tensor1, tensor2)
+    elif dim_tensor1 >= 1 and dim_tensor2 >= 1:
+        # We are multiplying b1 x n x m1 by x2 x m2 x p (where b1 can be a list);
+        # we track m1 vs m2 separately even though they must match for nicer error messages
+        n = tensor1[-2] if dim_tensor1 > 1 else 1
+        batch_tensor1: list[int] = []
+        # TODO: handling of slice
+        for i in range(dim_tensor1 - 2):
+            batch_tensor1.append(tensor1[i])
+        p = tensor2[-1]
+        batch_tensor2: list[int] = []
+        # TODO: handling of slice
+        for i in range(dim_tensor2 - 2):
+            batch_tensor2.append(tensor2[i])
+
+        # expand the batch portion (i.e. cut off matrix dimensions and expand rest)
+        expand_batch_portion = broadcast(batch_tensor1, batch_tensor2)
+
+        # todo: copy ?
+        output_shape = expand_batch_portion
+        if dim_tensor1 > 1:
+            output_shape.append(n)
+
+        if dim_tensor2 > 1:
+            output_shape.append(p)
+
+        return output_shape
+    else:
+        raise AssertionError("both arguments to matmul need to be at least 1D")
+
+
+def t(self: list[int]):
+    if len(self) > 2:
+        raise AssertionError(
+            f"Expected tensor to have <= 2 dimensions, but got {len(self)}"
+        )
+    self_len = len(self)
+    if self_len == 0:
+        out: list[int] = []
+        return out
+    elif self_len == 1:
+        return [self[0]]
+    else:
+        return [self[1], self[0]]
+
+
+def transpose(self: list[int], dim0: int, dim1: int):
+    ndims = len(self)
+    dim0 = maybe_wrap_dim(dim0, ndims)
+    dim1 = maybe_wrap_dim(dim1, ndims)
+    if dim0 == dim1:
+        return _copy(self)
+    out: list[int] = []
+    for i in range(ndims):
+        if i == dim0:
+            out.append(self[dim1])
+        elif i == dim1:
+            out.append(self[dim0])
+        else:
+            out.append(self[i])
+    return out
+
+
+def linear(input: list[int], weight: list[int], bias: Optional[list[int]]):
+    out = matmul(input, t(weight))
+    if bias is not None:
+        if broadcast(bias, out) != out:
+            raise AssertionError(
+                f"Bias shape {bias} is not broadcastable to output shape {out}"
+            )
+    return out
+
+
+def addmm(self: list[int], mat1: list[int], mat2: list[int], beta: Any, alpha: Any):
+    return broadcast(self, mm(mat1, mat2))
+
+
+def check_non_negative(array: list[int]) -> bool:
+    # TODO: look into rewriting with early return and getting loop unrolling to fire
+    non_negative = False
+    for val in array:
+        if val < 0:
+            non_negative = True
+    return non_negative
+
+
+def check_shape_forward(
+    input: list[int],
+    weight_sizes: list[int],
+    bias: Optional[list[int]],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    groups: int,
+):
+    k = len(input)
+    weight_dim = len(weight_sizes)
+
+    # TODO: assertions could be expanded with the error messages
+    if check_non_negative(padding):
+        raise AssertionError(f"Padding must be non-negative, got {padding}")
+    if check_non_negative(stride):
+        raise AssertionError(f"Stride must be non-negative, got {stride}")
+
+    if weight_dim != k:
+        raise AssertionError(f"Expected weight_dim ({weight_dim}) == k ({k})")
+    if weight_sizes[0] < groups:
+        raise AssertionError(
+            f"Expected weight_sizes[0] ({weight_sizes[0]}) >= groups ({groups})"
+        )
+    if (weight_sizes[0] % groups) != 0:
+        raise AssertionError(
+            f"Expected weight_sizes[0] ({weight_sizes[0]}) to be divisible by "
+            f"groups ({groups})"
+        )
+    # only handling not transposed
+    if input[1] != weight_sizes[1] * groups:
+        raise AssertionError(
+            f"Expected input[1] ({input[1]}) == weight_sizes[1] * groups "
+            f"({weight_sizes[1] * groups})"
+        )
+    if bias is not None and not (len(bias) == 1 and bias[0] == weight_sizes[0]):
+        raise AssertionError(
+            f"Expected bias to be None or have shape [1] with value "
+            f"weight_sizes[0]={weight_sizes[0]}, got {bias}"
+        )
+
+    for i in range(2, k):
+        if (input[i] + 2 * padding[i - 2]) < (
+            dilation[i - 2] * (weight_sizes[i] - 1) + 1
+        ):
+            raise AssertionError(
+                f"Calculated padded input size ({input[i] + 2 * padding[i - 2]}) "
+                f"is smaller than effective kernel size "
+                f"({dilation[i - 2] * (weight_sizes[i] - 1) + 1}) at dimension {i}"
+            )
+
+    # this is not handling transposed convolution yet
+
+
+def conv_output_size(
+    input_size: list[int],
+    weight_size: list[int],
+    bias: Optional[list[int]],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    groups: int,
+):
+    check_shape_forward(
+        input_size, weight_size, bias, stride, padding, dilation, groups
+    )
+
+    has_dilation = len(dilation) > 0
+    dim = len(input_size)
+    output_size: list[int] = []
+    input_batch_size_dim = 0
+    weight_output_channels_dim = 0
+    output_size.append(input_size[input_batch_size_dim])
+    output_size.append(weight_size[weight_output_channels_dim])
+
+    for d in range(2, dim):
+        dilation_ = dilation[d - 2] if has_dilation else 1
+        kernel = dilation_ * (weight_size[d] - 1) + 1
+        output_size.append(
+            (input_size[d] + (2 * padding[d - 2]) - kernel) // stride[d - 2] + 1
+        )
+    return output_size
+
+
+def conv1d(
+    input: list[int],
+    weight: list[int],
+    bias: Optional[list[int]],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    groups: int,
+):
+    if len(weight) != 3:
+        raise AssertionError(f"Expected 3D weight for conv1d, got {len(weight)}D")
+    if len(input) != 3:
+        raise AssertionError(f"Expected 3D input for conv1d, got {len(input)}D")
+    return conv_output_size(input, weight, bias, stride, padding, dilation, groups)
+
+
+def conv2d(
+    input: list[int],
+    weight: list[int],
+    bias: Optional[list[int]],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    groups: int,
+):
+    if len(weight) != 4:
+        raise AssertionError(f"Expected 4D weight for conv2d, got {len(weight)}D")
+    if len(input) != 4:
+        raise AssertionError(f"Expected 4D input for conv2d, got {len(input)}D")
+    return conv_output_size(input, weight, bias, stride, padding, dilation, groups)
+
+
+def conv_backwards(
+    grad_output: list[int],
+    input: list[int],
+    weight: list[int],
+    biases: Optional[list[int]],
+):
+    # Bias gradient is always generated regardess of if biases is supplied
+    return _copy(input), _copy(weight), [grad_output[1]]
+
+
+def conv_transpose2d_input(
+    input: list[int],
+    weight: list[int],
+    bias: Optional[list[int]] = None,
+    stride: Optional[list[int]] = None,
+    padding: Optional[list[int]] = None,
+    output_padding: Optional[list[int]] = None,
+    groups: int = 1,
+    dilation: Optional[list[int]] = None,
+) -> list[int]:
+    if stride is None:
+        stride = [1, 1]
+    if padding is None:
+        padding = [0, 0]
+    if output_padding is None:
+        output_padding = [0, 0]
+    if dilation is None:
+        dilation = [1, 1]
+    has_dilation = len(dilation) > 0
+    dim = len(input)
+    output_size: list[int] = []
+    input_batch_size_dim = 0
+    weight_output_channels_dim = 1
+    output_size.append(input[input_batch_size_dim])
+    output_size.append(weight[weight_output_channels_dim] * groups)
+
+    for d in range(2, dim):
+        dilation_ = dilation[d - 2] if has_dilation else 1
+        kernel = dilation_ * (weight[d] - 1)
+        output_size.append(
+            (input[d] - 1) * stride[d - 2]
+            - 2 * padding[d - 2]
+            + kernel
+            + output_padding[d - 2]
+            + 1
+        )
+    return output_size
+
+
+def conv_forwards(
+    input: list[int],
+    weight: list[int],
+    bias: Optional[list[int]],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    transposed: bool,
+    output_padding: list[int],
+    groups: int,
+) -> list[int]:
+    has_dilation = len(dilation) > 0
+    has_output_padding = len(output_padding) > 0
+    dim = len(input)
+    output_size: list[int] = []
+    input_batch_size_dim = 0
+    weight_output_channels_dim = 1 if transposed else 0
+    output_size.append(input[input_batch_size_dim])
+    if transposed:
+        output_size.append(weight[weight_output_channels_dim] * groups)
+    else:
+        output_size.append(weight[weight_output_channels_dim])
+
+    for d in range(2, dim):
+        dilation_ = dilation[d - 2] if has_dilation else 1
+        output_padding_ = output_padding[d - 2] if has_output_padding else 0
+        if transposed:
+            kernel = dilation_ * (weight[d] - 1)
+            output_size.append(
+                (input[d] - 1) * stride[d - 2]
+                - 2 * padding[d - 2]
+                + kernel
+                + output_padding_
+                + 1
+            )
+        else:
+            kernel = dilation_ * (weight[d] - 1) + 1
+            output_size.append(
+                (input[d] + (2 * padding[d - 2]) - kernel) // stride[d - 2] + 1
+            )
+    return output_size
+
+
+def _conv_forwards(
+    input: list[int],
+    weight: list[int],
+    bias: Optional[list[int]],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    transposed: bool,
+    output_padding: list[int],
+    groups: int,
+    benchmark: bool,
+    deterministic: bool,
+    cudnn_enabled: bool,
+    allow_tf32: bool,
+) -> list[int]:
+    return conv_forwards(
+        input,
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        transposed,
+        output_padding,
+        groups,
+    )
+
+
+def batch_norm(
+    input: list[int],
+    weight: Optional[list[int]],
+    bias: Optional[list[int]],
+    running_mean: Optional[list[int]],
+    running_var: Optional[list[int]],
+    training: bool,
+    momentum: float,
+    eps: float,
+    cudnn_enabled: bool,
+):
+    out: list[int] = []
+    for elem in input:
+        out.append(elem)
+    return out
+
+
+def conv3d(
+    input: list[int],
+    weight: list[int],
+    bias: Optional[list[int]],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    groups: int,
+):
+    if len(weight) != 5:
+        raise AssertionError(f"Expected 5D weight for conv3d, got {len(weight)}D")
+    if len(input) != 5:
+        raise AssertionError(f"Expected 5D input for conv3d, got {len(input)}D")
+    return conv_output_size(input, weight, bias, stride, padding, dilation, groups)
+
+
+def maybe_wrap_dim(dim: int, dim_post_expr: int, wrap_scalar: bool = True):
+    if dim_post_expr <= 0:
+        if not wrap_scalar:
+            raise AssertionError(
+                "Expected wrap_scalar to be True when dim_post_expr <= 0"
+            )
+        dim_post_expr = 1
+    min = -dim_post_expr
+    max = dim_post_expr - 1
+    if dim < min or dim > max:
+        raise AssertionError(
+            f"Dimension {dim} out of range (expected to be in range [{min}, {max}])"
+        )
+    if dim < 0:
+        dim += dim_post_expr
+    return dim
+
+
+def zero_dim_tensor(input: Any):
+    out: list[int] = []
+    return out
+
+
+def multiply_integers(li: list[int]):
+    out = 1
+    for elem in li:
+        out = out * elem
+    return out
+
+
+def arange_end(end: number, inp0: Any, inp1: Any, inp2: Any, inp3: Any):
+    if end < 0:
+        raise AssertionError(f"Expected end ({end}) >= 0")
+    return [int(math.ceil(end))]
+
+
+def arange_start(
+    start: number, end: number, inp0: Any, inp1: Any, inp2: Any, inp3: Any
+):
+    if end < 0:
+        raise AssertionError(f"Expected end ({end}) >= 0")
+    if end < start:
+        raise AssertionError(f"Expected end ({end}) >= start ({start})")
+    return [int(math.ceil(end - start))]
+
+
+def arange_start_step(
+    start: number, end: number, step: number, inp0: Any, inp1: Any, inp2: Any, inp3: Any
+):
+    if step == 0:
+        raise AssertionError("step must not be zero")
+    if step < 0:
+        if start < end:
+            raise AssertionError(
+                f"Expected start ({start}) >= end ({end}) when step < 0"
+            )
+    else:
+        if end < start:
+            raise AssertionError(
+                f"Expected end ({end}) >= start ({start}) when step > 0"
+            )
+    return [int(math.ceil((end - start) / step))]
+
+
+def permute(input: list[int], dims: list[int]):
+    if len(input) != len(dims):
+        raise AssertionError(
+            f"Expected len(input) ({len(input)}) == len(dims) ({len(dims)})"
+        )
+    ndim = len(dims)
+    seen_dims: list[int] = []
+    newSizes: list[int] = []
+    for i in range(ndim):
+        dim = maybe_wrap_dim(dims[i], ndim)
+        seen_dims.append(dim)
+        newSizes.append(input[dim])
+    for i in range(1, ndim):
+        for j in range(i):
+            if seen_dims[i] == seen_dims[j]:
+                raise AssertionError(
+                    f"Repeated dimension {seen_dims[i]} in permute dimensions"
+                )
+    return newSizes
+
+
+def movedim(self: list[int], source: list[int], destination: list[int]) -> list[int]:
+    self_dim = len(self)
+    if self_dim <= 1:
+        return self
+    normalized_src: list[int] = []
+    normalized_dst: list[int] = []
+    for i in range(len(source)):
+        normalized_src.append(maybe_wrap_dim(source[i], self_dim))
+        normalized_dst.append(maybe_wrap_dim(destination[i], self_dim))
+    order = [-1 for i in range(self_dim)]
+    src_dims = [i for i in range(self_dim)]
+    dst_dims = [i for i in range(self_dim)]
+
+    for i in range(len(source)):
+        order[normalized_dst[i]] = normalized_src[i]
+        src_dims[normalized_src[i]] = -1
+        dst_dims[normalized_dst[i]] = -1
+
+    source_dims: list[int] = []
+    destination_dims: list[int] = []
+    for ele in src_dims:
+        if ele != -1:
+            source_dims.append(ele)
+    for ele in dst_dims:
+        if ele != -1:
+            destination_dims.append(ele)
+
+    rest_dim = self_dim - len(source)
+    for i in range(rest_dim):
+        order[destination_dims[i]] = source_dims[i]
+    return permute(self, order)
+
+
+def flatten(input: list[int], start_dim: int, end_dim: int):
+    start_dim = maybe_wrap_dim(start_dim, len(input))
+    end_dim = maybe_wrap_dim(end_dim, len(input))
+    if start_dim > end_dim:
+        raise AssertionError(f"Expected start_dim ({start_dim}) <= end_dim ({end_dim})")
+    if len(input) == 0:
+        return [1]
+    if start_dim == end_dim:
+        # TODO: return self
+        out: list[int] = []
+        for elem in input:
+            out.append(elem)
+        return out
+    slice_numel = 1
+    for i in range(start_dim, end_dim + 1):
+        slice_numel *= input[i]
+    # TODO: use slicing when slice optimization has landed
+    # slice_numel = multiply_integers(input[start_dim:end_dim - start_dim + 1])
+    shape: list[int] = []
+    for i in range(start_dim):
+        shape.append(input[i])
+    shape.append(slice_numel)
+    for i in range(end_dim + 1, len(input)):
+        shape.append(input[i])
+    return shape
+
+
+def nonzero_lower_bound(input: list[int]):
+    return [0, len(input)]
+
+
+def nonzero_upper_bound(input: list[int]):
+    return [numel(input), len(input)]
+
+
+def _reduce_along_dim(self: list[int], dim: int, keepdim: bool):
+    dim = maybe_wrap_dim(dim, len(self))
+    out: list[int] = []
+    for i, self_dim in enumerate(self):
+        if i == dim:
+            if keepdim:
+                out.append(1)
+        else:
+            out.append(self_dim)
+    return out
+
+
+def argmax(
+    self: list[int], dim: Optional[int] = None, keepdim: bool = False
+) -> list[int]:
+    if dim is None:
+        return []
+    return _reduce_along_dim(self, dim, keepdim)
+
+
+def bmm(self: list[int], mat2: list[int]) -> list[int]:
+    if len(self) != 3:
+        raise AssertionError(f"bmm only supports 3D tensors, got {len(self)}D")
+    if len(mat2) != 3:
+        raise AssertionError(f"bmm only supports 3D tensors, got {len(mat2)}D")
+    if self[0] != mat2[0]:
+        raise AssertionError(
+            f"mismatching batch dimension: self[0]={self[0]}, mat2[0]={mat2[0]}"
+        )
+    if self[2] != mat2[1]:
+        raise AssertionError(
+            f"mismatching contracting dimension: self[2]={self[2]}, mat2[1]={mat2[1]}"
+        )
+    return [self[0], self[1], mat2[2]]
+
+
+def _shape_as_tensor(self: list[int]) -> list[int]:
+    return [len(self)]
+
+
+def topk(self: list[int], k: int, dim: int = -1) -> tuple[list[int], list[int]]:
+    if len(self) == 0:
+        result: list[int] = []
+    else:
+        if k > self[dim]:
+            raise AssertionError(
+                f"k ({k}) is too big for dimension {dim} of size {self[dim]}"
+            )
+        result = _copy(self)
+        result[dim] = k
+    return result, result
+
+
+def nll_loss_forward(
+    self: list[int], target: list[int], weight: Optional[list[int]], reduction: int
+) -> tuple[list[int], list[int]]:
+    # This is taken shamelessly from the meta function in LossNLL.cpp
+    self_dim = len(self)
+    target_dim = len(target)
+    if not (0 < self_dim <= 2):
+        raise AssertionError(f"Expected 0 < self_dim <= 2, but got self_dim={self_dim}")
+    if target_dim > 1:
+        raise AssertionError(f"Expected target_dim <= 1, but got {target_dim}")
+    no_batch_dim = self_dim == 1 and target_dim == 0
+    if not (no_batch_dim or (self[0] == target[0])):
+        raise AssertionError(
+            f"Batch size mismatch: self[0]={self[0]}, target[0]={target[0]}"
+        )
+    n_classes = self[-1]
+    scalar_shape: list[int] = []
+    if weight is not None and not (len(weight) == 1 and weight[0] == n_classes):
+        raise AssertionError(
+            f"Expected weight to be None or have shape [n_classes], "
+            f"got {weight} with n_classes={n_classes}"
+        )
+    if reduction == 0 and self_dim == 2:
+        reduction_shape = [self[0]]
+    else:
+        reduction_shape = scalar_shape
+    return reduction_shape, scalar_shape
+
+
+def native_layer_norm(
+    input: list[int], normalized_shape: list[int]
+) -> tuple[list[int], list[int], list[int]]:
+    reduction_shape: list[int] = []
+    num_unreduced_dimensions = len(input) - len(normalized_shape)
+    if num_unreduced_dimensions < 0:
+        raise AssertionError(
+            f"Expected len(input) ({len(input)}) >= len(normalized_shape) "
+            f"({len(normalized_shape)})"
+        )
+    for i in range(num_unreduced_dimensions):
+        reduction_shape.append(input[i])
+    for i in range(num_unreduced_dimensions, len(input)):
+        reduction_shape.append(1)
+    return _copy(input), reduction_shape, reduction_shape
+
+
+def native_batch_norm(
+    input: list[int],
+    weight: Optional[list[int]],
+    bias: Optional[list[int]],
+    running_mean: Optional[list[int]],
+    running_var: Optional[list[int]],
+    training: bool,
+) -> tuple[list[int], list[int], list[int]]:
+    if training:
+        _size = [input[1]]
+    else:
+        _size = [0]
+    return _copy(input), _size, _size
+
+
+def _batch_norm_with_update(
+    input: list[int],
+    weight: Optional[list[int]],
+    bias: Optional[list[int]],
+    running_mean: Optional[list[int]],
+    running_var: Optional[list[int]],
+) -> tuple[list[int], list[int], list[int], list[int]]:
+    _size = [input[1]]
+    return _copy(input), _size, _size, [0]
+
+
+def cross_entropy_loss(
+    self: list[int],
+    target: list[int],
+    weight: Optional[list[int]] = None,
+    reduction: int = 1,
+    ignore_index: int = -100,
+    label_smoothing: float = 0.0,
+) -> list[int]:
+    result_shape = nll_loss_forward(self, target, weight, reduction)[0]
+    return result_shape
+
+
+"""
+Currently deferring the enabling of this, as part of the propoasal to suspend
+adding ops.
+There are currently cases in the test case where this is being called
+in the SSA opinfo tests with with unexpected values (eg list of two ints, see the first
+opinfo test). The behavior of index is significantly dependent on the inputs.
+
+This could be an error with how we are matching up shape functions, or that this
+function needs to just implement everything.
+
+def index_Tensor(self: List[int], indices: List[Optional[List[int]]]) -> List[int]:
+    assert len(indices) <= len(self), "More indices than dimensions to index"
+    broadcasted_shape: List[int] = []
+    for index_tensor_shape in indices:
+        if index_tensor_shape is not None:
+            broadcasted_shape = broadcast(broadcasted_shape, index_tensor_shape)
+    return broadcasted_shape
+"""
+
+ScriptFn = torch._C.ScriptFunction
+shape_compute_graph_mapping: dict[str, ScriptFn] = {}
+bounded_compute_graph_mapping: dict[str, tuple[ScriptFn, ScriptFn]] = {}
+script_func_map: dict[Callable, ScriptFn] = {}
+
+
+def process_func(func: Callable):
+    if func not in script_func_map:
+        scripted_func = torch.jit.script(func)
+
+        torch._C._jit_pass_inline(scripted_func.graph)
+
+        for _ in range(2):
+            torch._C._jit_pass_peephole(scripted_func.graph)
+            torch._C._jit_pass_constant_propagation(scripted_func.graph)
+
+        script_func_map[func] = scripted_func
+    return script_func_map[func]
+
+
+def add_shape_compute_mapping(operator_schema: str, func: Callable):
+    global shape_compute_graph_mapping
+
+    shape_compute_graph_mapping[operator_schema] = process_func(func)
+
+
+def add_bounded_compute_mapping(
+    operator_schema: str, lower_bound_func: Callable, upper_bound_func: Callable
+):
+    # Adds a shape compute function for both upper and lower bounds
+    fns = (process_func(lower_bound_func), process_func(upper_bound_func))
+    bounded_compute_graph_mapping[operator_schema] = fns
+
+
+add_shape_compute_mapping(
+    "aten::contiguous(Tensor(a) self, *, MemoryFormat memory_format=contiguous_format) -> Tensor(a)",
+    unary,
+)
+add_shape_compute_mapping(
+    "aten::rsub.Tensor(Tensor self, Scalar other, Scalar alpha=1) -> Tensor", unary
+)
+add_shape_compute_mapping(
+    "aten::dropout(Tensor input, float p, bool train) -> Tensor", unary
+)
+add_shape_compute_mapping(
+    "aten::adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor",
+    adaptive_avg_pool2d,
+)
+add_shape_compute_mapping(
+    "prim::NumToTensor.Scalar(Scalar a) -> Tensor", zero_dim_tensor
+)
+add_shape_compute_mapping("prim::NumToTensor.bool(bool a) -> Tensor", zero_dim_tensor)
+add_shape_compute_mapping(
+    "aten::zeros(int[] size, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor)",
+    unary,
+)
+add_shape_compute_mapping(
+    "aten::to.dtype(Tensor(a) self, int dtype, bool non_blocking=False, bool copy=False, int? memory_format=None) -> (Tensor(a))",
+    unary,
+)
+add_shape_compute_mapping(
+    "aten::arange(Scalar end, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor)",
+    arange_end,
+)
+add_shape_compute_mapping(
+    "aten::arange.start(Scalar start, Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor",
+    arange_start,
+)
+add_shape_compute_mapping(
+    "aten::arange.start_step(Scalar start, Scalar end, Scalar step, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor",
+    arange_start_step,
+)
+add_shape_compute_mapping("aten::squeeze(Tensor(a) self) -> Tensor(a)", squeeze_nodim)
+add_shape_compute_mapping(
+    "aten::squeeze.dim(Tensor(a) self, int dim) -> Tensor(a)", squeeze
+)
+add_shape_compute_mapping(
+    "aten::squeeze.dims(Tensor(a) self, int[] dim) -> Tensor(a)", squeeze_dims
+)
+add_shape_compute_mapping(
+    "aten::unsqueeze(Tensor(a) self, int dim) -> Tensor(a)", unsqueeze
+)
+add_shape_compute_mapping(
+    "aten::slice.Tensor(Tensor(a) self, int dim=0, int? start=None, int? end=None, int step=1) -> Tensor(a)",
+    slice,
+)
+add_shape_compute_mapping(
+    "aten::select.int(Tensor(a) self, int dim, int index) -> Tensor(a)", select
+)
+add_shape_compute_mapping(
+    "aten::index_select(Tensor self, int dim, Tensor index) -> Tensor", index_select
+)
+add_shape_compute_mapping(
+    "aten::layer_norm(Tensor input, int[] normalized_shape, Tensor? weight=None, Tensor? bias=None, "
+    "float eps=1e-05, bool cudnn_enable=True) -> Tensor",
+    unary,
+)
+add_shape_compute_mapping(
+    "aten::softmax.int(Tensor self, int dim, ScalarType? dtype=None) -> Tensor", unary
+)
+add_shape_compute_mapping(
+    "aten::_no_grad_embedding_renorm_(Tensor weight, Tensor input, float max_norm, float norm_type) -> Tensor",
+    unary,
+)
+add_shape_compute_mapping(
+    "aten::embedding_renorm_(Tensor(a!) self, Tensor indices, float max_norm, float norm_type) -> Tensor(a!)",
+    unary,
+)
+add_shape_compute_mapping(
+    "aten::embedding(Tensor weight, Tensor indices, int padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor",
+    embedding,
+)
+add_shape_compute_mapping("aten::mm(Tensor self, Tensor mat2) -> Tensor", mm)
+add_shape_compute_mapping("aten::dot(Tensor self, Tensor tensor) -> Tensor", dot)
+add_shape_compute_mapping("aten::mv(Tensor self, Tensor vec) -> Tensor", mv)
+add_shape_compute_mapping("aten::matmul(Tensor self, Tensor other) -> Tensor", matmul)
+add_shape_compute_mapping(
+    "aten::linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor", linear
+)
+add_shape_compute_mapping(
+    "aten::max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor",
+    max_pool2d,
+)
+add_shape_compute_mapping(
+    "aten::max_pool2d_with_indices(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor)",
+    max_pool2d_with_indices,
+)
+add_shape_compute_mapping("aten::t(Tensor(a) self) -> Tensor(a)", t)
+add_shape_compute_mapping(
+    "aten::transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a)", transpose
+)
+add_shape_compute_mapping(
+    "aten::conv1d(Tensor input, Tensor weight, Tensor? bias=None, int[1] stride=1, int[1] padding=0, int[1] dilation=1, int groups=1) -> Tensor",
+    conv1d,
+)
+add_shape_compute_mapping(
+    "aten::conv2d(Tensor input, Tensor weight, Tensor? bias=None, int[2] stride=1, int[2] padding=0, int[2] dilation=1, int groups=1) -> Tensor",
+    conv2d,
+)
+add_shape_compute_mapping(
+    "aten::batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor",
+    batch_norm,
+)
+add_shape_compute_mapping(
+    "aten::conv3d(Tensor input, Tensor weight, Tensor? bias=None, int[3] stride=1, int[3] padding=0, int[3] dilation=1, int groups=1) -> Tensor",
+    conv3d,
+)
+add_shape_compute_mapping(
+    "aten::convolution_backward(Tensor grad_output, Tensor input, Tensor weight, int[]? bias_sizes, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups, bool[3] output_mask) -> (Tensor, Tensor, Tensor)",
+    conv_backwards,
+)
+add_shape_compute_mapping(
+    "aten::convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups) -> Tensor",
+    conv_forwards,
+)
+add_shape_compute_mapping(
+    "aten::_convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, bool transposed, int[] output_padding, int groups, bool benchmark, bool deterministic, bool cudnn_enabled, bool allow_tf32) -> Tensor",
+    _conv_forwards,
+)
+add_shape_compute_mapping(
+    "aten::conv_transpose2d.input(Tensor input, Tensor weight, Tensor? bias=None, int[2] stride=1, int[2] padding=0, int[2] output_padding=0, int groups=1, int[2] dilation=1) -> Tensor",
+    conv_transpose2d_input,
+)
+add_shape_compute_mapping(
+    "aten::flatten.using_ints(Tensor(a) self, int start_dim=0, int end_dim=-1) -> Tensor(a)",
+    flatten,
+)
+add_shape_compute_mapping("aten::cat(Tensor[] tensors, int dim=0) -> Tensor", cat)
+add_shape_compute_mapping("aten::stack(Tensor[] tensors, int dim=0) -> Tensor", stack)
+add_shape_compute_mapping(
+    "aten::permute(Tensor(a) self, int[] dims) -> Tensor(a)", permute
+)
+add_shape_compute_mapping(
+    "aten::movedim.intlist(Tensor(a) self, int[] source, int[] destination) -> Tensor(a)",
+    movedim,
+)
+add_shape_compute_mapping("aten::view(Tensor(a) self, int[] size) -> Tensor(a)", view)
+add_shape_compute_mapping(
+    "aten::expand_as(Tensor(a) self, Tensor other) -> Tensor(a)", expand
+)
+add_shape_compute_mapping(
+    "aten::expand(Tensor(a) self, int[] size, *, bool implicit=False) -> Tensor(a)",
+    expand_one_unused,
+)
+add_shape_compute_mapping(
+    "aten::mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor",
+    sum_mean_dim,
+)
+add_shape_compute_mapping(
+    "aten::sum.dim_IntList(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor",
+    sum_mean_dim,
+)
+add_shape_compute_mapping(
+    "aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)",
+    max_dim,
+)
+add_shape_compute_mapping(
+    "aten::mean(Tensor self, *, ScalarType? dtype=None) -> Tensor", zero_dim_tensor
+)
+add_shape_compute_mapping(
+    "aten::sum(Tensor self, *, ScalarType? dtype=None) -> Tensor", zero_dim_tensor
+)
+add_shape_compute_mapping(
+    "aten::addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor",
+    addmm,
+)
+add_shape_compute_mapping(
+    "aten::upsample_nearest2d.vec(Tensor input, int[]? output_size, float[]? scale_factors) -> (Tensor)",
+    upsample_nearest2d,
+)
+add_shape_compute_mapping(
+    "aten::quantize_per_tensor(Tensor self, float scale, int zero_point, ScalarType dtype) -> Tensor",
+    unary,
+)
+add_shape_compute_mapping(
+    "aten::quantize_per_tensor.tensor_qparams(Tensor self, Tensor scale, Tensor zero_point, ScalarType dtype) -> Tensor",
+    unary,
+)
+add_shape_compute_mapping("aten::dequantize(Tensor self) -> Tensor", unary)
+add_shape_compute_mapping(
+    "quantized::add(Tensor qa, Tensor qb, float scale, int zero_point) -> Tensor qc",
+    broadcast,
+)
+add_shape_compute_mapping(
+    "aten::argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor", argmax
+)
+add_shape_compute_mapping("aten::bmm(Tensor self, Tensor mat2) -> Tensor", bmm)
+add_shape_compute_mapping(
+    "aten::_shape_as_tensor(Tensor self) -> Tensor", _shape_as_tensor
+)
+add_shape_compute_mapping(
+    "aten::topk(Tensor self, int k, int dim=-1, bool largest=True, bool sorted=True) -> (Tensor values, Tensor indices)",
+    topk,
+)
+add_shape_compute_mapping(
+    "aten::nll_loss_forward(Tensor self, Tensor target, Tensor? weight, int reduction, int ignore_index) -> (Tensor output, Tensor total_weight)",
+    nll_loss_forward,
+)
+add_shape_compute_mapping(
+    "aten::native_layer_norm(Tensor input, int[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)",
+    native_layer_norm,
+)
+add_shape_compute_mapping(
+    "aten::native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor)",
+    native_batch_norm,
+)
+add_shape_compute_mapping(
+    "aten::_native_batch_norm_legit(Tensor input, Tensor? weight, Tensor? bias, Tensor running_mean, Tensor running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor)",
+    native_batch_norm,
+)
+add_shape_compute_mapping(
+    "aten::_native_batch_norm_legit.no_stats(Tensor input, Tensor? weight, Tensor? bias, Tensor running_mean, Tensor running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor)",
+    native_batch_norm,
+)
+add_shape_compute_mapping(
+    "_batch_norm_with_update(Tensor input, Tensor? weight, Tensor? bias, Tensor(a!) running_mean, Tensor(b!) running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor, Tensor)",
+    _batch_norm_with_update,
+)
+
+add_shape_compute_mapping(
+    "aten::cross_entropy_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, float label_smoothing=0.0) -> Tensor",
+    cross_entropy_loss,
+)
+# add_shape_compute_mapping("aten::index.Tensor(Tensor self, Tensor?[] indices) -> Tensor", index_Tensor)
+
+# TODO: migrate over all of symbolic_shape_registry_util.cpp
+# These are duplicated here so that the functions will be serialized
+add_shape_compute_mapping(
+    "aten::lerp.Tensor(Tensor self, Tensor end, Tensor weight) -> Tensor",
+    broadcast_three,
+)
+add_shape_compute_mapping(
+    "aten::where.ScalarSelf(Tensor condition, Scalar self, Tensor other) -> Tensor",
+    broadcast_one_three,
+)
+add_shape_compute_mapping(
+    "aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!)",
+    broadcast_inplace,
+)
+
+# quantized_conv_prepack TODO
+
+# Shape Compute Fn with upper and lower bounds
+add_bounded_compute_mapping(
+    "aten::nonzero(Tensor self) -> (Tensor)", nonzero_lower_bound, nonzero_upper_bound
+)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_state.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_state.py
new file mode 100644
index 0000000000000000000000000000000000000000..e811f14d370c031fa455451d64db21db21a6d2af
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_state.py
@@ -0,0 +1,129 @@
+# mypy: allow-untyped-defs
+"""JIT-related state.
+
+This module stores various pieces of Python-global state relating to the JIT.
+
+This is not intended to be imported directly; please the exposed
+functionalities in `torch.jit`.
+"""
+
+import os
+import weakref
+from typing import Any
+
+import torch
+
+
+class EnabledProxy:
+    """Stores whether the JIT is enabled or not.
+
+    This is just a wrapper for a bool, so that we get reference semantics
+    """
+
+    def __init__(self) -> None:
+        self.enabled = self.parse_env(
+            "PYTORCH_JIT", True, "> Using PyTorch JIT", "> PyTorch JIT DISABLED"
+        )
+
+    def parse_env(self, name, default, true_message, false_message):
+        value = os.environ.get(name)
+        if value is None:
+            return default
+        if value.lower() in {"1", "true", "yes"}:
+            return True
+        elif value.lower() in {"0", "false", "no"}:
+            return False
+        if value == "1v":
+            print(true_message)
+            return True
+        elif value == "0v":
+            print(false_message)
+            return False
+        raise ValueError(f"Unknown setting of {name}. Try using 0 or 1.")
+
+    def __bool__(self) -> bool:
+        return self.enabled
+
+
+_enabled = EnabledProxy()
+
+
+def disable() -> None:
+    _enabled.enabled = False
+
+
+def enable() -> None:
+    _enabled.enabled = True
+
+
+# The Python CompilationUnit. All functions and modules defined in Python will
+# live in here. It's defined in Python because doing in cpp creates static
+# destruction order issues.
+_python_cu = torch._C.CompilationUnit()
+
+
+# python class => ScriptClass mapping
+_script_classes: dict[type[Any], type[Any]] = {}
+_name_to_pyclass: dict[str, type[Any]] = {}
+
+
+def _add_script_class(python_class, script_class) -> None:
+    _script_classes[python_class] = script_class
+    _name_to_pyclass[script_class.qualified_name()] = python_class
+
+
+def _get_script_class(python_class):
+    override = getattr(python_class, "_jit_override_qualname", None)
+    if override is not None:
+        python_class = _get_python_class(override)
+    return _script_classes.get(python_class)
+
+
+def _get_python_class(qualified_name):
+    return _name_to_pyclass.get(qualified_name)
+
+
+def _clear_class_state() -> None:
+    _script_classes.clear()
+    _name_to_pyclass.clear()
+
+
+# Caching: we currently cache compilation of free functions and overloaded functions.
+# To cache free functions we hold a weak ref to the function object and
+# map to the compiled fn's qualified name.
+# To cache overloaded functions we hold a weak ref to the function obj and
+# map to all of its overloaded compiled fns.
+# In the future we could consider caching more types of objects so that
+# aliasing is preserved across separate compilations of the same object.
+
+_jit_caching_layer: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
+_jit_function_overload_caching: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
+
+
+def _try_get_jit_cached_overloads(key):
+    qual_names = _jit_function_overload_caching.get(key, None)
+    if qual_names:
+        return [_python_cu.find_function(qual_name) for qual_name in qual_names]
+    else:
+        return None
+
+
+def _set_jit_overload_cache(key, compiled_fns) -> None:
+    _jit_function_overload_caching[key] = [fn.qualified_name for fn in compiled_fns]
+
+
+def _try_get_jit_cached_function(key):
+    if getattr(key, "__disable_jit_function_caching__", False) is True:
+        return None
+    qual_name = _jit_caching_layer.get(key, None)
+    if qual_name:
+        return _python_cu.find_function(qual_name)
+    else:
+        return None
+
+
+def _set_jit_function_cache(key, value) -> None:
+    # only free functions currently supported
+    if not isinstance(value, torch.jit.ScriptFunction):
+        raise AssertionError(f"Expected ScriptFunction, got {type(value)}")
+    _jit_caching_layer[key] = value.qualified_name
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_trace.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_trace.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca2d85a76bc3a863c6d7d5d2ec687c8a50068fe8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/_trace.py
@@ -0,0 +1,1442 @@
+# mypy: allow-untyped-defs
+"""Tracing.
+
+This module contains functionality to support the JIT's tracing frontend, notably:
+    * torch.jit.trace
+    * torch.jit.trace_module
+
+This is not intended to be imported directly; please use the exposed
+functionalities in `torch.jit`.
+"""
+
+import contextlib
+import copy
+import functools
+import inspect
+import os
+import re
+import sys
+import warnings
+from collections.abc import Callable
+from enum import Enum
+from typing import Any, TypeVar
+from typing_extensions import ParamSpec
+
+import torch
+from torch._jit_internal import (
+    _get_model_id,
+    _qualified_name,
+    get_callable_argument_names,
+    is_scripting,
+)
+from torch.autograd import function
+from torch.jit._script import _CachedForward, script, ScriptModule
+from torch.jit._state import _enabled, _python_cu
+from torch.nn import Module
+from torch.testing._comparison import default_tolerances
+
+
+_flatten = torch._C._jit_flatten
+_unflatten = torch._C._jit_unflatten
+
+R = TypeVar("R", covariant=True)  # return type (always covariant)
+P = ParamSpec("P")
+
+
+def _create_interpreter_name_lookup_fn(frames_up=1):
+    def _get_interpreter_name_for_var(var):
+        frame = inspect.currentframe()
+        if not frame:
+            raise RuntimeError("failed to inspect frame")
+
+        i = 0
+        while i < frames_up + 1:
+            frame = frame.f_back
+            if not frame:
+                raise RuntimeError("failed to get frame")
+            i += 1
+
+        f_locals = frame.f_locals
+
+        for k, v in f_locals.items():
+            if isinstance(v, torch.Tensor) and var is v:
+                return k if k != "self" else ""
+        return ""
+
+    return _get_interpreter_name_for_var
+
+
+def _unique_state_dict(module, keep_vars=False):
+    # since Parameter.detach() always creates a new torch.Tensor instance,
+    # id(v) doesn't work with it. So we always get the Parameter or Buffer
+    # as values, and deduplicate the params using Parameters and Buffers
+    state_dict = module.state_dict(keep_vars=True)
+    filtered_dict = type(state_dict)()
+    seen_ids: set[int] = set()
+    for k, v in state_dict.items():
+        if id(v) in seen_ids:
+            continue
+        seen_ids.add(id(v))
+        if keep_vars:
+            filtered_dict[k] = v
+        else:
+            filtered_dict[k] = v.detach()
+    return filtered_dict
+
+
+class ONNXTracedModule(torch.nn.Module):
+    def __init__(
+        self,
+        inner,
+        strict=True,
+        force_outplace=False,
+        return_inputs=False,
+        return_inputs_states=False,
+    ):
+        super().__init__()
+        # inner may be a Module, or it may be an arbitrary callable
+        # If it's a Module, we get its parameters automatically, which lets
+        # us avoid a special casing functions versus modules.
+        self.inner = inner
+        self.strict = strict
+        self._force_outplace = force_outplace
+        self._return_inputs = return_inputs
+        self._return_inputs_states = return_inputs_states
+
+    def forward(self, *args: torch.Tensor):
+        in_vars, in_desc = _flatten(args)
+        # NOTE: use full state, because we need it for BatchNorm export
+        # This differs from the compiler path, which doesn't support it at the moment.
+        module_state = list(_unique_state_dict(self, keep_vars=True).values())
+
+        ret_inputs = []
+        inputs_states = []
+        outs = []
+
+        def wrapper(*args):
+            in_args: list[torch.Tensor] = []
+            for i in range(len(in_vars)):
+                if not isinstance(args[i], torch.Tensor):
+                    raise RuntimeError("Expected Tensor argument")
+                in_args.append(args[i])
+
+            trace_inputs = _unflatten(in_args, in_desc)
+
+            if self._return_inputs:
+                ret_inputs.append(
+                    tuple(x.clone(memory_format=torch.preserve_format) for x in args)
+                )
+            if self._return_inputs_states:
+                inputs_states.append(_unflatten(in_args, in_desc))
+            outs.append(self.inner(*trace_inputs))
+            if self._return_inputs_states:
+                inputs_states[0] = (inputs_states[0], trace_inputs)
+            out_vars, _ = _flatten(outs)
+            if len(out_vars) == 1:
+                return out_vars[0]
+            else:
+                return tuple(out_vars)
+
+        graph, _out = torch._C._create_graph_by_tracing(
+            wrapper,
+            in_vars + module_state,
+            _create_interpreter_name_lookup_fn(),
+            self.strict,
+            self._force_outplace,
+        )
+
+        if self._return_inputs:
+            return graph, outs[0], ret_inputs[0]
+        if self._return_inputs_states:
+            return graph, outs[0], inputs_states[0]
+        else:
+            return graph, outs[0]
+
+
+def _clone_inputs(args):
+    def clone_input(a):
+        if a is None:
+            return None
+        elif isinstance(a, torch.Tensor):
+            # TODO: figure out one liner to .clone() and set requires_grad
+            v = (
+                a.detach()
+                .clone(memory_format=None if a.is_mkldnn else torch.preserve_format)
+                .requires_grad_(a.requires_grad)
+            )
+            if a.grad is not None:
+                v.grad = clone_input(v.grad)
+            return v
+        else:
+            return a.clone(memory_format=torch.preserve_format)
+
+    # pyrefly: ignore [missing-attribute]
+    return function._nested_map(
+        lambda x: isinstance(x, torch.Tensor), clone_input, condition_msg="tensors"
+    )(args)
+
+
+# This is purely for developer debugging.  We are not going to advertise it.
+_JIT_TIME = os.environ.get("PYTORCH_JIT_TIME", False)  # CUDA-only timing
+_JIT_DISABLE = os.environ.get("PYTORCH_JIT_DISABLE", False)
+_JIT_STATS = os.environ.get("PYTORCH_JIT_STATS", False)
+
+
+@contextlib.contextmanager
+def _time(trace_name, name, time=True):
+    if (not _JIT_TIME and not time) or not torch.cuda.is_available():
+        yield
+        return
+    stream = torch.cuda.current_stream()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    stream.record_event(start)
+    try:
+        yield
+    finally:
+        stream.record_event(end)
+        end.synchronize()
+        print(f"{trace_name} {name} time: {start.elapsed_time(end)} ms")
+
+
+def verify(model, args, loss_fn=torch.sum, devices=None):
+    """
+    Verify that a JIT compiled model has the same behavior as its uncompiled version along with its backwards pass.
+
+    If your model returns multiple outputs,
+    you must also specify a `loss_fn` to produce a loss for which
+    the backwards will be computed.
+
+    This function has side-effects (e.g., it executes your model / saves and loads
+    parameters), so don't expect the model to come out exactly the same as what
+    you passed in.
+
+    Args:
+        model (compiled torch.nn.Module or function): the module/function to be
+            verified.  The module/function definition MUST have been decorated with
+            `@torch.jit.compile`.
+        args (tuple or Tensor): the positional arguments to pass to the
+            compiled function/module to be verified.  A non-tuple is assumed to
+            be a single positional argument to be passed to the model.
+        loss_fn (function, optional): the loss function to be applied to
+            the output of the model, before backwards is invoked.  By default,
+            we assume that a model returns a single result, and we :func:`torch.sum`
+            before calling backwards; if this is inappropriate, you can pass your
+            own loss function.  Note that if a model returns a tuple of results,
+            these are passed as separate positional arguments to `loss_fn`.
+        devices (iterable of device IDs, optional): the GPU devices which the
+            compiled module will be run on.  This determines the RNG state we
+            must save when running both compiled and uncompiled versions of the model.
+    """
+    # TODO: In principle, we track device information in our trace, so it
+    # should be possible to check if our execution actually obeyed the 'devices'
+    # the user provided.
+
+    # TODO: Consider adding a utility function to torch.jit to test
+    # for this case
+    if not isinstance(model, torch._C.CompiledFunction):  # type: ignore[attr-defined]
+        raise TypeError(
+            "Cannot verify an uncompiled module.  Add @torch.jit.compile to compile it"
+        )
+    is_module = isinstance(model, Module)
+
+    if not isinstance(args, tuple):
+        args = (args,)
+
+    if is_module:
+        saved_state = copy.deepcopy(model.state_dict())
+
+    def run_fwd_bwd(args, force_trace=False, assert_compiled=False):
+        params = list(model.parameters()) if is_module else []
+        in_vars, _ = _flatten((args, params))
+        # We use a special API to reset the trace and compile it from scratch.
+        compiled_fn = model
+        if force_trace:
+            compiled_fn.clear_cache()
+        if assert_compiled:
+            hits = compiled_fn.hits
+        out = model(*args)
+        if assert_compiled and compiled_fn.hits == hits:  # type: ignore[possibly-undefined]
+            raise RuntimeError("failed to use the compiled function")
+        if not isinstance(out, tuple):
+            out = (out,)
+        if loss_fn == torch.sum and len(out) != 1:
+            raise ValueError(
+                f"Model returns {len(out)} outputs, but default loss function "
+                "(torch.sum) can only handle a single output"
+            )
+        out_vars, _ = _flatten(out)
+        saved_outs = [
+            v.detach().clone(memory_format=torch.preserve_format) for v in out_vars
+        ]
+        loss = loss_fn(*out)
+        grads = torch.autograd.grad([loss], in_vars)
+        # TODO: I'm not sure if the clone here is necessary but it is safer
+        saved_grads = [
+            v.detach().clone(memory_format=torch.preserve_format) for v in grads
+        ]
+        return (saved_outs, saved_grads)
+
+    with torch.random.fork_rng(devices, _caller="torch.jit.verify"):
+        uncompiled_outs, uncompiled_grads = run_fwd_bwd(args, force_trace=True)
+        if not model.has_trace_for(*args):
+            raise AssertionError("Model should have trace for the given args")
+
+    if is_module:
+        model.load_state_dict(saved_state)  # type: ignore[possibly-undefined]
+    compiled_outs, compiled_grads = run_fwd_bwd(args, assert_compiled=True)
+
+    _verify_equal(uncompiled_outs, compiled_outs)
+    _verify_equal(uncompiled_grads, compiled_grads)
+
+
+def _verify_equal(xs, ys):
+    for x, y in zip(xs, ys):
+        if x.sub(y).abs().max() > 1e-6:
+            raise RuntimeError("JIT and real computation mismatch")
+
+
+def indent(s):
+    return "\n".join(["\t" + line for line in s.splitlines()])
+
+
+class TracingCheckError(Exception):
+    def __init__(self, graph_diff_error, tensor_compare_error, extra_msg=None):
+        self.message = "Tracing failed sanity checks!\n"
+        if extra_msg is not None:
+            self.message += extra_msg + "\n"
+        if graph_diff_error is not None:
+            self.message += "ERROR: Graphs differed across invocations!\n"
+            self.message += indent(graph_diff_error) + "\n"
+        if tensor_compare_error is not None:
+            self.message += (
+                "ERROR: Tensor-valued Constant nodes differed in value "
+                "across invocations. This often indicates that the tracer has"
+                " encountered untraceable code.\n"
+            )
+            self.message += indent(tensor_compare_error) + "\n"
+        super().__init__(self.message)
+
+
+# Check the traced module against a set of user-provided validation inputs
+@torch.no_grad()
+def _check_trace(
+    check_inputs,
+    func,
+    traced_func,
+    check_tolerance,
+    strict,
+    force_outplace,
+    is_trace_module,
+    _module_class,
+    example_inputs_is_kwarg=False,
+):
+    # Note: tracing is independent of optimizations, which consume the trace
+    for inputs in check_inputs:
+        if isinstance(inputs, torch.Tensor):
+            inputs = (inputs,)
+
+        if is_trace_module:
+            copied_dict = {}
+
+            for name, data in inputs.items():
+                copied_dict[name] = _clone_inputs(data)
+            check_mod = torch.jit.trace_module(
+                getattr(func, "__self__", func),
+                copied_dict,
+                check_trace=False,
+                strict=strict,
+                _force_outplace=force_outplace,
+                _module_class=_module_class,
+                _compilation_unit=torch._C.CompilationUnit(),
+                example_inputs_is_kwarg=example_inputs_is_kwarg,
+                _store_inputs=False,
+            )
+            check_mod_func = check_mod._c._get_method(traced_func.name)
+            inputs = inputs[traced_func.name]
+            if (
+                isinstance(inputs, (torch.Tensor))
+                or isinstance(inputs, dict)
+                and not example_inputs_is_kwarg
+            ):
+                inputs = (inputs,)
+        else:
+            if example_inputs_is_kwarg:
+                check_mod = torch.jit.trace(
+                    func,
+                    check_trace=False,
+                    strict=strict,
+                    _force_outplace=force_outplace,
+                    _module_class=_module_class,
+                    example_kwarg_inputs=_clone_inputs(inputs),
+                    _store_inputs=False,
+                )
+            else:
+                check_mod = torch.jit.trace(
+                    func,
+                    _clone_inputs(inputs),
+                    check_trace=False,
+                    strict=strict,
+                    _force_outplace=force_outplace,
+                    _module_class=_module_class,
+                    _store_inputs=False,
+                )
+            check_mod_func = check_mod
+
+        def graph_diagnostic_info():
+            mod_canonicalized = torch._C._jit_pass_canonicalize(traced_func.graph)
+            torch._C._jit_pass_inline(mod_canonicalized)
+            torch._C._jit_pass_erase_shape_information(mod_canonicalized)
+            mod_str = str(mod_canonicalized)
+            mod_str = re.sub(r"___torch_mangle_[0-9]+\.", "", mod_str)
+            check_canonicalized = torch._C._jit_pass_canonicalize(check_mod_func.graph)
+            torch._C._jit_pass_inline(check_canonicalized)
+            torch._C._jit_pass_erase_shape_information(check_canonicalized)
+            check_str = str(check_canonicalized)
+            check_str = re.sub(r"___torch_mangle_[0-9]+\.", "", check_str)
+
+            graph_diff_errors = None
+            if mod_str != check_str:
+                import difflib
+
+                graph_diff = difflib.ndiff(
+                    mod_str.splitlines(True), check_str.splitlines(True)
+                )
+                graph_diff_errors = "Graph diff:\n" + indent("".join(graph_diff)) + "\n"
+
+                for n_mod, n_check in zip(
+                    mod_canonicalized.nodes(), check_canonicalized.nodes()
+                ):
+                    if str(n_mod) != str(n_check):
+                        graph_diff_errors += "First diverging operator:\n"
+                        node_diff = difflib.ndiff(
+                            str(n_mod).splitlines(True), str(n_check).splitlines(True)
+                        )
+                        source_printout = (
+                            "Node diff:\n" + indent("".join(node_diff)) + "\n"
+                        )
+                        mod_stack = n_mod.sourceRange()
+                        if mod_stack:
+                            source_printout += (
+                                "Trace source location:\n" + indent(mod_stack) + "\n"
+                            )
+                        check_stack = n_check.sourceRange()
+                        if check_stack:
+                            source_printout += (
+                                "Check source location:\n" + indent(check_stack) + "\n"
+                            )
+                        graph_diff_errors += source_printout
+
+                        break  # For now, only print out the first pair of nodes that diverges
+
+            tensor_compare_errors = None
+            # Check Tensor-valued constant nodes
+            for n_mod, n_check in zip(
+                mod_canonicalized.nodes(), check_canonicalized.nodes()
+            ):
+                if n_mod.kind() != n_check.kind():
+                    break  # Graphs have already diverged
+
+                if n_mod.kind() == "prim::Constant" and not (
+                    n_mod.mustBeNone() or n_check.mustBeNone()
+                ):
+                    if not n_mod.hasAttribute("value"):
+                        continue
+                    if n_mod.kindOf("value") != "t" or n_check.kindOf("value") != "t":
+                        continue
+
+                    mod_tensor_val = n_mod.t("value")
+                    check_tensor_val = n_check.t("value")
+
+                    try:
+                        torch.testing.assert_close(
+                            mod_tensor_val, check_tensor_val, equal_nan=True
+                        )
+                    except (RuntimeError, AssertionError) as e:
+                        if tensor_compare_errors is None:
+                            tensor_compare_errors = ""
+                        tensor_compare_errors += "Node:\n" + indent(str(n_mod)) + "\n"
+                        compare_stack = n_mod.sourceRange()
+                        if compare_stack:
+                            tensor_compare_errors += (
+                                "Source Location:\n" + indent(compare_stack) + "\n"
+                            )
+                        tensor_compare_errors += "Comparison exception: " + indent(
+                            str(e)
+                        )
+
+                        break  # For now, only print the first diverging pair
+
+            return graph_diff_errors, tensor_compare_errors
+
+        def wrap_retval(x):
+            return x if isinstance(x, tuple) else (x,)
+
+        def run_mod_and_filter_tensor_outputs(mod, inputs, running_what):
+            try:
+                if isinstance(inputs, dict) and example_inputs_is_kwarg:
+                    outs = wrap_retval(mod(**inputs))
+                else:
+                    outs = wrap_retval(mod(*_clone_inputs(inputs)))
+                outs = [out for out in outs if isinstance(out, torch.Tensor)]
+                return outs
+            except Exception as e:
+                graph_diff_errors, tensor_compare_errors = graph_diagnostic_info()
+                msg = f"encountered an exception while running the {running_what} with test inputs.\nException:\n{indent(str(e))}"
+                raise TracingCheckError(
+                    graph_diff_errors,
+                    tensor_compare_errors,
+                    extra_msg=msg,
+                ) from e
+
+        has_warned = [False]
+
+        def maybe_warn_nondeterministic():
+            if has_warned[0]:
+                return
+            has_warned[0] = True
+            nondeterm_ops = [
+                op for op in traced_func.graph.nodes() if op.isNondeterministic()
+            ]
+            if len(nondeterm_ops) > 0:
+                nondeterministic_ops_warning = "Trace had nondeterministic nodes. "
+                nondeterministic_ops_warning += (
+                    "Did you forget call .eval() on your model? Nodes:\n"
+                )
+                nondeterministic_ops_warning += "\n".join(
+                    [indent(str(op)) for op in nondeterm_ops][:20]
+                )
+                nondeterministic_ops_warning += (
+                    "\nThis may cause errors in trace checking. To disable trace checking,"
+                    " pass check_trace=False to torch.jit.trace()"
+                )
+                warnings.warn(
+                    nondeterministic_ops_warning, category=TracerWarning, stacklevel=5
+                )
+
+        def compare_outputs(original, reference, match_what):
+            all_ok = True
+            for i, (orig, ref) in enumerate(zip(original, reference)):
+                try:
+                    if orig.is_quantized:
+                        orig = orig.dequantize()
+                    if ref.is_quantized:
+                        ref = ref.dequantize()
+                    if orig.is_mkldnn:
+                        orig = orig.to_dense()
+                    if ref.is_mkldnn:
+                        ref = ref.to_dense()
+                    if ref.is_complex() or orig.is_complex():
+                        torch.testing.assert_close(
+                            orig.to(torch.cdouble),
+                            ref.to(torch.cdouble),
+                            rtol=check_tolerance,
+                            atol=default_tolerances(orig, ref)[1],
+                            equal_nan=True,
+                        )
+                    else:
+                        if orig.is_mps or ref.is_mps:
+                            torch.testing.assert_close(
+                                orig.float(),
+                                ref.float(),
+                                rtol=check_tolerance,
+                                atol=default_tolerances(orig, ref)[1],
+                                equal_nan=True,
+                            )
+                        elif getattr(orig, "is_nested", None) or getattr(
+                            ref, "is_nested", None
+                        ):
+                            if getattr(orig, "is_nested", None) != getattr(
+                                ref, "is_nested", None
+                            ):
+                                raise AssertionError(
+                                    f"Nested tensor mismatch: orig.is_nested="
+                                    f"{getattr(orig, 'is_nested', None)}, "
+                                    f"ref.is_nested={getattr(ref, 'is_nested', None)}"
+                                )
+                            for t_orig, t_ref in zip(orig.unbind(), ref.unbind()):
+                                torch.testing.assert_close(
+                                    t_orig.double(),
+                                    t_ref.double(),
+                                    rtol=check_tolerance,
+                                    atol=default_tolerances(t_orig, t_ref)[1],
+                                    equal_nan=True,
+                                )
+                        else:
+                            torch.testing.assert_close(
+                                orig.double(),
+                                ref.double(),
+                                rtol=check_tolerance,
+                                atol=default_tolerances(orig, ref)[1],
+                                equal_nan=True,
+                            )
+                except AssertionError as e:
+                    maybe_warn_nondeterministic()
+                    warnings.warn(
+                        "Output nr "
+                        + str(i + 1)
+                        + ". of the traced function does not match "
+                        "the corresponding output of the "
+                        + match_what
+                        + ". Detailed error:\n"
+                        + str(e),
+                        category=TracerWarning,
+                        stacklevel=4,
+                    )
+                    all_ok = False
+
+            return all_ok
+
+        traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, "trace")
+        fn_outs = run_mod_and_filter_tensor_outputs(func, inputs, "Python function")
+        if compare_outputs(traced_outs, fn_outs, "Python function"):
+            check_outs = run_mod_and_filter_tensor_outputs(
+                check_mod_func, inputs, "repeated trace"
+            )
+            compare_outputs(traced_outs, check_outs, "repeated trace")
+
+        diag_info = graph_diagnostic_info()
+        if any(info is not None for info in diag_info):
+            raise TracingCheckError(*diag_info)
+
+
+class TracerWarning(Warning):
+    @staticmethod
+    def ignore_lib_warnings():
+        # We ignore warnings from all submodules excluding the JIT, because we need them e.g. for _check_trace
+        warnings.filterwarnings(
+            "ignore", category=TracerWarning, module="torch.(?!jit)"
+        )
+        warnings.filterwarnings("ignore", "torch::jit::fuser::cuda")
+
+
+# We ignore the tracer warnings coming form inside the library, because all our shape
+# checks in nn will trigger them.
+TracerWarning.ignore_lib_warnings()
+torch._C._tracer_warn_use_python()
+
+
+def make_tuple(example_inputs):
+    if isinstance(example_inputs, (torch.Tensor, dict)):
+        return (example_inputs,)
+    # done primarily so that weird iterables fail here and not pybind11 code
+    if not isinstance(example_inputs, tuple):
+        return tuple(example_inputs)
+    return example_inputs
+
+
+def make_module(mod, _module_class, _compilation_unit):
+    if isinstance(mod, ScriptModule):
+        return mod
+    elif torch._jit_internal.module_has_exports(mod):
+        infer_methods_stubs_fn = torch.jit._recursive.make_stubs_from_exported_methods
+        return torch.jit._recursive.create_script_module(
+            mod, infer_methods_stubs_fn, share_types=False, is_tracing=True
+        )
+    else:
+        if _module_class is None:
+            _module_class = TopLevelTracedModule
+        return _module_class(mod, _compilation_unit=_compilation_unit)
+
+
+def wrap_check_inputs(check_inputs):
+    if check_inputs is None:
+        return None
+
+    return [{"forward": c} for c in check_inputs]
+
+
+def analyze_ts_result_with_export_result(export, trace):
+    import torch.utils._pytree as pytree
+
+    flat_export = pytree.tree_leaves(export)
+    flat_trace = pytree.tree_leaves(trace)
+
+    for orig, loaded in zip(flat_export, flat_trace):
+        if orig.layout != loaded.layout:
+            return False
+        # mkldnn is not supported for torch.allclose
+        if orig.layout == torch._mkldnn:  # type: ignore[attr-defined]
+            return True
+        if type(orig) is not type(loaded):
+            return False
+
+        if isinstance(orig, torch._subclasses.FakeTensor):
+            # Skip for FakeTensor.
+            return True
+        elif isinstance(orig, torch.Tensor):
+            if orig.dtype != loaded.dtype:
+                return False
+            if not torch.allclose(orig, loaded):
+                return False
+        else:
+            if orig != loaded:
+                return False
+    return True
+
+
+def _trace_impl(
+    func,
+    example_inputs=None,
+    optimize=None,
+    check_trace=True,
+    check_inputs=None,
+    check_tolerance=1e-5,
+    strict=True,
+    _force_outplace=False,
+    _module_class=None,
+    _compilation_unit=_python_cu,
+    example_kwarg_inputs=None,
+    _store_inputs=True,
+):
+    if isinstance(func, torch.jit.ScriptModule):
+        # it is hard to trace it because the forward method on ScriptModule is already defined, so it
+        # would result in an error.
+        warnings.warn(
+            "The input to trace is already a ScriptModule, tracing it is a no-op. Returning the object as is.",
+            stacklevel=2,
+        )
+        return func
+
+    if isinstance(func, torch.nn.Module):
+        if example_inputs is None:
+            if isinstance(example_kwarg_inputs, dict):
+                example_inputs = example_kwarg_inputs
+            else:
+                raise RuntimeError("example_kwarg_inputs should be a dict")
+        return trace_module(
+            func,
+            {"forward": example_inputs},
+            None,
+            check_trace,
+            wrap_check_inputs(check_inputs),
+            check_tolerance,
+            strict,
+            _force_outplace,
+            _module_class,
+            example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),
+            _store_inputs=_store_inputs,
+        )
+    if (
+        hasattr(func, "__self__")
+        and isinstance(func.__self__, torch.nn.Module)
+        and func.__name__ == "forward"
+    ):
+        if example_inputs is None:
+            if isinstance(example_kwarg_inputs, dict):
+                example_inputs = example_kwarg_inputs
+            else:
+                raise RuntimeError("example_kwarg_inputs should be a dict")
+        return trace_module(
+            func.__self__,
+            {"forward": example_inputs},
+            None,
+            check_trace,
+            wrap_check_inputs(check_inputs),
+            check_tolerance,
+            strict,
+            _force_outplace,
+            _module_class,
+            example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),
+            _store_inputs=_store_inputs,
+        )
+
+    # Special case for common case of passing a single Tensor
+    if (
+        isinstance(example_inputs, (torch.Tensor, dict))
+        and example_kwarg_inputs is None
+    ):
+        example_inputs = (example_inputs,)
+    # done primarily so that weird iterables fail here and not pybind11 code
+    elif example_kwarg_inputs is None and not isinstance(example_inputs, tuple):
+        # pyrefly: ignore [bad-argument-type]
+        example_inputs = tuple(example_inputs)
+
+    var_lookup_fn = _create_interpreter_name_lookup_fn(0)
+
+    if hasattr(func, "__self__") and isinstance(func.__self__, torch.nn.Module):
+        raise AttributeError(
+            "trace doesn't support compiling individual module's functions.\n"
+            "Please use trace_module"
+        )
+
+    name = _qualified_name(func)
+    if isinstance(example_kwarg_inputs, dict):
+        example_inputs = example_kwarg_inputs
+        traced = torch._C._create_function_from_trace_with_dict(
+            name,
+            func,
+            example_kwarg_inputs,
+            var_lookup_fn,
+            strict,
+            _force_outplace,
+            get_callable_argument_names(func),
+        )
+    else:
+        traced = torch._C._create_function_from_trace(
+            name,
+            func,
+            # pyrefly: ignore [bad-argument-type]
+            example_inputs,
+            var_lookup_fn,
+            strict,
+            _force_outplace,
+            get_callable_argument_names(func),
+        )
+
+    # Check the trace against new traces created from user-specified inputs
+    if check_trace:
+        if check_inputs is not None:
+            _check_trace(
+                check_inputs,
+                func,
+                traced,
+                check_tolerance,
+                strict,
+                _force_outplace,
+                False,
+                _module_class,
+                example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),
+            )
+        else:
+            _check_trace(
+                [example_inputs],
+                func,
+                traced,
+                check_tolerance,
+                strict,
+                _force_outplace,
+                False,
+                _module_class,
+                example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),
+            )
+
+    # Allow torch.compile() to inline
+    traced._torchdynamo_inline = func  # type: ignore[attr-defined]
+    return traced
+
+
+class _ExportType(str, Enum):
+    DIRECT_EXPORT = "DIRECT_EXPORT"
+    TRACE_AND_EXPORT = "TRACE_AND_EXPORT"
+    SOURCE_TO_SOURCE = "SOURCE_TO_SOURCE"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+class _ExportOutcome(str, Enum):
+    SUCCESS = "SUCCESS"
+    FAILED_TO_EXPORT = "FAILED_TO_EXPORT"
+    FAILED_TO_RUN = "FAILED_TO_RUN"
+    ACCURACY_ERROR = "ACCURACY_ERROR"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+def trace(
+    func,
+    example_inputs=None,
+    optimize=None,
+    check_trace=True,
+    check_inputs=None,
+    check_tolerance=1e-5,
+    strict=True,
+    _force_outplace=False,
+    _module_class=None,
+    _compilation_unit=_python_cu,
+    example_kwarg_inputs=None,
+    _store_inputs=True,
+):
+    r"""
+    Trace a function and return an executable  or :class:`ScriptFunction` that will be optimized using just-in-time compilation.
+
+    Tracing is ideal for code that operates only on
+    ``Tensor``\\s and lists, dictionaries, and
+    tuples of ``Tensor``\\s.
+
+    Using `torch.jit.trace` and `torch.jit.trace_module`, you can turn an
+    existing module or Python function into a TorchScript
+    :class:`ScriptFunction` or :class:`ScriptModule`. You must provide example
+    inputs, and we run the function, recording the operations performed on all
+    the tensors.
+
+    * The resulting recording of a standalone function produces `ScriptFunction`.
+    * The resulting recording of `nn.Module.forward` or `nn.Module` produces
+      `ScriptModule`.
+
+    This module also contains any parameters that the original
+    module had as well.
+
+    Warning:
+        Tracing only correctly records functions and modules which are not data
+        dependent (e.g., do not have conditionals on data in tensors) and do not have
+        any untracked external dependencies (e.g., perform input/output or
+        access global variables). Tracing only records operations done when the given
+        function is run on the given tensors. Therefore, the returned
+        `ScriptModule` will always run the same traced graph on any input. This
+        has some important implications when your module is expected to run
+        different sets of operations, depending on the input and/or the module
+        state. For example,
+
+        * Tracing will not record any control-flow like if-statements or loops.
+          When this control-flow is constant across your module, this is fine
+          and it often inlines the control-flow decisions. But sometimes the
+          control-flow is actually part of the model itself. For instance, a
+          recurrent network is a loop over the (possibly dynamic) length of an
+          input sequence.
+        * In the returned :class:`ScriptModule`, operations that have different
+          behaviors in ``training`` and ``eval`` modes will always behave as if
+          it is in the mode it was in during tracing, no matter which mode the
+          `ScriptModule` is in.
+
+        In cases like these, tracing would not be appropriate and
+        :func:`scripting <torch.jit.script>` is a better choice. If you trace
+        such models, you may silently get incorrect results on subsequent
+        invocations of the model. The tracer will try to emit warnings when
+        doing something that may cause an incorrect trace to be produced.
+
+    Args:
+        func (callable or torch.nn.Module):  A Python function or `torch.nn.Module`
+            that will be run with `example_inputs`. `func` arguments and return
+            values  must be tensors or (possibly nested) tuples that contain
+            tensors. When a module is passed `torch.jit.trace`, only the
+            ``forward`` method is run and traced (see :func:`torch.jit.trace
+            <torch.jit.trace_module>` for details).
+
+    Keyword arguments:
+        example_inputs (tuple or torch.Tensor or None, optional): A tuple of example
+            inputs that will be passed to the function while tracing.
+            Default: ``None``. Either this argument or ``example_kwarg_inputs``
+            should be specified. The resulting trace can be run with inputs of
+            different types and shapes assuming the traced operations support those
+            types and shapes. `example_inputs` may also be a single Tensor in which
+            case it is automatically wrapped in a tuple. When the value is None,
+            ``example_kwarg_inputs`` should be specified.
+
+        check_trace (``bool``, optional): Check if the same inputs run through
+            traced code produce the same outputs. Default: ``True``. You might want
+            to disable this if, for example, your network contains non-
+            deterministic ops or if you are sure that the network is correct despite
+            a checker failure.
+
+        check_inputs (list of tuples, optional): A list of tuples of input
+            arguments that should be used to check the trace against what is
+            expected. Each tuple is equivalent to a set of input arguments that
+            would be specified in ``example_inputs``. For best results, pass in
+            a set of checking inputs representative of the space of shapes and
+            types of inputs you expect the network to see.  If not specified,
+            the original ``example_inputs`` are used for checking
+        check_tolerance (float, optional): Floating-point comparison tolerance
+            to use in the checker procedure.  This can be used to relax the
+            checker strictness in the event that results diverge numerically
+            for a known reason, such as operator fusion.
+        strict (``bool``, optional): run the tracer in a strict mode or not
+            (default: ``True``). Only turn this off when you want the tracer to
+            record your mutable container types (currently ``list``/``dict``)
+            and you are sure that the container you are using in your
+            problem is a ``constant`` structure and does not get used as
+            control flow (if, for) conditions.
+        example_kwarg_inputs (dict, optional): This parameter is a pack of keyword
+            arguments of example inputs that will be passed to the function while
+            tracing. Default: ``None``. Either this argument or ``example_inputs``
+            should be specified. The dict will be unpacking by the arguments name
+            of the traced function. If the keys of the dict don't not match with
+            the traced function's arguments name, a runtime exception will be raised.
+
+    Returns:
+        If `func` is `nn.Module` or ``forward`` of `nn.Module`, `trace` returns
+        a :class:`ScriptModule` object with a single ``forward`` method
+        containing the traced code.  The returned `ScriptModule` will
+        have the same set of sub-modules and parameters as the original
+        ``nn.Module``.  If ``func`` is a standalone function, ``trace``
+        returns `ScriptFunction`.
+
+    Example (tracing a function):
+
+    .. testcode::
+
+        import torch
+
+        def foo(x, y):
+            return 2 * x + y
+
+        # Run `foo` with the provided inputs and record the tensor operations
+        traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))
+
+        # `traced_foo` can now be run with the TorchScript interpreter or saved
+        # and loaded in a Python-free environment
+
+    Example (tracing an existing module)::
+
+        import torch
+        import torch.nn as nn
+
+
+        class Net(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = nn.Conv2d(1, 1, 3)
+
+            def forward(self, x):
+                return self.conv(x)
+
+
+        n = Net()
+        example_weight = torch.rand(1, 1, 3, 3)
+        example_forward_input = torch.rand(1, 1, 3, 3)
+
+        # Trace a specific method and construct `ScriptModule` with
+        # a single `forward` method
+        module = torch.jit.trace(n.forward, example_forward_input)
+
+        # Trace a module (implicitly traces `forward`) and construct a
+        # `ScriptModule` with a single `forward` method
+        module = torch.jit.trace(n, example_forward_input)
+
+    """
+    if sys.version_info >= (3, 14):
+        warnings.warn(
+            "`torch.jit.trace` is not supported in Python 3.14+ and may break. "
+            "Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    else:
+        warnings.warn(
+            "`torch.jit.trace` is deprecated. Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    if not _enabled:
+        return func
+    if optimize is not None:
+        warnings.warn(
+            "`optimize` is deprecated and has no effect. "
+            "Use `with torch.jit.optimized_execution()` instead",
+            FutureWarning,
+            stacklevel=2,
+        )
+
+    from torch._utils_internal import log_torchscript_usage
+
+    traced_func = _trace_impl(
+        func,
+        example_inputs,
+        optimize,
+        check_trace,
+        check_inputs,
+        check_tolerance,
+        strict,
+        _force_outplace,
+        _module_class,
+        _compilation_unit,
+        example_kwarg_inputs,
+        _store_inputs,
+    )
+    log_torchscript_usage("trace", model_id=_get_model_id(traced_func))
+    return traced_func
+
+
+_trace_module_map: dict[Any, Any] | None = None
+
+
+def trace_module(
+    mod,
+    inputs,
+    optimize=None,
+    check_trace=True,
+    check_inputs=None,
+    check_tolerance=1e-5,
+    strict=True,
+    _force_outplace=False,
+    _module_class=None,
+    _compilation_unit=_python_cu,
+    example_inputs_is_kwarg=False,
+    _store_inputs=True,
+):
+    """
+    Trace a module and return an executable :class:`ScriptModule` that will be optimized using just-in-time compilation.
+
+    When a module is passed to :func:`torch.jit.trace <torch.jit.trace>`, only
+    the ``forward`` method is run and traced. With ``trace_module``, you can specify a dictionary of
+    method names to example inputs to trace (see the ``inputs``) argument below.
+
+    See :func:`torch.jit.trace <torch.jit.trace>` for more information on tracing.
+
+    Args:
+        mod (torch.nn.Module):  A ``torch.nn.Module`` containing methods whose names are
+                                specified in ``inputs``. The given methods will be compiled
+                                as a part of a single `ScriptModule`.
+        inputs (dict):  A dict containing sample inputs indexed by method names in ``mod``.
+                                The inputs will be passed to methods whose names correspond to inputs'
+                                keys while tracing.
+                                ``{ 'forward' : example_forward_input, 'method2': example_method2_input}``
+    Keyword arguments:
+        check_trace (``bool``, optional): Check if the same inputs run through
+                                      traced code produce the same outputs. Default: ``True``. You might want
+                                      to disable this if, for example, your network contains non-
+                                      deterministic ops or if you are sure that the network is correct despite
+                                      a checker failure.
+
+        check_inputs (list of dicts, optional): A list of dicts of input arguments that should be used
+                                                 to check the trace against what is expected. Each tuple
+                                                 is equivalent to a set of input arguments that would
+                                                 be specified in ``inputs``. For best results, pass in a
+                                                 set of checking inputs representative of the space of
+                                                 shapes and types of inputs you expect the network to see.
+                                                 If not specified, the original ``inputs`` are used for checking
+        check_tolerance (float, optional): Floating-point comparison tolerance to use in the checker procedure.
+                                           This can be used to relax the checker strictness in the event that
+                                           results diverge numerically for a known reason, such as operator fusion.
+        example_inputs_is_kwarg (``bool``, optional): This parameter indicate whether the example inputs is a pack
+                                           pack of keyword arguments. Default: ``False``.
+
+    Returns:
+        A :class:`ScriptModule` object with a single ``forward`` method containing the traced code.
+        When ``func`` is a ``torch.nn.Module``, the returned :class:`ScriptModule` will have the same set of
+        sub-modules and parameters as ``func``.
+
+    Example (tracing a module with multiple methods)::
+
+        import torch
+        import torch.nn as nn
+
+
+        class Net(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = nn.Conv2d(1, 1, 3)
+
+            def forward(self, x):
+                return self.conv(x)
+
+            def weighted_kernel_sum(self, weight):
+                return weight * self.conv.weight
+
+
+        n = Net()
+        example_weight = torch.rand(1, 1, 3, 3)
+        example_forward_input = torch.rand(1, 1, 3, 3)
+
+        # Trace a specific method and construct `ScriptModule` with
+        # a single `forward` method
+        module = torch.jit.trace(n.forward, example_forward_input)
+
+        # Trace a module (implicitly traces `forward`) and construct a
+        # `ScriptModule` with a single `forward` method
+        module = torch.jit.trace(n, example_forward_input)
+
+        # Trace specific methods on a module (specified in `inputs`), constructs
+        # a `ScriptModule` with `forward` and `weighted_kernel_sum` methods
+        inputs = {
+            "forward": example_forward_input,
+            "weighted_kernel_sum": example_weight,
+        }
+        module = torch.jit.trace_module(n, inputs)
+
+    """
+    if sys.version_info >= (3, 14):
+        warnings.warn(
+            "`torch.jit.trace_method` is not supported in Python 3.14+ and may break. "
+            "Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    else:
+        warnings.warn(
+            "`torch.jit.trace_method` is deprecated. Please switch to `torch.compile` or `torch.export`.",
+            DeprecationWarning,
+        )
+    if not _enabled:
+        return mod
+    if optimize is not None:
+        warnings.warn(
+            "`optimize` is deprecated and has no effect. "
+            "Use `with torch.jit.optimized_execution()` instead",
+            FutureWarning,
+            stacklevel=2,
+        )
+
+    var_lookup_fn = _create_interpreter_name_lookup_fn(0)
+
+    if not isinstance(mod, torch.nn.Module):
+        raise AttributeError("expected torch.nn.Module as the first argument")
+
+    if not isinstance(inputs, dict):
+        raise AttributeError("expected a dictionary of (method_name, input) pairs")
+
+    old_module_map = torch.jit._trace._trace_module_map
+    try:
+        trace_module_map: dict[Any, Any] = {}
+
+        def register_submods(mod, prefix):
+            for name, child in mod.named_children():
+                submod_qualname = prefix + "." + name
+                trace_module_map[child] = submod_qualname
+                register_submods(child, submod_qualname)
+
+        trace_module_map["__module"] = mod
+        torch.jit._trace._trace_module_map = trace_module_map
+        register_submods(mod, "__module")
+
+        module = make_module(mod, _module_class, _compilation_unit)
+
+        for method_name, example_inputs in inputs.items():
+            if method_name == "forward":
+                # "forward" is a special case because we need to trace
+                # `Module.__call__`, which sets up some extra tracing, but uses
+                # argument names of the real `Module.forward` method.
+                func = mod
+                forward_method = getattr(mod, method_name)
+                argument_names = get_callable_argument_names(forward_method)
+            else:
+                func = getattr(mod, method_name)
+                argument_names = get_callable_argument_names(func)
+
+            if isinstance(example_inputs, dict) and example_inputs_is_kwarg:
+                # Raise exception when the user provided key names are not aligned with forward() method's arguments' name/
+                for key in example_inputs:
+                    if key not in argument_names:
+                        valid_arguments = "[" + ",".join(argument_names) + "]"
+                        raise NameError(
+                            f"""'{key}' is not in forward() method's arguments,
+                         valid arguments name are {valid_arguments}"""
+                        )
+                module._c._create_method_from_trace_with_dict(
+                    method_name,
+                    func,
+                    example_inputs,
+                    var_lookup_fn,
+                    strict,
+                    _force_outplace,
+                    argument_names,
+                    _store_inputs,
+                )
+            else:
+                example_inputs = make_tuple(example_inputs)
+                module._c._create_method_from_trace(
+                    method_name,
+                    func,
+                    example_inputs,
+                    var_lookup_fn,
+                    strict,
+                    _force_outplace,
+                    argument_names,
+                    _store_inputs,
+                )
+
+            check_trace_method = module._c._get_method(method_name)
+
+            # Check the trace against new traces created from user-specified inputs
+            if check_trace:
+                if check_inputs is not None:
+                    _check_trace(
+                        check_inputs,
+                        func,
+                        check_trace_method,
+                        check_tolerance,
+                        strict,
+                        _force_outplace,
+                        True,
+                        _module_class,
+                        example_inputs_is_kwarg=example_inputs_is_kwarg,
+                    )
+                else:
+                    _check_trace(
+                        [inputs],
+                        func,
+                        check_trace_method,
+                        check_tolerance,
+                        strict,
+                        _force_outplace,
+                        True,
+                        _module_class,
+                        example_inputs_is_kwarg=example_inputs_is_kwarg,
+                    )
+    finally:
+        torch.jit._trace._trace_module_map = old_module_map
+
+    return module
+
+
+def is_tracing():
+    """Return a boolean value.
+
+    Returns ``True`` in tracing (if a function is called during the
+    tracing of code with ``torch.jit.trace``) and ``False`` otherwise.
+    """
+    if is_scripting():
+        return False
+    return torch._C._is_tracing()
+
+
+class TracedModule(ScriptModule):
+    _disable_script_meta = True
+
+    def __init__(self, orig, id_set=None, _compilation_unit=None):
+        # XXX: orig can be a nn.Module or a function!
+        super().__init__()
+        if not isinstance(orig, torch.nn.Module):
+            raise AssertionError(f"Expected nn.Module, got {type(orig)}")
+
+        # Copy a subset of `orig` to a temporary nn.Module.
+        # This is a way to customize what will actually get compiled by create_script_module
+        id_set = set()
+
+        # This allows us to preserve the original module's qualified name by defining a new
+        # type with the attribute _jit_override_qualname. In torch._jit_internal._qualified_name
+        # we have a special case that will look up this attribute to override whatever qualname
+        # we would get from the python type system
+        class QualnameWrapper(torch.nn.Module):
+            pass
+
+        QualnameWrapper._jit_override_qualname = torch._jit_internal._qualified_name(  # type: ignore[attr-defined]
+            type(orig)
+        )
+
+        tmp_module = QualnameWrapper()
+
+        def check_unique(param):
+            if param in id_set:
+                raise ValueError(
+                    "TracedModules don't support parameter sharing between modules"
+                )
+            id_set.add(param)
+
+        tmp_module.training = orig.training
+
+        for name, param in orig._parameters.items():
+            if param is not None:
+                tmp_module._parameters[name] = param
+                check_unique(param)
+        for name, buf in orig._buffers.items():
+            if buf is not None:
+                tmp_module._buffers[name] = buf
+                check_unique(buf)
+        for name, val in orig.__dict__.items():
+            if (
+                torch._C._jit_is_script_object(val)
+                and name not in orig._parameters
+                and name not in orig._buffers
+            ):
+                setattr(tmp_module, name, val)
+
+        if orig._backward_hooks:
+            raise ValueError(
+                "Modules that have backward hooks assigned can't be compiled: "
+                + str(orig)
+            )
+
+        for name, submodule in orig._modules.items():
+            if submodule is None:
+                continue
+            tmp_module._modules[name] = make_module(
+                submodule, TracedModule, _compilation_unit=None
+            )
+
+        script_module = torch.jit._recursive.create_script_module(
+            tmp_module, lambda module: (), share_types=False, is_tracing=True
+        )
+
+        self.__dict__["_name"] = type(orig).__name__
+        self.__dict__["_actual_script_module"] = script_module
+        for name in ("_parameters", "_buffers", "_modules", "training"):
+            delattr(self, name)
+
+    def forward(self, *args, **kwargs):
+        raise RuntimeError("Trace submodules cannot be called.")
+
+    def __getattr__(self, attr):
+        if "_actual_script_module" not in self.__dict__:
+            return super().__getattr__(attr)
+        return getattr(self._actual_script_module, attr)
+
+    def __setattr__(self, attr, value):
+        if "_actual_script_module" not in self.__dict__:
+            return super().__setattr__(attr, value)
+        setattr(self._actual_script_module, attr, value)
+
+    def _get_name(self):
+        return self._name
+
+    def extra_repr(self):
+        return f"original_name={self._name}"
+
+
+class TopLevelTracedModule(TracedModule):
+    forward: Callable[..., Any] = _CachedForward()  # type: ignore[assignment]
+
+    def _reconstruct(self, cpp_module):
+        """
+        Re-construct an instance of TopLevelTracedModule using an instance of a C++ module.
+
+        Args:
+            cpp_module: The C++ module that this TopLevelTracedModule will be rebuilt around.
+        """
+        self.__dict__["_actual_script_module"]._reconstruct(cpp_module)
+
+
+def _script_if_tracing(fn: Callable[P, R]) -> Callable[P, R]:
+    @functools.wraps(fn)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        if not is_tracing():
+            # Not tracing, don't do anything
+            return fn(*args, **kwargs)
+
+        compiled_fn: Callable[P, R] = script(wrapper.__original_fn)  # type: ignore[attr-defined]
+        return compiled_fn(*args, **kwargs)
+
+    wrapper.__original_fn = fn  # type: ignore[attr-defined]
+    wrapper.__script_if_tracing_wrapper = True  # type: ignore[attr-defined]
+
+    return wrapper
+
+
+def _get_trace_graph(
+    f,
+    args=(),
+    kwargs=None,
+    strict=True,
+    _force_outplace=False,
+    return_inputs=False,
+    _return_inputs_states=False,
+):
+    """Return a tuple on tracing a function or model.
+
+    .. warning::
+        This function is internal-only and should only be used by the ONNX
+        exporter. If you are trying to get a graph through tracing, please go
+        through the public API instead::
+
+            trace = torch.jit.trace(nn.LSTMCell(), (input, hidden))
+            trace_graph = trace.graph
+
+    Trace a function or model, returning a tuple consisting of the both the
+    *trace* of an execution, as well as the original return value. If return_inputs,
+    also returns the trace inputs as part of the tuple
+
+    Tracing is guaranteed not to change the semantics of the function/module
+    that is traced.
+
+    Args:
+        f (torch.nn.Module or function): the function or module
+            to be traced.
+        args (tuple or Tensor): the positional arguments to pass to the
+            function/module to be traced.  A non-tuple is assumed to
+            be a single positional argument to be passed to the model.
+        kwargs (dict): the keyword arguments to pass to the function/module
+            to be traced.
+
+    Example (trace a cell):
+
+    .. testcode::
+
+        trace = torch.jit.trace(nn.LSTMCell(), (input, hidden))
+    """
+    if kwargs is None:
+        kwargs = {}
+    if not isinstance(args, tuple):
+        args = (args,)
+    outs = ONNXTracedModule(
+        f, strict, _force_outplace, return_inputs, _return_inputs_states
+    )(*args, **kwargs)
+    return outs
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/annotations.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/annotations.py
new file mode 100644
index 0000000000000000000000000000000000000000..c09de1971a4c423d916585edf44f4be57efd67bc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/annotations.py
@@ -0,0 +1,562 @@
+# mypy: allow-untyped-defs
+import ast
+import builtins
+import dis
+import enum
+import inspect
+import re
+import typing
+import warnings
+from textwrap import dedent
+
+import torch
+from torch._C import (
+    _GeneratorType,
+    AnyType,
+    AwaitType,
+    BoolType,
+    ComplexType,
+    DeviceObjType,
+    DictType,
+    EnumType,
+    FloatType,
+    FutureType,
+    InterfaceType,
+    IntType,
+    ListType,
+    NoneType,
+    NumberType,
+    OptionalType,
+    StreamObjType,
+    StringType,
+    TensorType,
+    TupleType,
+    UnionType,
+)
+from torch._jit_internal import (
+    _Await,
+    _qualified_name,
+    Any,
+    BroadcastingList1,
+    BroadcastingList2,  # pyrefly: ignore [missing-module-attribute]
+    BroadcastingList3,  # pyrefly: ignore [missing-module-attribute]
+    Dict,
+    Future,
+    is_await,
+    is_dict,
+    is_future,
+    is_ignored_fn,
+    is_list,
+    is_optional,
+    is_tuple,
+    is_union,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+from torch._sources import get_source_lines_and_file
+
+from ._state import _get_script_class
+
+
+if torch.distributed.rpc.is_available():
+    from torch._C import RRefType
+    from torch._jit_internal import is_rref, RRef
+
+from torch._ops import OpOverloadPacket
+
+
+class Module:
+    def __init__(self, name, members) -> None:
+        self.name = name
+        self.members = members
+
+    def __getattr__(self, name):
+        try:
+            return self.members[name]
+        except KeyError:
+            raise RuntimeError(
+                f"Module {self.name} has no member called {name}"
+            ) from None
+
+
+class EvalEnv:
+    env = {
+        "torch": Module("torch", {"Tensor": torch.Tensor}),
+        "Tensor": torch.Tensor,
+        "typing": Module("typing", {"Tuple": Tuple}),
+        "Tuple": Tuple,
+        "List": List,
+        "Dict": Dict,
+        "Optional": Optional,
+        "Union": Union,
+        "Future": Future,
+        "Await": _Await,
+    }
+
+    def __init__(self, rcb) -> None:
+        self.rcb = rcb
+        if torch.distributed.rpc.is_available():
+            # pyrefly: ignore [bad-typed-dict-key, unsupported-operation]
+            self.env["RRef"] = RRef
+
+    def __getitem__(self, name):
+        if name in self.env:
+            return self.env[name]
+        if self.rcb is not None:
+            return self.rcb(name)
+        return getattr(builtins, name, None)
+
+
+def get_signature(fn, rcb, loc, is_method):
+    if isinstance(fn, OpOverloadPacket):
+        signature = try_real_annotations(fn.op, loc)
+    else:
+        signature = try_real_annotations(fn, loc)
+    if signature is not None and is_method:
+        # If this is a method, then the signature will include a type for
+        # `self`, but type comments do not contain a `self`. So strip it
+        # away here so everything is consistent (`inspect.ismethod` does
+        # not work here since `fn` is unbound at this point)
+        param_types, return_type = signature
+        param_types = param_types[1:]
+        signature = (param_types, return_type)
+
+    if signature is None:
+        type_line, source = None, None
+        try:
+            source = dedent("".join(get_source_lines_and_file(fn)[0]))
+            type_line = get_type_line(source)
+        except TypeError:
+            pass
+        # This might happen both because we failed to get the source of fn, or
+        # because it didn't have any annotations.
+        if type_line is not None:
+            signature = parse_type_line(type_line, rcb, loc)
+
+    return signature
+
+
+def is_function_or_method(the_callable):
+    # A stricter version of `inspect.isroutine` that does not pass for built-in
+    # functions
+    return inspect.isfunction(the_callable) or inspect.ismethod(the_callable)
+
+
+def is_vararg(the_callable):
+    if not is_function_or_method(the_callable) and callable(the_callable):  # noqa: B004
+        # If `the_callable` is a class, de-sugar the call so we can still get
+        # the signature
+        the_callable = the_callable.__call__
+
+    if is_function_or_method(the_callable):
+        return inspect.getfullargspec(the_callable).varargs is not None
+    else:
+        return False
+
+
+def get_param_names(fn, n_args):
+    if isinstance(fn, OpOverloadPacket):
+        fn = fn.op
+
+    if (
+        not is_function_or_method(fn)
+        and callable(fn)
+        and is_function_or_method(fn.__call__)
+    ):  # noqa: B004
+        # De-sugar calls to classes
+        fn = fn.__call__
+
+    if is_function_or_method(fn):
+        if is_ignored_fn(fn):
+            fn = inspect.unwrap(fn)
+        return inspect.getfullargspec(fn).args
+    else:
+        # The `fn` was not a method or function (maybe a class with a __call__
+        # method, so use a default param name list)
+        return [str(i) for i in range(n_args)]
+
+
+def check_fn(fn, loc) -> None:
+    # Make sure the function definition is not a class instantiation
+    try:
+        source = dedent("".join(get_source_lines_and_file(fn)[0]))
+    except (OSError, TypeError):
+        return
+    if source is None:
+        return
+
+    py_ast = ast.parse(source)
+    if len(py_ast.body) == 1 and isinstance(py_ast.body[0], ast.ClassDef):
+        raise torch.jit.frontend.FrontendError(
+            loc,
+            f"Cannot instantiate class '{py_ast.body[0].name}' in a script function",
+        )
+    if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef):
+        raise torch.jit.frontend.FrontendError(
+            loc, "Expected a single top-level function"
+        )
+
+
+def _eval_no_call(stmt, glob, loc):
+    """Evaluate statement as long as it does not contain any method/function calls."""
+    bytecode = compile(stmt, "", mode="eval")
+    for insn in dis.get_instructions(bytecode):
+        if "CALL" in insn.opname:
+            raise RuntimeError(
+                f"Type annotation should not contain calls, but '{stmt}' does"
+            )
+    return eval(bytecode, glob, loc)  # type: ignore[arg-type] # noqa: P204
+
+
+def parse_type_line(type_line, rcb, loc):
+    """Parse a type annotation specified as a comment.
+
+    Example inputs:
+        # type: (Tensor, torch.Tensor) -> Tuple[Tensor]
+        # type: (Tensor, Tuple[Tensor, Tensor]) -> Tensor
+    """
+    arg_ann_str, ret_ann_str = split_type_line(type_line)
+
+    try:
+        arg_ann = _eval_no_call(arg_ann_str, {}, EvalEnv(rcb))
+    except (NameError, SyntaxError) as e:
+        raise RuntimeError(
+            "Failed to parse the argument list of a type annotation"
+        ) from e
+
+    if not isinstance(arg_ann, tuple):
+        arg_ann = (arg_ann,)
+
+    try:
+        ret_ann = _eval_no_call(ret_ann_str, {}, EvalEnv(rcb))
+    except (NameError, SyntaxError) as e:
+        raise RuntimeError(
+            "Failed to parse the return type of a type annotation"
+        ) from e
+
+    arg_types = [ann_to_type(ann, loc) for ann in arg_ann]
+    return arg_types, ann_to_type(ret_ann, loc)
+
+
+def get_type_line(source):
+    """Try to find the line containing a comment with the type annotation."""
+    type_comment = "# type:"
+
+    lines = source.split("\n")
+    lines = list(enumerate(lines))
+    type_lines = list(filter(lambda line: type_comment in line[1], lines))
+    # `type: ignore` comments may be needed in JIT'ed functions for mypy, due
+    # to the hack in torch/_VF.py.
+
+    # An ignore type comment can be of following format:
+    #   1) type: ignore
+    #   2) type: ignore[rule-code]
+    # This ignore statement must be at the end of the line
+
+    # adding an extra backslash before the space, to avoid triggering
+    # one of the checks in .github/workflows/lint.yml
+    type_pattern = re.compile("# type:\\ ignore(\\[[a-zA-Z-]+\\])?$")
+    type_lines = list(filter(lambda line: not type_pattern.search(line[1]), type_lines))
+
+    if len(type_lines) == 0:
+        # Catch common typo patterns like extra spaces, typo in 'ignore', etc.
+        wrong_type_pattern = re.compile("#[\t ]*type[\t ]*(?!: ignore(\\[.*\\])?$):")
+        wrong_type_lines = list(
+            filter(lambda line: wrong_type_pattern.search(line[1]), lines)
+        )
+        if len(wrong_type_lines) > 0:
+            raise RuntimeError(
+                "The annotation prefix in line "
+                + str(wrong_type_lines[0][0])
+                + " is probably invalid.\nIt must be '# type:'"
+                + "\nSee PEP 484 (https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)"  # noqa: B950
+                + "\nfor examples"
+            )
+        return None
+    elif len(type_lines) == 1:
+        # Only 1 type line, quit now
+        return type_lines[0][1].strip()
+
+    # Parse split up argument types according to PEP 484
+    # https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code
+    return_line = None
+    parameter_type_lines = []
+    for line_num, line in type_lines:
+        if "# type: (...) -> " in line:
+            return_line = (line_num, line)
+            break
+        elif type_comment in line:
+            parameter_type_lines.append(line)
+    if return_line is None:
+        raise RuntimeError(
+            "Return type line '# type: (...) -> ...' not found on multiline "
+            "type annotation\nfor type lines:\n"
+            + "\n".join([line[1] for line in type_lines])
+            + "\n(See PEP 484 https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)"
+        )
+
+    def get_parameter_type(line):
+        item_type = line[line.find(type_comment) + len(type_comment) :]
+        return item_type.strip()
+
+    types = map(get_parameter_type, parameter_type_lines)
+    parameter_types = ", ".join(types)
+
+    return return_line[1].replace("...", parameter_types)
+
+
+def split_type_line(type_line):
+    """Split the comment with the type annotation into parts for argument and return types.
+
+    For example, for an input of:
+        # type: (Tensor, torch.Tensor) -> Tuple[Tensor, Tensor]
+
+    This function will return:
+        ("(Tensor, torch.Tensor)", "Tuple[Tensor, Tensor]")
+
+    """
+    start_offset = len("# type:")
+    try:
+        arrow_pos = type_line.index("->")
+    except ValueError:
+        raise RuntimeError(
+            "Syntax error in type annotation (couldn't find `->`)"
+        ) from None
+    return type_line[start_offset:arrow_pos].strip(), type_line[arrow_pos + 2 :].strip()
+
+
+def try_real_annotations(fn, loc):
+    """Try to use the Py3.5+ annotation syntax to get the type."""
+    try:
+        # Note: anything annotated as `Optional[T]` will automatically
+        # be returned as `Union[T, None]` per
+        # https://github.com/python/cpython/blob/main/Lib/typing.py#L732
+        sig = inspect.signature(fn)
+    except ValueError:
+        return None
+
+    all_annots = [sig.return_annotation] + [
+        p.annotation for p in sig.parameters.values()
+    ]
+    if all(ann is sig.empty for ann in all_annots):
+        return None
+
+    arg_types = [ann_to_type(p.annotation, loc) for p in sig.parameters.values()]
+    return_type = ann_to_type(sig.return_annotation, loc)
+    return arg_types, return_type
+
+
+# Finds common type for enum values belonging to an Enum class. If not all
+# values have the same type, AnyType is returned.
+def get_enum_value_type(e: type[enum.Enum], loc):
+    enum_values: List[enum.Enum] = list(e)
+    if not enum_values:
+        raise ValueError(f"No enum values defined for: '{e.__class__}'")
+
+    types = {type(v.value) for v in enum_values}
+    ir_types = [try_ann_to_type(t, loc) for t in types]
+
+    # If Enum values are of different types, an exception will be raised here.
+    # Even though Python supports this case, we chose to not implement it to
+    # avoid overcomplicate logic here for a rare use case. Please report a
+    # feature request if you find it necessary.
+    res = torch._C.unify_type_list(ir_types)
+    if not res:
+        return AnyType.get()
+    return res
+
+
+def is_tensor(ann) -> bool:
+    if issubclass(ann, torch.Tensor):
+        return True
+
+    if issubclass(
+        ann,
+        (
+            torch.LongTensor,
+            torch.DoubleTensor,
+            torch.FloatTensor,
+            torch.IntTensor,
+            torch.ShortTensor,
+            torch.HalfTensor,
+            torch.CharTensor,
+            torch.ByteTensor,
+            torch.BoolTensor,
+        ),
+    ):
+        warnings.warn(
+            "TorchScript will treat type annotations of Tensor "
+            "dtype-specific subtypes as if they are normal Tensors. "
+            "dtype constraints are not enforced in compilation either.",
+            stacklevel=2,
+        )
+        return True
+
+    return False
+
+
+def _fake_rcb(inp) -> None:
+    return None
+
+
+def try_ann_to_type(ann, loc, rcb=None):
+    ann_args = typing.get_args(ann)  # always returns a tuple!
+
+    if ann is inspect.Signature.empty:
+        return TensorType.getInferred()
+    if ann is None:
+        return NoneType.get()
+    if inspect.isclass(ann) and is_tensor(ann):
+        return TensorType.get()
+    if is_tuple(ann):
+        # Special case for the empty Tuple type annotation `Tuple[()]`
+        if len(ann_args) == 1 and ann_args[0] == ():
+            return TupleType([])
+        return TupleType([try_ann_to_type(a, loc) for a in ann_args])
+    if is_list(ann):
+        elem_type = try_ann_to_type(ann_args[0], loc)
+        if elem_type:
+            return ListType(elem_type)
+    if is_dict(ann):
+        key = try_ann_to_type(ann_args[0], loc)
+        value = try_ann_to_type(ann_args[1], loc)
+        # Raise error if key or value is None
+        if key is None:
+            raise ValueError(
+                f"Unknown type annotation: '{ann_args[0]}' at {loc.highlight()}"
+            )
+        if value is None:
+            raise ValueError(
+                f"Unknown type annotation: '{ann_args[1]}' at {loc.highlight()}"
+            )
+        return DictType(key, value)
+    if is_optional(ann):
+        if issubclass(ann_args[1], type(None)):
+            contained = ann_args[0]
+        else:
+            contained = ann_args[1]
+        valid_type = try_ann_to_type(contained, loc)
+        if not valid_type:
+            raise AssertionError(
+                f"Unsupported annotation {repr(ann)} could not be resolved because "
+                f"{repr(contained)} could not be resolved. At\n{repr(loc)}"
+            )
+        return OptionalType(valid_type)
+    if is_union(ann):
+        # TODO: this is hack to recognize NumberType
+        if set(ann_args) == {int, float, complex}:
+            return NumberType.get()
+        inner: List = []
+        # We need these extra checks because both `None` and invalid
+        # values will return `None`
+        # TODO: Determine if the other cases need to be fixed as well
+        for a in typing.get_args(ann):
+            if a is None:
+                inner.append(NoneType.get())
+            maybe_type = try_ann_to_type(a, loc)
+            if not maybe_type:
+                raise AssertionError(
+                    f"Unsupported annotation {repr(ann)} could not be resolved because "
+                    f"{repr(a)} could not be resolved. At\n{repr(loc)}"
+                )
+            inner.append(maybe_type)
+        return UnionType(inner)  # type: ignore[arg-type]
+    if torch.distributed.rpc.is_available() and is_rref(ann):
+        return RRefType(try_ann_to_type(ann_args[0], loc))
+    if is_future(ann):
+        return FutureType(try_ann_to_type(ann_args[0], loc))
+    if is_await(ann):
+        elementType = try_ann_to_type(ann_args[0], loc) if ann_args else AnyType.get()
+        return AwaitType(elementType)
+    if ann is float:
+        return FloatType.get()
+    if ann is complex:
+        return ComplexType.get()
+    if ann is int or ann is torch.SymInt:
+        return IntType.get()
+    if ann is str:
+        return StringType.get()
+    if ann is bool:
+        return BoolType.get()
+    if ann is Any:
+        return AnyType.get()
+    if ann is type(None):
+        return NoneType.get()
+    if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
+        return InterfaceType(ann.__torch_script_interface__)
+    if ann is torch.device:
+        return DeviceObjType.get()
+    if ann is torch.Generator:
+        return _GeneratorType.get()
+    if ann is torch.Stream:
+        return StreamObjType.get()
+    if ann is torch.dtype:
+        return IntType.get()  # dtype not yet bound in as its own type
+    if ann is torch.qscheme:
+        return IntType.get()  # qscheme not yet bound in as its own type
+
+    if inspect.isclass(ann) and issubclass(ann, enum.Enum):
+        if _get_script_class(ann) is None:
+            scripted_class = torch.jit._script._recursive_compile_class(ann, loc)
+            name = scripted_class.qualified_name()
+        else:
+            name = _qualified_name(ann)
+        return EnumType(name, get_enum_value_type(ann, loc), list(ann))
+    if inspect.isclass(ann):
+        maybe_script_class = _get_script_class(ann)
+        if maybe_script_class is not None:
+            return maybe_script_class
+        if torch._jit_internal.can_compile_class(ann):
+            return torch.jit._script._recursive_compile_class(ann, loc)
+
+    # Maybe resolve a NamedTuple to a Tuple Type
+    if rcb is None:
+        rcb = _fake_rcb
+    # pyrefly: ignore [bad-argument-type]
+    return torch._C._resolve_type_from_object(ann, loc, rcb)
+
+
+def ann_to_type(ann, loc, rcb=None):
+    the_type = try_ann_to_type(ann, loc, rcb)
+    if the_type is not None:
+        return the_type
+    raise ValueError(f"Unknown type annotation: '{ann}' at {loc.highlight()}")
+
+
+__all__ = [
+    "Any",
+    "List",
+    "BroadcastingList1",
+    "BroadcastingList2",
+    "BroadcastingList3",
+    "Tuple",
+    "is_tuple",
+    "is_list",
+    "Dict",
+    "is_dict",
+    "is_optional",
+    "is_union",
+    "TensorType",
+    "TupleType",
+    "FloatType",
+    "ComplexType",
+    "IntType",
+    "ListType",
+    "StringType",
+    "DictType",
+    "AnyType",
+    "Module",
+    # TODO: Consider not exporting these during wildcard import (reserve
+    # that for the types; for idiomatic typing code.)
+    "get_signature",
+    "check_fn",
+    "get_param_names",
+    "parse_type_line",
+    "get_type_line",
+    "split_type_line",
+    "try_real_annotations",
+    "try_ann_to_type",
+    "ann_to_type",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/frontend.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/frontend.py
new file mode 100644
index 0000000000000000000000000000000000000000..4164c6570ffbc72c4505deac9d34da6b2f3fe6da
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/frontend.py
@@ -0,0 +1,1288 @@
+# mypy: allow-untyped-defs
+import ast
+import copy
+import dataclasses
+import inspect
+import re
+import string
+from collections import namedtuple
+from textwrap import dedent
+
+import torch
+import torch.jit.annotations
+from torch import _jit_internal
+from torch._C._jit_tree_views import (
+    Apply,
+    Assert,
+    Assign,
+    Attribute,
+    AugAssign,
+    BinOp,
+    Break,
+    ClassDef,
+    Const,
+    Continue,
+    Decl,
+    Def,
+    Delete,
+    DictComp,
+    DictLiteral,
+    Dots,
+    EmptyTypeAnnotation,
+    ExprStmt,
+    FalseLiteral,
+    For,
+    Ident,
+    If,
+    ListComp,
+    ListLiteral,
+    NoneLiteral,
+    Param,
+    Pass,
+    Property,
+    Raise,
+    Return,
+    Select,
+    SliceExpr,
+    Starred,
+    Stmt,
+    StringLiteral,
+    Subscript,
+    TernaryIf,
+    TrueLiteral,
+    TupleLiteral,
+    UnaryOp,
+    Var,
+    While,
+    With,
+    WithItem,
+)
+from torch._jit_internal import (  # noqa: F401
+    _is_drop_fn,
+    FunctionModifiers,
+    is_static_fn,
+    should_drop,
+)
+from torch._sources import (
+    get_source_lines_and_file,
+    make_source_context,
+    parse_def,
+    ParsedDef as _ParsedDef,
+)
+from torch.jit._dataclass_impls import DATACLASS_MAGIC_METHODS
+from torch.jit._monkeytype_config import get_qualified_name, monkeytype_trace
+
+
+# Borrowed from cPython implementation
+# https://github.com/python/cpython/blob/561612d8456cfab5672c9b445521113b847bd6b3/Lib/textwrap.py#L411#
+
+_reserved_prefix = "__jit"
+_reserved_names = {"print"}
+_identifier_chars = set(string.ascii_lowercase + string.ascii_uppercase + string.digits)
+
+
+def is_reserved_name(name):
+    return name.startswith(_reserved_prefix) or name in _reserved_names
+
+
+pretty_node_names = {
+    ast.FunctionDef: "function definitions",
+    ast.For: "for loops",
+    ast.Delete: "del statements",
+    ast.ClassDef: "class definitions",
+    ast.With: "with statements",
+    ast.Raise: "raise statements",
+    ast.Assert: "assertions",
+    ast.Import: "import statements",
+    ast.ImportFrom: "import statements",
+    ast.Global: "global variables",
+    ast.Break: "break statements",
+    ast.Continue: "continue statements",
+}
+
+node_start_tokens = {
+    ast.FunctionDef: "def",
+    ast.For: "for",
+    ast.Delete: "del",
+    ast.ClassDef: "class",
+    ast.With: "with",
+    ast.Raise: "raise",
+    ast.Assert: "assert",
+    ast.Import: "import",
+    ast.ImportFrom: "from",
+    ast.Global: "global",
+    ast.Break: "break",
+    ast.Continue: "continue",
+}
+
+# pyrefly: ignore [no-matching-overload]
+pretty_node_names.update(
+    {
+        ast.AsyncFunctionDef: "async function definitions",
+        ast.AsyncFor: "async for loops",
+        ast.AsyncWith: "async with statements",
+        ast.Try: "try blocks",
+        ast.Nonlocal: "nonlocal variables",
+    }
+)
+
+# pyrefly: ignore [no-matching-overload]
+node_start_tokens.update(
+    {
+        ast.AsyncFunctionDef: "async def",
+        ast.AsyncFor: "async for",
+        ast.AsyncWith: "async with",
+        ast.Try: "try",
+        ast.Nonlocal: "nonlocal",
+    }
+)
+
+# pyrefly: ignore [no-matching-overload]
+pretty_node_names.update(
+    {
+        ast.AnnAssign: "annotated assignments",
+    }
+)
+# NB: no specific token for AnnAssign
+
+
+class FrontendError(Exception):
+    def __init__(self, source_range, msg) -> None:
+        self.source_range = source_range
+        self.msg = msg
+
+        # This has to be instantiated here so the ErrorReport is accurate to the
+        # call stack when the FrontendError was raised
+        self.error_report = torch._C.ErrorReport(self.source_range)
+
+    def __str__(self) -> str:
+        return self.msg + self.error_report.what().lstrip()
+
+
+class NotSupportedError(FrontendError):
+    pass
+
+
+class UnsupportedNodeError(NotSupportedError):
+    def __init__(self, ctx, offending_node, reason="") -> None:
+        # If we don't have a specific token, we default to length of 1
+        node_type = type(offending_node)
+        range_len = len(node_start_tokens.get(node_type, " "))
+        source_range = ctx.make_range(
+            offending_node.lineno,
+            offending_node.col_offset,
+            offending_node.col_offset + range_len,
+        )
+        feature_name = pretty_node_names.get(node_type, node_type.__name__)
+        msg = f"{feature_name} {reason + ' ' if reason else ''}aren't supported"
+        super().__init__(source_range, msg)
+
+
+class FrontendTypeError(FrontendError):
+    pass
+
+
+def build_withitems(ctx, items):
+    items = [build_withitem(ctx, i) for i in items]
+    return list(items)
+
+
+def build_stmts(ctx, stmts):
+    stmts = [build_stmt(ctx, s) for s in stmts]
+    return list(filter(None, stmts))
+
+
+def get_class_properties(cls, self_name):
+    """
+    Get a list of Property objects representing the properties of a class.
+
+    Args:
+        cls:  The class to get properties of.
+        self_name: The name of the class that the properties should belong to.
+    Returns:
+        A list of Property objects corresponding to the properties of cls. Property
+        here refers to the subclass of TreeView.
+    """
+    props = inspect.getmembers(cls, predicate=lambda m: isinstance(m, property))
+    # Any property that should not compiled must be in this list on the Module.
+    unused_properties = getattr(cls, "__jit_unused_properties__", [])
+
+    # Create Property TreeView objects from inspected property objects.
+    properties = []
+    for prop in props:
+        if prop[0] not in unused_properties and not should_drop(prop[1].fget):
+            getter = get_jit_def(
+                prop[1].fget, f"__{prop[0]}_getter", self_name=self_name
+            )
+            setter = (
+                get_jit_def(prop[1].fset, f"__{prop[0]}_setter", self_name=self_name)
+                if prop[1].fset
+                else None
+            )
+            properties.append(
+                Property(getter.range(), Ident(getter.range(), prop[0]), getter, setter)
+            )
+
+    return properties
+
+
+def get_class_assigns(ctx, cls_ast):
+    assigns = []
+
+    def maybe_build_assign(builder, entry) -> None:
+        nonlocal assigns
+        try:
+            assigns.append(builder(ctx, entry))
+        except NotSupportedError:
+            pass
+
+    for entry in cls_ast.body:
+        if isinstance(entry, ast.Assign):
+            maybe_build_assign(StmtBuilder.build_Assign, entry)
+        elif isinstance(entry, ast.AnnAssign):
+            maybe_build_assign(StmtBuilder.build_AnnAssign, entry)
+    return assigns
+
+
+def get_jit_class_def(cls, self_name):
+    """Get definitions for each method within the current class independently.
+
+    Args:
+        cls: The class to get definition of.
+        self_name: The name of the class that the properties should belong to.
+
+    Returns:
+        torch._C._jit_tree_views.ClassDef: A representation of the class,
+            the methods in the class and their definition as a tree.
+    """
+    # TODO: proper overriding analysis when implementing class inheritance
+    methods = inspect.getmembers(
+        cls,
+        predicate=lambda m: (inspect.ismethod(m) or inspect.isfunction(m))
+        and not is_static_fn(cls, m.__name__)
+        and m.__name__ in cls.__dict__
+        and not _is_drop_fn(m),
+    )
+
+    def is_classmethod(fn):
+        return inspect.ismethod(fn) and getattr(fn, "__self__", None) == cls
+
+    # Get and parse the source code for this class
+    sourcelines, file_lineno, filename = get_source_lines_and_file(
+        cls, torch._C.ErrorReport.call_stack()
+    )
+    source = "".join(sourcelines)
+
+    dedent_src = dedent(source)
+    py_ast = ast.parse(dedent_src)
+
+    class_ast = py_ast.body[0]
+    if not isinstance(class_ast, ast.ClassDef):
+        raise AssertionError(
+            f"Expected class definition, got {type(class_ast).__name__}"
+        )
+
+    # Special case for dataclasses. In general we need access to the source code for
+    # an object in order to JIT compile it. But the dataclasses module dynamically synthesizes
+    # magic methods for classes, and we can't get the source code for these methods. As a
+    # workaround, we synthesize TorchScript-friendly implementations ourselves.
+    if dataclasses.is_dataclass(cls):
+        # Detect whether the user manually implemented any of the magic methods. If they did,
+        # we don't want to synthesize/override them.
+        overrides = {
+            method.name
+            for method in class_ast.body
+            if isinstance(method, ast.FunctionDef)
+            and method.name in DATACLASS_MAGIC_METHODS
+        }
+        for i, (name, _) in enumerate(methods):
+            # Is this a magic method we can synthesize?
+            synthesizer_fn = DATACLASS_MAGIC_METHODS.get(name)
+            if synthesizer_fn and name not in overrides:
+                parsed_def = synthesizer_fn(cls)
+                methods[i] = name, parsed_def
+                func = getattr(cls, name)
+                _jit_internal.loader.cache(func, parsed_def.source)
+
+    method_defs = [
+        get_jit_def(obj, name, self_name=self_name, is_classmethod=is_classmethod(obj))
+        for (name, obj) in methods
+    ]
+    properties = get_class_properties(cls, self_name)
+
+    leading_whitespace_len = len(source.split("\n", 1)[0]) - len(
+        dedent_src.split("\n", 1)[0]
+    )
+    ctx = make_source_context(
+        source, filename, file_lineno, leading_whitespace_len, False
+    )
+    assigns = get_class_assigns(ctx, class_ast)
+
+    return build_class_def(ctx, class_ast, method_defs, properties, self_name, assigns)
+
+
+def get_jit_def(fn, def_name, self_name=None, is_classmethod=False):
+    """
+    Build a JIT AST (TreeView) from the given function.
+
+    Args:
+        fn: A function object to compile or a pre-parsed ParsedDef object
+        def_name: The name to give to the resulting AST object. This is not
+            always the same as `fn.__name__`, for example:
+                def _forward(self):
+                    ...
+                forward = _forward
+            In this case, the `__name__` attribute of the function object is "_forward",
+            but we want the result AST to have the name "forward".
+        self_name: If this function is a method, what the type name of `self` is.
+    """
+    parsed_def = parse_def(fn) if not isinstance(fn, _ParsedDef) else fn
+    type_line = torch.jit.annotations.get_type_line(parsed_def.source)
+    fn_def = parsed_def.ast.body[0]
+
+    if is_classmethod:
+        arg_name = fn_def.args.args[0].arg  # type:ignore[union-attr]
+        # Insert a statement that assigns the first argument to the class
+        assign_stmt = ast.parse(f"{arg_name} = {self_name}").body[0]
+        fn_def.body.insert(0, assign_stmt)  # type:ignore[union-attr]
+
+    # Swap out the function signature and body if it is unused
+    if should_drop(fn):
+        unused_fn_def = ast.parse(
+            'def unused_fn(self: Any):\n\traise RuntimeError("Cannot call @unused methods")'
+        )
+        if len(unused_fn_def.body) != 1 or not isinstance(
+            unused_fn_def.body[0], ast.FunctionDef
+        ):
+            raise RuntimeError(
+                f"Expected a single top-level function: {parsed_def.filename}:{parsed_def.file_lineno}"
+            )
+        unused_def = unused_fn_def.body[0]
+        fn_def.body = unused_def.body  # type:ignore[union-attr]
+        # kwarg/vararg not supported by `build_def`
+        fn_def.args.kwarg = fn_def.args.vararg = None  # type:ignore[union-attr]
+        for arg in fn_def.args.args + fn_def.args.kwonlyargs:  # type:ignore[union-attr]
+            # Replace potentially unsupported type annotations by "Any"
+            arg.annotation = unused_def.args.args[0].annotation
+        if _is_drop_fn(fn):
+            # Dropping potentially unsupported return type annotation for jit._drop
+            fn_def.returns = None  # type:ignore[union-attr]
+            fn_def.type_comment = None  # type:ignore[union-attr]
+
+    # If MonkeyType is installed, get all the consolidated type traces
+    # for the arguments from type_trace_db
+    type_trace_db = torch.jit._script._get_type_trace_db()
+    pdt_arg_types = None
+    if monkeytype_trace and not isinstance(fn, _ParsedDef):  # type: ignore[truthy-function]
+        qualname = get_qualified_name(fn)
+        pdt_arg_types = type_trace_db.get_args_types(qualname)
+
+    return build_def(
+        parsed_def.ctx,
+        fn_def,
+        type_line,
+        def_name,
+        self_name=self_name,
+        pdt_arg_types=pdt_arg_types,
+    )
+
+
+# TODO: more robust handling of recognizing ignore context manager
+def is_torch_jit_ignore_context_manager(stmt) -> bool:
+    # checks if the statement is torch.jit.ignore context manager
+    if isinstance(stmt.items[0].context_expr, ast.Call):
+        # extract torch part
+        function = stmt.items[0].context_expr.func
+        if isinstance(function, ast.Attribute):
+            attr_name = function.attr
+            attr_value = function.value
+            if attr_name == "_IgnoreContextManager" and isinstance(
+                attr_value, ast.Attribute
+            ):
+                # there should be at most two nested attributes (e.g torch.jit._IgnoreContextManager)
+                if attr_value.attr == "jit" and isinstance(attr_value.value, ast.Name):
+                    if attr_value.value.id == "torch":
+                        return True
+    return False
+
+
+class Builder:
+    def __call__(self, ctx, node):
+        method = getattr(self, "build_" + node.__class__.__name__, None)
+        if method is None:
+            raise UnsupportedNodeError(ctx, node)
+        return method(ctx, node)
+
+
+def build_class_def(ctx, py_def, methods, properties, self_name, assigns):
+    r = ctx.make_range(
+        py_def.lineno, py_def.col_offset, py_def.col_offset + len("class")
+    )
+    return ClassDef(
+        Ident(r, self_name), [Stmt(method) for method in methods], properties, assigns
+    )
+
+
+def build_def(ctx, py_def, type_line, def_name, self_name=None, pdt_arg_types=None):
+    body = py_def.body
+    r = ctx.make_range(py_def.lineno, py_def.col_offset, py_def.col_offset + len("def"))
+
+    param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types)
+    return_type = None
+    if getattr(py_def, "returns", None) is not None:
+        return_type = build_expr(ctx, py_def.returns)
+
+    decl = Decl(r, param_list, return_type)
+    is_method = self_name is not None
+    if type_line is not None:
+        type_comment_decl = torch._C.parse_type_comment(type_line)
+        decl = torch._C.merge_type_from_type_comment(
+            decl,  # type: ignore[arg-type]
+            type_comment_decl,
+            is_method,  # type: ignore[assignment]
+        )
+
+    return Def(Ident(r, def_name), decl, build_stmts(ctx, body))
+
+
+_vararg_kwarg_err = (
+    "Compiled functions can't take variable number of arguments "
+    "or use keyword-only arguments with defaults"
+)
+
+
+def build_param_list(ctx, py_args, self_name, pdt_arg_types=None):
+    if py_args.kwarg is not None:
+        expr = py_args.kwarg
+        ctx_range = ctx.make_range(
+            expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg)
+        )
+        raise NotSupportedError(ctx_range, _vararg_kwarg_err)
+    if py_args.vararg is not None:
+        expr = py_args.vararg
+        ctx_range = ctx.make_range(
+            expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg)
+        )
+        raise NotSupportedError(ctx_range, _vararg_kwarg_err)
+    if len(py_args.kw_defaults) > 0:
+        # kw_defaults is a list of the values for the kwargs (which default to None),
+        # so they don't actually have line numbers.
+        for arg in py_args.kw_defaults:
+            if arg is not None:
+                ctx_range = build_expr(ctx, arg).range()
+                raise NotSupportedError(ctx_range, _vararg_kwarg_err)
+
+    # List of Tuple of args and type as inferred by profile directed typing
+    arg_and_types = [
+        (
+            arg,
+            pdt_arg_types[arg.arg]
+            if pdt_arg_types and bool(pdt_arg_types[arg.arg])
+            else None,
+        )
+        for arg in py_args.args
+    ]
+    arg_and_types_kwonlyargs = [
+        (
+            arg,
+            pdt_arg_types[arg.arg]
+            if pdt_arg_types and bool(pdt_arg_types[arg.arg])
+            else None,
+        )
+        for arg in py_args.kwonlyargs
+    ]
+
+    result = [
+        build_param(ctx, arg, self_name, kwarg_only=False, pdt_arg_type=arg_type)
+        for arg, arg_type in arg_and_types
+    ]
+    result += [
+        build_param(ctx, arg, self_name, kwarg_only=True, pdt_arg_type=arg_type)
+        for arg, arg_type in arg_and_types_kwonlyargs
+    ]
+    return result
+
+
+def build_param(ctx, py_arg, self_name, kwarg_only, pdt_arg_type=None):
+    # NB: In Python3 py_arg is a pair of (str arg, expr? annotation)
+    name = py_arg.arg
+    r = ctx.make_range(py_arg.lineno, py_arg.col_offset, py_arg.col_offset + len(name))
+    if getattr(py_arg, "annotation", None) is not None:
+        annotation_expr = build_expr(ctx, py_arg.annotation)
+    elif pdt_arg_type:
+        annotation_expr = Var(Ident(r, pdt_arg_type))
+    elif self_name is not None and name == "self":
+        annotation_expr = Var(Ident(r, self_name))
+    else:
+        annotation_expr = EmptyTypeAnnotation(r)
+    return Param(annotation_expr, Ident(r, name), kwarg_only)
+
+
+def build_ignore_context_manager(ctx, stmt):
+    InputType = namedtuple("InputType", ["name", "ann"])
+    OutputType = namedtuple("OutputType", ["name", "ann"])
+
+    def process_ins_outs(args):
+        # parse the context manager to figure out inputs and outputs
+        # with their annotated types
+        # TODO: add input, output validator
+        inputs = []
+        outputs = []
+        for arg in args:
+            var_name = arg.arg
+            var_ann = arg.value.value
+            var_decl_type, var_ann = var_ann.split(":")
+            if var_decl_type == "inp":
+                inputs.append(InputType(var_name, var_ann))
+            if var_decl_type == "out":
+                outputs.append(OutputType(var_name, var_ann))
+        return inputs, outputs
+
+    def create_unique_name_ext(ctx, stmt) -> str:
+        # extension will be based on the full path filename plus
+        # the line number of original context manager
+        fn = re.sub(r"[^a-zA-Z0-9_]", "_", ctx.filename)
+        return f"{fn}_{stmt.lineno}"
+
+    def build_return_ann_stmt(outputs):
+        return_type_ann = ""
+        return_statement_str = "return "
+        if len(outputs) == 0:
+            return_type_ann += " -> None"
+        if len(outputs) == 1:
+            return_type_ann = " -> " + outputs[0].ann
+            return_statement_str += outputs[0].name
+        if len(outputs) > 1:
+            return_type_ann = " -> tuple"
+            return_type_ann += "[" + ", ".join([var.ann for var in outputs]) + "]"
+            return_statement_str += ", ".join([var.name for var in outputs])
+        return return_type_ann, return_statement_str
+
+    def build_args(args):
+        return ", ".join([arg.name for arg in args])
+
+    inputs, outputs = process_ins_outs(stmt.items[0].context_expr.keywords)
+
+    # build the replacement function str with given inputs and outputs
+    ignore_function_name = "func_ignore_" + create_unique_name_ext(ctx, stmt)
+    ignore_function_str = "\ndef " + ignore_function_name
+    ignore_function_str += (
+        "(" + ", ".join([var.name + " :" + var.ann for var in inputs]) + ")"
+    )
+
+    return_ann, return_stmt = build_return_ann_stmt(outputs)
+    ignore_function_str += return_ann + ": pass"
+
+    # first create the functionDef object from just declaration
+    ignore_function = ast.parse(ignore_function_str).body[0]
+
+    # dump the body of context manager to dummy function
+    ignore_function.body = stmt.body  # type: ignore[attr-defined]
+
+    # insert return statement to the function
+    return_stmt = ast.parse(return_stmt).body[0]
+    ignore_function.body.append(return_stmt)  # type: ignore[attr-defined]
+
+    ignore_func_str = f"""\
+# Backward compat: These used to be imported into the outer global scope so some
+# code may still expect them.
+from typing import List, Dict, Tuple
+
+@torch.jit.ignore
+{ast.unparse(ignore_function)}
+"""
+    g = copy.copy(globals())
+    exec(ignore_func_str, g)  # noqa: P204
+    # registers the custom function in the global context
+    globals()[ignore_function_name] = g[ignore_function_name]
+
+    # build the statements as:
+    # <out_1>, <out_2>, ... = torch.jit.frontend.<func>(<in_1>, <in_2>)
+    assign_str_lhs = build_args(outputs)
+    # this function will be registered in torch.jit.frontend module by default
+    assign_str_rhs = (
+        f"torch.jit.frontend.{ignore_function_name}(" + build_args(inputs) + ")"
+    )
+
+    if len(outputs) > 0:
+        assign_str = assign_str_lhs + " = " + assign_str_rhs
+    else:
+        assign_str = assign_str_rhs
+    assign_ast = ast.parse(assign_str).body[0]
+    return assign_ast
+
+
+def get_default_args(fn):
+    """
+    Get a dictionary of default arguments for a function.
+
+    Args:
+        fn: Callable - The function to inspect for default arguments.
+    Returns:
+        (Dict[str, Any]): mapping argument names to their default values if
+        :attr:`fn` is not None, else empty dictionary.
+    """
+    if fn is None:
+        return {}
+
+    signature = inspect.signature(fn)
+
+    return {
+        k: v.default
+        for k, v in signature.parameters.items()
+        if v.default is not inspect.Parameter.empty
+    }
+
+
+def get_default_args_for_class(cls):
+    """
+    Get default arguments for all methods in a class (except for static methods).
+
+    Args:
+        cls: type - The class type to inspect for default arguments.
+    Returns:
+        A Dict[str, Dict[str, Any]] which maps each method name to a Dict[str, Any]
+        that maps each argument name to its default value.
+    """
+    # Get methods (except static methods because those are compiled separately as
+    # if they were independent script functions).
+    methods = inspect.getmembers(
+        cls,
+        predicate=lambda m: (inspect.ismethod(m) or inspect.isfunction(m))
+        and not is_static_fn(cls, m.__name__)
+        and m.__name__ in cls.__dict__,
+    )
+
+    # Get method defaults. Property defaults do not need to be considered
+    # because setters cannot be invoked without a value.
+    defaults = {
+        method_name: get_default_args(method_impl)
+        for method_name, method_impl in methods
+    }
+
+    return defaults
+
+
+class WithItemBuilder(Builder):
+    @staticmethod
+    def build_withitem(ctx, item):
+        lineno = item.context_expr.lineno
+        start = item.context_expr.col_offset
+        end = start + len(pretty_node_names[ast.With])
+        op_vars = item.optional_vars
+        r = ctx.make_range(lineno, start, end)
+
+        return WithItem(
+            r,
+            build_expr(ctx, item.context_expr),
+            build_expr(ctx, op_vars) if op_vars else None,
+        )
+
+
+class StmtBuilder(Builder):
+    augassign_map = {
+        ast.Add: "+",
+        ast.Sub: "-",
+        ast.Mult: "*",
+        ast.Div: "/",
+        ast.Mod: "%",
+        ast.BitOr: "|",
+        ast.BitAnd: "&",
+        ast.BitXor: "^",
+        ast.LShift: "<<",
+        ast.RShift: ">>",
+        ast.Pow: "**",
+    }
+
+    @staticmethod
+    def build_Expr(ctx, stmt):
+        value = stmt.value
+        if value.__class__.__name__ == "Str":
+            # If a statement is a string literal expression,
+            # then it is a docstring. Just ignore it.
+            return None
+        else:
+            return ExprStmt(build_expr(ctx, value))
+
+    @staticmethod
+    def build_Assign(ctx, stmt):
+        rhs = build_expr(ctx, stmt.value)
+        lhs = [build_expr(ctx, x) for x in stmt.targets]
+        return Assign(lhs, rhs)
+
+    @staticmethod
+    def build_AnnAssign(ctx, stmt):
+        if stmt.value is None:
+            raise UnsupportedNodeError(ctx, stmt, reason="without assigned value")
+
+        # Disallow type annotations on instance attributes outside of __init__
+        if (
+            type(stmt.target) is ast.Attribute
+            and stmt.target.value.id == "self"  # type: ignore[attr-defined]
+            and ctx.funcname != "__init__"
+        ):
+            start = stmt.col_offset
+            end = start + len(f"self.{stmt.target.attr}")
+            if hasattr(stmt.annotation, "id"):
+                end += len(f": {stmt.annotation.id}")
+            sr = ctx.make_range(stmt.lineno, start, end)
+            raise ValueError(
+                "Type annotations on instance attributes must be declared in "
+                f"__init__, not '{ctx.funcname}': {sr}"
+            )
+
+        rhs = build_expr(ctx, stmt.value)
+        lhs = build_expr(ctx, stmt.target)
+        the_type = build_expr(ctx, stmt.annotation)
+        return Assign([lhs], rhs, the_type)
+
+    @staticmethod
+    def build_Delete(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("del"))
+
+        return Delete(r, [build_expr(ctx, target) for target in stmt.targets])
+
+    @staticmethod
+    def build_Return(ctx, stmt):
+        r = ctx.make_range(
+            stmt.lineno, stmt.col_offset, stmt.col_offset + len("return")
+        )
+        return Return(r, None if stmt.value is None else build_expr(ctx, stmt.value))
+
+    @staticmethod
+    def build_Raise(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("raise"))
+        expr = build_expr(ctx, stmt.exc)
+        return Raise(r, expr)
+
+    @staticmethod
+    def build_Assert(ctx, stmt):
+        r = ctx.make_range(
+            stmt.lineno, stmt.col_offset, stmt.col_offset + len("assert")
+        )
+        test = build_expr(ctx, stmt.test)
+        msg = build_expr(ctx, stmt.msg) if stmt.msg is not None else None
+        return Assert(r, test, msg)
+
+    @staticmethod
+    def build_AugAssign(ctx, stmt):
+        lhs = build_expr(ctx, stmt.target)
+        rhs = build_expr(ctx, stmt.value)
+        op = type(stmt.op)
+        if op in StmtBuilder.augassign_map:
+            op_token = StmtBuilder.augassign_map[op]
+        else:
+            raise NotSupportedError(
+                find_before(ctx, rhs.range().start, "=", offsets=(-1, 0)),
+                "unsupported kind of augmented assignment: " + op.__name__,
+            )
+        return AugAssign(lhs, op_token, rhs)
+
+    @staticmethod
+    def build_While(ctx, stmt):
+        if stmt.orelse:
+            # TODO: try to recover the location of else:? Python doesn't give us useful
+            # annotations in this case
+            raise NotSupportedError(
+                None, "else branches of while loops aren't supported"
+            )
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("while"))
+        return While(r, build_expr(ctx, stmt.test), build_stmts(ctx, stmt.body))
+
+    @staticmethod
+    def build_For(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("for"))
+        if stmt.orelse:
+            raise NotSupportedError(r, "else branches of for loops aren't supported")
+
+        return For(
+            r,
+            [build_expr(ctx, stmt.target)],
+            [build_expr(ctx, stmt.iter)],
+            build_stmts(ctx, stmt.body),
+        )
+
+    @staticmethod
+    def build_If(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("if"))
+        return If(
+            r,
+            build_expr(ctx, stmt.test),
+            build_stmts(ctx, stmt.body),
+            build_stmts(ctx, stmt.orelse),
+        )
+
+    @staticmethod
+    def build_Print(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("print"))
+        if stmt.dest:
+            raise NotSupportedError(
+                r, "print statements with non-default destinations aren't supported"
+            )
+        args = [build_expr(ctx, val) for val in stmt.values]
+        return ExprStmt(Apply(Var(Ident(r, "print")), args, []))
+
+    @staticmethod
+    def build_Pass(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("pass"))
+        return Pass(r)
+
+    @staticmethod
+    def build_Break(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("break"))
+        return Break(r)
+
+    @staticmethod
+    def build_Continue(ctx, stmt):
+        r = ctx.make_range(
+            stmt.lineno, stmt.col_offset, stmt.col_offset + len("continue")
+        )
+        return Continue(r)
+
+    @staticmethod
+    def build_With(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("with"))
+        # Handle ignore context manager
+        if is_torch_jit_ignore_context_manager(stmt):
+            assign_ast = build_ignore_context_manager(ctx, stmt)
+            return build_stmt(ctx, assign_ast)
+        return With(r, build_withitems(ctx, stmt.items), build_stmts(ctx, stmt.body))
+
+
+class ExprBuilder(Builder):
+    binop_map = {
+        ast.Add: "+",
+        ast.Sub: "-",
+        ast.Mult: "*",
+        ast.Div: "/",
+        ast.Pow: "**",
+        ast.Mod: "%",
+        ast.FloorDiv: "//",
+        ast.BitAnd: "&",
+        ast.BitXor: "^",
+        ast.BitOr: "|",
+        ast.LShift: "<<",
+        ast.RShift: ">>",
+    }
+
+    # pyrefly: ignore [unsupported-operation]
+    binop_map[ast.MatMult] = "@"
+
+    unop_map = {
+        ast.Not: "not",
+        ast.USub: "-",
+        ast.Invert: "~",
+    }
+
+    boolop_map = {
+        ast.And: "and",
+        ast.Or: "or",
+    }
+
+    cmpop_map = {
+        ast.Eq: "==",
+        ast.NotEq: "!=",
+        ast.LtE: "<=",
+        ast.Lt: "<",
+        ast.GtE: ">=",
+        ast.Gt: ">",
+        ast.Is: "is",
+        ast.IsNot: "is not",
+        ast.In: "in",
+        ast.NotIn: "not in",
+    }
+
+    @staticmethod
+    def build_Attribute(ctx, expr):
+        base = build_expr(ctx, expr.value)
+        # expr.attr is just a string, so it's not annotated in any way, so we have
+        # to build the range manually
+        source = ctx.source.encode("utf-8")
+
+        def get_char(index):
+            return chr(source[index])
+
+        start_pos = base.range().end + 1
+        while get_char(start_pos) in string.whitespace:  # Skip whitespace
+            start_pos += 1
+        end_pos = start_pos + len(expr.attr)
+        name_range = ctx.make_raw_range(start_pos, end_pos)
+        return Select(base, Ident(name_range, expr.attr))
+
+    @staticmethod
+    def build_Call(ctx, expr):
+        func = build_expr(ctx, expr.func)
+        args = [build_expr(ctx, py_arg) for py_arg in expr.args]
+        if hasattr(expr, "starargs") and expr.starargs:
+            stararg_expr = build_expr(ctx, expr.starargs)
+            args += [Starred(stararg_expr.range(), stararg_expr)]
+        kwargs = []
+        for kw in expr.keywords:
+            kw_expr = build_expr(ctx, kw.value)
+            # XXX: we could do a better job at figuring out the range for the name here
+            if not kw.arg:
+                raise NotSupportedError(
+                    kw_expr.range(), "keyword-arg expansion is not supported"
+                )
+            kwargs.append(Attribute(Ident(kw_expr.range(), kw.arg), kw_expr))
+        return Apply(func, args, kwargs)
+
+    @staticmethod
+    def build_Ellipsis(ctx, expr):
+        r = ctx.make_range(
+            expr.lineno, expr.col_offset, expr.col_offset + 3
+        )  # len("...") == 3
+        return Dots(r)
+
+    @staticmethod
+    def build_Name(ctx, expr):
+        r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(expr.id))
+        if expr.id.startswith(_reserved_prefix):
+            raise NotSupportedError(
+                r,
+                "names of variables used in JIT-ed functions "
+                "can't start with " + _reserved_prefix,
+            )
+        if expr.id == "True":
+            return TrueLiteral(r)
+        elif expr.id == "False":
+            return FalseLiteral(r)
+        elif expr.id == "None":
+            return NoneLiteral(r)
+        elif expr.id == "Ellipsis":
+            return Dots(r)
+        return Var(Ident(r, expr.id))
+
+    @staticmethod
+    def build_NameConstant(ctx, expr):
+        r = ctx.make_range(
+            expr.lineno, expr.col_offset, expr.col_offset + len(str(expr.value))
+        )
+        if expr.value is True:
+            return TrueLiteral(r)
+        elif expr.value is False:
+            return FalseLiteral(r)
+        elif expr.value is None:
+            return NoneLiteral(r)
+        elif expr.value == Ellipsis:
+            return Dots(r)
+        else:
+            raise ValueError("Name constant value unsupported: " + str(expr.value))
+
+    @staticmethod
+    def build_BinOp(ctx, expr):
+        lhs = build_expr(ctx, expr.left)
+        rhs = build_expr(ctx, expr.right)
+        op = type(expr.op)
+
+        if op == ast.Div and not ctx.uses_true_division:
+            err_range = ctx.make_raw_range(lhs.range().end, rhs.range().start)
+            raise FrontendError(
+                err_range,
+                "Division of ints in TorchScript uses Python 3 true "
+                "division semantics. Please put `from __future__ "
+                "import division` at the top of your file",
+            )
+        op_token = ExprBuilder.binop_map.get(op)
+        if op_token is None:
+            err_range = ctx.make_raw_range(lhs.range().end, rhs.range().start)
+            raise NotSupportedError(
+                err_range, "unsupported binary operator: " + op.__name__
+            )
+        return BinOp(op_token, lhs, rhs)
+
+    @staticmethod
+    def build_UnaryOp(ctx, expr):
+        sub_expr = build_expr(ctx, expr.operand)
+        op = type(expr.op)
+        op_token = ExprBuilder.unop_map.get(op)
+        if op_token is None:
+            raise NotSupportedError(
+                expr.range(), "unsupported unary operator: " + op.__name__
+            )
+        r = ctx.make_range(
+            expr.lineno, expr.col_offset, expr.col_offset + len(op_token)
+        )
+        return UnaryOp(r, op_token, sub_expr)
+
+    @staticmethod
+    def build_BoolOp(ctx, expr):
+        if len(expr.values) < 2:
+            raise AssertionError(
+                "expected at least 2 values in BoolOp, but got " + str(len(expr.values))
+            )
+        sub_exprs = [build_expr(ctx, sub_expr) for sub_expr in expr.values]
+        op = type(expr.op)
+        op_token = ExprBuilder.boolop_map.get(op)
+        if op_token is None:
+            err_range = ctx.make_raw_range(
+                sub_exprs[0].range().end, sub_exprs[1].range().start
+            )
+            raise NotSupportedError(
+                err_range, "unsupported boolean operator: " + op.__name__
+            )
+        lhs = sub_exprs[0]
+        for rhs in sub_exprs[1:]:
+            lhs = BinOp(op_token, lhs, rhs)
+        return lhs
+
+    @staticmethod
+    def build_IfExp(ctx, expr):
+        return TernaryIf(
+            build_expr(ctx, expr.test),
+            build_expr(ctx, expr.body),
+            build_expr(ctx, expr.orelse),
+        )
+
+    @staticmethod
+    def build_Compare(ctx, expr):
+        operands = [build_expr(ctx, e) for e in [expr.left] + list(expr.comparators)]
+        result = None
+        # pyrefly: ignore [bad-assignment]
+        for lhs, op_, rhs in zip(operands, expr.ops, operands[1:]):
+            op = type(op_)
+            op_token = ExprBuilder.cmpop_map.get(op)
+            r = ctx.make_raw_range(lhs.range().end, rhs.range().start)
+            if op_token is None:
+                raise NotSupportedError(
+                    r, "unsupported comparison operator: " + op.__name__
+                )
+
+            if op == ast.NotIn:
+                # NB: `not in` is just `not( in )`, so we don't introduce new tree view
+                # but just make it a nested call in our tree view structure
+                in_expr = BinOp("in", lhs, rhs)
+                cmp_expr = UnaryOp(r, "not", in_expr)
+            else:
+                cmp_expr = BinOp(op_token, lhs, rhs)  # type: ignore[assignment]
+
+            if result is None:
+                result = cmp_expr
+            else:
+                result = BinOp("and", result, cmp_expr)  # type: ignore[assignment]
+        return result
+
+    @staticmethod
+    def build_Subscript(ctx, expr):
+        def build_SliceExpr(ctx, base, slice_expr):
+            lower = (
+                build_expr(ctx, slice_expr.lower)
+                if slice_expr.lower is not None
+                else None
+            )
+            upper = (
+                build_expr(ctx, slice_expr.upper)
+                if slice_expr.upper is not None
+                else None
+            )
+            step = (
+                build_expr(ctx, slice_expr.step)
+                if slice_expr.step is not None
+                else None
+            )
+            return SliceExpr(base.range(), lower, upper, step)
+
+        def build_Index(ctx, base, index_expr):
+            if isinstance(index_expr.value, ast.Tuple):
+                raise NotSupportedError(
+                    base.range(),
+                    "slicing multiple dimensions with tuples not supported yet",
+                )
+            return build_expr(ctx, index_expr.value)
+
+        def build_ExtSlice(ctx, base, extslice):
+            sub_exprs = []
+            for expr in extslice.dims:
+                sub_type = type(expr)
+                if sub_type is ast.Index:
+                    sub_exprs.append(build_Index(ctx, base, expr))
+                elif sub_type is ast.Slice:
+                    sub_exprs.append(build_SliceExpr(ctx, base, expr))
+                elif sub_type is ast.Constant and expr.value is Ellipsis:
+                    sub_exprs.append(Dots(base.range()))
+                else:
+                    raise NotSupportedError(
+                        base.range(),
+                        f"slicing multiple dimensions with {sub_type} not supported",
+                    )
+            return sub_exprs
+
+        base = build_expr(ctx, expr.value)
+        sub_type = type(expr.slice)
+        if sub_type is ast.Index:
+            if isinstance(expr.slice.value, ast.Tuple):
+                # N-dimensional indexing using Tuple: x[(i, j, k)] is equivalent to x[i, j, k]
+                # XXX: Indexing using a list is **different**! It triggers advanced indexing.
+                indices = [
+                    build_expr(ctx, index_expr) for index_expr in expr.slice.value.elts
+                ]
+                if not indices:
+                    # `col_offset` is an int, but `end_col_offset` is
+                    # `Optional[int]`. The magic number is here to make
+                    # sure we can parse `()` on any machine
+                    r = ctx.make_range(
+                        expr.lineno,
+                        expr.slice.value.col_offset,
+                        expr.slice.value.col_offset + 2,
+                    )
+                    tup = TupleLiteral(r, [])
+                    indices.append(tup)
+                return Subscript(base, indices)
+            else:
+                return Subscript(base, [build_expr(ctx, expr.slice.value)])
+        elif sub_type is ast.Slice:
+            return Subscript(base, [build_SliceExpr(ctx, base, expr.slice)])
+        elif sub_type is ast.ExtSlice:
+            return Subscript(base, build_ExtSlice(ctx, base, expr.slice))
+        else:  # In Python3.9 array indices are not wrapped in ast.Index
+            if sub_type is ast.Tuple:
+                # N-dimensional indexing using Tuple: x[(i, j, k)] is equivalent to x[i, j, k]
+                indices = []
+                for index_expr in expr.slice.elts:
+                    if isinstance(index_expr, ast.Slice):
+                        indices.append(build_SliceExpr(ctx, base, index_expr))
+                    else:
+                        indices.append(build_expr(ctx, index_expr))
+                # Special-case logic for `typing.Tuple[()]`
+                if not indices:
+                    # See note above r.e. magic number
+                    r = ctx.make_range(
+                        expr.lineno, expr.slice.col_offset, expr.slice.col_offset + 2
+                    )
+                    tup = TupleLiteral(r, [])
+                    indices.append(tup)
+                return Subscript(base, indices)
+            return Subscript(base, [build_expr(ctx, expr.slice)])
+
+    @staticmethod
+    def build_List(ctx, expr):
+        return ListLiteral(
+            ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1),
+            [build_expr(ctx, e) for e in expr.elts],
+        )
+
+    @staticmethod
+    def build_Tuple(ctx, expr):
+        return TupleLiteral(
+            ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1),
+            [build_expr(ctx, e) for e in expr.elts],
+        )
+
+    @staticmethod
+    def build_Dict(ctx, expr):
+        range = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1)
+        if expr.keys and not expr.keys[0]:
+            raise NotSupportedError(
+                range, "Dict expansion (e.g. `{**dict}`) is not supported"
+            )
+        return DictLiteral(
+            range,
+            [build_expr(ctx, e) for e in expr.keys],
+            [build_expr(ctx, e) for e in expr.values],
+        )
+
+    @staticmethod
+    def build_Num(ctx, expr):
+        value = str(expr.value)
+        r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(value))
+        return Const(r, value)
+
+    @staticmethod
+    def build_Constant(ctx, expr):
+        value = expr.value
+        if value is None or isinstance(value, bool):
+            # NB: this check has to happen before the int check because bool is
+            # a subclass of int
+            return ExprBuilder.build_NameConstant(ctx, expr)
+        if isinstance(value, (int, float, complex)):
+            return ExprBuilder.build_Num(ctx, expr)
+        elif isinstance(value, str):
+            return ExprBuilder.build_Str(ctx, expr)
+        elif isinstance(value, type(Ellipsis)):
+            return ExprBuilder.build_Ellipsis(ctx, expr)
+        else:
+            error_range = ctx.make_range(
+                expr.lineno, expr.col_offset, expr.col_offset + len(str(value))
+            )
+            raise FrontendError(error_range, "Unknown Constant expression type")
+
+    @staticmethod
+    def build_Str(ctx, expr):
+        value = str(expr.value)
+        r = ctx.make_range(
+            expr.lineno, expr.col_offset, expr.col_offset + len(value) + 1
+        )
+        return StringLiteral(r, value)
+
+    @staticmethod
+    def build_JoinedStr(ctx, expr):
+        s = ""
+        args = []
+        for value in expr.values:
+            r = ctx.make_range(value.lineno, value.col_offset, value.col_offset + 1)
+            if isinstance(value, ast.FormattedValue):
+                if value.conversion != -1:
+                    raise NotSupportedError(r, "Don't support conversion in JoinedStr")
+                if value.format_spec is not None:
+                    raise NotSupportedError(r, "Don't support formatting in JoinedStr")
+                s += "{}"
+                args.append(build_expr(ctx, value.value))
+            elif isinstance(value, ast.Constant):
+                # pyrefly: ignore [unsupported-operation]
+                s += value.value
+            else:
+                raise NotSupportedError(r, "Unsupported value in JoinedStr")
+
+        r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1)
+        return Apply(Select(StringLiteral(r, s), Ident(r, "format")), args, [])
+
+    @staticmethod
+    def build_ListComp(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset)
+        if len(stmt.generators) != 1:
+            raise NotSupportedError(r, "Only a single generator is currently supported")
+
+        if len(stmt.generators[0].ifs) != 0:
+            raise NotSupportedError(r, "Comprehension ifs are not supported yet")
+
+        elt_expr = build_expr(ctx, stmt.elt)
+        target_expr = build_expr(ctx, stmt.generators[0].target)
+        iter_expr = build_expr(ctx, stmt.generators[0].iter)
+
+        return ListComp(r, elt_expr, target_expr, iter_expr)
+
+    @staticmethod
+    def build_GeneratorExp(ctx, stmt):
+        # Convert Generator expression to ListComp
+        return ExprBuilder.build_ListComp(ctx, stmt)
+
+    @staticmethod
+    def build_DictComp(ctx, stmt):
+        r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset)
+        if len(stmt.generators) != 1:
+            raise NotSupportedError(r, "Only a single generator is currently supported")
+
+        if len(stmt.generators[0].ifs) != 0:
+            raise NotSupportedError(r, "Comprehension ifs are not supported yet")
+
+        key_expr = build_expr(ctx, stmt.key)
+        value_expr = build_expr(ctx, stmt.value)
+        target_expr = build_expr(ctx, stmt.generators[0].target)
+        iter_expr = build_expr(ctx, stmt.generators[0].iter)
+
+        return DictComp(r, key_expr, value_expr, target_expr, iter_expr)
+
+    @staticmethod
+    def build_Starred(ctx, expr):
+        r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1)
+        return Starred(r, build_expr(ctx, expr.value))
+
+
+build_expr = ExprBuilder()
+build_stmt = StmtBuilder()
+build_withitem = WithItemBuilder()
+
+
+def find_before(ctx, pos, substr, offsets=(0, 0)):
+    new_pos = ctx.source[:pos].rindex(substr)
+    return ctx.make_raw_range(new_pos + offsets[0], new_pos + len(substr) + offsets[1])
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/generate_bytecode.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/generate_bytecode.py
new file mode 100644
index 0000000000000000000000000000000000000000..ebdceaddca92ad8ac40dfdebb50f9eaae596d42c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/generate_bytecode.py
@@ -0,0 +1,33 @@
+# mypy: allow-untyped-defs
+
+from torch._C import _compile_graph_to_code_table, _generate_upgraders_graph
+
+
+def format_bytecode(table):
+    # given a nested tuple, convert it to nested list
+    def listify(content):
+        if not isinstance(content, tuple):
+            return content
+        return [listify(i) for i in content]
+
+    formatted_table = {}
+    for entry in table:
+        identifier = entry[0]
+        content = entry[1]
+        content = listify(content)
+        formatted_table[identifier] = content
+    return formatted_table
+
+
+def generate_upgraders_bytecode() -> list:
+    yaml_content = []
+    upgraders_graph_map = _generate_upgraders_graph()
+    for upgrader_name, upgrader_graph in upgraders_graph_map.items():
+        bytecode_table = _compile_graph_to_code_table(upgrader_name, upgrader_graph)
+        entry = {upgrader_name: format_bytecode(bytecode_table)}
+        yaml_content.append(entry)
+    return yaml_content
+
+
+if __name__ == "__main__":
+    raise RuntimeError("This file is not meant to be run directly")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/mobile/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/mobile/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fbfc1c36480f9938a03e0fda34cc0300d0b7fb77
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/mobile/__init__.py
@@ -0,0 +1,238 @@
+# mypy: allow-untyped-defs
+import os
+
+import torch
+from torch.jit._serialization import validate_map_location
+
+
+def _load_for_lite_interpreter(f, map_location=None):
+    r"""
+    Load a :class:`LiteScriptModule` saved with :func:`torch.jit._save_for_lite_interpreter`.
+
+    Args:
+        f: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+        map_location: a string or torch.device used to dynamically remap
+            storages to an alternative set of devices.
+
+    Returns:
+        A :class:`LiteScriptModule` object.
+
+    Example:
+
+    .. testcode::
+
+        import torch
+        import io
+
+        # Load LiteScriptModule from saved file path
+        torch.jit._load_for_lite_interpreter('lite_script_module.pt')
+
+        # Load LiteScriptModule from io.BytesIO object
+        with open('lite_script_module.pt', 'rb') as f:
+            buffer = io.BytesIO(f.read())
+
+        # Load all tensors to the original device
+        torch.jit.mobile._load_for_lite_interpreter(buffer)
+    """
+    if isinstance(f, (str, os.PathLike)):
+        if not os.path.exists(f):
+            raise ValueError(f"The provided filename {f} does not exist")
+        if os.path.isdir(f):
+            raise ValueError(f"The provided filename {f} is a directory")
+
+    map_location = validate_map_location(map_location)
+
+    if isinstance(f, (str, os.PathLike)):
+        cpp_module = torch._C._load_for_lite_interpreter(os.fspath(f), map_location)
+    else:
+        cpp_module = torch._C._load_for_lite_interpreter_from_buffer(
+            f.read(),
+            map_location,
+        )
+
+    return LiteScriptModule(cpp_module)
+
+
+class LiteScriptModule:
+    def __init__(self, cpp_module) -> None:
+        self._c = cpp_module
+        super().__init__()
+
+    def __call__(self, *input):
+        return self._c.forward(input)
+
+    def find_method(self, method_name):
+        return self._c.find_method(method_name)
+
+    def forward(self, *input):
+        return self._c.forward(input)
+
+    def run_method(self, method_name, *input):
+        return self._c.run_method(method_name, input)
+
+
+def _export_operator_list(module: LiteScriptModule):
+    r"""Return a set of root operator names (with overload name) that are used by any method in this mobile module."""
+    return torch._C._export_operator_list(module._c)
+
+
+def _get_model_bytecode_version(f_input) -> int:
+    r"""Take a file-like object to return an integer.
+
+    Args:
+        f_input: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+
+    Returns:
+        version: An integer. If the integer is -1, the version is invalid. A warning
+            will show in the log.
+
+    Example:
+    .. testcode::
+
+        from torch.jit.mobile import _get_model_bytecode_version
+
+        # Get bytecode version from a saved file path
+        version = _get_model_bytecode_version("path/to/model.ptl")
+
+    """
+    if isinstance(f_input, (str, os.PathLike)):
+        if not os.path.exists(f_input):
+            raise ValueError(f"The provided filename {f_input} does not exist")
+        if os.path.isdir(f_input):
+            raise ValueError(f"The provided filename {f_input} is a directory")
+
+    if isinstance(f_input, (str, os.PathLike)):
+        return torch._C._get_model_bytecode_version(os.fspath(f_input))
+    else:
+        return torch._C._get_model_bytecode_version_from_buffer(f_input.read())
+
+
+def _get_mobile_model_contained_types(f_input) -> int:
+    r"""Take a file-like object and return a set of string, like ("int", "Optional").
+
+    Args:
+        f_input: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+
+    Returns:
+        type_list: A set of string, like ("int", "Optional"). These are types used in bytecode.
+
+    Example:
+
+    .. testcode::
+
+        from torch.jit.mobile import _get_mobile_model_contained_types
+
+        # Get type list from a saved file path
+        type_list = _get_mobile_model_contained_types("path/to/model.ptl")
+
+    """
+    if isinstance(f_input, (str, os.PathLike)):
+        if not os.path.exists(f_input):
+            raise ValueError(f"The provided filename {f_input} does not exist")
+        if os.path.isdir(f_input):
+            raise ValueError(f"The provided filename {f_input} is a directory")
+
+    if isinstance(f_input, (str, os.PathLike)):
+        return torch._C._get_mobile_model_contained_types(os.fspath(f_input))
+    else:
+        return torch._C._get_mobile_model_contained_types_from_buffer(f_input.read())
+
+
+def _backport_for_mobile(f_input, f_output, to_version):
+    r"""Take a input string containing a file name (file-like object) and a new destination to return a boolean.
+
+    Args:
+        f_input: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+        f_output: path to new model destination
+        to_version: the expected output model bytecode version
+    Returns:
+        success: A boolean. If backport success, return true, otherwise false
+    """
+    if isinstance(f_input, (str, os.PathLike)):
+        if not os.path.exists(f_input):
+            raise ValueError(f"The provided filename {f_input} does not exist")
+        if os.path.isdir(f_input):
+            raise ValueError(f"The provided filename {f_input} is a directory")
+
+    if (isinstance(f_input, (str, os.PathLike))) and (
+        isinstance(f_output, (str, os.PathLike))
+    ):
+        return torch._C._backport_for_mobile(
+            os.fspath(f_input),
+            os.fspath(f_output),
+            to_version,
+        )
+    else:
+        return torch._C._backport_for_mobile_from_buffer(
+            f_input.read(),
+            str(f_output),
+            to_version,
+        )
+
+
+def _backport_for_mobile_to_buffer(f_input, to_version):
+    r"""Take a string containing a file name (file-like object).
+
+    Args:
+        f_input: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+
+    """
+    if isinstance(f_input, (str, os.PathLike)):
+        if not os.path.exists(f_input):
+            raise ValueError(f"The provided filename {f_input} does not exist")
+        if os.path.isdir(f_input):
+            raise ValueError(f"The provided filename {f_input} is a directory")
+
+    if isinstance(f_input, (str, os.PathLike)):
+        return torch._C._backport_for_mobile_to_buffer(os.fspath(f_input), to_version)
+    else:
+        return torch._C._backport_for_mobile_from_buffer_to_buffer(
+            f_input.read(),
+            to_version,
+        )
+
+
+def _get_model_ops_and_info(f_input):
+    r"""Retrieve the root (top level) operators of a model and their corresponding compatibility info.
+
+    These root operators can call other operators within them (traced ops), and
+    a root op can call many different traced ops depending on internal code paths in the root op.
+    These traced ops are not returned by this function. Those operators are abstracted into the
+    runtime as an implementation detail (and the traced ops themselves can also call other operators)
+    making retrieving them difficult and their value from this api negligible since they will differ
+    between which runtime version the model is run on. Because of this, there is a false positive this
+    api can't prevent in a compatibility usecase. All the root ops of a model are present in a
+    target runtime, but not all the traced ops are which prevents a model from being able to run.
+    Args:
+        f_input: a file-like object (has to implement read, readline, tell, and seek),
+            or a string containing a file name
+
+    Returns:
+        Operators and info: A Dictionary mapping strings (the qualified names of the root operators)
+        of the model to their OperatorInfo structs.
+
+    Example:
+
+    .. testcode::
+
+        from torch.jit.mobile import _get_model_ops_and_info
+
+        # Get bytecode version from a saved file path
+        ops_and_info = _get_model_ops_and_info("path/to/model.ptl")
+
+    """
+    if isinstance(f_input, (str, os.PathLike)):
+        if not os.path.exists(f_input):
+            raise ValueError(f"The provided filename {f_input} does not exist")
+        if os.path.isdir(f_input):
+            raise ValueError(f"The provided filename {f_input} is a directory")
+
+    if isinstance(f_input, (str, os.PathLike)):
+        return torch._C._get_model_ops_and_info(os.fspath(f_input))
+    else:
+        return torch._C._get_model_ops_and_info(f_input.read())
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/quantized.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/quantized.py
new file mode 100644
index 0000000000000000000000000000000000000000..a2500c1f1b9fed481ab713953746294dc52668af
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/quantized.py
@@ -0,0 +1,100 @@
+# mypy: allow-untyped-defs
+import torch
+
+
+class QuantizedLinear(torch.jit.ScriptModule):
+    def __init__(self, other):
+        raise RuntimeError(
+            "torch.jit.QuantizedLinear is no longer supported. Please use "
+            "torch.ao.nn.quantized.dynamic.Linear instead."
+        )
+
+
+# FP16 weights
+class QuantizedLinearFP16(torch.jit.ScriptModule):
+    def __init__(self, other):
+        super().__init__()
+        raise RuntimeError(
+            "torch.jit.QuantizedLinearFP16 is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic.Linear instead."
+        )
+
+
+# Quantized RNN cell implementations
+class QuantizedRNNCellBase(torch.jit.ScriptModule):
+    def __init__(self, other):
+        raise RuntimeError(
+            "torch.jit.QuantizedRNNCellBase is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic.RNNCell instead."
+        )
+
+
+class QuantizedRNNCell(QuantizedRNNCellBase):
+    def __init__(self, other):
+        raise RuntimeError(
+            "torch.jit.QuantizedRNNCell is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic.RNNCell instead."
+        )
+
+
+class QuantizedLSTMCell(QuantizedRNNCellBase):
+    def __init__(self, other):
+        super().__init__(other)
+        raise RuntimeError(
+            "torch.jit.QuantizedLSTMCell is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic.LSTMCell instead."
+        )
+
+
+class QuantizedGRUCell(QuantizedRNNCellBase):
+    def __init__(self, other):
+        super().__init__(other)
+        raise RuntimeError(
+            "torch.jit.QuantizedGRUCell is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic.GRUCell instead."
+        )
+
+
+class QuantizedRNNBase(torch.jit.ScriptModule):
+    def __init__(self, other, dtype=torch.int8):
+        raise RuntimeError(
+            "torch.jit.QuantizedRNNBase is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic instead."
+        )
+
+
+class QuantizedLSTM(QuantizedRNNBase):
+    def __init__(self, other, dtype):
+        raise RuntimeError(
+            "torch.jit.QuantizedLSTM is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic.LSTM instead."
+        )
+
+
+class QuantizedGRU(QuantizedRNNBase):
+    def __init__(self, *args, **kwargs):
+        raise RuntimeError(
+            "torch.jit.QuantizedGRU is no longer supported. "
+            "Please use the torch.ao.nn.quantized.dynamic.GRU instead."
+        )
+
+
+def quantize_rnn_cell_modules(module):
+    raise RuntimeError(
+        "quantize_rnn_cell_modules function is no longer supported. "
+        "Please use torch.ao.quantization.quantize_dynamic API instead."
+    )
+
+
+def quantize_linear_modules(module, dtype=torch.int8):
+    raise RuntimeError(
+        "quantize_linear_modules function is no longer supported. "
+        "Please use torch.ao.quantization.quantize_dynamic API instead."
+    )
+
+
+def quantize_rnn_modules(module, dtype=torch.int8):
+    raise RuntimeError(
+        "quantize_rnn_modules function is no longer supported. "
+        "Please use torch.ao.quantization.quantize_dynamic API instead."
+    )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/supported_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/supported_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a258280ea352d993b5a56b739c28d70b0273e8a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/supported_ops.py
@@ -0,0 +1,346 @@
+# mypy: allow-untyped-defs
+import inspect
+import textwrap
+
+import torch.jit
+from torch.jit._builtins import _find_builtin
+
+
+# this file is for generating documentation using sphinx autodoc
+# > help(torch.jit.supported_ops) will also give a nice listed of the
+# supported ops programmatically
+
+
+def _hidden(name):
+    return name.startswith("_") and not name.startswith("__")
+
+
+def _emit_type(type):
+    return str(type)
+
+
+def _emit_arg(indent, i, arg):
+    v = f"{arg.name} : {_emit_type(arg.type)}"
+    default = arg.default_value
+    if default is not None:
+        v = f"{v}={str(default)}"
+    if i > 0:
+        v = f"\n{' ' * indent}{v}"
+    return v
+
+
+def _emit_args(indent, arguments):
+    return ",".join(_emit_arg(indent, i, arg) for i, arg in enumerate(arguments))
+
+
+def _emit_ret(ret):
+    return _emit_type(ret.type)
+
+
+def _emit_rets(returns):
+    if len(returns) == 1:
+        return _emit_ret(returns[0])
+    return f"Tuple[{', '.join(_emit_ret(r) for r in returns)}]"
+
+
+def _emit_schema(mod, name, schema, arg_start=0, padding=4):
+    if mod is None:
+        qualified_name = name
+    else:
+        qualified_name = f"{mod}.{name}"
+    schema_str = (
+        f"{qualified_name}"
+        f"({_emit_args(len(qualified_name) + 1 + padding, schema.arguments[arg_start:])}) "
+        f"-> {_emit_rets(schema.returns)}"
+    )
+    return schema_str
+
+
+def _get_tensor_ops():
+    def is_tensor_method(schema) -> bool:
+        if len(schema.arguments) == 0:
+            return False
+        self = schema.arguments[0]
+        if self.name != "self":
+            return False
+        if not self.type.isSubtypeOf(torch._C.TensorType.get()):
+            return False
+        return True
+
+    methods = []
+    # discover methods
+    for elem in dir(torch.Tensor):
+        if not _hidden(elem):
+            schemas = torch._C._jit_get_schemas_for_operator("aten::" + elem)
+            for schema in schemas:
+                if is_tensor_method(schema):
+                    methods.append(_emit_schema("Tensor", elem, schema, arg_start=1))
+
+    return "Supported Tensor Methods", methods
+
+
+def _get_nn_functional_ops():
+    functions = []
+
+    # Iterate over torch.nn.functional
+    mod = torch.nn.functional
+    name = mod.__name__
+    for elem in dir(torch.nn.functional):
+        attr = getattr(mod, elem)
+        if not inspect.isfunction(attr) or _hidden(elem[0]):
+            # Ignore non-functions and internal methods
+            continue
+
+        attr_module = inspect.getmodule(attr)
+        if not attr_module:
+            raise RuntimeError(f"Module for {attr} not found")
+
+        if "torch.nn.functional" not in attr_module.__name__:
+            # Ignore functions from outside torch.nn.functional
+            continue
+
+        try:
+            # compile fn, get schema
+            scripted = torch.jit.script(attr)
+            scripted_schema = scripted.schema
+            functions.append(_emit_schema(name, elem, scripted_schema))
+        except:  # noqa: B001,E722
+            # Skip interpolate / boolean dispatched things
+            pass
+
+    # Iterate over modules that we know contain a lot of builtins
+    for mod in torch.jit._builtins._modules_containing_builtins:
+        name = mod.__name__
+        for elem in dir(mod):
+            builtin = _find_builtin(getattr(mod, elem))
+            if builtin is not None:
+                schemas = torch._C._jit_get_schemas_for_operator(builtin)
+                for schema in schemas:
+                    # remove _tan but not __and__
+                    if not _hidden(elem):
+                        functions.append(_emit_schema(name, elem, schema))
+    return "Supported PyTorch Functions", functions
+
+
+def _get_builtins_helper():
+    builtins = []
+    for fn, _builtin_name in torch.jit._builtins._builtin_ops:
+        mod = inspect.getmodule(fn)
+
+        if not hasattr(fn, "__name__"):
+            # typing classes
+            continue
+        if not mod:
+            continue
+        if _hidden(fn.__name__) or _hidden(fn.__qualname__) or _hidden(mod.__name__):
+            # skip internal-only methods
+            continue
+
+        if "torch._C" in mod.__name__:
+            continue
+
+        builtins.append((fn, _builtin_name))
+
+    return builtins
+
+
+def _is_math_fn(fn):
+    mod = inspect.getmodule(fn)
+    if not mod:
+        raise RuntimeError(f"Module for {fn} not found")
+
+    return mod.__name__ == "math"
+
+
+def _get_torchscript_builtins():
+    functions = []
+    builtins = filter(lambda fn: not _is_math_fn(fn[0]), _get_builtins_helper())
+    builtins_list = list(builtins)
+    # Iterate over the specially added builtins
+    for fn, _builtin_name in builtins_list:
+        mod = inspect.getmodule(fn)
+        if not mod:
+            raise RuntimeError(f"Module for {fn} not found")
+        builtin = _find_builtin(fn)
+        if builtin is not None:
+            schemas = torch._C._jit_get_schemas_for_operator(builtin)
+            for schema in schemas:
+                functions.append(_emit_schema(mod.__name__, fn.__name__, schema))
+
+    return "TorchScript Builtin Functions", functions
+
+
+def _get_math_builtins():
+    functions = []
+    builtins = filter(lambda fn: _is_math_fn(fn[0]), _get_builtins_helper())
+    builtins_list = list(builtins)
+    # Iterate over the specially added builtins
+    for fn, _builtin_name in builtins_list:
+        mod = inspect.getmodule(fn)
+        if not mod:
+            raise RuntimeError(f"Module for {fn} not found")
+        builtin = _find_builtin(fn)
+        if builtin is not None:
+            schemas = torch._C._jit_get_schemas_for_operator(builtin)
+            for schema in schemas:
+                schema_str = _emit_schema(mod.__name__, fn.__name__, schema)
+                if "Tensor" in schema_str:
+                    # Skip Tensor ops that have the same name as math functions
+                    # (they will show up in the tensor methods section)
+                    continue
+                functions.append(schema)
+
+    return "``math`` Module", functions
+
+
+def _get_global_builtins():
+    # Taken from the 'globals' map in torch/csrc/jit/frontend/ir_emitter.cpp
+    supported_builtins = [
+        "print",
+        "tuple",
+        "float",
+        "complex",
+        "int",
+        "bool",
+        "str",
+        "getattr",
+        "hasattr",
+        "isinstance",
+        "len",
+        "hex",
+        "oct",
+        "round",
+        "hash",
+        "min",
+        "max",
+        "abs",
+        "all",
+        "divmod",
+        "list",
+        "ord",
+        "chr",
+        "bin",
+        "range",
+        "zip",
+        "enumerate",
+        "sorted",
+    ]
+
+    op_renames = {
+        "bool": "aten::Bool",
+        "int": "aten::Int",
+        "float": "aten::Float",
+        "complex": "aten::Complex",
+        "abs": "prim::abs",
+        "max": "prim::max",
+        "min": "prim::min",
+        "range": "fake::does_not_exist",
+    }
+
+    schemaless_op_explanations = {
+        "print": "Print any value",
+        "tuple": "Lists cannot be converted to tuples with this method since their size is not statically known",
+        "getattr": "Attribute name must be a literal string",
+        "hasattr": "Attribute name must be a literal string",
+        "isinstance": "Result is static",
+        "zip": "Arguments must be iterable.",
+        "enumerate": "Arguments must be iterable.",
+        "range": "Can only be used as an iterator in a for loop",
+    }
+
+    magic_methods = [
+        ("complex", "__complex__"),
+        ("float", "__float__"),
+        ("int", "__int__"),
+        ("bool", "__bool__"),
+        ("str", "__str__"),
+        ("len", "__len__"),
+        ("hex", "__hex__"),
+        ("oct", "__oct__"),
+    ]
+
+    magic_methods_rows = []
+    for fn, magic_method in magic_methods:
+        # pyrefly: ignore [bad-argument-type]
+        magic_methods_rows.append(f'"{fn}", "``{magic_method}``"')
+
+    schematized_ops = []
+    schemaless_ops = []
+
+    for fn in supported_builtins:
+        op_name = f"aten::{fn}"
+        if fn in op_renames:
+            op_name = op_renames[fn]
+        schemas = torch._C._jit_get_schemas_for_operator(op_name)
+        for s in schemas:
+            schematized_ops.append(_emit_schema(None, fn, s, padding=0))
+        if len(schemas) > 0:
+            schematized_ops.append("")
+        else:
+            table_row = (
+                f'":external+python:py:obj:`{fn}`", "{schemaless_op_explanations[fn]}"'
+            )
+            # pyrefly: ignore [bad-argument-type]
+            schemaless_ops.append(table_row)
+
+    schematized_ops_str = "\n".join(schematized_ops)
+    schemaless_ops_str = "\n".join(schemaless_ops)
+    magic_methods_rows_str = "\n".join(magic_methods_rows)
+    schematized_ops_str = textwrap.indent(schematized_ops_str, "\t")
+    schemaless_ops_str = textwrap.indent(schemaless_ops_str, "\t")
+    magic_methods_rows_str = textwrap.indent(magic_methods_rows_str, "\t")
+    section = f"""
+The functions in the following table are supported but do not have a static schema
+
+.. csv-table::
+    :header: "Function", "Note"
+
+{schemaless_ops_str}
+
+The following functions will use the corresponding magic method on TorchScript classes
+
+.. csv-table::
+    :header: "Function", "Magic Method"
+
+{magic_methods_rows_str}
+
+These built-in functions use the schema
+
+.. rst-class:: codeblock-height-limiter
+
+::
+
+{schematized_ops_str}
+    """
+
+    return "Python Built-in Functions", section
+
+
+def _list_supported_ops():
+    def emit_block(decls):
+        return "\n.. rst-class:: codeblock-height-limiter\n\n::\n\n{}\n".format(
+            "".join(f"    {d}\n\n" for d in decls)
+        )
+
+    body = ""
+    op_gathering_fns = (
+        _get_tensor_ops,
+        _get_nn_functional_ops,
+        _get_torchscript_builtins,
+        _get_global_builtins,
+        _get_math_builtins,
+    )
+    for fn in op_gathering_fns:
+        header, items = fn()
+        link_target = header.replace("`", "").replace("-", "").lower().replace(" ", "-")
+        if isinstance(items, str):
+            section = f"{header}\n{'~' * len(header)}\n{items}\n"
+        else:
+            section = f"{header}\n{'~' * len(header)}\n{emit_block(items)}"
+        section = f".. _{link_target}:" + "\n\n" + section
+        body += section
+
+    return body
+
+
+__doc__ = _list_supported_ops()
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/unsupported_tensor_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/unsupported_tensor_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..12bca0fce337d291a6e1c75861f0b3f020cc7a6c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/jit/unsupported_tensor_ops.py
@@ -0,0 +1,78 @@
+# mypy: allow-untyped-defs
+from textwrap import dedent
+from typing import Any
+
+import torch.jit
+
+
+def execWrapper(code, glob, loc) -> None:
+    exec(code, glob, loc)
+
+
+def _gen_unsupported_methods_properties():
+    tensor_attrs = set(filter(lambda x: x[0] != "_", dir(torch.Tensor)))
+    tensor = torch.tensor([2])
+    funcs_template = dedent(
+        """
+    def func(x):
+        return x.{op}()
+    """
+    )
+
+    deprecated_apis = {
+        "volatile",
+        "resize",
+        "reinforce",
+        "new",
+        "name",
+        "map2_",
+        "has_names",
+        "grad_fn",
+        "resize_as",
+    }
+    tensor_attrs = tensor_attrs - deprecated_apis
+
+    properties = []
+    methods = []
+    sorted_tensor_attrs = sorted(tensor_attrs, key=lambda x: x.lower())
+    for attr in sorted_tensor_attrs:
+        funcs_str = funcs_template.format(op=attr)
+        scope: dict[str, Any] = {}
+        execWrapper(funcs_str, globals(), scope)
+        try:
+            torch.jit.CompilationUnit(funcs_str)
+        except Exception as e:
+            if "nonexistent attribute" not in repr(e):
+                continue
+            attr_repr = repr(getattr(tensor, attr))
+            if "bound method" in attr_repr or "built-in method" in attr_repr:
+                methods.append(attr)
+            else:
+                properties.append(attr)
+
+    mapped_methods = ("\t*  :meth:`~torch.Tensor." + x + r"`" for x in methods)
+    mapped_properties = ("\t*  :attr:`~torch.Tensor." + x + r"`" for x in properties)
+    return "\n".join(mapped_methods), "\n".join(mapped_properties)
+
+
+def _list_unsupported_tensor_ops():
+    header = """\n\n
+Unsupported Tensor Methods
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    """
+    methods, properties = _gen_unsupported_methods_properties()
+    return (
+        header
+        + "\n"
+        + methods
+        + """
+
+Unsupported Tensor Properties
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    """
+        + "\n"
+        + properties
+    )
+
+
+__doc__ = _list_unsupported_tensor_ops()
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10.so
new file mode 100644
index 0000000000000000000000000000000000000000..9ff3def1157f527caba7c00fa508765a2508c747
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d72c6da442fdaf1b4b3916b2efb9c5b1ba3564f5b54eefb3e9319907dbdfcefc
+size 1423328
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10_cuda.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10_cuda.so
new file mode 100644
index 0000000000000000000000000000000000000000..20c2a6d50ed5b069aad67e829a5367fa33853b91
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libc10_cuda.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8dd4112d02b9d479f3d81955ef9453ce04c36b93c21aabad46e9296385f931d3
+size 589080
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libcaffe2_nvrtc.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libcaffe2_nvrtc.so
new file mode 100644
index 0000000000000000000000000000000000000000..94548e95b327848d0f19b649c17ce65c06c952d3
Binary files /dev/null and b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libcaffe2_nvrtc.so differ
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm.so
new file mode 100644
index 0000000000000000000000000000000000000000..8338a1fd037d854b0fff6236cfaf150c6861ee64
Binary files /dev/null and b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm.so differ
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/alloc_info.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/alloc_info.h
new file mode 100644
index 0000000000000000000000000000000000000000..e441ff5a28936d8ca999fcb61ddc8dbbb2c8c12b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/alloc_info.h
@@ -0,0 +1,9 @@
+#pragma once
+
+#include <unistd.h>
+
+struct AllocInfo {
+  pid_t pid;
+  char free;
+  char filename[60];
+};
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/err.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/err.h
new file mode 100644
index 0000000000000000000000000000000000000000..e1e6aa4e277c3a94dd642ff2a27e6cd564322e46
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/err.h
@@ -0,0 +1,25 @@
+#pragma once
+
+#include <cerrno>
+#include <system_error>
+
+// `errno` is only meaningful when it fails. E.g., a  successful `fork()` sets
+// `errno` to `EINVAL` in child process on some macos
+// (https://stackoverflow.com/a/20295079), and thus `errno` should really only
+// be inspected if an error occurred.
+//
+// All functions used in `libshm` (so far) indicate error by returning `-1`. If
+// you want to use a function with a different error reporting mechanism, you
+// need to port `SYSCHECK` from `torch/lib/c10d/Utils.hpp`.
+#define SYSCHECK_ERR_RETURN_NEG1(expr)                          \
+  while (true) {                                                \
+    if ((expr) == -1) {                                         \
+      if (errno == EINTR) {                                     \
+        continue;                                               \
+      } else {                                                  \
+        throw std::system_error(errno, std::system_category()); \
+      }                                                         \
+    } else {                                                    \
+      break;                                                    \
+    }                                                           \
+  }
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/libshm.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/libshm.h
new file mode 100644
index 0000000000000000000000000000000000000000..d3f7c7061abc9e56b7147fad7e85d1bcdacc61c8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/libshm.h
@@ -0,0 +1,46 @@
+#pragma once
+
+#include <ATen/MapAllocator.h>
+
+#ifdef __cplusplus
+
+void libshm_init(const char* manager_exec_path);
+
+// Superclass to run a constructor before at::RefcountedMapAllocator
+class THManagedMapAllocatorInit {
+ protected:
+  THManagedMapAllocatorInit(const char* manager_handle, const char* filename);
+  std::string manager_handle_;
+};
+
+// Like a at::RefcountedMapAllocator, but it also makes use of an external
+// shared memory manager process to ensure that shared memory regions actually
+// get freed in the end (even if processes lose the memory).
+class THManagedMapAllocator : private THManagedMapAllocatorInit,
+                              public at::RefcountedMapAllocator {
+ public:
+  THManagedMapAllocator(
+      const char* manager_handle,
+      const char* filename,
+      int flags,
+      size_t size);
+
+  void close() override;
+
+  ~THManagedMapAllocator() override {
+    close();
+  }
+
+  static at::DataPtr makeDataPtr(
+      const char* manager_handle,
+      const char* filename,
+      int flags,
+      size_t size);
+  static THManagedMapAllocator* fromDataPtr(const at::DataPtr& /*dptr*/);
+
+  const char* manager_handle() const {
+    return manager_handle_.c_str();
+  }
+};
+
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/socket.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/socket.h
new file mode 100644
index 0000000000000000000000000000000000000000..e048098b94efac3360d4d72835d60b346fab4842
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm/socket.h
@@ -0,0 +1,164 @@
+#pragma once
+
+#include <poll.h>
+#include <sys/socket.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <sys/un.h>
+#include <unistd.h>
+#include <cstddef>
+#include <cstdio>
+#include <cstring>
+#include <string>
+
+#include <libshm/alloc_info.h>
+#include <libshm/err.h>
+
+class Socket {
+ public:
+  int socket_fd;
+  Socket(const Socket& other) = delete;
+
+ protected:
+  Socket() {
+    SYSCHECK_ERR_RETURN_NEG1(socket_fd = socket(AF_UNIX, SOCK_STREAM, 0));
+  }
+  Socket(Socket&& other) noexcept : socket_fd(other.socket_fd) {
+    other.socket_fd = -1;
+  };
+  explicit Socket(int fd) : socket_fd(fd) {}
+
+  virtual ~Socket() {
+    if (socket_fd != -1)
+      close(socket_fd);
+  }
+
+  struct sockaddr_un prepare_address(const char* path) {
+    struct sockaddr_un address;
+    address.sun_family = AF_UNIX;
+    strcpy(address.sun_path, path);
+    return address;
+  }
+
+  // Implemented based on https://man7.org/linux/man-pages/man7/unix.7.html
+  size_t address_length(struct sockaddr_un address) {
+    return offsetof(sockaddr_un, sun_path) + strlen(address.sun_path) + 1;
+  }
+
+  void recv(void* _buffer, size_t num_bytes) {
+    char* buffer = (char*)_buffer;
+    size_t bytes_received = 0;
+    ssize_t step_received;
+    struct pollfd pfd = {};
+    pfd.fd = socket_fd;
+    pfd.events = POLLIN;
+    while (bytes_received < num_bytes) {
+      SYSCHECK_ERR_RETURN_NEG1(poll(&pfd, 1, 1000));
+      if (pfd.revents & POLLIN) {
+        SYSCHECK_ERR_RETURN_NEG1(
+            step_received =
+                ::read(socket_fd, buffer, num_bytes - bytes_received));
+        TORCH_CHECK(step_received != 0, "Other end has closed the connection");
+        bytes_received += step_received;
+        buffer += step_received;
+      } else if (pfd.revents & (POLLERR | POLLHUP)) {
+        TORCH_CHECK(false, "An error occurred while waiting for the data");
+      } else {
+        TORCH_CHECK(false, "Shared memory manager connection has timed out");
+      }
+    }
+  }
+
+  void send(const void* _buffer, size_t num_bytes) {
+    const char* buffer = (const char*)_buffer;
+    size_t bytes_sent = 0;
+    ssize_t step_sent;
+    while (bytes_sent < num_bytes) {
+      SYSCHECK_ERR_RETURN_NEG1(
+          step_sent = ::write(socket_fd, buffer, num_bytes));
+      bytes_sent += step_sent;
+      buffer += step_sent;
+    }
+  }
+};
+
+class ManagerSocket : public Socket {
+ public:
+  explicit ManagerSocket(int fd) : Socket(fd) {}
+
+  AllocInfo receive() {
+    AllocInfo info;
+    recv(&info, sizeof(info));
+    return info;
+  }
+
+  void confirm() {
+    send("OK", 2);
+  }
+};
+
+class ManagerServerSocket : public Socket {
+ public:
+  explicit ManagerServerSocket(const std::string& path) {
+    socket_path = path;
+    try {
+      struct sockaddr_un address = prepare_address(path.c_str());
+      size_t len = address_length(address);
+      SYSCHECK_ERR_RETURN_NEG1(
+          bind(socket_fd, (struct sockaddr*)&address, len));
+      SYSCHECK_ERR_RETURN_NEG1(listen(socket_fd, 10));
+    } catch (std::exception&) {
+      SYSCHECK_ERR_RETURN_NEG1(close(socket_fd));
+      throw;
+    }
+  }
+
+  void remove() {
+    struct stat file_stat;
+    if (fstat(socket_fd, &file_stat) == 0)
+      SYSCHECK_ERR_RETURN_NEG1(unlink(socket_path.c_str()));
+  }
+
+  ~ManagerServerSocket() override {
+    unlink(socket_path.c_str());
+  }
+
+  ManagerSocket accept() {
+    int client_fd;
+    struct sockaddr_un addr;
+    socklen_t addr_len = sizeof(addr);
+    SYSCHECK_ERR_RETURN_NEG1(
+        client_fd = ::accept(socket_fd, (struct sockaddr*)&addr, &addr_len));
+    return ManagerSocket(client_fd);
+  }
+
+  std::string socket_path;
+};
+
+class ClientSocket : public Socket {
+ public:
+  explicit ClientSocket(const std::string& path) {
+    try {
+      struct sockaddr_un address = prepare_address(path.c_str());
+      size_t len = address_length(address);
+      SYSCHECK_ERR_RETURN_NEG1(
+          connect(socket_fd, (struct sockaddr*)&address, len));
+    } catch (std::exception&) {
+      SYSCHECK_ERR_RETURN_NEG1(close(socket_fd));
+      throw;
+    }
+  }
+
+  void register_allocation(AllocInfo& info) {
+    char buffer[3] = {0, 0, 0};
+    send(&info, sizeof(info));
+    recv(buffer, 2);
+    TORCH_CHECK(
+        strcmp(buffer, "OK") == 0,
+        "Shared memory manager didn't respond with an OK");
+  }
+
+  void register_deallocation(AllocInfo& info) {
+    send(&info, sizeof(info));
+  }
+};
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm_windows/libshm.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm_windows/libshm.h
new file mode 100644
index 0000000000000000000000000000000000000000..4dd193df93d110e3a04d33a3f9d3e3ec24948277
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libshm_windows/libshm.h
@@ -0,0 +1,36 @@
+#pragma once
+
+#include <ATen/MapAllocator.h>
+
+#ifdef __cplusplus
+
+#ifdef SHM_EXPORTS
+#define SHM_API __declspec(dllexport)
+#else
+#define SHM_API __declspec(dllimport)
+#endif
+
+SHM_API void libshm_init(const char* manager_exec_path);
+
+class SHM_API THManagedMapAllocator : public at::RefcountedMapAllocator {
+ public:
+  THManagedMapAllocator(
+      const char* manager_handle,
+      const char* filename,
+      int flags,
+      size_t size)
+      : at::RefcountedMapAllocator(filename, flags, size) {}
+
+  static at::DataPtr makeDataPtr(
+      const char* manager_handle,
+      const char* filename,
+      int flags,
+      size_t size);
+  static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
+
+  const char* manager_handle() const {
+    return "no_manager";
+  }
+};
+
+#endif
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch.so
new file mode 100644
index 0000000000000000000000000000000000000000..3ca0a5eeb4b3d10434d0410e9a7897acd6ba2301
Binary files /dev/null and b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch.so differ
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
new file mode 100644
index 0000000000000000000000000000000000000000..223af0e5bc5d0d2b3699be0fb34e6ffea1ecc685
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb7d16aa6557493538ec9f41030d6e29cb37103e5db9abb31213365f9285eff3
+size 320251880
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
new file mode 100644
index 0000000000000000000000000000000000000000..a650f6e52af5b11d3e5f669cefd9a40271ae89e2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebb958be5a08cadec613ae9a87f1dc3d823909b56a319758482ec8e04391bf7b
+size 539185872
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda_linalg.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda_linalg.so
new file mode 100644
index 0000000000000000000000000000000000000000..1f6b00f1f62b4c0e23ed5b43b7e7e9e21be22c19
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda_linalg.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03c403dab767c9747e406e61224c7ffe6dde7bbd179fb2414f6aa51a4609004b
+size 574312
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_global_deps.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_global_deps.so
new file mode 100644
index 0000000000000000000000000000000000000000..7045d9b20475378bb18ee6eb8754b082f6fee692
Binary files /dev/null and b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_global_deps.so differ
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_python.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_python.so
new file mode 100644
index 0000000000000000000000000000000000000000..74271415ebe38d5cb15aacc194e5f16619abdb79
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_python.so
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfe3670763f5b340e1db3214d5651cb599803ddd2f18c80705d1aa65cdeb9a68
+size 31291720
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/library.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/library.py
new file mode 100644
index 0000000000000000000000000000000000000000..377c464e76a16c17a5462a0a4b22c8b0bc407440
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/library.py
@@ -0,0 +1,1757 @@
+# mypy: allow-untyped-defs
+import contextlib
+import functools
+import inspect
+import re
+import sys
+import weakref
+from collections.abc import Callable, Sequence
+from typing import Any, overload, TYPE_CHECKING, TypeVar, Union
+from typing_extensions import deprecated, ParamSpec
+
+import torch
+import torch._library as _library
+from torch._library.custom_ops import (
+    _cast,
+    _maybe_get_opdef,
+    custom_op,
+    CustomOpDef,
+    device_types_t,
+)
+from torch._library.effects import EffectType
+from torch._library.infer_schema import infer_schema  # noqa: F401
+from torch._library.triton import triton_op, wrap_triton
+from torch._ops import OpOverload
+from torch.types import _dtype
+
+
+__all__ = [
+    "Library",
+    "impl",
+    "define",
+    "fallthrough_kernel",
+    "impl_abstract",
+    "register_autocast",
+    "register_fake",
+    "register_torch_dispatch",
+    "register_vmap",
+    "get_ctx",
+    "get_kernel",
+    "custom_op",
+    "triton_op",
+    "wrap_triton",
+    "infer_schema",
+]
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+# Set containing the combination of (namespace, operator, DispatchKey) for which a new kernel has been registered
+# The keys in the set are of the form `namespace + "/" + op_name + "/" + dispatch_key`.
+# This set is maintained to ensure that two libraries don't try to override the exact same functionality to avoid
+# libraries calling into kernels not intended to be called.
+_impls: set[str] = set()
+_defs: set[str] = set()
+
+# prim is reserved by TorchScript interpreter
+_reserved_namespaces = ["prim"]
+
+
+def fallthrough_kernel():
+    """
+    A dummy function to pass to ``Library.impl`` in order to register a fallthrough.
+    """
+    raise NotImplementedError("fallthrough_kernel() should never be called.")
+
+
+class Library:
+    """
+    A class to create libraries that can be used to register new operators or
+    override operators in existing libraries from Python.
+    A user can optionally pass in a dispatch keyname if they only want to register
+    kernels corresponding to only one specific dispatch key.
+
+    To create a library to override operators in an existing library (with name ns), set the kind to "IMPL".
+    To create a new library (with name ns) to register new operators, set the kind to "DEF".
+    To create a fragment of a possibly existing library to register operators (and bypass
+    the limitation that there is only one library for a given namespace), set the kind to
+    "FRAGMENT".
+
+    Args:
+        ns: library name
+        kind: "DEF", "IMPL", "FRAGMENT"
+        dispatch_key: PyTorch dispatch key (default: "")
+    """
+
+    def __init__(self, ns, kind, dispatch_key=""):
+        from torch.fx.operator_schemas import _SCHEMA_TO_SIGNATURE_CACHE
+
+        if kind not in ("IMPL", "DEF", "FRAGMENT"):
+            raise ValueError(f"Unsupported kind: {kind}")
+
+        if ns in _reserved_namespaces and (kind == "DEF" or kind == "FRAGMENT"):
+            raise ValueError(
+                f"{ns} is a reserved namespace. Please try creating a library with another name."
+            )
+
+        f = sys._getframe(1)
+        filename, lineno = f.f_code.co_filename, f.f_lineno
+        self.m: Any | None = torch._C._dispatch_library(
+            kind, ns, dispatch_key, filename, lineno
+        )
+        self.ns = ns
+        self._op_defs: set[str] = set()
+        self._op_impls: set[str] = set()
+        self._registration_handles: list[torch._library.utils.RegistrationHandle] = []
+        self.kind = kind
+        self.dispatch_key = dispatch_key
+        # Use a finalizer to setup the "destructor" instead of __del__.
+        # Python __del__ can lead to weird things (globals and locals may already
+        # be gone when __del__ actually gets called!). finalizers help the
+        # situation because it lets us capture references and keeps them alive
+        weakref.finalize(
+            self,
+            _del_library,
+            _impls,
+            self._op_impls,
+            _defs,
+            self._op_defs,
+            self._registration_handles,
+            self.m,
+            _SCHEMA_TO_SIGNATURE_CACHE,
+        )
+
+    def __repr__(self):
+        return f"Library(kind={self.kind}, ns={self.ns}, dispatch_key={self.dispatch_key})>"
+
+    def define(self, schema, alias_analysis="", *, tags=()):
+        r"""Defines a new operator and its semantics in the ns namespace.
+
+        Args:
+            schema: function schema to define a new operator.
+            alias_analysis (optional): Indicates if the aliasing properties of the operator arguments can be
+                                       inferred from the schema (default behavior) or not ("CONSERVATIVE").
+            tags (Tag | Sequence[Tag]): one or more torch.Tag to apply to this
+                                       operator. Tagging an operator changes the operator's behavior
+                                       under various PyTorch subsystems; please read the docs for the
+                                       torch.Tag carefully before applying it.
+
+        Returns:
+            name of the operator as inferred from the schema.
+
+        Example::
+
+            >>> my_lib = Library("mylib", "DEF")
+            >>> my_lib.define("sum(Tensor self) -> Tensor")
+        """
+
+        # This is added because we also want to disallow PURE_FUNCTION alias analysis which is a valid
+        # AliasAnalysis type in C++
+        if alias_analysis not in ["", "FROM_SCHEMA", "CONSERVATIVE"]:
+            raise RuntimeError(f"Invalid alias_analysis type {alias_analysis}")
+        if self.m is None:
+            raise AssertionError("Library object has been destroyed")
+        if isinstance(tags, torch.Tag):
+            tags = (tags,)
+
+        name = schema.split("(")[0]
+        packet_name = name.split(".")[0] if "." in name else name
+        has_preexisting_packet = hasattr(torch.ops, self.ns) and hasattr(
+            getattr(torch.ops, self.ns), packet_name
+        )
+
+        result = self.m.define(schema, alias_analysis, tuple(tags))
+        name = schema.split("(")[0]
+        qualname = self.ns + "::" + name
+
+        # If the OpOverloadPacket exists already, then this means we're adding a
+        # new OpOverload for it. Refresh the packet to include the new OpOverload.
+        if has_preexisting_packet:
+            ns = getattr(torch.ops, self.ns)
+            packet = getattr(ns, packet_name)
+            torch._ops._refresh_packet(packet)
+
+        self._op_defs.add(qualname)
+        _defs.add(qualname)
+        return result
+
+    def _register_fake(self, op_name, fn, _stacklevel=1, *, allow_override=False):
+        r"""Registers the fake impl for an operator defined in the library."""
+
+        source = torch._library.utils.get_source(_stacklevel + 1)
+        frame = sys._getframe(_stacklevel)
+        caller_module = inspect.getmodule(frame)
+        # Can be none if you call register_fake from somewhere there isn't a module
+        # (e.g. __main__)
+        caller_module_name = None if caller_module is None else caller_module.__name__
+
+        # TODO(rzou): We're gonna need to stage this change with torchvision,
+        # since torchvision is github first.
+        if caller_module_name is not None and caller_module_name.startswith(
+            "torchvision."
+        ):
+            caller_module_name = None
+
+        qualname = f"{self.ns}::{op_name}"
+        entry = torch._library.simple_registry.singleton.find(qualname)
+        if caller_module_name is not None:
+            func_to_register = _check_pystubs_once(fn, qualname, caller_module_name)
+        else:
+            func_to_register = fn
+
+        handle = entry.fake_impl.register(
+            func_to_register, source, lib=self, allow_override=allow_override
+        )
+        self._registration_handles.append(handle)
+
+    def _register_torch_dispatch_rule(self, op_name, torch_dispatch_class, fn):
+        r"""Registers a torch_dispatch rule for the given operator and torch_dispatch_class.
+
+        This allows for open registration to specify the behavior between the operator
+        and the torch_dispatch_class without needing to modify the torch_dispatch_class
+        or the operator directly.
+
+        The torch_dispatch_class is either a Tensor subclass with `__torch_dispatch__` or a
+        TorchDispatchMode.
+
+        If it is a Tensor subclass, we expect fn to have the following signature:
+        (cls, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any
+
+        If it is a TorchDispatchMode, we expect fn to have the following signature:
+        (mode, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any
+        """
+
+        qualname = f"{self.ns}::{op_name}"
+        entry = torch._library.simple_registry.singleton.find(qualname)
+        handle = entry.torch_dispatch_rules.register(torch_dispatch_class, fn)
+        self._registration_handles.append(handle)
+
+    def _impl_with_aoti_compile(self, op_name, dispatch_key=""):
+        r"""Register the operator to use the AOTI-compiled implementation.
+
+        Args:
+            op_name: operator name (along with the overload) or OpOverload object.
+            dispatch_key: dispatch key that the input function should be registered for. By default, it uses
+                          the dispatch key that the library was created with.
+
+        Example::
+
+            >>> my_lib = Library("aten", "IMPL")
+            >>> my_lib._impl_with_aoti_compile("div.Tensor", "CPU")
+        """
+
+        if dispatch_key == "":
+            dispatch_key = self.dispatch_key
+        # pyrefly: ignore [bad-argument-type]
+        if not torch.DispatchKeySet(dispatch_key).has(torch._C.DispatchKey.Dense):
+            raise AssertionError(
+                f"dispatch_key {dispatch_key} does not have Dense in its keyset"
+            )
+
+        if isinstance(op_name, str):
+            name = op_name
+        elif isinstance(op_name, OpOverload):
+            name = op_name._schema.name
+            overload_name = op_name._schema.overload_name
+            if overload_name != "":
+                name = name + "." + overload_name
+        else:
+            raise RuntimeError(
+                "_impl_with_aoti_compile should be passed either a name or an OpOverload object "
+                "as the first argument"
+            )
+
+        key = self.ns + "/" + name.split("::")[-1] + "/" + dispatch_key
+        if key in _impls:
+            # TODO: in future, add more info about where the existing function is registered (this info is
+            # today already returned by the C++ warning when _impl_with_aoti_compile is called but we error out before that)
+            raise RuntimeError(
+                "This is not allowed since there's already a kernel registered from python overriding {}"
+                "'s behavior for {} dispatch key and {} namespace.".format(
+                    name.split("::")[-1], dispatch_key, self.ns
+                )
+            )
+
+        if self.m is None:
+            raise AssertionError("Library object has been destroyed")
+        impl_fn: Callable = self.m.impl_with_aoti_compile
+        impl_fn(self.ns, name.split("::")[-1], dispatch_key)
+
+        _impls.add(key)
+        self._op_impls.add(key)
+
+    def impl(
+        self, op_name, fn, dispatch_key="", *, with_keyset=False, allow_override=False
+    ):
+        r"""Registers the function implementation for an operator defined in the library.
+
+        Args:
+            op_name: operator name (along with the overload) or OpOverload object.
+            fn: function that's the operator implementation for the input dispatch key or :func:`~fallthrough_kernel`
+                to register a fallthrough.
+            dispatch_key: dispatch key that the input function should be registered for. By default, it uses
+                          the dispatch key that the library was created with.
+            with_keyset: flag controlling if the current dispatcher call keyset should be passed as the first argument
+                         to :attr:`fn` when calling. This should be used to create the appropriate keyset for redispatch calls.
+            allow_override: Flag controlling if we want to override an
+                         existing registered kernel implementation. This is by
+                         default off, and will error you're trying to register a
+                         kernel to a dispatch key with a kernel already
+                         registered.
+
+        Example::
+            >>> # xdoctest: +SKIP("Requires Python <= 3.11")
+            >>> my_lib = Library("aten", "IMPL")
+            >>> def div_cpu(self, other):
+            >>>     return self * (1 / other)
+            >>> my_lib.impl("div.Tensor", div_cpu, "CPU")
+        """
+
+        if not callable(fn):
+            raise TypeError(
+                f"Input function is required to be a callable but found type {type(fn)}"
+            )
+        if dispatch_key == "":
+            dispatch_key = self.dispatch_key
+
+        if isinstance(op_name, str):
+            name = op_name
+        elif isinstance(op_name, OpOverload):
+            name = op_name._schema.name
+            overload_name = op_name._schema.overload_name
+            if overload_name != "":
+                name = name + "." + overload_name
+        else:
+            raise RuntimeError(
+                "impl should be passed either a name or an OpOverload object as the first argument"
+            )
+
+        key = self.ns + "/" + name.split("::")[-1] + "/" + dispatch_key
+        if (not allow_override) and key in _impls:
+            # TODO: in future, add more info about where the existing function is registered (this info is
+            # today already returned by the C++ warning when impl is called but we error out before that)
+            raise RuntimeError(
+                "This is not allowed since there's already a kernel registered from python overriding {}"
+                "'s behavior for {} dispatch key and {} namespace.".format(
+                    name.split("::")[-1], dispatch_key, self.ns
+                )
+            )
+
+        if dispatch_key == "Meta":
+            dispatcher_op_name = name
+            if "::" not in dispatcher_op_name:
+                dispatcher_op_name = f"{self.ns}::{dispatcher_op_name}"
+
+            # Internally, we shouldn't be registering meta kernels for any operators that
+            # have CompositeImplicitAutograd kernels.
+            # Instead, we should be letting those decompositions run, and writing meta kernels
+            # only for the base operators.
+            if torch._C._dispatch_has_kernel_for_dispatch_key(
+                dispatcher_op_name, "CompositeImplicitAutograd"
+            ):
+                raise RuntimeError(
+                    f"We should not register a meta kernel directly to the operator '{name}',"
+                    " because it has a CompositeImplicitAutograd kernel in core."
+                    " Instead we should let the operator decompose, and ensure that we have meta kernels"
+                    " for the base ops that it decomposes into."
+                )
+
+        if self.m is None:
+            raise AssertionError("Library object has been destroyed")
+        self.m.impl(
+            name,
+            dispatch_key if dispatch_key != "" else "CompositeImplicitAutograd",
+            fn,
+            with_keyset,
+        )
+
+        _impls.add(key)
+        self._op_impls.add(key)
+
+    def fallback(self, fn, dispatch_key="", *, with_keyset=False):
+        r"""Registers the function implementation as the fallback for the given key.
+
+        This function only works for a library with global namespace ("_").
+
+        Args:
+            fn: function used as fallback for the given dispatch key or :func:`~fallthrough_kernel`
+                to register a fallthrough.
+            dispatch_key: dispatch key that the input function should be registered for. By default, it uses
+                          the dispatch key that the library was created with.
+            with_keyset: flag controlling if the current dispatcher call keyset should be passed as the first argument
+                         to :attr:`fn` when calling. This should be used to create the appropriate keyset for redispatch calls.
+
+        Example::
+
+            >>> my_lib = Library("_", "IMPL")
+            >>> def fallback_kernel(op, *args, **kwargs):
+            >>>     # Handle all autocast ops generically
+            >>>     # ...
+            >>> my_lib.fallback(fallback_kernel, "Autocast")
+        """
+
+        if dispatch_key == "":
+            dispatch_key = self.dispatch_key
+
+        if self.ns != "_":
+            raise RuntimeError(
+                f"""Fallback can only be registered using library fragment on the global namespace "_" but it is {self.ns}"""
+            )
+
+        if dispatch_key == "":
+            raise AssertionError("dispatch_key must not be empty for fallback")
+        if self.m is None:
+            raise AssertionError("Library object has been destroyed")
+
+        self.m.fallback(dispatch_key, fn, with_keyset)
+
+    def _register_effectful_op(self, op_name: str, effect: EffectType | None):
+        """
+        Registers an effect to an operator. This is used to register an op that
+        has side effects that is not capturable by the schema.
+
+        Args:
+            op_name: operator name (along with the overload) or OpOverload object.
+            effect: The effect of the op.
+        """
+        from torch._higher_order_ops.effects import (
+            _register_effectful_op as hoo_register_effect,
+        )
+
+        handle = hoo_register_effect(op_name, effect)
+        self._registration_handles.append(handle)
+
+    def _destroy(self):
+        if self.m is not None:
+            self.m.reset()
+        self.m = None
+        for handle in self._registration_handles:
+            handle.destroy()
+        self._registration_handles.clear()
+        global _impls
+        _impls -= self._op_impls
+        for name in self._op_defs:
+            # Delete the cached torch.ops.ns.foo if it was registered.
+            # Otherwise, accessing it leads to a segfault.
+            # It's possible that we only registered an overload in this Library
+            # and another library owns an alive overload.
+            # That's OK - the next time torch.ops.ns.foo gets called, it'll be
+            # recomputed to point at the right collection of overloads.
+            ns, name_with_overload = name.split("::")
+            name = name_with_overload.split(".")[0]
+            if not hasattr(torch.ops, ns):
+                continue
+            namespace = getattr(torch.ops, ns)
+            if not hasattr(namespace, name):
+                continue
+            delattr(namespace, name)
+            namespace._dir.remove(name)
+
+
+def _del_library(
+    captured_impls,
+    op_impls,
+    captured_defs,
+    op_defs,
+    registration_handles,
+    m,
+    schema_to_signature_cache,
+):
+    for op_def in op_defs:
+        name = op_def
+        overload_name = ""
+        if "." in op_def:
+            name, overload_name = op_def.split(".")
+        if (
+            name,
+            overload_name,
+        ) in schema_to_signature_cache:
+            del schema_to_signature_cache[(name, overload_name)]
+
+    captured_impls -= op_impls
+    captured_defs -= op_defs
+    for handle in registration_handles:
+        handle.destroy()
+
+    if m is not None:
+        m.reset()
+
+
+@contextlib.contextmanager
+def _scoped_library(*args, **kwargs):
+    try:
+        lib = Library(*args, **kwargs)
+        yield lib
+    finally:
+        lib._destroy()
+
+
+_keep_alive: list[Library] = []
+
+
+NAMELESS_SCHEMA = re.compile(r"\(.*\) -> .*")
+
+
+@functools.singledispatch
+def define(qualname, schema, *, lib=None, tags=()):
+    r"""Defines a new operator.
+
+    In PyTorch, defining an op (short for "operator") is a two step-process:
+    - we need to define the op (by providing an operator name and schema)
+    - we need to implement behavior for how the operator interacts with
+    various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc.
+
+    This entrypoint defines the custom operator (the first step)
+    you must then perform the second step by calling various
+    ``impl_*`` APIs, like :func:`torch.library.impl` or
+    :func:`torch.library.register_fake`.
+
+    Args:
+        qualname (str): The qualified name for the operator. Should be
+            a string that looks like "namespace::name", e.g. "aten::sin".
+            Operators in PyTorch need a namespace to
+            avoid name collisions; a given operator may only be created once.
+            If you are writing a Python library, we recommend the namespace to
+            be the name of your top-level module.
+        schema (str): The schema of the operator. E.g. "(Tensor x) -> Tensor"
+            for an op that accepts one Tensor and returns one Tensor. It does
+            not contain the operator name (that is passed in ``qualname``).
+        lib (Optional[Library]): If provided, the lifetime of this operator
+            will be tied to the lifetime of the Library object.
+        tags (Tag | Sequence[Tag]): one or more torch.Tag to apply to this
+            operator. Tagging an operator changes the operator's behavior
+            under various PyTorch subsystems; please read the docs for the
+            torch.Tag carefully before applying it.
+
+    Example::
+        >>> import torch
+        >>> import numpy as np
+        >>>
+        >>> # Define the operator
+        >>> torch.library.define("mylib::sin", "(Tensor x) -> Tensor")
+        >>>
+        >>> # Add implementations for the operator
+        >>> @torch.library.impl("mylib::sin", "cpu")
+        >>> def f(x):
+        >>>     return torch.from_numpy(np.sin(x.numpy()))
+        >>>
+        >>> # Call the new operator from torch.ops.
+        >>> x = torch.randn(3)
+        >>> y = torch.ops.mylib.sin(x)
+        >>> assert torch.allclose(y, x.sin())
+
+    """
+    if not isinstance(qualname, str):
+        raise ValueError(
+            f"define(qualname, schema): expected qualname "
+            f"to be instance of str, got {type(qualname)}"
+        )
+    namespace, name = torch._library.utils.parse_namespace(qualname)
+    if lib is None:
+        lib = Library(namespace, "FRAGMENT")
+        _keep_alive.append(lib)
+    if not NAMELESS_SCHEMA.fullmatch(schema):
+        raise ValueError(
+            f"define(qualname, schema, ...): expected schema "
+            f'to look like e.g. "(Tensor x) -> Tensor" but '
+            f'got "{schema}"'
+        )
+    lib.define(name + schema, alias_analysis="", tags=tags)
+
+
+@define.register
+def _(lib: Library, schema, alias_analysis=""):
+    """The old torch.library.define.
+    We're keeping this around for BC reasons
+    """
+
+    def wrap(f):
+        name = lib.define(schema, alias_analysis)
+        lib.impl(name, f)
+        return f
+
+    return wrap
+
+
+@overload
+def impl(
+    qualname: str,
+    types: str | Sequence[str],
+    func: None = None,
+    *,
+    lib: Library | None = None,
+) -> Callable[[Callable[..., object]], None]: ...
+
+
+@overload
+def impl(
+    qualname: str,
+    types: str | Sequence[str],
+    func: Callable[..., object],
+    *,
+    lib: Library | None = None,
+) -> None: ...
+
+
+# Deprecated BC API
+@overload
+def impl(
+    lib: Library,
+    name: str,
+    dispatch_key: str = "",
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: ...
+
+
+@functools.singledispatch
+def impl(
+    qualname: str,
+    types: str | Sequence[str],
+    func: Callable[_P, _T] | None = None,
+    *,
+    lib: Library | None = None,
+) -> object:
+    """Register an implementation for a device type for this operator.
+
+    You may pass "default" for ``types`` to register this implementation as the
+    default implementation for ALL device types.
+    Please only use this if the implementation truly supports all device types;
+    for example, this is true if it is a composition of built-in PyTorch operators.
+
+    This API may be used as a decorator. You can use nested decorators
+    with this API provided they return a function and are placed inside
+    this API (see Example 2).
+
+    Some valid types are: "cpu", "cuda", "xla", "mps", "ipu", "xpu".
+
+    Args:
+        qualname (str): Should be a string that looks like "namespace::operator_name".
+        types (str | Sequence[str]): The device types to register an impl to.
+        lib (Optional[Library]): If provided, the lifetime of this registration
+            will be tied to the lifetime of the Library object.
+
+    Examples:
+        >>> # xdoctest: +SKIP("Requires Python <= 3.11")
+        >>> import torch
+        >>> import numpy as np
+        >>> # Example 1: Register function.
+        >>> # Define the operator
+        >>> torch.library.define("mylib::mysin", "(Tensor x) -> Tensor")
+        >>>
+        >>> # Add implementations for the cpu device
+        >>> @torch.library.impl("mylib::mysin", "cpu")
+        >>> def f(x):
+        >>>     return torch.from_numpy(np.sin(x.numpy()))
+        >>>
+        >>> x = torch.randn(3)
+        >>> y = torch.ops.mylib.mysin(x)
+        >>> assert torch.allclose(y, x.sin())
+        >>>
+        >>> # Example 2: Register function with decorator.
+        >>> def custom_decorator(func):
+        >>>     def wrapper(*args, **kwargs):
+        >>>         return func(*args, **kwargs) + 1
+        >>>     return wrapper
+        >>>
+        >>> # Define the operator
+        >>> torch.library.define("mylib::sin_plus_one", "(Tensor x) -> Tensor")
+        >>>
+        >>> # Add implementations for the operator
+        >>> @torch.library.impl("mylib::sin_plus_one", "cpu")
+        >>> @custom_decorator
+        >>> def f(x):
+        >>>     return torch.from_numpy(np.sin(x.numpy()))
+        >>>
+        >>> # Call the new operator from torch.ops.
+        >>> x = torch.randn(3)
+        >>>
+        >>> y1 = torch.ops.mylib.sin_plus_one(x)
+        >>> y2 = torch.sin(x) + 1
+        >>> assert torch.allclose(y1, y2)
+    """
+
+    return _impl(qualname, types, func, lib=lib, disable_dynamo=False)
+
+
+if not TYPE_CHECKING:
+
+    @impl.register
+    def _(
+        lib: Library, name: str, dispatch_key: str = ""
+    ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
+        """Legacy torch.library.impl API. Kept around for BC"""
+
+        def wrap(f: Callable[_P, _T]) -> Callable[_P, _T]:
+            lib.impl(name, f, dispatch_key)
+            return f
+
+        return wrap
+
+
+@overload
+def _impl(
+    qualname: str,
+    types: str | Sequence[str],
+    func: None = None,
+    *,
+    lib: Library | None = None,
+    disable_dynamo: bool = False,
+) -> Callable[[Callable[..., object]], None]: ...
+
+
+@overload
+def _impl(
+    qualname: str,
+    types: str | Sequence[str],
+    func: Callable[..., object],
+    *,
+    lib: Library | None = None,
+    disable_dynamo: bool = False,
+) -> None: ...
+
+
+def _impl(
+    qualname: str,
+    types: str | Sequence[str],
+    func: Callable[..., object] | None = None,
+    *,
+    lib: Library | None = None,
+    disable_dynamo: bool = False,
+) -> Callable[[Callable[..., object]], None] | None:
+    # See impl()
+    if isinstance(types, str):
+        types = (types,)
+    keys = set({})
+    for typ in types:
+        is_dispatch_key = torch._C._parse_dispatch_key(typ)
+        if is_dispatch_key:
+            # We also support passing a DispatchKey to impl. Please prefer using
+            # the higher-level torch.library APIs and only pass DispatchKey to
+            # torch.library.impl with caution (or even better, don't use this
+            # option and file an issue on GitHub for what you need).
+            # We don't advertise this to users because
+            # it is very easy to shoot yourself in the foot.
+            keys.add(typ)
+        else:
+            keys.add(_device_type_to_key(typ))
+
+    def register_(func: Callable[..., object]) -> None:
+        namespace, _ = torch._library.utils.parse_namespace(qualname)
+
+        if lib is None:
+            use_lib = Library(namespace, "FRAGMENT")
+            _keep_alive.append(use_lib)
+        else:
+            use_lib = lib
+        if disable_dynamo:
+
+            @torch._disable_dynamo
+            def func_no_dynamo(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            for key in keys:
+                use_lib.impl(qualname, func_no_dynamo, key)
+        else:
+            for key in keys:
+                use_lib.impl(qualname, func, key)
+
+    if func is None:
+        return register_
+    else:
+        register_(func)
+        return None
+
+
+def _device_type_to_key(device_type: str) -> str:
+    if device_type == "default":
+        # This is technically not correct, because although all device_type
+        # DispatchKeys are included in CompositeExplicitAutograd,
+        # not everything in CompositeExplicitAutograd is associated with a
+        # device_type. I don't really care that much about the difference.
+        return "CompositeExplicitAutograd"
+    return torch._C._dispatch_key_for_device(device_type)
+
+
+@deprecated(
+    "`torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that "
+    "instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.",
+    category=FutureWarning,
+)
+def impl_abstract(qualname, func=None, *, lib=None, _stacklevel=1):
+    r"""This API was renamed to :func:`torch.library.register_fake` in PyTorch 2.4.
+    Please use that instead.
+    """
+    if func is not None:
+        _stacklevel = _stacklevel + 1
+    return register_fake(qualname, func, lib=lib, _stacklevel=_stacklevel)
+
+
+_op_identifier = Union[
+    str, "torch._ops.OpOverload", "torch._library.custom_ops.CustomOpDef"
+]
+
+
+def register_kernel(
+    op: _op_identifier,
+    device_types: device_types_t,
+    func: Callable | None = None,
+    /,
+    *,
+    lib: Library | None = None,
+):
+    """Register an implementation for a device type for this operator.
+
+    Some valid device_types are: "cpu", "cuda", "xla", "mps", "ipu", "xpu".
+    This API may be used as a decorator.
+
+    Args:
+        op (str | OpOverload): The operator to register an impl to.
+        device_types (str | None | Sequence[str]): The device_types to register an impl to.
+            If None, we will register to all device types -- please only use
+            this option if your implementation is truly device-type-agnostic.
+        func (Callable): The function to register as the implementation for
+            the given device types.
+        lib (Optional[Library]): If provided, the lifetime of this registration
+
+    Examples::
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+        >>> import torch
+        >>> from torch import Tensor
+        >>> from torch.library import custom_op
+        >>> import numpy as np
+        >>>
+        >>> # Create a custom op that works on cpu
+        >>> @custom_op("mylib::numpy_sin", mutates_args=(), device_types="cpu")
+        >>> def numpy_sin(x: Tensor) -> Tensor:
+        >>>     x_np = x.numpy()
+        >>>     y_np = np.sin(x_np)
+        >>>     return torch.from_numpy(y_np)
+        >>>
+        >>> # Add implementations for the cuda device
+        >>> @torch.library.register_kernel("mylib::numpy_sin", "cuda")
+        >>> def _(x):
+        >>>     x_np = x.cpu().numpy()
+        >>>     y_np = np.sin(x_np)
+        >>>     return torch.from_numpy(y_np).to(device=x.device)
+        >>>
+        >>> x_cpu = torch.randn(3)
+        >>> x_cuda = x_cpu.cuda()
+        >>> assert torch.allclose(numpy_sin(x_cpu), x_cpu.sin())
+        >>> assert torch.allclose(numpy_sin(x_cuda), x_cuda.sin())
+
+    """
+
+    if not isinstance(
+        op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef)
+    ):
+        raise ValueError(
+            f"register_kernel({op}): got unexpected type for op: {type(op)}"
+        )
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+    opdef = _maybe_get_opdef(op)
+    if opdef is not None:
+        return opdef.register_kernel(device_types, func)
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str at this point, got {type(op).__name__}")
+    if device_types is None:
+        device_types = "CompositeExplicitAutograd"
+
+    return _impl(op, device_types, func, lib=lib, disable_dynamo=True)
+
+
+def register_autocast(
+    op: _op_identifier,
+    device_type: str,
+    cast_inputs: _dtype,
+    /,
+    *,
+    lib: Library | None = None,
+):
+    r"""Register an autocast dispatch rule for this custom op.
+
+    Valid `device_type` include: "cpu" and "cuda".
+
+    Args:
+        op (str | OpOverload): The operator to register an autocast dispatch rule to.
+        device_type(str):  Device type to use. 'cuda' or 'cpu'.
+            The type is the same as the `type` attribute of a :class:`torch.device`.
+            Thus, you may obtain the device type of a tensor using `Tensor.device.type`.
+        cast_inputs (:class:`torch.dtype`): When custom op runs in an autocast-enabled region,
+            casts incoming floating-point Tensors to the target dtype (non-floating-point Tensors
+            are not affected), then executes custom op with autocast disabled.
+        lib (Optional[Library]): If provided, the lifetime of this registration
+
+    Examples::
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+        >>> import torch
+        >>> from torch import Tensor
+        >>> from torch.library import custom_op
+        >>>
+        >>> # Create a custom op that works on cuda
+        >>> @torch.library.custom_op("mylib::my_sin", mutates_args=())
+        >>> def my_sin(x: Tensor) -> Tensor:
+        >>>     return torch.sin(x)
+        >>>
+        >>> # Register autocast dispatch rule for the cuda device
+        >>> torch.library.register_autocast("mylib::my_sin", "cuda", torch.float16)
+        >>>
+        >>> x = torch.randn(3, dtype=torch.float32, device="cuda")
+        >>> with torch.autocast("cuda", dtype=torch.float16):
+        >>>     y = torch.ops.mylib.my_sin(x)
+        >>> assert y.dtype == torch.float16
+
+    """
+    if not isinstance(
+        op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef)
+    ):
+        raise ValueError(
+            f"register_autocast({op}): got unexpected type for op: {type(op)}"
+        )
+    if device_type not in ["cpu", "cuda"]:
+        raise ValueError(f"Unknown device type: {device_type}")
+
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+    opdef = _maybe_get_opdef(op)
+    if opdef is not None:
+        return opdef.register_autocast(device_type, cast_inputs)
+
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str at this point, got {type(op).__name__}")
+    qualname = op
+    _op = torch._library.utils.lookup_op(qualname)
+
+    namespace, opname = torch._library.utils.parse_namespace(qualname)
+    if lib is None:
+        lib = Library(namespace, "FRAGMENT")
+        _keep_alive.append(lib)
+
+    def _maybe_override_py_impl(op: torch._ops.OpOverload, dispatch_key):
+        def inner(kernel):
+            if op.has_kernel_for_dispatch_key(dispatch_key):
+                op.py_kernels.pop(dispatch_key)
+            return op.py_impl(dispatch_key)(kernel)
+
+        return inner
+
+    @_maybe_override_py_impl(_op, torch._C.DispatchKey.AutocastCPU)
+    @_maybe_override_py_impl(_op, torch._C.DispatchKey.AutocastCUDA)
+    def _autocast_py_impl(*args, **kwargs):
+        if len(kwargs) != 0:
+            raise AssertionError("Custom ops do not support kwargs yet.")
+        autocast_keyset = torch._C.DispatchKeySet(
+            torch._C.DispatchKey.AutocastCPU
+        ) | torch._C.DispatchKeySet(torch._C.DispatchKey.AutocastCUDA)
+        with torch._C._ExcludeDispatchKeyGuard(autocast_keyset):
+            return _op(*_cast(args, device_type, cast_inputs))
+
+    def kernel(_, *args, **kwargs):
+        if len(kwargs) != 0:
+            raise AssertionError("Custom ops do not support kwargs yet.")
+        return _autocast_py_impl(*args, **kwargs)
+
+    if device_type == "cuda":
+        return lib.impl(opname, kernel, "AutocastCUDA", with_keyset=True)
+    else:
+        # device_type is "cpu"
+        return lib.impl(opname, kernel, "AutocastCPU", with_keyset=True)
+
+
+def register_fake(
+    op: _op_identifier,
+    func: Callable | None = None,
+    /,
+    *,
+    lib: Library | None = None,
+    _stacklevel: int = 1,
+    allow_override: bool = False,
+):
+    r"""Register a FakeTensor implementation ("fake impl") for this operator.
+
+    Also sometimes known as a "meta kernel", "abstract impl".
+
+    An "FakeTensor implementation" specifies the behavior of this operator on
+    Tensors that carry no data ("FakeTensor"). Given some input Tensors with
+    certain properties (sizes/strides/storage_offset/device), it specifies
+    what the properties of the output Tensors are.
+
+    The FakeTensor implementation has the same signature as the operator.
+    It is run for both FakeTensors and meta tensors. To write a FakeTensor
+    implementation, assume that all Tensor inputs to the operator are
+    regular CPU/CUDA/Meta tensors, but they do not have storage, and
+    you are trying to return regular CPU/CUDA/Meta tensor(s) as output.
+    The FakeTensor implementation must consist of only PyTorch operations
+    (and may not directly access the storage or data of any input or
+    intermediate Tensors).
+
+    This API may be used as a decorator (see examples).
+
+    For a detailed guide on custom ops, please see
+    https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html
+
+    Args:
+        op_name: Operator name (along with the overload) or OpOverload object.
+        func: Fake tensor implementation.
+        lib (Optional[Library]): Library to register the fake tensor to.
+        allow_override: Flag controlling if we want to override an
+                        existing registered fake impl. This is by default off,
+                        and will error you're trying to register a fake impl to
+                        an operator that already has a fake impl. This also only
+                        applies if the custom operator was not created via
+                        torch.library.custom_op, as overriding and existing fake
+                        impl is already allowed.
+
+    Examples:
+        >>> import torch
+        >>> import numpy as np
+        >>> from torch import Tensor
+        >>>
+        >>> # Example 1: an operator without data-dependent output shape
+        >>> @torch.library.custom_op("mylib::custom_linear", mutates_args=())
+        >>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
+        >>>     raise NotImplementedError("Implementation goes here")
+        >>>
+        >>> @torch.library.register_fake("mylib::custom_linear")
+        >>> def _(x, weight, bias):
+        >>>     assert x.dim() == 2
+        >>>     assert weight.dim() == 2
+        >>>     assert bias.dim() == 1
+        >>>     assert x.shape[1] == weight.shape[1]
+        >>>     assert weight.shape[0] == bias.shape[0]
+        >>>     assert x.device == weight.device
+        >>>
+        >>>     return (x @ weight.t()) + bias
+        >>>
+        >>> with torch._subclasses.fake_tensor.FakeTensorMode():
+        >>>     x = torch.randn(2, 3)
+        >>>     w = torch.randn(3, 3)
+        >>>     b = torch.randn(3)
+        >>>     y = torch.ops.mylib.custom_linear(x, w, b)
+        >>>
+        >>> assert y.shape == (2, 3)
+        >>>
+        >>> # Example 2: an operator with data-dependent output shape
+        >>> @torch.library.custom_op("mylib::custom_nonzero", mutates_args=())
+        >>> def custom_nonzero(x: Tensor) -> Tensor:
+        >>>     x_np = x.numpy(force=True)
+        >>>     res = np.stack(np.nonzero(x_np), axis=1)
+        >>>     return torch.tensor(res, device=x.device)
+        >>>
+        >>> @torch.library.register_fake("mylib::custom_nonzero")
+        >>> def _(x):
+        >>> # Number of nonzero-elements is data-dependent.
+        >>> # Since we cannot peek at the data in an fake impl,
+        >>> # we use the ctx object to construct a new symint that
+        >>> # represents the data-dependent size.
+        >>>     ctx = torch.library.get_ctx()
+        >>>     nnz = ctx.new_dynamic_size()
+        >>>     shape = [nnz, x.dim()]
+        >>>     result = x.new_empty(shape, dtype=torch.int64)
+        >>>     return result
+        >>>
+        >>> from torch.fx.experimental.proxy_tensor import make_fx
+        >>>
+        >>> x = torch.tensor([0, 1, 2, 3, 4, 0])
+        >>> trace = make_fx(torch.ops.mylib.custom_nonzero, tracing_mode="symbolic")(x)
+        >>> trace.print_readable()
+        >>>
+        >>> assert torch.allclose(trace(x), torch.ops.mylib.custom_nonzero(x))
+
+    """
+    if not isinstance(
+        op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef)
+    ):
+        raise ValueError(f"register_fake({op}): got unexpected type for op: {type(op)}")
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+    opdef = _maybe_get_opdef(op)
+    if opdef is not None:
+        if func is None:
+            return opdef.register_fake
+        else:
+            return opdef.register_fake(func)
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str at this point, got {type(op).__name__}")
+
+    stacklevel = _stacklevel
+
+    def register(func):
+        namespace, op_name = torch._library.utils.parse_namespace(op)
+        if lib is None:
+            use_lib = Library(namespace, "FRAGMENT")
+            _keep_alive.append(use_lib)
+        else:
+            use_lib = lib
+        use_lib._register_fake(
+            op_name, func, _stacklevel=stacklevel + 1, allow_override=allow_override
+        )
+        return func
+
+    if func is None:
+        return register
+    else:
+        stacklevel += 1
+        return register(func)
+
+
+def _register_effectful_op(
+    op: _op_identifier,
+    effect: EffectType | None,
+    *,
+    lib: Library | None = None,
+) -> None:
+    r"""
+    To specify that an operator has side-effects, we must register an effect
+    type for the operator. This will prevent graph passes in torch.compile from
+    reordering operations with the same effect type.
+
+    Args:
+        op_name: Operator name (along with the overload) or OpOverload object.
+        effect: Effect type to register. None means the operator is not effectful.
+    """
+    if not isinstance(
+        op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef)
+    ):
+        raise ValueError(
+            f"register_effectful_op({op}): got unexpected type for op: {type(op)}"
+        )
+
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+    opdef = _maybe_get_opdef(op)
+    if opdef is not None:
+        opdef.register_effect(effect)
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str at this point, got {type(op).__name__}")
+
+    namespace, _ = torch._library.utils.parse_namespace(op)
+    if lib is None:
+        use_lib = Library(namespace, "FRAGMENT")
+        _keep_alive.append(use_lib)
+    else:
+        use_lib = lib
+    use_lib._register_effectful_op(op, effect)
+
+
+def register_autograd(
+    op: _op_identifier,
+    backward: Callable,
+    /,
+    *,
+    setup_context: Callable | None = None,
+    lib=None,
+) -> None:
+    r"""Register a backward formula for this custom op.
+
+    In order for an operator to work with autograd, you need to register
+    a backward formula:
+    1. You must tell us how to compute gradients during the backward pass
+    by providing us a "backward" function.
+    2. If you need any values from the forward to compute gradients, you can
+    use `setup_context` to save values for backward.
+
+    ``backward`` runs during the backward pass. It accepts ``(ctx, *grads)``:
+    - ``grads`` is one or more gradients. The number of gradients matches
+    the number of outputs of the operator.
+    The ``ctx`` object is `the same ctx object <context_method_mixins>`_ used by
+    :class:`torch.autograd.Function`. The semantics of ``backward_fn`` are the
+    same as :meth:`torch.autograd.Function.backward`.
+
+    ``setup_context(ctx, inputs, output)`` runs during the forward pass.
+    Please save quantities needed for backward onto the ``ctx`` object via
+    either :meth:`torch.autograd.function.FunctionCtx.save_for_backward`
+    or assigning them as attributes of ``ctx``. If your custom op has
+    kwarg-only arguments, we expect the signature of ``setup_context``
+    to be ``setup_context(ctx, inputs, keyword_only_inputs, output)``.
+
+    Both ``setup_context_fn`` and ``backward_fn`` must be traceable. That is,
+    they may not directly access :meth:`torch.Tensor.data_ptr` and they must
+    not depend on or mutate global state. If you need a non-traceable backward,
+    you can make it a separate custom_op that you call inside ``backward_fn``.
+
+    If you need different autograd behavior on different devices, then we
+    recommend creating two different custom operators, one for each device
+    that needs different behavior, and switching between them at runtime.
+
+    Examples:
+        >>> import torch
+        >>> import numpy as np
+        >>> from torch import Tensor
+        >>>
+        >>> @torch.library.custom_op("mylib::numpy_sin", mutates_args=())
+        >>> def numpy_sin(x: Tensor) -> Tensor:
+        >>>     x_np = x.cpu().numpy()
+        >>>     y_np = np.sin(x_np)
+        >>>     return torch.from_numpy(y_np).to(device=x.device)
+        >>>
+        >>> def setup_context(ctx, inputs, output) -> Tensor:
+        >>>     x, = inputs
+        >>>     ctx.save_for_backward(x)
+        >>>
+        >>> def backward(ctx, grad):
+        >>>     x, = ctx.saved_tensors
+        >>>     return grad * x.cos()
+        >>>
+        >>> torch.library.register_autograd(
+        ...     "mylib::numpy_sin", backward, setup_context=setup_context
+        ... )
+        >>>
+        >>> x = torch.randn(3, requires_grad=True)
+        >>> y = numpy_sin(x)
+        >>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
+        >>> assert torch.allclose(grad_x, x.cos())
+        >>>
+        >>> # Example with a keyword-only arg
+        >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
+        >>> def numpy_mul(x: Tensor, *, val: float) -> Tensor:
+        >>>     x_np = x.cpu().numpy()
+        >>>     y_np = x_np * val
+        >>>     return torch.from_numpy(y_np).to(device=x.device)
+        >>>
+        >>> def setup_context(ctx, inputs, keyword_only_inputs, output) -> Tensor:
+        >>>     ctx.val = keyword_only_inputs["val"]
+        >>>
+        >>> def backward(ctx, grad):
+        >>>     return grad * ctx.val
+        >>>
+        >>> torch.library.register_autograd(
+        ...     "mylib::numpy_mul", backward, setup_context=setup_context
+        ... )
+        >>>
+        >>> x = torch.randn(3, requires_grad=True)
+        >>> y = numpy_mul(x, val=3.14)
+        >>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y))
+        >>> assert torch.allclose(grad_x, torch.full_like(x, 3.14))
+
+    """
+    if not isinstance(
+        op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef)
+    ):
+        raise ValueError(
+            f"register_autograd({op}): got unexpected type for op: {type(op)}"
+        )
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+    opdef = _maybe_get_opdef(op)
+    if opdef is not None:
+        opdef.register_autograd(backward, setup_context=setup_context)
+        return
+
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str at this point, got {type(op).__name__}")
+    qualname = op
+    op = torch._library.utils.lookup_op(qualname)
+    schema = op._schema
+    if not _library.utils.is_functional_schema(schema):
+        raise RuntimeError(
+            f"Cannot register autograd formula for non-functional operator "
+            f"{op} with schema {schema}. Please create "
+            f"a functional operator and register an autograd formula for that."
+        )
+    if _library.utils.has_kwarg_only_tensors(schema):
+        raise NotImplementedError(
+            f"register_autograd with kwarg-only Tensor args. In the original "
+            f"definition of the op, please make your tensors not kwarg-only. "
+            f"Got: {schema}"
+        )
+
+    info = _library.autograd.Info(backward, setup_context)
+    autograd_kernel = _library.autograd.make_autograd_impl(op, info)
+    namespace, opname = torch._library.utils.parse_namespace(qualname)
+    if lib is None:
+        lib = Library(namespace, "FRAGMENT")
+        _keep_alive.append(lib)
+    lib.impl(opname, autograd_kernel, "Autograd", with_keyset=True)
+
+
+def register_torch_dispatch(
+    op: _op_identifier,
+    torch_dispatch_class: Any,
+    func: Callable | None = None,
+    /,
+    *,
+    lib: Library | None = None,
+):
+    r"""Registers a torch_dispatch rule for the given operator and ``torch_dispatch_class``.
+
+    This allows for open registration to specify the behavior between the operator
+    and the ``torch_dispatch_class`` without needing to modify the ``torch_dispatch_class``
+    or the operator directly.
+
+    The ``torch_dispatch_class`` is either a Tensor subclass with ``__torch_dispatch__`` or a
+    TorchDispatchMode.
+
+    If it is a Tensor subclass, we expect ``func`` to have the following signature:
+    ``(cls, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any``
+
+    If it is a TorchDispatchMode, we expect ``func`` to have the following signature:
+    ``(mode, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any``
+
+    ``args`` and ``kwargs`` will have been normalized the same way they are
+    in ``__torch_dispatch__`` (see :ref:`torch-dispatch-calling-convention`).
+
+    Examples:
+
+        >>> import torch
+        >>>
+        >>> @torch.library.custom_op("mylib::foo", mutates_args={})
+        >>> def foo(x: torch.Tensor) -> torch.Tensor:
+        >>>     return x.clone()
+        >>>
+        >>> class MyMode(torch.utils._python_dispatch.TorchDispatchMode):
+        >>>     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        >>>         return func(*args, **kwargs)
+        >>>
+        >>> @torch.library.register_torch_dispatch("mylib::foo", MyMode)
+        >>> def _(mode, func, types, args, kwargs):
+        >>>     x, = args
+        >>>     return x + 1
+        >>>
+        >>> x = torch.randn(3)
+        >>> y = foo(x)
+        >>> assert torch.allclose(y, x)
+        >>>
+        >>> with MyMode():
+        >>>     y = foo(x)
+        >>> assert torch.allclose(y, x + 1)
+
+    """
+    if not isinstance(
+        op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef)
+    ):
+        raise ValueError(
+            f"register_torch_dispatch({op}): got unexpected type for op: {type(op)}"
+        )
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+    opdef = _maybe_get_opdef(op)
+    if opdef is not None:
+        return opdef.register_torch_dispatch(torch_dispatch_class, func)
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str at this point, got {type(op).__name__}")
+
+    def register(func):
+        namespace, op_name = torch._library.utils.parse_namespace(op)
+        if lib is None:
+            use_lib = Library(namespace, "FRAGMENT")
+            _keep_alive.append(use_lib)
+        else:
+            use_lib = lib
+        use_lib._register_torch_dispatch_rule(op_name, torch_dispatch_class, func)
+        return func
+
+    if func is None:
+        return register
+    else:
+        return register(func)
+
+
+def register_vmap(
+    op: _op_identifier,
+    func: Callable | None = None,
+    /,
+    *,
+    lib=None,
+):
+    r"""Register a vmap implementation to support :func:`torch.vmap` for this custom op.
+
+    This API may be used as a decorator (see examples).
+
+    In order for an operator to work with :func:`torch.vmap`, you may need to register a
+    vmap implementation in the following signature:
+
+        ``vmap_func(info, in_dims: Tuple[Optional[int]], *args, **kwargs)``,
+
+    where ``*args`` and ``**kwargs`` are the arguments and kwargs for ``op``.
+    We do not support kwarg-only Tensor args.
+
+    It specifies how do we compute the batched version of ``op`` given inputs with an additional
+    dimension (specified by ``in_dims``).
+
+    For each arg in ``args``, ``in_dims`` has a corresponding ``Optional[int]``. It is ``None``
+    if the arg is not a Tensor or if the arg is not being vmapped over, otherwise, it is an integer
+    specifying what dimension of the Tensor is being vmapped over.
+
+    ``info`` is a collection of additional metadata that may be helpful:
+    ``info.batch_size`` specifies the size of the dimension being vmapped over, while
+    ``info.randomness`` is the ``randomness`` option that was passed to :func:`torch.vmap`.
+
+    The return of the function ``func`` is a tuple of ``(output, out_dims)``. Similar to ``in_dims``,
+    ``out_dims`` should be of the same structure as ``output`` and contain one ``out_dim``
+    per output that specifies if the output has the vmapped dimension and what index it is in.
+
+    Examples:
+        >>> import torch
+        >>> import numpy as np
+        >>> from torch import Tensor
+        >>> from typing import Tuple
+        >>>
+        >>> def to_numpy(tensor):
+        >>>     return tensor.cpu().numpy()
+        >>>
+        >>> lib = torch.library.Library("mylib", "FRAGMENT")
+        >>> @torch.library.custom_op("mylib::numpy_cube", mutates_args=())
+        >>> def numpy_cube(x: Tensor) -> Tuple[Tensor, Tensor]:
+        >>>     x_np = to_numpy(x)
+        >>>     dx = torch.tensor(3 * x_np ** 2, device=x.device)
+        >>>     return torch.tensor(x_np ** 3, device=x.device), dx
+        >>>
+        >>> def numpy_cube_vmap(info, in_dims, x):
+        >>>     result = numpy_cube(x)
+        >>>     return result, (in_dims[0], in_dims[0])
+        >>>
+        >>> torch.library.register_vmap(numpy_cube, numpy_cube_vmap)
+        >>>
+        >>> x = torch.randn(3)
+        >>> torch.vmap(numpy_cube)(x)
+        >>>
+        >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
+        >>> def numpy_mul(x: Tensor, y: Tensor) -> Tensor:
+        >>>     return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device)
+        >>>
+        >>> @torch.library.register_vmap("mylib::numpy_mul")
+        >>> def numpy_mul_vmap(info, in_dims, x, y):
+        >>>     x_bdim, y_bdim = in_dims
+        >>>     x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1)
+        >>>     y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1)
+        >>>     result = x * y
+        >>>     result = result.movedim(-1, 0)
+        >>>     return result, 0
+        >>>
+        >>>
+        >>> x = torch.randn(3)
+        >>> y = torch.randn(3)
+        >>> torch.vmap(numpy_mul)(x, y)
+
+    .. note::
+        The vmap function should aim to preserve the semantics of the entire custom operator.
+        That is, ``grad(vmap(op))`` should be replaceable with a ``grad(map(op))``.
+
+        If your custom operator has any custom behavior in the backward pass, please
+        keep this in mind.
+
+    """
+    if not isinstance(
+        op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef)
+    ):
+        raise ValueError(f"register_vmap({op}): got unexpected type for op: {type(op)}")
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+    opdef = _maybe_get_opdef(op)
+    if opdef is not None:
+        return opdef.register_vmap(func)
+    if not isinstance(op, str):
+        raise AssertionError(f"op must be str at this point, got {type(op).__name__}")
+    qualname = op
+    op = torch._library.utils.lookup_op(qualname)
+    schema = op._schema
+    if _library.utils.has_kwarg_only_tensors(schema):
+        raise NotImplementedError(
+            f"register_vmap with kwarg-only Tensor args. In the original "
+            f"definition of the op, please make your tensors not kwarg-only. "
+            f"Got: {schema}"
+        )
+
+    def register(func):
+        nonlocal op, lib
+
+        namespace, opname = torch._library.utils.parse_namespace(qualname)
+        if lib is None:
+            lib = Library(namespace, "FRAGMENT")
+            _keep_alive.append(lib)
+
+        from torch._functorch.autograd_function import custom_function_call_vmap_helper
+        from torch._functorch.pyfunctorch import retrieve_current_functorch_interpreter
+
+        def wrapped_func(keyset, *args, **kwargs):
+            interpreter = retrieve_current_functorch_interpreter()
+            return custom_function_call_vmap_helper(
+                # pyrefly: ignore[bad-argument-type]
+                interpreter,
+                func,
+                op,
+                *args,
+                **kwargs,
+            )
+
+        lib.impl(opname, wrapped_func, "FuncTorchBatched", with_keyset=True)
+
+    if func is None:
+        return register
+    else:
+        return register(func)
+
+
+# If the op was defined in C++, then we want to make sure there was an
+# m.set_python_module(module, ...) call and that the module is the
+# same as the module that called torch.library.register_fake.
+def _check_pystubs_once(func, qualname, actual_module_name):
+    checked = False
+
+    def inner(*args, **kwargs):
+        nonlocal checked
+        if checked:
+            return func(*args, **kwargs)
+
+        op = torch._library.utils.lookup_op(qualname)
+        if op._defined_in_python:
+            checked = True
+            return func(*args, **kwargs)
+
+        maybe_pystub = torch._C._dispatch_pystub(
+            op._schema.name, op._schema.overload_name
+        )
+        if maybe_pystub is None:
+            if torch._library.utils.requires_set_python_module():
+                namespace = op.namespace
+                cpp_filename = op._handle.debug()
+                raise RuntimeError(
+                    f"Operator '{qualname}' was defined in C++ and has a Python "
+                    f"fake impl. In this situation, we require there to also be a "
+                    f'companion C++ `m.set_python_module("{actual_module_name}")` '
+                    f"call, but we could not find one. Please add that to "
+                    f"to the top of the C++ TORCH_LIBRARY({namespace}, ...) block the "
+                    f"operator was registered in ({cpp_filename})"
+                )
+        else:
+            pystub_module = maybe_pystub[0]
+            if actual_module_name != pystub_module:
+                cpp_filename = op._handle.debug()
+                raise RuntimeError(
+                    f"Operator '{qualname}' specified that its python fake impl "
+                    f"is in the Python module '{pystub_module}' but it was actually found "
+                    f"in '{actual_module_name}'. Please either move the fake impl "
+                    f"or correct the m.set_python_module call ({cpp_filename})"
+                )
+        checked = True
+        return func(*args, **kwargs)
+
+    return inner
+
+
+# NOTE [ctx inside the fake implementation]
+# If a user has an operator with data-dependent output shape, then when writing
+# a fake implementation they must query the current ctx and use methods on the
+# ctx to construct a new unbacked symint.
+#
+# This is done via us setting the global_ctx_getter function every time a fake
+# implementation is invoked.
+def get_ctx() -> "torch._library.fake_impl.FakeImplCtx":
+    """get_ctx() returns the current AbstractImplCtx object.
+
+    Calling ``get_ctx()`` is only valid inside of an fake impl
+    (see :func:`torch.library.register_fake` for more usage details.
+    """
+    return torch._library.fake_impl.global_ctx_getter()
+
+
+def get_kernel(
+    op: _op_identifier, dispatch_key: str | torch.DispatchKey
+) -> torch._C._SafeKernelFunction:
+    """Returns the computed kernel for a given operator and dispatch key.
+
+    This function retrieves the kernel that would be executed for a given
+    operator and dispatch key combination. The returned SafeKernelFunction
+    can be used to call the kernel in a boxed fashion. The intended use
+    case for this function is to retrieve the original kernel for a given
+    dispatch key and then register another kernel to the same dispatch key
+    that calls into the original kernel for certain cases.
+
+    Args:
+        op: Operator name (along with the overload) or OpOverload object
+            Can be a string (e.g., "aten::add.Tensor"), an OpOverload, or a CustomOpDef.
+        dispatch_key (str | torch.DispatchKey): The dispatch key to get the kernel for.
+            Can be a string (e.g., "CPU", "CUDA") or a DispatchKey enum value.
+
+    Returns:
+        torch._C._SafeKernelFunction: A safe kernel function that can be used to
+            call the kernel.
+
+    Raises:
+        RuntimeError: If the operator does not exist.
+
+    Example:
+        >>> # Get the CPU kernel for torch.add
+        >>> kernel = torch.library.get_kernel("aten::add.Tensor", "CPU")
+        >>>
+        >>> # You can also use DispatchKey enum
+        >>> kernel = torch.library.get_kernel("aten::add.Tensor", torch.DispatchKey.CPU)
+        >>>
+        >>> # Or use an OpOverload directly
+        >>> kernel = torch.library.get_kernel(torch.ops.aten.add.Tensor, "CPU")
+        >>>
+        >>> # Example: Using get_kernel in a custom op with conditional dispatch
+        >>> # Get the original kernel for torch.sin
+        >>> original_sin_kernel = torch.library.get_kernel("aten::sin", "CPU")
+        >>>
+        >>> # If input has negative values, use original sin, otherwise return zeros
+        >>> def conditional_sin_impl(dispatch_keys, x):
+        >>>     if (x < 0).any():
+        >>>         return original_sin_kernel.call_boxed(dispatch_keys, x)
+        >>>     else:
+        >>>         return torch.zeros_like(x)
+        >>>
+        >>> lib = torch.library.Library("aten", "IMPL")
+        >>> # with_keyset=True so the first argument to the impl is the current DispatchKeySet
+        >>> which needs to be the first argument to ``kernel.call_boxed``
+        >>> lib.impl("sin", conditional_sin_impl, "CPU", with_keyset=True)
+        >>>
+        >>> # Test the conditional behavior
+        >>> x_positive = torch.tensor([1.0, 2.0])
+        >>> x_mixed = torch.tensor([-1.0, 2.0])
+        >>> torch.sin(x_positive)
+        tensor([0., 0.])
+        >>> torch.sin(x_mixed)
+        tensor([-0.8415, 0.9093])
+    """
+    if not isinstance(op, (str, torch._ops.OpOverload)):
+        raise ValueError(f"get_kernel({op}): got unexpected type for op: {type(op)}")
+
+    if isinstance(op, torch._ops.OpOverload):
+        op = op._name
+
+    if isinstance(dispatch_key, str):
+        try:
+            dispatch_key = torch._C.DispatchKey.__members__[dispatch_key]
+        except KeyError:
+            raise ValueError(f"Invalid dispatch key: {dispatch_key}") from None
+
+    return torch._C._dispatch_get_computed_kernel_for_dispatch_key(op, dispatch_key)
+
+
+_OPCHECK_DEFAULT_UTILS = (
+    "test_schema",
+    "test_autograd_registration",
+    "test_faketensor",
+    "test_aot_dispatch_dynamic",
+)
+
+
+def opcheck(
+    op: torch._ops.OpOverload | torch._ops.OpOverloadPacket | CustomOpDef,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None = None,
+    *,
+    test_utils: str | Sequence[str] = _OPCHECK_DEFAULT_UTILS,
+    raise_exception: bool = True,
+    atol=None,
+    rtol=None,
+) -> dict[str, str]:
+    """Given an operator and some sample arguments, tests if the operator is
+    registered correctly.
+
+    That is, when you use the torch.library/TORCH_LIBRARY APIs to create a
+    custom op, you specified metadata (e.g. mutability info) about the custom op
+    and these APIs require that the functions you pass them satisfy certain
+    properties (e.g. no data pointer access in the fake/meta/abstract kernel)
+    ``opcheck`` tests these metadata and properties.
+
+    Concretely, we test the following:
+
+    - test_schema: If the schema matches the implementation of
+      the operator. For example: if the schema specifies a Tensor is mutated,
+      then we check the implementation mutates the Tensor. If the schema
+      specifies that we return a new Tensor, then we check that the
+      implementation returns a new Tensor (instead of an existing one or
+      a view of an existing one).
+    - test_autograd_registration: If the operator supports training
+      (autograd): we check that its autograd formula is registered via
+      torch.library.register_autograd or a manual registration to one
+      or more DispatchKey::Autograd keys. Any other DispatchKey-based
+      registrations may lead to undefined behavior.
+    - test_faketensor: If the operator has a FakeTensor kernel
+      (and if it is correct). The FakeTensor kernel is necessary (
+      but not sufficient) for the operator to work with PyTorch compilation
+      APIs (torch.compile/export/FX). We check that a FakeTensor kernel
+      (also sometimes known as a meta kernel) was registered for the
+      operator and that it is correct. This test takes the result of
+      running the operator on real tensors and the result of running
+      the operator on FakeTensors and checks that they have the same
+      Tensor metadata (sizes/strides/dtype/device/etc).
+    - test_aot_dispatch_dynamic: If the operator has correct behavior
+      with PyTorch compilation APIs (torch.compile/export/FX).
+      This checks that the outputs (and gradients, if applicable) are the
+      same under eager-mode PyTorch and torch.compile.
+      This test is a superset of ``test_faketensor`` and is an e2e test;
+      other things it tests are that the operator supports
+      functionalization and that the backward pass (if it exists) also
+      supports FakeTensor and functionalization.
+
+    For best results, please call ``opcheck`` multiple times with a
+    representative set of inputs. If your operator supports
+    autograd, please use ``opcheck`` with inputs with ``requires_grad = True``;
+    if your operator supports multiple devices (e.g. CPU and CUDA), please
+    use ``opcheck`` with inputs on all supported devices.
+
+    Args:
+        op: The operator. Must either be a function decorated with
+            :func:`torch.library.custom_op` or an OpOverload/OpOverloadPacket
+            found in torch.ops.* (e.g. torch.ops.aten.sin, torch.ops.mylib.foo)
+        args: The args to the operator
+        kwargs: The kwargs to the operator
+        test_utils: Tests that we should run. Default: all of them.
+            Example: ("test_schema", "test_faketensor")
+        raise_exception: If we should raise an exception on the first
+            error. If False, we will return a dict with information
+            on if each test passed or not.
+        rtol (Optional[float]): Relative tolerance for floating point comparisons.
+            If specified ``atol`` must also be specified.
+            If omitted, default values based on the ``dtype`` are selected
+            (see the table in :func:`torch.testing.assert_close`).
+        atol (Optional[float]): Absolute tolerance for floating point comparisons.
+            If specified ``rtol`` must also be specified.
+            If omitted, default values based on the ``dtype`` are selected
+            (see the table in :func:`torch.testing.assert_close`).
+
+    .. warning::
+
+        opcheck and :func:`torch.autograd.gradcheck` test different things;
+        opcheck tests if your usage of torch.library APIs is correct while
+        :func:`torch.autograd.gradcheck` tests if your autograd formula is
+        mathematically correct. Use both to test custom ops that support
+        gradient computation.
+
+    Example:
+
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+        >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=())
+        >>> def numpy_mul(x: Tensor, y: float) -> Tensor:
+        >>>     x_np = x.numpy(force=True)
+        >>>     z_np = x_np * y
+        >>>     return torch.from_numpy(z_np).to(x.device)
+        >>>
+        >>> @numpy_mul.register_fake
+        >>> def _(x, y):
+        >>>     return torch.empty_like(x)
+        >>>
+        >>> def setup_context(ctx, inputs, output):
+        >>>     y, = inputs
+        >>>     ctx.y = y
+        >>>
+        >>> def backward(ctx, grad):
+        >>>     return grad * ctx.y, None
+        >>>
+        >>> numpy_mul.register_autograd(backward, setup_context=setup_context)
+        >>>
+        >>> sample_inputs = [
+        >>>     (torch.randn(3), 3.14),
+        >>>     (torch.randn(2, 3, device='cuda'), 2.718),
+        >>>     (torch.randn(1, 10, requires_grad=True), 1.234),
+        >>>     (torch.randn(64, 64, device='cuda', requires_grad=True), 90.18),
+        >>> ]
+        >>>
+        >>> for args in sample_inputs:
+        >>>     torch.library.opcheck(numpy_mul, args)
+
+    """
+    import torch.testing._internal.optests as optests
+
+    return optests.opcheck(
+        op,
+        args,
+        kwargs,
+        test_utils=test_utils,
+        raise_exception=raise_exception,
+        rtol=rtol,
+        atol=atol,
+    )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/linalg/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/linalg/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ce47d60eebcadfa0beb33122918a8d94413528e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/linalg/__init__.py
@@ -0,0 +1,3058 @@
+from torch._C import (
+    _add_docstr,
+    _linalg,  # pyrefly: ignore [missing-module-attribute]
+    _LinAlgError as LinAlgError,  # pyrefly: ignore  # missing-module-attribute
+)
+
+
+common_notes = {
+    "experimental_warning": """This function is "experimental" and it may change in a future PyTorch release.""",
+    "sync_note": "When inputs are on a CUDA device, this function synchronizes that device with the CPU.",
+    "sync_note_ex": r"When the inputs are on a CUDA device, this function synchronizes only when :attr:`check_errors`\ `= True`.",
+    "sync_note_has_ex": (
+        "When inputs are on a CUDA device, this function synchronizes that device with the CPU. "
+        "For a version of this function that does not synchronize, see :func:`{}`."
+    ),
+}
+
+
+# Note: This not only adds doc strings for functions in the linalg namespace, but
+# also connects the torch.linalg Python namespace to the torch._C._linalg builtins.
+
+cross = _add_docstr(
+    _linalg.linalg_cross,
+    r"""
+linalg.cross(input, other, *, dim=-1, out=None) -> Tensor
+
+
+Computes the cross product of two 3-dimensional vectors.
+
+Supports input of float, double, cfloat and cdouble dtypes. Also supports batches
+of vectors, for which it computes the product along the dimension :attr:`dim`.
+It broadcasts over the batch dimensions.
+
+Args:
+    input (Tensor): the first input tensor.
+    other (Tensor): the second input tensor.
+    dim  (int, optional): the dimension along which to take the cross-product. Default: `-1`.
+
+Keyword args:
+    out (Tensor, optional): the output tensor. Ignored if `None`. Default: `None`.
+
+Example:
+    >>> a = torch.randn(4, 3)
+    >>> a
+    tensor([[-0.3956,  1.1455,  1.6895],
+            [-0.5849,  1.3672,  0.3599],
+            [-1.1626,  0.7180, -0.0521],
+            [-0.1339,  0.9902, -2.0225]])
+    >>> b = torch.randn(4, 3)
+    >>> b
+    tensor([[-0.0257, -1.4725, -1.2251],
+            [-1.1479, -0.7005, -1.9757],
+            [-1.3904,  0.3726, -1.1836],
+            [-0.9688, -0.7153,  0.2159]])
+    >>> torch.linalg.cross(a, b)
+    tensor([[ 1.0844, -0.5281,  0.6120],
+            [-2.4490, -1.5687,  1.9792],
+            [-0.8304, -1.3037,  0.5650],
+            [-1.2329,  1.9883,  1.0551]])
+    >>> a = torch.randn(1, 3)  # a is broadcast to match shape of b
+    >>> a
+    tensor([[-0.9941, -0.5132,  0.5681]])
+    >>> torch.linalg.cross(a, b)
+    tensor([[ 1.4653, -1.2325,  1.4507],
+            [ 1.4119, -2.6163,  0.1073],
+            [ 0.3957, -1.9666, -1.0840],
+            [ 0.2956, -0.3357,  0.2139]])
+""",
+)
+
+cholesky = _add_docstr(
+    _linalg.linalg_cholesky,
+    r"""
+linalg.cholesky(A, *, upper=False, out=None) -> Tensor
+
+Computes the Cholesky decomposition of a complex Hermitian or real symmetric positive-definite matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **Cholesky decomposition** of a complex Hermitian or real symmetric positive-definite matrix
+:math:`A \in \mathbb{K}^{n \times n}` is defined as
+
+.. math::
+
+    A = LL^{\text{H}}\mathrlap{\qquad L \in \mathbb{K}^{n \times n}}
+
+where :math:`L` is a lower triangular matrix with real positive diagonal (even in the complex case) and
+:math:`L^{\text{H}}` is the conjugate transpose when :math:`L` is complex, and the transpose when :math:`L` is real-valued.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.cholesky_ex")}
+"""
+    + r"""
+
+.. seealso::
+
+        :func:`torch.linalg.cholesky_ex` for a version of this operation that
+        skips the (slow) error checking by default and instead returns the debug
+        information. This makes it a faster way to check if a matrix is
+        positive-definite.
+
+        :func:`torch.linalg.eigh` for a different decomposition of a Hermitian matrix.
+        The eigenvalue decomposition gives more information about the matrix but it
+        slower to compute than the Cholesky decomposition.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                consisting of symmetric or Hermitian positive-definite matrices.
+
+Keyword args:
+    upper (bool, optional): whether to return an upper triangular matrix.
+        The tensor returned with upper=True is the conjugate transpose of the tensor
+        returned with upper=False.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Raises:
+    RuntimeError: if the :attr:`A` matrix or any matrix in a batched :attr:`A` is not Hermitian
+                  (resp. symmetric) positive-definite. If :attr:`A` is a batch of matrices,
+                  the error message will include the batch index of the first matrix that fails
+                  to meet this condition.
+
+Examples::
+
+    >>> A = torch.randn(2, 2, dtype=torch.complex128)
+    >>> A = A @ A.T.conj() + torch.eye(2) # creates a Hermitian positive-definite matrix
+    >>> A
+    tensor([[2.5266+0.0000j, 1.9586-2.0626j],
+            [1.9586+2.0626j, 9.4160+0.0000j]], dtype=torch.complex128)
+    >>> L = torch.linalg.cholesky(A)
+    >>> L
+    tensor([[1.5895+0.0000j, 0.0000+0.0000j],
+            [1.2322+1.2976j, 2.4928+0.0000j]], dtype=torch.complex128)
+    >>> torch.dist(L @ L.T.conj(), A)
+    tensor(4.4692e-16, dtype=torch.float64)
+
+    >>> A = torch.randn(3, 2, 2, dtype=torch.float64)
+    >>> A = A @ A.mT + torch.eye(2)  # batch of symmetric positive-definite matrices
+    >>> L = torch.linalg.cholesky(A)
+    >>> torch.dist(L @ L.mT, A)
+    tensor(5.8747e-16, dtype=torch.float64)
+""",
+)
+
+cholesky_ex = _add_docstr(
+    _linalg.linalg_cholesky_ex,
+    r"""
+linalg.cholesky_ex(A, *, upper=False, check_errors=False, out=None) -> (Tensor, Tensor)
+
+Computes the Cholesky decomposition of a complex Hermitian or real
+symmetric positive-definite matrix.
+
+This function skips the (slow) error checking and error message construction
+of :func:`torch.linalg.cholesky`, instead directly returning the LAPACK
+error codes as part of a named tuple ``(L, info)``. This makes this function
+a faster way to check if a matrix is positive-definite, and it provides an
+opportunity to handle decomposition errors more gracefully or performantly
+than :func:`torch.linalg.cholesky` does.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+If :attr:`A` is not a Hermitian positive-definite matrix, or if it's a batch of matrices
+and one or more of them is not a Hermitian positive-definite matrix,
+then ``info`` stores a positive integer for the corresponding matrix.
+The positive integer indicates the order of the leading minor that is not positive-definite,
+and the decomposition could not be completed.
+``info`` filled with zeros indicates that the decomposition was successful.
+If ``check_errors=True`` and ``info`` contains positive integers, then a RuntimeError is thrown.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_ex"]}
+
+.. warning:: {common_notes["experimental_warning"]}
+"""
+    + r"""
+
+.. seealso::
+        :func:`torch.linalg.cholesky` is a NumPy compatible variant that always checks for errors.
+
+Args:
+    A (Tensor): the Hermitian `n \times n` matrix or the batch of such matrices of size
+                    `(*, n, n)` where `*` is one or more batch dimensions.
+
+Keyword args:
+    upper (bool, optional): whether to return an upper triangular matrix.
+        The tensor returned with upper=True is the conjugate transpose of the tensor
+        returned with upper=False.
+    check_errors (bool, optional): controls whether to check the content of ``infos``. Default: `False`.
+    out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.randn(2, 2, dtype=torch.complex128)
+    >>> A = A @ A.t().conj()  # creates a Hermitian positive-definite matrix
+    >>> L, info = torch.linalg.cholesky_ex(A)
+    >>> A
+    tensor([[ 2.3792+0.0000j, -0.9023+0.9831j],
+            [-0.9023-0.9831j,  0.8757+0.0000j]], dtype=torch.complex128)
+    >>> L
+    tensor([[ 1.5425+0.0000j,  0.0000+0.0000j],
+            [-0.5850-0.6374j,  0.3567+0.0000j]], dtype=torch.complex128)
+    >>> info
+    tensor(0, dtype=torch.int32)
+
+""",
+)
+
+inv = _add_docstr(
+    _linalg.linalg_inv,
+    r"""
+linalg.inv(A, *, out=None) -> Tensor
+
+Computes the inverse of a square matrix if it exists.
+Throws a `RuntimeError` if the matrix is not invertible.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+for a matrix :math:`A \in \mathbb{K}^{n \times n}`,
+its **inverse matrix** :math:`A^{-1} \in \mathbb{K}^{n \times n}` (if it exists) is defined as
+
+.. math::
+
+    A^{-1}A = AA^{-1} = \mathrm{I}_n
+
+where :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix.
+
+The inverse matrix exists if and only if :math:`A` is `invertible`_. In this case,
+the inverse is unique.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices
+then the output has the same batch dimensions.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.inv_ex")}
+"""
+    + r"""
+
+.. note::
+    Consider using :func:`torch.linalg.solve` if possible for multiplying a matrix on the left by
+    the inverse, as::
+
+        linalg.solve(A, B) == linalg.inv(A) @ B  # When B is a matrix
+
+    It is always preferred to use :func:`~solve` when possible, as it is faster and more
+    numerically stable than computing the inverse explicitly.
+
+.. seealso::
+
+        :func:`torch.linalg.pinv` computes the pseudoinverse (Moore-Penrose inverse) of matrices
+        of any shape.
+
+        :func:`torch.linalg.solve` computes :attr:`A`\ `.inv() @ \ `:attr:`B` with a
+        numerically stable algorithm.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                consisting of invertible matrices.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Raises:
+    RuntimeError: if the matrix :attr:`A` or any matrix in the batch of matrices :attr:`A` is not invertible.
+
+Examples::
+
+    >>> A = torch.randn(4, 4)
+    >>> Ainv = torch.linalg.inv(A)
+    >>> torch.dist(A @ Ainv, torch.eye(4))
+    tensor(1.1921e-07)
+
+    >>> A = torch.randn(2, 3, 4, 4)  # Batch of matrices
+    >>> Ainv = torch.linalg.inv(A)
+    >>> torch.dist(A @ Ainv, torch.eye(4))
+    tensor(1.9073e-06)
+
+    >>> A = torch.randn(4, 4, dtype=torch.complex128)  # Complex matrix
+    >>> Ainv = torch.linalg.inv(A)
+    >>> torch.dist(A @ Ainv, torch.eye(4))
+    tensor(7.5107e-16, dtype=torch.float64)
+
+.. _invertible:
+    https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem
+""",
+)
+
+solve_ex = _add_docstr(
+    _linalg.linalg_solve_ex,
+    r"""
+linalg.solve_ex(A, B, *, left=True, check_errors=False, out=None) -> (Tensor, Tensor)
+
+A version of :func:`~solve` that does not perform error checks unless :attr:`check_errors`\ `= True`.
+It also returns the :attr:`info` tensor returned by `LAPACK's getrf`_.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_ex"]}
+
+.. warning:: {common_notes["experimental_warning"]}
+"""
+    + r"""
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`.
+    check_errors (bool, optional): controls whether to check the content of ``infos`` and raise
+                                   an error if it is non-zero. Default: `False`.
+    out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(result, info)`.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> Ainv, info = torch.linalg.solve_ex(A)
+    >>> torch.dist(torch.linalg.inv(A), Ainv)
+    tensor(0.)
+    >>> info
+    tensor(0, dtype=torch.int32)
+
+.. _LAPACK's getrf:
+    https://www.netlib.org/lapack/explore-html-3.6.1/dd/d9a/group__double_g_ecomputational_ga0019443faea08275ca60a734d0593e60.html
+""",
+)
+
+inv_ex = _add_docstr(
+    _linalg.linalg_inv_ex,
+    r"""
+linalg.inv_ex(A, *, check_errors=False, out=None) -> (Tensor, Tensor)
+
+Computes the inverse of a square matrix if it is invertible.
+
+Returns a namedtuple ``(inverse, info)``. ``inverse`` contains the result of
+inverting :attr:`A` and ``info`` stores the LAPACK error codes.
+
+If :attr:`A` is not an invertible matrix, or if it's a batch of matrices
+and one or more of them is not an invertible matrix,
+then ``info`` stores a positive integer for the corresponding matrix.
+The positive integer indicates the diagonal element of the LU decomposition of
+the input matrix that is exactly zero.
+``info`` filled with zeros indicates that the inversion was successful.
+If ``check_errors=True`` and ``info`` contains positive integers, then a RuntimeError is thrown.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_ex"]}
+
+.. warning:: {common_notes["experimental_warning"]}
+"""
+    + r"""
+
+.. seealso::
+
+        :func:`torch.linalg.inv` is a NumPy compatible variant that always checks for errors.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                    consisting of square matrices.
+    check_errors (bool, optional): controls whether to check the content of ``info``. Default: `False`.
+
+Keyword args:
+    out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> Ainv, info = torch.linalg.inv_ex(A)
+    >>> torch.dist(torch.linalg.inv(A), Ainv)
+    tensor(0.)
+    >>> info
+    tensor(0, dtype=torch.int32)
+
+""",
+)
+
+det = _add_docstr(
+    _linalg.linalg_det,
+    r"""
+linalg.det(A, *, out=None) -> Tensor
+
+Computes the determinant of a square matrix.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+.. seealso::
+
+        :func:`torch.linalg.slogdet` computes the sign and natural logarithm of the absolute
+        value of the determinant of square matrices.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> torch.linalg.det(A)
+    tensor(0.0934)
+
+    >>> A = torch.randn(3, 2, 2)
+    >>> torch.linalg.det(A)
+    tensor([1.1990, 0.4099, 0.7386])
+""",
+)
+
+slogdet = _add_docstr(
+    _linalg.linalg_slogdet,
+    r"""
+linalg.slogdet(A, *, out=None) -> (Tensor, Tensor)
+
+Computes the sign and natural logarithm of the absolute value of the determinant of a square matrix.
+
+For complex :attr:`A`, it returns the sign and the natural logarithm of the modulus of the
+determinant, that is, a logarithmic polar decomposition of the determinant.
+
+The determinant can be recovered as `sign * exp(logabsdet)`.
+When a matrix has a determinant of zero, it returns `(0, -inf)`.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+.. seealso::
+
+        :func:`torch.linalg.det` computes the determinant of square matrices.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(sign, logabsdet)`.
+
+    `sign` will have the same dtype as :attr:`A`.
+
+    `logabsdet` will always be real-valued, even when :attr:`A` is complex.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> A
+    tensor([[ 0.0032, -0.2239, -1.1219],
+            [-0.6690,  0.1161,  0.4053],
+            [-1.6218, -0.9273, -0.0082]])
+    >>> torch.linalg.det(A)
+    tensor(-0.7576)
+    >>> torch.logdet(A)
+    tensor(nan)
+    >>> torch.linalg.slogdet(A)
+    torch.return_types.linalg_slogdet(sign=tensor(-1.), logabsdet=tensor(-0.2776))
+""",
+)
+
+eig = _add_docstr(
+    _linalg.linalg_eig,
+    r"""
+linalg.eig(A, *, out=None) -> (Tensor, Tensor)
+
+Computes the eigenvalue decomposition of a square matrix if it exists.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **eigenvalue decomposition** of a square matrix
+:math:`A \in \mathbb{K}^{n \times n}` (if it exists) is defined as
+
+.. math::
+
+    A = V \operatorname{diag}(\Lambda) V^{-1}\mathrlap{\qquad V \in \mathbb{C}^{n \times n}, \Lambda \in \mathbb{C}^n}
+
+This decomposition exists if and only if :math:`A` is `diagonalizable`_.
+This is the case when all its eigenvalues are different.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+The returned eigenvalues are not guaranteed to be in any specific order.
+
+.. note:: The eigenvalues and eigenvectors of a real matrix may be complex.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note"]}
+"""
+    + r"""
+
+.. warning:: This function assumes that :attr:`A` is `diagonalizable`_ (for example, when all the
+             eigenvalues are different). If it is not diagonalizable, the returned
+             eigenvalues will be correct but :math:`A \neq V \operatorname{diag}(\Lambda)V^{-1}`.
+
+.. warning:: The returned eigenvectors are normalized to have norm `1`.
+             Even then, the eigenvectors of a matrix are not unique, nor are they continuous with respect to
+             :attr:`A`. Due to this lack of uniqueness, different hardware and software may compute
+             different eigenvectors.
+
+             This non-uniqueness is caused by the fact that multiplying an eigenvector by
+             by :math:`e^{i \phi}, \phi \in \mathbb{R}` produces another set of valid eigenvectors
+             of the matrix.  For this reason, the loss function shall not depend on the phase of the
+             eigenvectors, as this quantity is not well-defined.
+             This is checked when computing the gradients of this function. As such,
+             when inputs are on a CUDA device, the computation of the gradients
+             of this function synchronizes that device with the CPU.
+
+
+.. warning:: Gradients computed using the `eigenvectors` tensor will only be finite when
+             :attr:`A` has distinct eigenvalues.
+             Furthermore, if the distance between any two eigenvalues is close to zero,
+             the gradient will be numerically unstable, as it depends on the eigenvalues
+             :math:`\lambda_i` through the computation of
+             :math:`\frac{1}{\min_{i \neq j} \lambda_i - \lambda_j}`.
+
+.. seealso::
+
+        :func:`torch.linalg.eigvals` computes only the eigenvalues.
+        Unlike :func:`torch.linalg.eig`, the gradients of :func:`~eigvals` are always
+        numerically stable.
+
+        :func:`torch.linalg.eigh` for a (faster) function that computes the eigenvalue decomposition
+        for Hermitian and symmetric matrices.
+
+        :func:`torch.linalg.svd` for a function that computes another type of spectral
+        decomposition that works on matrices of any shape.
+
+        :func:`torch.linalg.qr` for another (much faster) decomposition that works on matrices of
+        any shape.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                consisting of diagonalizable matrices.
+
+Keyword args:
+    out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(eigenvalues, eigenvectors)` which corresponds to :math:`\Lambda` and :math:`V` above.
+
+    `eigenvalues` and `eigenvectors` will always be complex-valued, even when :attr:`A` is real. The eigenvectors
+    will be given by the columns of `eigenvectors`.
+
+Examples::
+
+    >>> A = torch.randn(2, 2, dtype=torch.complex128)
+    >>> A
+    tensor([[ 0.9828+0.3889j, -0.4617+0.3010j],
+            [ 0.1662-0.7435j, -0.6139+0.0562j]], dtype=torch.complex128)
+    >>> L, V = torch.linalg.eig(A)
+    >>> L
+    tensor([ 1.1226+0.5738j, -0.7537-0.1286j], dtype=torch.complex128)
+    >>> V
+    tensor([[ 0.9218+0.0000j,  0.1882-0.2220j],
+            [-0.0270-0.3867j,  0.9567+0.0000j]], dtype=torch.complex128)
+    >>> torch.dist(V @ torch.diag(L) @ torch.linalg.inv(V), A)
+    tensor(7.7119e-16, dtype=torch.float64)
+
+    >>> A = torch.randn(3, 2, 2, dtype=torch.float64)
+    >>> L, V = torch.linalg.eig(A)
+    >>> torch.dist(V @ torch.diag_embed(L) @ torch.linalg.inv(V), A)
+    tensor(3.2841e-16, dtype=torch.float64)
+
+.. _diagonalizable:
+    https://en.wikipedia.org/wiki/Diagonalizable_matrix#Definition
+""",
+)
+
+eigvals = _add_docstr(
+    _linalg.linalg_eigvals,
+    r"""
+linalg.eigvals(A, *, out=None) -> Tensor
+
+Computes the eigenvalues of a square matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **eigenvalues** of a square matrix :math:`A \in \mathbb{K}^{n \times n}` are defined
+as the roots (counted with multiplicity) of the polynomial `p` of degree `n` given by
+
+.. math::
+
+    p(\lambda) = \operatorname{det}(A - \lambda \mathrm{I}_n)\mathrlap{\qquad \lambda \in \mathbb{C}}
+
+where :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+The returned eigenvalues are not guaranteed to be in any specific order.
+
+.. note:: The eigenvalues of a real matrix may be complex, as the roots of a real polynomial may be complex.
+
+          The eigenvalues of a matrix are always well-defined, even when the matrix is not diagonalizable.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note"]}
+"""
+    + r"""
+
+.. seealso::
+
+        :func:`torch.linalg.eig` computes the full eigenvalue decomposition.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Returns:
+    A complex-valued tensor containing the eigenvalues even when :attr:`A` is real.
+
+Examples::
+
+    >>> A = torch.randn(2, 2, dtype=torch.complex128)
+    >>> L = torch.linalg.eigvals(A)
+    >>> L
+    tensor([ 1.1226+0.5738j, -0.7537-0.1286j], dtype=torch.complex128)
+
+    >>> torch.dist(L, torch.linalg.eig(A).eigenvalues)
+    tensor(2.4576e-07)
+""",
+)
+
+eigh = _add_docstr(
+    _linalg.linalg_eigh,
+    r"""
+linalg.eigh(A, UPLO='L', *, out=None) -> (Tensor, Tensor)
+
+Computes the eigenvalue decomposition of a complex Hermitian or real symmetric matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **eigenvalue decomposition** of a complex Hermitian or real symmetric matrix
+:math:`A \in \mathbb{K}^{n \times n}` is defined as
+
+.. math::
+
+    A = Q \operatorname{diag}(\Lambda) Q^{\text{H}}\mathrlap{\qquad Q \in \mathbb{K}^{n \times n}, \Lambda \in \mathbb{R}^n}
+
+where :math:`Q^{\text{H}}` is the conjugate transpose when :math:`Q` is complex, and the transpose when :math:`Q` is real-valued.
+:math:`Q` is orthogonal in the real case and unitary in the complex case.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+:attr:`A` is assumed to be Hermitian (resp. symmetric), but this is not checked internally, instead:
+
+- If :attr:`UPLO`\ `= 'L'` (default), only the lower triangular part of the matrix is used in the computation.
+- If :attr:`UPLO`\ `= 'U'`, only the upper triangular part of the matrix is used.
+
+The eigenvalues are returned in ascending order.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note"]}
+"""
+    + r"""
+
+.. note:: The eigenvalues of real symmetric or complex Hermitian matrices are always real.
+
+.. warning:: The eigenvectors of a symmetric matrix are not unique, nor are they continuous with
+             respect to :attr:`A`. Due to this lack of uniqueness, different hardware and
+             software may compute different eigenvectors.
+
+             This non-uniqueness is caused by the fact that multiplying an eigenvector by
+             `-1` in the real case or by :math:`e^{i \phi}, \phi \in \mathbb{R}` in the complex
+             case produces another set of valid eigenvectors of the matrix.
+             For this reason, the loss function shall not depend on the phase of the eigenvectors, as
+             this quantity is not well-defined.
+             This is checked for complex inputs when computing the gradients of this function. As such,
+             when inputs are complex and are on a CUDA device, the computation of the gradients
+             of this function synchronizes that device with the CPU.
+
+.. warning:: Gradients computed using the `eigenvectors` tensor will only be finite when
+             :attr:`A` has distinct eigenvalues.
+             Furthermore, if the distance between any two eigenvalues is close to zero,
+             the gradient will be numerically unstable, as it depends on the eigenvalues
+             :math:`\lambda_i` through the computation of
+             :math:`\frac{1}{\min_{i \neq j} \lambda_i - \lambda_j}`.
+
+.. warning:: User may see pytorch crashes if running `eigh` on CUDA devices with CUDA versions before 12.1 update 1
+             with large ill-conditioned matrices as inputs.
+             Refer to :ref:`Linear Algebra Numerical Stability<Linear Algebra Stability>` for more details.
+             If this is the case, user may (1) tune their matrix inputs to be less ill-conditioned,
+             or (2) use :func:`torch.backends.cuda.preferred_linalg_library` to
+             try other supported backends.
+
+.. seealso::
+
+        :func:`torch.linalg.eigvalsh` computes only the eigenvalues of a Hermitian matrix.
+        Unlike :func:`torch.linalg.eigh`, the gradients of :func:`~eigvalsh` are always
+        numerically stable.
+
+        :func:`torch.linalg.cholesky` for a different decomposition of a Hermitian matrix.
+        The Cholesky decomposition gives less information about the matrix but is much faster
+        to compute than the eigenvalue decomposition.
+
+        :func:`torch.linalg.eig` for a (slower) function that computes the eigenvalue decomposition
+        of a not necessarily Hermitian square matrix.
+
+        :func:`torch.linalg.svd` for a (slower) function that computes the more general SVD
+        decomposition of matrices of any shape.
+
+        :func:`torch.linalg.qr` for another (much faster) decomposition that works on general
+        matrices.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                consisting of symmetric or Hermitian matrices.
+    UPLO ('L', 'U', optional): controls whether to use the upper or lower triangular part
+                               of :attr:`A` in the computations. Default: `'L'`.
+
+Keyword args:
+    out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(eigenvalues, eigenvectors)` which corresponds to :math:`\Lambda` and :math:`Q` above.
+
+    `eigenvalues` will always be real-valued, even when :attr:`A` is complex.
+    It will also be ordered in ascending order.
+
+    `eigenvectors` will have the same dtype as :attr:`A` and will contain the eigenvectors as its columns.
+
+Examples::
+    >>> A = torch.randn(2, 2, dtype=torch.complex128)
+    >>> A = A + A.T.conj()  # creates a Hermitian matrix
+    >>> A
+    tensor([[2.9228+0.0000j, 0.2029-0.0862j],
+            [0.2029+0.0862j, 0.3464+0.0000j]], dtype=torch.complex128)
+    >>> L, Q = torch.linalg.eigh(A)
+    >>> L
+    tensor([0.3277, 2.9415], dtype=torch.float64)
+    >>> Q
+    tensor([[-0.0846+-0.0000j, -0.9964+0.0000j],
+            [ 0.9170+0.3898j, -0.0779-0.0331j]], dtype=torch.complex128)
+    >>> torch.dist(Q @ torch.diag(L.cdouble()) @ Q.T.conj(), A)
+    tensor(6.1062e-16, dtype=torch.float64)
+
+    >>> A = torch.randn(3, 2, 2, dtype=torch.float64)
+    >>> A = A + A.mT  # creates a batch of symmetric matrices
+    >>> L, Q = torch.linalg.eigh(A)
+    >>> torch.dist(Q @ torch.diag_embed(L) @ Q.mH, A)
+    tensor(1.5423e-15, dtype=torch.float64)
+""",
+)
+
+eigvalsh = _add_docstr(
+    _linalg.linalg_eigvalsh,
+    r"""
+linalg.eigvalsh(A, UPLO='L', *, out=None) -> Tensor
+
+Computes the eigenvalues of a complex Hermitian or real symmetric matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **eigenvalues** of a complex Hermitian or real symmetric  matrix :math:`A \in \mathbb{K}^{n \times n}`
+are defined as the roots (counted with multiplicity) of the polynomial `p` of degree `n` given by
+
+.. math::
+
+    p(\lambda) = \operatorname{det}(A - \lambda \mathrm{I}_n)\mathrlap{\qquad \lambda \in \mathbb{R}}
+
+where :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix.
+The eigenvalues of a real symmetric or complex Hermitian matrix are always real.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+The eigenvalues are returned in ascending order.
+
+:attr:`A` is assumed to be Hermitian (resp. symmetric), but this is not checked internally, instead:
+
+- If :attr:`UPLO`\ `= 'L'` (default), only the lower triangular part of the matrix is used in the computation.
+- If :attr:`UPLO`\ `= 'U'`, only the upper triangular part of the matrix is used.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note"]}
+"""
+    + r"""
+
+.. seealso::
+
+        :func:`torch.linalg.eigh` computes the full eigenvalue decomposition.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                consisting of symmetric or Hermitian matrices.
+    UPLO ('L', 'U', optional): controls whether to use the upper or lower triangular part
+                               of :attr:`A` in the computations. Default: `'L'`.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Returns:
+    A real-valued tensor containing the eigenvalues even when :attr:`A` is complex.
+    The eigenvalues are returned in ascending order.
+
+Examples::
+
+    >>> A = torch.randn(2, 2, dtype=torch.complex128)
+    >>> A = A + A.T.conj()  # creates a Hermitian matrix
+    >>> A
+    tensor([[2.9228+0.0000j, 0.2029-0.0862j],
+            [0.2029+0.0862j, 0.3464+0.0000j]], dtype=torch.complex128)
+    >>> torch.linalg.eigvalsh(A)
+    tensor([0.3277, 2.9415], dtype=torch.float64)
+
+    >>> A = torch.randn(3, 2, 2, dtype=torch.float64)
+    >>> A = A + A.mT  # creates a batch of symmetric matrices
+    >>> torch.linalg.eigvalsh(A)
+    tensor([[ 2.5797,  3.4629],
+            [-4.1605,  1.3780],
+            [-3.1113,  2.7381]], dtype=torch.float64)
+""",
+)
+
+householder_product = _add_docstr(
+    _linalg.linalg_householder_product,
+    r"""
+householder_product(A, tau, *, out=None) -> Tensor
+
+Computes the first `n` columns of a product of Householder matrices.
+
+Let :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, and
+let :math:`A \in \mathbb{K}^{m \times n}` be a matrix with columns :math:`a_i \in \mathbb{K}^m`
+for :math:`i=1,\ldots,m` with :math:`m \geq n`. Denote by :math:`b_i` the vector resulting from
+zeroing out the first :math:`i-1` components of :math:`a_i` and setting to `1` the :math:`i`-th.
+For a vector :math:`\tau \in \mathbb{K}^k` with :math:`k \leq n`, this function computes the
+first :math:`n` columns of the matrix
+
+.. math::
+
+    H_1H_2 ... H_k \qquad\text{with}\qquad H_i = \mathrm{I}_m - \tau_i b_i b_i^{\text{H}}
+
+where :math:`\mathrm{I}_m` is the `m`-dimensional identity matrix and :math:`b^{\text{H}}` is the
+conjugate transpose when :math:`b` is complex, and the transpose when :math:`b` is real-valued.
+The output matrix is the same size as the input matrix :attr:`A`.
+
+See `Representation of Orthogonal or Unitary Matrices`_ for further details.
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if the inputs are batches of matrices then
+the output has the same batch dimensions.
+
+.. seealso::
+
+        :func:`torch.geqrf` can be used together with this function to form the `Q` from the
+        :func:`~qr` decomposition.
+
+        :func:`torch.ormqr` is a related function that computes the matrix multiplication
+        of a product of Householder matrices with another matrix.
+        However, that function is not supported by autograd.
+
+.. warning::
+    Gradient computations are only well-defined if :math:`\tau_i \neq \frac{1}{||a_i||^2}`.
+    If this condition is not met, no error will be thrown, but the gradient produced may contain `NaN`.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+    tau (Tensor): tensor of shape `(*, k)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Raises:
+    RuntimeError: if :attr:`A` doesn't satisfy the requirement `m >= n`,
+                  or :attr:`tau` doesn't satisfy the requirement `n >= k`.
+
+Examples::
+
+    >>> A = torch.randn(2, 2)
+    >>> h, tau = torch.geqrf(A)
+    >>> Q = torch.linalg.householder_product(h, tau)
+    >>> torch.dist(Q, torch.linalg.qr(A).Q)
+    tensor(0.)
+
+    >>> h = torch.randn(3, 2, 2, dtype=torch.complex128)
+    >>> tau = torch.randn(3, 1, dtype=torch.complex128)
+    >>> Q = torch.linalg.householder_product(h, tau)
+    >>> Q
+    tensor([[[ 1.8034+0.4184j,  0.2588-1.0174j],
+            [-0.6853+0.7953j,  2.0790+0.5620j]],
+
+            [[ 1.4581+1.6989j, -1.5360+0.1193j],
+            [ 1.3877-0.6691j,  1.3512+1.3024j]],
+
+            [[ 1.4766+0.5783j,  0.0361+0.6587j],
+            [ 0.6396+0.1612j,  1.3693+0.4481j]]], dtype=torch.complex128)
+
+.. _Representation of Orthogonal or Unitary Matrices:
+    https://www.netlib.org/lapack/lug/node128.html
+""",
+)
+
+ldl_factor = _add_docstr(
+    _linalg.linalg_ldl_factor,
+    r"""
+linalg.ldl_factor(A, *, hermitian=False, out=None) -> (Tensor, Tensor)
+
+Computes a compact representation of the LDL factorization of a Hermitian or symmetric (possibly indefinite) matrix.
+
+When :attr:`A` is complex valued it can be Hermitian (:attr:`hermitian`\ `= True`)
+or symmetric (:attr:`hermitian`\ `= False`).
+
+The factorization is of the form the form :math:`A = L D L^T`.
+If :attr:`hermitian` is `True` then transpose operation is the conjugate transpose.
+
+:math:`L` (or :math:`U`) and :math:`D` are stored in compact form in ``LD``.
+They follow the format specified by `LAPACK's sytrf`_ function.
+These tensors may be used in :func:`torch.linalg.ldl_solve` to solve linear systems.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.ldl_factor_ex")}
+"""
+    + r"""
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                consisting of symmetric or Hermitian matrices.
+
+Keyword args:
+    hermitian (bool, optional): whether to consider the input to be Hermitian or symmetric.
+                                For real-valued matrices, this switch has no effect. Default: `False`.
+    out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(LD, pivots)`.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> A = A @ A.mT # make symmetric
+    >>> A
+    tensor([[7.2079, 4.2414, 1.9428],
+            [4.2414, 3.4554, 0.3264],
+            [1.9428, 0.3264, 1.3823]])
+    >>> LD, pivots = torch.linalg.ldl_factor(A)
+    >>> LD
+    tensor([[ 7.2079,  0.0000,  0.0000],
+            [ 0.5884,  0.9595,  0.0000],
+            [ 0.2695, -0.8513,  0.1633]])
+    >>> pivots
+    tensor([1, 2, 3], dtype=torch.int32)
+
+.. _LAPACK's sytrf:
+    https://www.netlib.org/lapack/explore-html-3.6.1/d3/db6/group__double_s_ycomputational_gad91bde1212277b3e909eb6af7f64858a.html
+""",
+)
+
+ldl_factor_ex = _add_docstr(
+    _linalg.linalg_ldl_factor_ex,
+    r"""
+linalg.ldl_factor_ex(A, *, hermitian=False, check_errors=False, out=None) -> (Tensor, Tensor, Tensor)
+
+This is a version of :func:`~ldl_factor` that does not perform error checks unless :attr:`check_errors`\ `= True`.
+It also returns the :attr:`info` tensor returned by `LAPACK's sytrf`_.
+``info`` stores integer error codes from the backend library.
+A positive integer indicates the diagonal element of :math:`D` that is zero.
+Division by 0 will occur if the result is used for solving a system of linear equations.
+``info`` filled with zeros indicates that the factorization was successful.
+If ``check_errors=True`` and ``info`` contains positive integers, then a `RuntimeError` is thrown.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_ex"]}
+
+.. warning:: {common_notes["experimental_warning"]}
+"""
+    + r"""
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions
+                consisting of symmetric or Hermitian matrices.
+
+Keyword args:
+    hermitian (bool, optional): whether to consider the input to be Hermitian or symmetric.
+                                For real-valued matrices, this switch has no effect. Default: `False`.
+    check_errors (bool, optional): controls whether to check the content of ``info`` and raise
+                                   an error if it is non-zero. Default: `False`.
+    out (tuple, optional): tuple of three tensors to write the output to. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(LD, pivots, info)`.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> A = A @ A.mT # make symmetric
+    >>> A
+    tensor([[7.2079, 4.2414, 1.9428],
+            [4.2414, 3.4554, 0.3264],
+            [1.9428, 0.3264, 1.3823]])
+    >>> LD, pivots, info = torch.linalg.ldl_factor_ex(A)
+    >>> LD
+    tensor([[ 7.2079,  0.0000,  0.0000],
+            [ 0.5884,  0.9595,  0.0000],
+            [ 0.2695, -0.8513,  0.1633]])
+    >>> pivots
+    tensor([1, 2, 3], dtype=torch.int32)
+    >>> info
+    tensor(0, dtype=torch.int32)
+
+.. _LAPACK's sytrf:
+    https://www.netlib.org/lapack/explore-html-3.6.1/d3/db6/group__double_s_ycomputational_gad91bde1212277b3e909eb6af7f64858a.html
+""",
+)
+
+ldl_solve = _add_docstr(
+    _linalg.linalg_ldl_solve,
+    r"""
+linalg.ldl_solve(LD, pivots, B, *, hermitian=False, out=None) -> Tensor
+
+Computes the solution of a system of linear equations using the LDL factorization.
+
+:attr:`LD` and :attr:`pivots` are the compact representation of the LDL factorization and
+are expected to be computed by :func:`torch.linalg.ldl_factor_ex`.
+:attr:`hermitian` argument to this function should be the same
+as the corresponding arguments in :func:`torch.linalg.ldl_factor_ex`.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+"""
+    + rf"""
+.. warning:: {common_notes["experimental_warning"]}
+"""
+    + r"""
+
+Args:
+    LD (Tensor): the `n \times n` matrix or the batch of such matrices of size
+                      `(*, n, n)` where `*` is one or more batch dimensions.
+    pivots (Tensor): the pivots corresponding to the LDL factorization of :attr:`LD`.
+    B (Tensor): right-hand side tensor of shape `(*, n, k)`.
+
+Keyword args:
+    hermitian (bool, optional): whether to consider the decomposed matrix to be Hermitian or symmetric.
+                                For real-valued matrices, this switch has no effect. Default: `False`.
+    out (tuple, optional): output tensor. `B` may be passed as `out` and the result is computed in-place on `B`.
+                           Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.randn(2, 3, 3)
+    >>> A = A @ A.mT # make symmetric
+    >>> LD, pivots, info = torch.linalg.ldl_factor_ex(A)
+    >>> B = torch.randn(2, 3, 4)
+    >>> X = torch.linalg.ldl_solve(LD, pivots, B)
+    >>> torch.linalg.norm(A @ X - B)
+    >>> tensor(0.0001)
+""",
+)
+
+lstsq = _add_docstr(
+    _linalg.linalg_lstsq,
+    r"""
+torch.linalg.lstsq(A, B, rcond=None, *, driver=None) -> (Tensor, Tensor, Tensor, Tensor)
+
+Computes a solution to the least squares problem of a system of linear equations.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **least squares problem** for a linear system :math:`AX = B` with
+:math:`A \in \mathbb{K}^{m \times n}, B \in \mathbb{K}^{m \times k}` is defined as
+
+.. math::
+
+    \min_{X \in \mathbb{K}^{n \times k}} \|AX - B\|_F
+
+where :math:`\|-\|_F` denotes the Frobenius norm.
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if the inputs are batches of matrices then
+the output has the same batch dimensions.
+
+:attr:`driver` chooses the backend function that will be used.
+For CPU inputs the valid values are `'gels'`, `'gelsy'`, `'gelsd`, `'gelss'`.
+To choose the best driver on CPU consider:
+
+- If :attr:`A` is well-conditioned (its `condition number`_ is not too large), or you do not mind some precision loss.
+
+  - For a general matrix: `'gelsy'` (QR with pivoting) (default)
+  - If :attr:`A` is full-rank: `'gels'` (QR)
+
+- If :attr:`A` is not well-conditioned.
+
+  - `'gelsd'` (tridiagonal reduction and SVD)
+  - But if you run into memory issues: `'gelss'` (full SVD).
+
+For CUDA input, the only valid driver is `'gels'`, which assumes that :attr:`A` is full-rank.
+
+See also the `full description of these drivers`_
+
+:attr:`rcond` is used to determine the effective rank of the matrices in :attr:`A`
+when :attr:`driver` is one of (`'gelsy'`, `'gelsd'`, `'gelss'`).
+In this case, if :math:`\sigma_i` are the singular values of `A` in decreasing order,
+:math:`\sigma_i` will be rounded down to zero if :math:`\sigma_i \leq \text{rcond} \cdot \sigma_1`.
+If :attr:`rcond`\ `= None` (default), :attr:`rcond` is set to the machine precision of the dtype of :attr:`A` times `max(m, n)`.
+
+This function returns the solution to the problem and some extra information in a named tuple of
+four tensors `(solution, residuals, rank, singular_values)`. For inputs :attr:`A`, :attr:`B`
+of shape `(*, m, n)`, `(*, m, k)` respectively, it contains
+
+- `solution`: the least squares solution. It has shape `(*, n, k)`.
+- `residuals`: the squared residuals of the solutions, that is, :math:`\|AX - B\|_F^2`.
+  It has shape `(*, k)`.
+  It is computed when `m > n` and every matrix in :attr:`A` is full-rank,
+  otherwise, it is an empty tensor.
+  If :attr:`A` is a batch of matrices and any matrix in the batch is not full rank,
+  then an empty tensor is returned. This behavior may change in a future PyTorch release.
+- `rank`: tensor of ranks of the matrices in :attr:`A`.
+  It has shape equal to the batch dimensions of :attr:`A`.
+  It is computed when :attr:`driver` is one of (`'gelsy'`, `'gelsd'`, `'gelss'`),
+  otherwise it is an empty tensor.
+- `singular_values`: tensor of singular values of the matrices in :attr:`A`.
+  It has shape `(*, min(m, n))`.
+  It is computed when :attr:`driver` is one of (`'gelsd'`, `'gelss'`),
+  otherwise it is an empty tensor.
+
+.. note::
+    This function computes `X = \ `:attr:`A`\ `.pinverse() @ \ `:attr:`B` in a faster and
+    more numerically stable way than performing the computations separately.
+
+.. warning::
+    The default value of :attr:`rcond` may change in a future PyTorch release.
+    It is therefore recommended to use a fixed value to avoid potential
+    breaking changes.
+
+Args:
+    A (Tensor): lhs tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+    B (Tensor): rhs tensor of shape `(*, m, k)` where `*` is zero or more batch dimensions.
+    rcond (float, optional): used to determine the effective rank of :attr:`A`.
+                             If :attr:`rcond`\ `= None`, :attr:`rcond` is set to the machine
+                             precision of the dtype of :attr:`A` times `max(m, n)`. Default: `None`.
+
+Keyword args:
+    driver (str, optional): name of the LAPACK/MAGMA method to be used.
+        If `None`, `'gelsy'` is used for CPU inputs and `'gels'` for CUDA inputs.
+        Default: `None`.
+
+Returns:
+    A named tuple `(solution, residuals, rank, singular_values)`.
+
+Examples::
+
+    >>> A = torch.randn(1,3,3)
+    >>> A
+    tensor([[[-1.0838,  0.0225,  0.2275],
+         [ 0.2438,  0.3844,  0.5499],
+         [ 0.1175, -0.9102,  2.0870]]])
+    >>> B = torch.randn(2,3,3)
+    >>> B
+    tensor([[[-0.6772,  0.7758,  0.5109],
+         [-1.4382,  1.3769,  1.1818],
+         [-0.3450,  0.0806,  0.3967]],
+        [[-1.3994, -0.1521, -0.1473],
+         [ 1.9194,  1.0458,  0.6705],
+         [-1.1802, -0.9796,  1.4086]]])
+    >>> X = torch.linalg.lstsq(A, B).solution # A is broadcasted to shape (2, 3, 3)
+    >>> torch.dist(X, torch.linalg.pinv(A) @ B)
+    tensor(1.5152e-06)
+
+    >>> S = torch.linalg.lstsq(A, B, driver='gelsd').singular_values
+    >>> torch.dist(S, torch.linalg.svdvals(A))
+    tensor(2.3842e-07)
+
+    >>> A[:, 0].zero_()  # Decrease the rank of A
+    >>> rank = torch.linalg.lstsq(A, B).rank
+    >>> rank
+    tensor([2])
+
+.. _condition number:
+    https://pytorch.org/docs/main/linalg.html#torch.linalg.cond
+.. _full description of these drivers:
+    https://www.netlib.org/lapack/lug/node27.html
+""",
+)
+
+matrix_power = _add_docstr(
+    _linalg.linalg_matrix_power,
+    r"""
+matrix_power(A, n, *, out=None) -> Tensor
+
+Computes the `n`-th power of a square matrix for an integer `n`.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+If :attr:`n`\ `= 0`, it returns the identity matrix (or batch) of the same shape
+as :attr:`A`. If :attr:`n` is negative, it returns the inverse of each matrix
+(if invertible) raised to the power of `abs(n)`.
+
+.. note::
+    Consider using :func:`torch.linalg.solve` if possible for multiplying a matrix on the left by
+    a negative power as, if :attr:`n`\ `> 0`::
+
+        torch.linalg.solve(matrix_power(A, n), B) == matrix_power(A, -n)  @ B
+
+    It is always preferred to use :func:`~solve` when possible, as it is faster and more
+    numerically stable than computing :math:`A^{-n}` explicitly.
+
+.. seealso::
+
+        :func:`torch.linalg.solve` computes :attr:`A`\ `.inverse() @ \ `:attr:`B` with a
+        numerically stable algorithm.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, m)` where `*` is zero or more batch dimensions.
+    n (int): the exponent.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Raises:
+    RuntimeError: if :attr:`n`\ `< 0` and the matrix :attr:`A` or any matrix in the
+                  batch of matrices :attr:`A` is not invertible.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> torch.linalg.matrix_power(A, 0)
+    tensor([[1., 0., 0.],
+            [0., 1., 0.],
+            [0., 0., 1.]])
+    >>> torch.linalg.matrix_power(A, 3)
+    tensor([[ 1.0756,  0.4980,  0.0100],
+            [-1.6617,  1.4994, -1.9980],
+            [-0.4509,  0.2731,  0.8001]])
+    >>> torch.linalg.matrix_power(A.expand(2, -1, -1), -2)
+    tensor([[[ 0.2640,  0.4571, -0.5511],
+            [-1.0163,  0.3491, -1.5292],
+            [-0.4899,  0.0822,  0.2773]],
+            [[ 0.2640,  0.4571, -0.5511],
+            [-1.0163,  0.3491, -1.5292],
+            [-0.4899,  0.0822,  0.2773]]])
+""",
+)
+
+matrix_rank = _add_docstr(
+    _linalg.linalg_matrix_rank,
+    r"""
+linalg.matrix_rank(A, *, atol=None, rtol=None, hermitian=False, out=None) -> Tensor
+
+Computes the numerical rank of a matrix.
+
+The matrix rank is computed as the number of singular values
+(or eigenvalues in absolute value when :attr:`hermitian`\ `= True`)
+that are greater than :math:`\max(\text{atol}, \sigma_1 * \text{rtol})` threshold,
+where :math:`\sigma_1` is the largest singular value (or eigenvalue).
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+If :attr:`hermitian`\ `= True`, :attr:`A` is assumed to be Hermitian if complex or
+symmetric if real, but this is not checked internally. Instead, just the lower
+triangular part of the matrix is used in the computations.
+
+If :attr:`rtol` is not specified and :attr:`A` is a matrix of dimensions `(m, n)`,
+the relative tolerance is set to be :math:`\text{rtol} = \max(m, n) \varepsilon`
+and :math:`\varepsilon` is the epsilon value for the dtype of :attr:`A` (see :class:`.finfo`).
+If :attr:`rtol` is not specified and :attr:`atol` is specified to be larger than zero then
+:attr:`rtol` is set to zero.
+
+If :attr:`atol` or :attr:`rtol` is a :class:`torch.Tensor`, its shape must be broadcastable to that
+of the singular values of :attr:`A` as returned by :func:`torch.linalg.svdvals`.
+
+.. note::
+    This function has NumPy compatible variant `linalg.matrix_rank(A, tol, hermitian=False)`.
+    However, use of the positional argument :attr:`tol` is deprecated in favor of :attr:`atol` and :attr:`rtol`.
+
+"""
+    + rf"""
+.. note:: The matrix rank is computed using a singular value decomposition
+          :func:`torch.linalg.svdvals` if :attr:`hermitian`\ `= False` (default) and the eigenvalue
+          decomposition :func:`torch.linalg.eigvalsh` when :attr:`hermitian`\ `= True`.
+          {common_notes["sync_note"]}
+"""
+    + r"""
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+    tol (float, Tensor, optional): [NumPy Compat] Alias for :attr:`atol`. Default: `None`.
+
+Keyword args:
+    atol (float, Tensor, optional): the absolute tolerance value. When `None` it's considered to be zero.
+                                    Default: `None`.
+    rtol (float, Tensor, optional): the relative tolerance value. See above for the value it takes when `None`.
+                                    Default: `None`.
+    hermitian(bool): indicates whether :attr:`A` is Hermitian if complex
+                     or symmetric if real. Default: `False`.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.eye(10)
+    >>> torch.linalg.matrix_rank(A)
+    tensor(10)
+    >>> B = torch.eye(10)
+    >>> B[0, 0] = 0
+    >>> torch.linalg.matrix_rank(B)
+    tensor(9)
+
+    >>> A = torch.randn(4, 3, 2)
+    >>> torch.linalg.matrix_rank(A)
+    tensor([2, 2, 2, 2])
+
+    >>> A = torch.randn(2, 4, 2, 3)
+    >>> torch.linalg.matrix_rank(A)
+    tensor([[2, 2, 2, 2],
+            [2, 2, 2, 2]])
+
+    >>> A = torch.randn(2, 4, 3, 3, dtype=torch.complex64)
+    >>> torch.linalg.matrix_rank(A)
+    tensor([[3, 3, 3, 3],
+            [3, 3, 3, 3]])
+    >>> torch.linalg.matrix_rank(A, hermitian=True)
+    tensor([[3, 3, 3, 3],
+            [3, 3, 3, 3]])
+    >>> torch.linalg.matrix_rank(A, atol=1.0, rtol=0.0)
+    tensor([[3, 2, 2, 2],
+            [1, 2, 1, 2]])
+    >>> torch.linalg.matrix_rank(A, atol=1.0, rtol=0.0, hermitian=True)
+    tensor([[2, 2, 2, 1],
+            [1, 2, 2, 2]])
+""",
+)
+
+norm = _add_docstr(
+    _linalg.linalg_norm,
+    r"""
+linalg.norm(A, ord=None, dim=None, keepdim=False, *, out=None, dtype=None) -> Tensor
+
+Computes a vector or matrix norm.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+
+Whether this function computes a vector or matrix norm is determined as follows:
+
+- If :attr:`dim` is an `int`, the vector norm will be computed.
+- If :attr:`dim` is a `2`-`tuple`, the matrix norm will be computed.
+- If :attr:`dim`\ `= None` and :attr:`ord`\ `= None`,
+  :attr:`A` will be flattened to 1D and the `2`-norm of the resulting vector will be computed.
+- If :attr:`dim`\ `= None` and :attr:`ord` `!= None`, :attr:`A` must be 1D or 2D.
+
+:attr:`ord` defines the norm that is computed. The following norms are supported:
+
+======================     ==========================   ======================================================
+:attr:`ord`                norm for matrices            norm for vectors
+======================     ==========================   ======================================================
+`None` (default)           Frobenius norm               `2`-norm (see below)
+`'fro'`                    Frobenius norm               -- not supported --
+`'nuc'`                    nuclear norm                 -- not supported --
+`inf`                      `max(sum(abs(x), dim=1))`    `max(abs(x))`
+`-inf`                     `min(sum(abs(x), dim=1))`    `min(abs(x))`
+`0`                        -- not supported --          `sum(x != 0)`
+`1`                        `max(sum(abs(x), dim=0))`    as below
+`-1`                       `min(sum(abs(x), dim=0))`    as below
+`2`                        largest `singular value`_    as below
+`-2`                       smallest `singular value`_   as below
+other `int` or `float`     -- not supported --          `sum(abs(x)^{ord})^{(1 / ord)}`
+======================     ==========================   ======================================================
+
+where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object.
+
+.. seealso::
+
+        :func:`torch.linalg.vector_norm` computes a vector norm.
+
+        :func:`torch.linalg.matrix_norm` computes a matrix norm.
+
+        The above functions are often clearer and more flexible than using :func:`torch.linalg.norm`.
+        For example, `torch.linalg.norm(A, ord=1, dim=(0, 1))` always
+        computes a matrix norm, but with `torch.linalg.vector_norm(A, ord=1, dim=(0, 1))` it is possible
+        to compute a vector norm over the two dimensions.
+
+Args:
+    A (Tensor): tensor of shape `(*, n)` or `(*, m, n)` where `*` is zero or more batch dimensions
+    ord (int, float, inf, -inf, 'fro', 'nuc', optional): order of norm. Default: `None`
+    dim (int, Tuple[int], optional): dimensions over which to compute
+        the vector or matrix norm. See above for the behavior when :attr:`dim`\ `= None`.
+        Default: `None`
+    keepdim (bool, optional): If set to `True`, the reduced dimensions are retained
+        in the result as dimensions with size one. Default: `False`
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+    dtype (:class:`torch.dtype`, optional): If specified, the input tensor is cast to
+        :attr:`dtype` before performing the operation, and the returned tensor's type
+        will be :attr:`dtype`. Default: `None`
+
+Returns:
+    A real-valued tensor, even when :attr:`A` is complex.
+
+Examples::
+
+    >>> from torch import linalg as LA
+    >>> a = torch.arange(9, dtype=torch.float) - 4
+    >>> a
+    tensor([-4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.])
+    >>> B = a.reshape((3, 3))
+    >>> B
+    tensor([[-4., -3., -2.],
+            [-1.,  0.,  1.],
+            [ 2.,  3.,  4.]])
+
+    >>> LA.norm(a)
+    tensor(7.7460)
+    >>> LA.norm(B)
+    tensor(7.7460)
+    >>> LA.norm(B, 'fro')
+    tensor(7.7460)
+    >>> LA.norm(a, float('inf'))
+    tensor(4.)
+    >>> LA.norm(B, float('inf'))
+    tensor(9.)
+    >>> LA.norm(a, -float('inf'))
+    tensor(0.)
+    >>> LA.norm(B, -float('inf'))
+    tensor(2.)
+
+    >>> LA.norm(a, 1)
+    tensor(20.)
+    >>> LA.norm(B, 1)
+    tensor(7.)
+    >>> LA.norm(a, -1)
+    tensor(0.)
+    >>> LA.norm(B, -1)
+    tensor(6.)
+    >>> LA.norm(a, 2)
+    tensor(7.7460)
+    >>> LA.norm(B, 2)
+    tensor(7.3485)
+
+    >>> LA.norm(a, -2)
+    tensor(0.)
+    >>> LA.norm(B.double(), -2)
+    tensor(1.8570e-16, dtype=torch.float64)
+    >>> LA.norm(a, 3)
+    tensor(5.8480)
+    >>> LA.norm(a, -3)
+    tensor(0.)
+
+Using the :attr:`dim` argument to compute vector norms::
+
+    >>> c = torch.tensor([[1., 2., 3.],
+    ...                   [-1, 1, 4]])
+    >>> LA.norm(c, dim=0)
+    tensor([1.4142, 2.2361, 5.0000])
+    >>> LA.norm(c, dim=1)
+    tensor([3.7417, 4.2426])
+    >>> LA.norm(c, ord=1, dim=1)
+    tensor([6., 6.])
+
+Using the :attr:`dim` argument to compute matrix norms::
+
+    >>> A = torch.arange(8, dtype=torch.float).reshape(2, 2, 2)
+    >>> LA.norm(A, dim=(1,2))
+    tensor([ 3.7417, 11.2250])
+    >>> LA.norm(A[0, :, :]), LA.norm(A[1, :, :])
+    (tensor(3.7417), tensor(11.2250))
+
+.. _singular value:
+    https://en.wikipedia.org/wiki/Singular_value_decomposition#Singular_values,_singular_vectors,_and_their_relation_to_the_SVD
+""",
+)
+
+vector_norm = _add_docstr(
+    _linalg.linalg_vector_norm,
+    r"""
+linalg.vector_norm(x, ord=2, dim=None, keepdim=False, *, dtype=None, out=None) -> Tensor
+
+Computes a vector norm.
+
+If :attr:`x` is complex valued, it computes the norm of :attr:`x`\ `.abs()`
+
+Supports input of float, double, cfloat and cdouble dtypes.
+
+This function does not necessarily treat multidimensional :attr:`x` as a batch of
+vectors, instead:
+
+- If :attr:`dim`\ `= None`, :attr:`x` will be flattened before the norm is computed.
+- If :attr:`dim` is an `int` or a `tuple`, the norm will be computed over these dimensions
+  and the other dimensions will be treated as batch dimensions.
+
+This behavior is for consistency with :func:`torch.linalg.norm`.
+
+:attr:`ord` defines the vector norm that is computed. The following norms are supported:
+
+======================   ===============================
+:attr:`ord`              vector norm
+======================   ===============================
+`2` (default)            `2`-norm (see below)
+`inf`                    `max(abs(x))`
+`-inf`                   `min(abs(x))`
+`0`                      `sum(x != 0)`
+other `int` or `float`   `sum(abs(x)^{ord})^{(1 / ord)}`
+======================   ===============================
+
+where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object.
+
+:attr:`dtype` may be used to perform the computation in a more precise dtype.
+It is semantically equivalent to calling ``linalg.vector_norm(x.to(dtype))``
+but it is faster in some cases.
+
+.. seealso::
+
+        :func:`torch.linalg.matrix_norm` computes a matrix norm.
+
+Args:
+    x (Tensor): tensor, flattened by default, but this behavior can be
+        controlled using :attr:`dim`.  (Note: the keyword argument
+        `input` can also be used as an alias for `x`.)
+    ord (int, float, inf, -inf, 'fro', 'nuc', optional): order of norm. Default: `2`
+    dim (int, Tuple[int], optional): dimensions over which to compute
+        the norm. See above for the behavior when :attr:`dim`\ `= None`.
+        Default: `None`
+    keepdim (bool, optional): If set to `True`, the reduced dimensions are retained
+        in the result as dimensions with size one. Default: `False`
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+    dtype (:class:`torch.dtype`, optional): type used to perform the accumulation and the return.
+        If specified, :attr:`x` is cast to :attr:`dtype` before performing the operation,
+        and the returned tensor's type will be :attr:`dtype` if real and of its real counterpart if complex.
+        :attr:`dtype` may be complex if :attr:`x` is complex, otherwise it must be real.
+        :attr:`x` should be convertible without narrowing to :attr:`dtype`. Default: None
+
+Returns:
+    A real-valued tensor, even when :attr:`x` is complex.
+
+Examples::
+
+    >>> from torch import linalg as LA
+    >>> a = torch.arange(9, dtype=torch.float) - 4
+    >>> a
+    tensor([-4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.])
+    >>> B = a.reshape((3, 3))
+    >>> B
+    tensor([[-4., -3., -2.],
+            [-1.,  0.,  1.],
+            [ 2.,  3.,  4.]])
+    >>> LA.vector_norm(a, ord=3.5)
+    tensor(5.4345)
+    >>> LA.vector_norm(B, ord=3.5)
+    tensor(5.4345)
+""",
+)
+
+matrix_norm = _add_docstr(
+    _linalg.linalg_matrix_norm,
+    r"""
+linalg.matrix_norm(A, ord='fro', dim=(-2, -1), keepdim=False, *, dtype=None, out=None) -> Tensor
+
+Computes a matrix norm.
+
+If :attr:`A` is complex valued, it computes the norm of :attr:`A`\ `.abs()`
+
+Support input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices: the norm will be computed over the
+dimensions specified by the 2-tuple :attr:`dim` and the other dimensions will
+be treated as batch dimensions. The output will have the same batch dimensions.
+
+:attr:`ord` defines the matrix norm that is computed. The following norms are supported:
+
+======================   ========================================================
+:attr:`ord`              matrix norm
+======================   ========================================================
+`'fro'` (default)        Frobenius norm
+`'nuc'`                  nuclear norm
+`inf`                    `max(sum(abs(x), dim=1))`
+`-inf`                   `min(sum(abs(x), dim=1))`
+`1`                      `max(sum(abs(x), dim=0))`
+`-1`                     `min(sum(abs(x), dim=0))`
+`2`                      largest singular value
+`-2`                     smallest singular value
+======================   ========================================================
+
+where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object.
+
+Args:
+    A (Tensor): tensor with two or more dimensions. By default its
+        shape is interpreted as `(*, m, n)` where `*` is zero or more
+        batch dimensions, but this behavior can be controlled using :attr:`dim`.
+    ord (int, inf, -inf, 'fro', 'nuc', optional): order of norm. Default: `'fro'`
+    dim (Tuple[int, int], optional): dimensions over which to compute the norm. Default: `(-2, -1)`
+    keepdim (bool, optional): If set to `True`, the reduced dimensions are retained
+        in the result as dimensions with size one. Default: `False`
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+    dtype (:class:`torch.dtype`, optional): If specified, the input tensor is cast to
+        :attr:`dtype` before performing the operation, and the returned tensor's type
+        will be :attr:`dtype`. Default: `None`
+
+Returns:
+    A real-valued tensor, even when :attr:`A` is complex.
+
+Examples::
+
+    >>> from torch import linalg as LA
+    >>> A = torch.arange(9, dtype=torch.float).reshape(3, 3)
+    >>> A
+    tensor([[0., 1., 2.],
+            [3., 4., 5.],
+            [6., 7., 8.]])
+    >>> LA.matrix_norm(A)
+    tensor(14.2829)
+    >>> LA.matrix_norm(A, ord=-1)
+    tensor(9.)
+    >>> B = A.expand(2, -1, -1)
+    >>> B
+    tensor([[[0., 1., 2.],
+            [3., 4., 5.],
+            [6., 7., 8.]],
+
+            [[0., 1., 2.],
+            [3., 4., 5.],
+            [6., 7., 8.]]])
+    >>> LA.matrix_norm(B)
+    tensor([14.2829, 14.2829])
+    >>> LA.matrix_norm(B, dim=(0, 2))
+    tensor([ 3.1623, 10.0000, 17.2627])
+""",
+)
+
+matmul = _add_docstr(
+    _linalg.linalg_matmul,
+    r"""
+linalg.matmul(input, other, *, out=None) -> Tensor
+
+Alias for :func:`torch.matmul`
+""",
+)
+
+diagonal = _add_docstr(
+    _linalg.linalg_diagonal,
+    r"""
+linalg.diagonal(A, *, offset=0, dim1=-2, dim2=-1) -> Tensor
+
+Alias for :func:`torch.diagonal` with defaults :attr:`dim1`\ `= -2`, :attr:`dim2`\ `= -1`.
+""",
+)
+
+multi_dot = _add_docstr(
+    _linalg.linalg_multi_dot,
+    r"""
+linalg.multi_dot(tensors, *, out=None)
+
+Efficiently multiplies two or more matrices by reordering the multiplications so that
+the fewest arithmetic operations are performed.
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+This function does not support batched inputs.
+
+Every tensor in :attr:`tensors` must be 2D, except for the first and last which
+may be 1D. If the first tensor is a 1D vector of shape `(n,)` it is treated as a row vector
+of shape `(1, n)`, similarly if the last tensor is a 1D vector of shape `(n,)` it is treated
+as a column vector of shape `(n, 1)`.
+
+If the first and last tensors are matrices, the output will be a matrix.
+However, if either is a 1D vector, then the output will be a 1D vector.
+
+Differences with `numpy.linalg.multi_dot`:
+
+- Unlike `numpy.linalg.multi_dot`, the first and last tensors must either be 1D or 2D
+  whereas NumPy allows them to be nD
+
+.. warning:: This function does not broadcast.
+
+.. note:: This function is implemented by chaining :func:`torch.mm` calls after
+          computing the optimal matrix multiplication order.
+
+.. note:: The cost of multiplying two matrices with shapes `(a, b)` and `(b, c)` is
+          `a * b * c`. Given matrices `A`, `B`, `C` with shapes `(10, 100)`,
+          `(100, 5)`, `(5, 50)` respectively, we can calculate the cost of different
+          multiplication orders as follows:
+
+          .. math::
+
+             \begin{align*}
+             \operatorname{cost}((AB)C) &= 10 \times 100 \times 5 + 10 \times 5 \times 50 = 7500 \\
+             \operatorname{cost}(A(BC)) &= 10 \times 100 \times 50 + 100 \times 5 \times 50 = 75000
+             \end{align*}
+
+          In this case, multiplying `A` and `B` first followed by `C` is 10 times faster.
+
+Args:
+    tensors (Sequence[Tensor]): two or more tensors to multiply. The first and last
+        tensors may be 1D or 2D. Every other tensor must be 2D.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> from torch.linalg import multi_dot
+
+    >>> multi_dot([torch.tensor([1, 2]), torch.tensor([2, 3])])
+    tensor(8)
+    >>> multi_dot([torch.tensor([[1, 2]]), torch.tensor([2, 3])])
+    tensor([8])
+    >>> multi_dot([torch.tensor([[1, 2]]), torch.tensor([[2], [3]])])
+    tensor([[8]])
+
+    >>> A = torch.arange(2 * 3).view(2, 3)
+    >>> B = torch.arange(3 * 2).view(3, 2)
+    >>> C = torch.arange(2 * 2).view(2, 2)
+    >>> multi_dot((A, B, C))
+    tensor([[ 26,  49],
+            [ 80, 148]])
+""",
+)
+
+svd = _add_docstr(
+    _linalg.linalg_svd,
+    r"""
+linalg.svd(A, full_matrices=True, *, driver=None, out=None) -> (Tensor, Tensor, Tensor)
+
+Computes the singular value decomposition (SVD) of a matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **full SVD** of a matrix
+:math:`A \in \mathbb{K}^{m \times n}`, if `k = min(m,n)`, is defined as
+
+.. math::
+
+    A = U \operatorname{diag}(S) V^{\text{H}}
+    \mathrlap{\qquad U \in \mathbb{K}^{m \times m}, S \in \mathbb{R}^k, V \in \mathbb{K}^{n \times n}}
+
+where :math:`\operatorname{diag}(S) \in \mathbb{K}^{m \times n}`,
+:math:`V^{\text{H}}` is the conjugate transpose when :math:`V` is complex, and the transpose when :math:`V` is real-valued.
+The matrices  :math:`U`, :math:`V` (and thus :math:`V^{\text{H}}`) are orthogonal in the real case, and unitary in the complex case.
+
+When `m > n` (resp. `m < n`) we can drop the last `m - n` (resp. `n - m`) columns of `U` (resp. `V`) to form the **reduced SVD**:
+
+.. math::
+
+    A = U \operatorname{diag}(S) V^{\text{H}}
+    \mathrlap{\qquad U \in \mathbb{K}^{m \times k}, S \in \mathbb{R}^k, V \in \mathbb{K}^{n \times k}}
+
+where :math:`\operatorname{diag}(S) \in \mathbb{K}^{k \times k}`.
+In this case, :math:`U` and :math:`V` also have orthonormal columns.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+The returned decomposition is a named tuple `(U, S, Vh)`
+which corresponds to :math:`U`, :math:`S`, :math:`V^{\text{H}}` above.
+
+The singular values are returned in descending order.
+
+The parameter :attr:`full_matrices` chooses between the full (default) and reduced SVD.
+
+The :attr:`driver` kwarg may be used in CUDA with a cuSOLVER backend to choose the algorithm used to compute the SVD.
+The choice of a driver is a trade-off between accuracy and speed.
+
+- If :attr:`A` is well-conditioned (its `condition number`_ is not too large), or you do not mind some precision loss.
+
+  - For a general matrix: `'gesvdj'` (Jacobi method)
+  - If :attr:`A` is tall or wide (`m >> n` or `m << n`): `'gesvda'` (Approximate method)
+
+- If :attr:`A` is not well-conditioned or precision is relevant: `'gesvd'` (QR based)
+
+By default (:attr:`driver`\ `= None`), we call `'gesvdj'` and, if it fails, we fallback to `'gesvd'`.
+
+Differences with `numpy.linalg.svd`:
+
+- Unlike `numpy.linalg.svd`, this function always returns a tuple of three tensors
+  and it doesn't support `compute_uv` argument.
+  Please use :func:`torch.linalg.svdvals`, which computes only the singular values,
+  instead of `compute_uv=False`.
+
+.. note:: When :attr:`full_matrices`\ `= True`, the gradients with respect to `U[..., :, min(m, n):]`
+          and `Vh[..., min(m, n):, :]` will be ignored, as those vectors can be arbitrary bases
+          of the corresponding subspaces.
+
+.. warning:: The returned tensors `U` and `V` are not unique, nor are they continuous with
+             respect to :attr:`A`.
+             Due to this lack of uniqueness, different hardware and software may compute
+             different singular vectors.
+
+             This non-uniqueness is caused by the fact that multiplying any pair of singular
+             vectors :math:`u_k, v_k` by `-1` in the real case or by
+             :math:`e^{i \phi}, \phi \in \mathbb{R}` in the complex case produces another two
+             valid singular vectors of the matrix.
+             For this reason, the loss function shall not depend on this :math:`e^{i \phi}` quantity,
+             as it is not well-defined.
+             This is checked for complex inputs when computing the gradients of this function. As such,
+             when inputs are complex and are on a CUDA device, the computation of the gradients
+             of this function synchronizes that device with the CPU.
+
+.. warning:: Gradients computed using `U` or `Vh` will only be finite when
+             :attr:`A` does not have repeated singular values. If :attr:`A` is rectangular,
+             additionally, zero must also not be one of its singular values.
+             Furthermore, if the distance between any two singular values is close to zero,
+             the gradient will be numerically unstable, as it depends on the singular values
+             :math:`\sigma_i` through the computation of
+             :math:`\frac{1}{\min_{i \neq j} \sigma_i^2 - \sigma_j^2}`.
+             In the rectangular case, the gradient will also be numerically unstable when
+             :attr:`A` has small singular values, as it also depends on the computation of
+             :math:`\frac{1}{\sigma_i}`.
+
+.. seealso::
+
+        :func:`torch.linalg.svdvals` computes only the singular values.
+        Unlike :func:`torch.linalg.svd`, the gradients of :func:`~svdvals` are always
+        numerically stable.
+
+        :func:`torch.linalg.eig` for a function that computes another type of spectral
+        decomposition of a matrix. The eigendecomposition works just on square matrices.
+
+        :func:`torch.linalg.eigh` for a (faster) function that computes the eigenvalue decomposition
+        for Hermitian and symmetric matrices.
+
+        :func:`torch.linalg.qr` for another (much faster) decomposition that works on general
+        matrices.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+    full_matrices (bool, optional): controls whether to compute the full or reduced
+                                    SVD, and consequently,
+                                    the shape of the returned tensors
+                                    `U` and `Vh`. Default: `True`.
+
+Keyword args:
+    driver (str, optional): name of the cuSOLVER method to be used. This keyword argument only works on CUDA inputs.
+        Available options are: `None`, `gesvd`, `gesvdj`, and `gesvda`.
+        Default: `None`.
+    out (tuple, optional): output tuple of three tensors. Ignored if `None`.
+
+Returns:
+    A named tuple `(U, S, Vh)` which corresponds to :math:`U`, :math:`S`, :math:`V^{\text{H}}` above.
+
+    `S` will always be real-valued, even when :attr:`A` is complex.
+    It will also be ordered in descending order.
+
+    `U` and `Vh` will have the same dtype as :attr:`A`. The left / right singular vectors will be given by
+    the columns of `U` and the rows of `Vh` respectively.
+
+Examples::
+
+    >>> A = torch.randn(5, 3)
+    >>> U, S, Vh = torch.linalg.svd(A, full_matrices=False)
+    >>> U.shape, S.shape, Vh.shape
+    (torch.Size([5, 3]), torch.Size([3]), torch.Size([3, 3]))
+    >>> torch.dist(A, U @ torch.diag(S) @ Vh)
+    tensor(1.0486e-06)
+
+    >>> U, S, Vh = torch.linalg.svd(A)
+    >>> U.shape, S.shape, Vh.shape
+    (torch.Size([5, 5]), torch.Size([3]), torch.Size([3, 3]))
+    >>> torch.dist(A, U[:, :3] @ torch.diag(S) @ Vh)
+    tensor(1.0486e-06)
+
+    >>> A = torch.randn(7, 5, 3)
+    >>> U, S, Vh = torch.linalg.svd(A, full_matrices=False)
+    >>> torch.dist(A, U @ torch.diag_embed(S) @ Vh)
+    tensor(3.0957e-06)
+
+.. _condition number:
+    https://pytorch.org/docs/main/linalg.html#torch.linalg.cond
+.. _the resulting vectors will span the same subspace:
+    https://en.wikipedia.org/wiki/Singular_value_decomposition#Singular_values,_singular_vectors,_and_their_relation_to_the_SVD
+""",
+)
+
+svdvals = _add_docstr(
+    _linalg.linalg_svdvals,
+    r"""
+linalg.svdvals(A, *, driver=None, out=None) -> Tensor
+
+Computes the singular values of a matrix.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+The singular values are returned in descending order.
+
+.. note:: This function is equivalent to NumPy's `linalg.svd(A, compute_uv=False)`.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note"]}
+"""
+    + r"""
+
+.. seealso::
+
+        :func:`torch.linalg.svd` computes the full singular value decomposition.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    driver (str, optional): name of the cuSOLVER method to be used. This keyword argument only works on CUDA inputs.
+        Available options are: `None`, `gesvd`, `gesvdj`, and `gesvda`.
+        Check :func:`torch.linalg.svd` for details.
+        Default: `None`.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Returns:
+    A real-valued tensor, even when :attr:`A` is complex.
+
+Examples::
+
+    >>> A = torch.randn(5, 3)
+    >>> S = torch.linalg.svdvals(A)
+    >>> S
+    tensor([2.5139, 2.1087, 1.1066])
+
+    >>> torch.dist(S, torch.linalg.svd(A, full_matrices=False).S)
+    tensor(2.4576e-07)
+""",
+)
+
+cond = _add_docstr(
+    _linalg.linalg_cond,
+    r"""
+linalg.cond(A, p=None, *, out=None) -> Tensor
+
+Computes the condition number of a matrix with respect to a matrix norm.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **condition number** :math:`\kappa` of a matrix
+:math:`A \in \mathbb{K}^{n \times n}` is defined as
+
+.. math::
+
+    \kappa(A) = \|A\|_p\|A^{-1}\|_p
+
+The condition number of :attr:`A` measures the numerical stability of the linear system `AX = B`
+with respect to a matrix norm.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+:attr:`p` defines the matrix norm that is computed. The following norms are supported:
+
+=========    =================================
+:attr:`p`    matrix norm
+=========    =================================
+`None`       `2`-norm (largest singular value)
+`'fro'`      Frobenius norm
+`'nuc'`      nuclear norm
+`inf`        `max(sum(abs(x), dim=1))`
+`-inf`       `min(sum(abs(x), dim=1))`
+`1`          `max(sum(abs(x), dim=0))`
+`-1`         `min(sum(abs(x), dim=0))`
+`2`          largest singular value
+`-2`         smallest singular value
+=========    =================================
+
+where `inf` refers to `float('inf')`, NumPy's `inf` object, or any equivalent object.
+
+For :attr:`p` is one of `('fro', 'nuc', inf, -inf, 1, -1)`, this function uses
+:func:`torch.linalg.norm` and :func:`torch.linalg.inv`.
+As such, in this case, the matrix (or every matrix in the batch) :attr:`A` has to be square
+and invertible.
+
+For :attr:`p` in `(2, -2)`, this function can be computed in terms of the singular values
+:math:`\sigma_1 \geq \ldots \geq \sigma_n`
+
+.. math::
+
+    \kappa_2(A) = \frac{\sigma_1}{\sigma_n}\qquad \kappa_{-2}(A) = \frac{\sigma_n}{\sigma_1}
+
+In these cases, it is computed using :func:`torch.linalg.svdvals`. For these norms, the matrix
+(or every matrix in the batch) :attr:`A` may have any shape.
+
+.. note :: When inputs are on a CUDA device, this function synchronizes that device with the CPU
+           if :attr:`p` is one of `('fro', 'nuc', inf, -inf, 1, -1)`.
+
+.. seealso::
+
+        :func:`torch.linalg.solve` for a function that solves linear systems of square matrices.
+
+        :func:`torch.linalg.lstsq` for a function that solves linear systems of general matrices.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions
+                    for :attr:`p` in `(2, -2)`, and of shape `(*, n, n)` where every matrix
+                    is invertible for :attr:`p` in `('fro', 'nuc', inf, -inf, 1, -1)`.
+    p (int, inf, -inf, 'fro', 'nuc', optional):
+        the type of the matrix norm to use in the computations (see above). Default: `None`
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Returns:
+    A real-valued tensor, even when :attr:`A` is complex.
+
+Raises:
+    RuntimeError:
+        if :attr:`p` is one of `('fro', 'nuc', inf, -inf, 1, -1)`
+        and the :attr:`A` matrix or any matrix in the batch :attr:`A` is not square
+        or invertible.
+
+Examples::
+
+    >>> A = torch.randn(3, 4, 4, dtype=torch.complex64)
+    >>> torch.linalg.cond(A)
+    >>> A = torch.tensor([[1., 0, -1], [0, 1, 0], [1, 0, 1]])
+    >>> torch.linalg.cond(A)
+    tensor([1.4142])
+    >>> torch.linalg.cond(A, 'fro')
+    tensor(3.1623)
+    >>> torch.linalg.cond(A, 'nuc')
+    tensor(9.2426)
+    >>> torch.linalg.cond(A, float('inf'))
+    tensor(2.)
+    >>> torch.linalg.cond(A, float('-inf'))
+    tensor(1.)
+    >>> torch.linalg.cond(A, 1)
+    tensor(2.)
+    >>> torch.linalg.cond(A, -1)
+    tensor(1.)
+    >>> torch.linalg.cond(A, 2)
+    tensor([1.4142])
+    >>> torch.linalg.cond(A, -2)
+    tensor([0.7071])
+
+    >>> A = torch.randn(2, 3, 3)
+    >>> torch.linalg.cond(A)
+    tensor([[9.5917],
+            [3.2538]])
+    >>> A = torch.randn(2, 3, 3, dtype=torch.complex64)
+    >>> torch.linalg.cond(A)
+    tensor([[4.6245],
+            [4.5671]])
+""",
+)
+
+pinv = _add_docstr(
+    _linalg.linalg_pinv,
+    r"""
+linalg.pinv(A, *, atol=None, rtol=None, hermitian=False, out=None) -> Tensor
+
+Computes the pseudoinverse (Moore-Penrose inverse) of a matrix.
+
+The pseudoinverse may be `defined algebraically`_
+but it is more computationally convenient to understand it `through the SVD`_
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+If :attr:`hermitian`\ `= True`, :attr:`A` is assumed to be Hermitian if complex or
+symmetric if real, but this is not checked internally. Instead, just the lower
+triangular part of the matrix is used in the computations.
+
+The singular values (or the norm of the eigenvalues when :attr:`hermitian`\ `= True`)
+that are below :math:`\max(\text{atol}, \sigma_1 \cdot \text{rtol})` threshold are
+treated as zero and discarded in the computation,
+where :math:`\sigma_1` is the largest singular value (or eigenvalue).
+
+If :attr:`rtol` is not specified and :attr:`A` is a matrix of dimensions `(m, n)`,
+the relative tolerance is set to be :math:`\text{rtol} = \max(m, n) \varepsilon`
+and :math:`\varepsilon` is the epsilon value for the dtype of :attr:`A` (see :class:`.finfo`).
+If :attr:`rtol` is not specified and :attr:`atol` is specified to be larger than zero then
+:attr:`rtol` is set to zero.
+
+If :attr:`atol` or :attr:`rtol` is a :class:`torch.Tensor`, its shape must be broadcastable to that
+of the singular values of :attr:`A` as returned by :func:`torch.linalg.svd`.
+
+.. note:: This function uses :func:`torch.linalg.svd` if :attr:`hermitian`\ `= False` and
+          :func:`torch.linalg.eigh` if :attr:`hermitian`\ `= True`.
+          For CUDA inputs, this function synchronizes that device with the CPU.
+
+.. note::
+    Consider using :func:`torch.linalg.lstsq` if possible for multiplying a matrix on the left by
+    the pseudoinverse, as::
+
+        torch.linalg.lstsq(A, B).solution == A.pinv() @ B
+
+    It is always preferred to use :func:`~lstsq` when possible, as it is faster and more
+    numerically stable than computing the pseudoinverse explicitly.
+
+.. note::
+    This function has NumPy compatible variant `linalg.pinv(A, rcond, hermitian=False)`.
+    However, use of the positional argument :attr:`rcond` is deprecated in favor of :attr:`rtol`.
+
+.. warning::
+    This function uses internally :func:`torch.linalg.svd` (or :func:`torch.linalg.eigh`
+    when :attr:`hermitian`\ `= True`), so its derivative has the same problems as those of these
+    functions. See the warnings in :func:`torch.linalg.svd` and :func:`torch.linalg.eigh` for
+    more details.
+
+.. seealso::
+
+        :func:`torch.linalg.inv` computes the inverse of a square matrix.
+
+        :func:`torch.linalg.lstsq` computes :attr:`A`\ `.pinv() @ \ `:attr:`B` with a
+        numerically stable algorithm.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+    rcond (float, Tensor, optional): [NumPy Compat]. Alias for :attr:`rtol`. Default: `None`.
+
+Keyword args:
+    atol (float, Tensor, optional): the absolute tolerance value. When `None` it's considered to be zero.
+                                    Default: `None`.
+    rtol (float, Tensor, optional): the relative tolerance value. See above for the value it takes when `None`.
+                                    Default: `None`.
+    hermitian(bool, optional): indicates whether :attr:`A` is Hermitian if complex
+                               or symmetric if real. Default: `False`.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.randn(3, 5)
+    >>> A
+    tensor([[ 0.5495,  0.0979, -1.4092, -0.1128,  0.4132],
+            [-1.1143, -0.3662,  0.3042,  1.6374, -0.9294],
+            [-0.3269, -0.5745, -0.0382, -0.5922, -0.6759]])
+    >>> torch.linalg.pinv(A)
+    tensor([[ 0.0600, -0.1933, -0.2090],
+            [-0.0903, -0.0817, -0.4752],
+            [-0.7124, -0.1631, -0.2272],
+            [ 0.1356,  0.3933, -0.5023],
+            [-0.0308, -0.1725, -0.5216]])
+
+    >>> A = torch.randn(2, 6, 3)
+    >>> Apinv = torch.linalg.pinv(A)
+    >>> torch.dist(Apinv @ A, torch.eye(3))
+    tensor(8.5633e-07)
+
+    >>> A = torch.randn(3, 3, dtype=torch.complex64)
+    >>> A = A + A.T.conj()  # creates a Hermitian matrix
+    >>> Apinv = torch.linalg.pinv(A, hermitian=True)
+    >>> torch.dist(Apinv @ A, torch.eye(3))
+    tensor(1.0830e-06)
+
+.. _defined algebraically:
+    https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse#Existence_and_uniqueness
+.. _through the SVD:
+    https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse#Singular_value_decomposition_(SVD)
+""",
+)
+
+matrix_exp = _add_docstr(
+    _linalg.linalg_matrix_exp,
+    r"""
+linalg.matrix_exp(A) -> Tensor
+
+Computes the matrix exponential of a square matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+this function computes the **matrix exponential** of :math:`A \in \mathbb{K}^{n \times n}`, which is defined as
+
+.. math::
+    \mathrm{matrix\_exp}(A) = \sum_{k=0}^\infty \frac{1}{k!}A^k \in \mathbb{K}^{n \times n}.
+
+If the matrix :math:`A` has eigenvalues :math:`\lambda_i \in \mathbb{C}`,
+the matrix :math:`\mathrm{matrix\_exp}(A)` has eigenvalues :math:`e^{\lambda_i} \in \mathbb{C}`.
+
+Supports input of bfloat16, float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions.
+
+Example::
+
+    >>> A = torch.empty(2, 2, 2)
+    >>> A[0, :, :] = torch.eye(2, 2)
+    >>> A[1, :, :] = 2 * torch.eye(2, 2)
+    >>> A
+    tensor([[[1., 0.],
+             [0., 1.]],
+
+            [[2., 0.],
+             [0., 2.]]])
+    >>> torch.linalg.matrix_exp(A)
+    tensor([[[2.7183, 0.0000],
+             [0.0000, 2.7183]],
+
+             [[7.3891, 0.0000],
+              [0.0000, 7.3891]]])
+
+    >>> import math
+    >>> A = torch.tensor([[0, math.pi/3], [-math.pi/3, 0]]) # A is skew-symmetric
+    >>> torch.linalg.matrix_exp(A) # matrix_exp(A) = [[cos(pi/3), sin(pi/3)], [-sin(pi/3), cos(pi/3)]]
+    tensor([[ 0.5000,  0.8660],
+            [-0.8660,  0.5000]])
+""",
+)
+
+
+solve = _add_docstr(
+    _linalg.linalg_solve,
+    r"""
+linalg.solve(A, B, *, left=True, out=None) -> Tensor
+
+Computes the solution of a square system of linear equations with a unique solution.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+this function computes the solution :math:`X \in \mathbb{K}^{n \times k}` of the **linear system** associated to
+:math:`A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{n \times k}`, which is defined as
+
+.. math:: AX = B
+
+If :attr:`left`\ `= False`, this function returns the matrix :math:`X \in \mathbb{K}^{n \times k}` that solves the system
+
+.. math::
+
+    XA = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.}
+
+This system of linear equations has one solution if and only if :math:`A` is `invertible`_.
+This function assumes that :math:`A` is invertible.
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if the inputs are batches of matrices then
+the output has the same batch dimensions.
+
+Letting `*` be zero or more batch dimensions,
+
+- If :attr:`A` has shape `(*, n, n)` and :attr:`B` has shape `(*, n)` (a batch of vectors) or shape
+  `(*, n, k)` (a batch of matrices or "multiple right-hand sides"), this function returns `X` of shape
+  `(*, n)` or `(*, n, k)` respectively.
+- Otherwise, if :attr:`A` has shape `(*, n, n)` and  :attr:`B` has shape `(n,)`  or `(n, k)`, :attr:`B`
+  is broadcasted to have shape `(*, n)` or `(*, n, k)` respectively.
+  This function then returns the solution of the resulting batch of systems of linear equations.
+
+.. note::
+    This function computes `X = \ `:attr:`A`\ `.inverse() @ \ `:attr:`B` in a faster and
+    more numerically stable way than performing the computations separately.
+
+.. note::
+    It is possible to compute the solution of the system :math:`XA = B` by passing the inputs
+    :attr:`A` and :attr:`B` transposed and transposing the output returned by this function.
+
+.. note::
+    :attr:`A` is allowed to be a non-batched `torch.sparse_csr_tensor`, but only with `left=True`.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.solve_ex")}
+"""
+    + r"""
+
+.. seealso::
+
+        :func:`torch.linalg.solve_triangular` computes the solution of a triangular system of linear
+        equations with a unique solution.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions.
+    B (Tensor): right-hand side tensor of shape `(*, n)` or  `(*, n, k)` or `(n,)` or `(n, k)`
+                according to the rules described above
+
+Keyword args:
+    left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Raises:
+    RuntimeError: if the :attr:`A` matrix is not invertible or any matrix in a batched :attr:`A`
+                  is not invertible.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> b = torch.randn(3)
+    >>> x = torch.linalg.solve(A, b)
+    >>> torch.allclose(A @ x, b)
+    True
+    >>> A = torch.randn(2, 3, 3)
+    >>> B = torch.randn(2, 3, 4)
+    >>> X = torch.linalg.solve(A, B)
+    >>> X.shape
+    torch.Size([2, 3, 4])
+    >>> torch.allclose(A @ X, B)
+    True
+
+    >>> A = torch.randn(2, 3, 3)
+    >>> b = torch.randn(3, 1)
+    >>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3, 1)
+    >>> x.shape
+    torch.Size([2, 3, 1])
+    >>> torch.allclose(A @ x, b)
+    True
+    >>> b = torch.randn(3)
+    >>> x = torch.linalg.solve(A, b) # b is broadcasted to size (2, 3)
+    >>> x.shape
+    torch.Size([2, 3])
+    >>> Ax = A @ x.unsqueeze(-1)
+    >>> torch.allclose(Ax, b.unsqueeze(-1).expand_as(Ax))
+    True
+
+.. _invertible:
+    https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem
+""",
+)
+
+solve_triangular = _add_docstr(
+    _linalg.linalg_solve_triangular,
+    r"""
+linalg.solve_triangular(A, B, *, upper, left=True, unitriangular=False, out=None) -> Tensor
+
+Computes the solution of a triangular system of linear equations with a unique solution.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+this function computes the solution :math:`X \in \mathbb{K}^{n \times k}` of the **linear system**
+associated to the triangular matrix :math:`A \in \mathbb{K}^{n \times n}` without zeros on the diagonal
+(that is, it is `invertible`_) and the rectangular matrix , :math:`B \in \mathbb{K}^{n \times k}`,
+which is defined as
+
+.. math:: AX = B
+
+The argument :attr:`upper` signals whether :math:`A` is upper or lower triangular.
+
+If :attr:`left`\ `= False`, this function returns the matrix :math:`X \in \mathbb{K}^{n \times k}` that
+solves the system
+
+.. math::
+
+    XA = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.}
+
+If :attr:`upper`\ `= True` (resp. `False`) just the upper (resp. lower) triangular half of :attr:`A`
+will be accessed. The elements below the main diagonal will be considered to be zero and will not be accessed.
+
+If :attr:`unitriangular`\ `= True`, the diagonal of :attr:`A` is assumed to be ones and will not be accessed.
+
+The result may contain `NaN` s if the diagonal of :attr:`A` contains zeros or elements that
+are very close to zero and :attr:`unitriangular`\ `= False` (default) or if the input matrix
+has very small eigenvalues.
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if the inputs are batches of matrices then
+the output has the same batch dimensions.
+
+.. seealso::
+
+        :func:`torch.linalg.solve` computes the solution of a general square system of linear
+        equations with a unique solution.
+
+Args:
+    A (Tensor): tensor of shape `(*, n, n)` (or `(*, k, k)` if :attr:`left`\ `= False`)
+                where `*` is zero or more batch dimensions.
+    B (Tensor): right-hand side tensor of shape `(*, n, k)`.
+
+Keyword args:
+    upper (bool): whether :attr:`A` is an upper or lower triangular matrix.
+    left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`.
+    unitriangular (bool, optional): if `True`, the diagonal elements of :attr:`A` are assumed to be
+                                    all equal to `1`. Default: `False`.
+    out (Tensor, optional): output tensor. `B` may be passed as `out` and the result is computed in-place on `B`.
+                            Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.randn(3, 3).triu_()
+    >>> B = torch.randn(3, 4)
+    >>> X = torch.linalg.solve_triangular(A, B, upper=True)
+    >>> torch.allclose(A @ X, B)
+    True
+
+    >>> A = torch.randn(2, 3, 3).tril_()
+    >>> B = torch.randn(2, 3, 4)
+    >>> X = torch.linalg.solve_triangular(A, B, upper=False)
+    >>> torch.allclose(A @ X, B)
+    True
+
+    >>> A = torch.randn(2, 4, 4).tril_()
+    >>> B = torch.randn(2, 3, 4)
+    >>> X = torch.linalg.solve_triangular(A, B, upper=False, left=False)
+    >>> torch.allclose(X @ A, B)
+    True
+
+.. _invertible:
+    https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem
+""",
+)
+
+lu_factor = _add_docstr(
+    _linalg.linalg_lu_factor,
+    r"""
+linalg.lu_factor(A, *, bool pivot=True, out=None) -> (Tensor, Tensor)
+
+Computes a compact representation of the LU factorization with partial pivoting of a matrix.
+
+This function computes a compact representation of the decomposition given by :func:`torch.linalg.lu`.
+If the matrix is square, this representation may be used in :func:`torch.linalg.lu_solve`
+to solve system of linear equations that share the matrix :attr:`A`.
+
+The returned decomposition is represented as a named tuple `(LU, pivots)`.
+The ``LU`` matrix has the same shape as the input matrix ``A``. Its upper and lower triangular
+parts encode the non-constant elements of ``L`` and ``U`` of the LU decomposition of ``A``.
+
+The returned permutation matrix is represented by a 1-indexed vector. `pivots[i] == j` represents
+that in the `i`-th step of the algorithm, the `i`-th row was permuted with the `j-1`-th row.
+
+On CUDA, one may use :attr:`pivot`\ `= False`. In this case, this function returns the LU
+decomposition without pivoting if it exists.
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if the inputs are batches of matrices then
+the output has the same batch dimensions.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_has_ex"].format("torch.linalg.lu_factor_ex")}
+"""
+    + r"""
+.. warning:: The LU decomposition is almost never unique, as often there are different permutation
+             matrices that can yield different LU decompositions.
+             As such, different platforms, like SciPy, or inputs on different devices,
+             may produce different valid decompositions.
+
+             Gradient computations are only supported if the input matrix is full-rank.
+             If this condition is not met, no error will be thrown, but the gradient may not be finite.
+             This is because the LU decomposition with pivoting is not differentiable at these points.
+
+.. seealso::
+
+        :func:`torch.linalg.lu_solve` solves a system of linear equations given the output of this
+        function provided the input matrix was square and invertible.
+
+        :func:`torch.lu_unpack` unpacks the tensors returned by :func:`~lu_factor` into the three
+        matrices `P, L, U` that form the decomposition.
+
+        :func:`torch.linalg.lu` computes the LU decomposition with partial pivoting of a possibly
+        non-square matrix. It is a composition of :func:`~lu_factor` and :func:`torch.lu_unpack`.
+
+        :func:`torch.linalg.solve` solves a system of linear equations. It is a composition
+        of :func:`~lu_factor` and :func:`~lu_solve`.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    pivot (bool, optional): Whether to compute the LU decomposition with partial pivoting, or the regular LU
+                            decomposition. :attr:`pivot`\ `= False` not supported on CPU. Default: `True`.
+    out (tuple, optional): tuple of two tensors to write the output to. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(LU, pivots)`.
+
+Raises:
+    RuntimeError: if the :attr:`A` matrix is not invertible or any matrix in a batched :attr:`A`
+                  is not invertible.
+
+Examples::
+
+    >>> A = torch.randn(2, 3, 3)
+    >>> B1 = torch.randn(2, 3, 4)
+    >>> B2 = torch.randn(2, 3, 7)
+    >>> LU, pivots = torch.linalg.lu_factor(A)
+    >>> X1 = torch.linalg.lu_solve(LU, pivots, B1)
+    >>> X2 = torch.linalg.lu_solve(LU, pivots, B2)
+    >>> torch.allclose(A @ X1, B1)
+    True
+    >>> torch.allclose(A @ X2, B2)
+    True
+
+.. _invertible:
+    https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem
+""",
+)
+
+lu_factor_ex = _add_docstr(
+    _linalg.linalg_lu_factor_ex,
+    r"""
+linalg.lu_factor_ex(A, *, pivot=True, check_errors=False, out=None) -> (Tensor, Tensor, Tensor)
+
+This is a version of :func:`~lu_factor` that does not perform error checks unless :attr:`check_errors`\ `= True`.
+It also returns the :attr:`info` tensor returned by `LAPACK's getrf`_.
+
+"""
+    + rf"""
+.. note:: {common_notes["sync_note_ex"]}
+
+.. warning:: {common_notes["experimental_warning"]}
+"""
+    + r"""
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+
+Keyword args:
+    pivot (bool, optional): Whether to compute the LU decomposition with partial pivoting, or the regular LU
+                            decomposition. :attr:`pivot`\ `= False` not supported on CPU. Default: `True`.
+    check_errors (bool, optional): controls whether to check the content of ``infos`` and raise
+                                   an error if it is non-zero. Default: `False`.
+    out (tuple, optional): tuple of three tensors to write the output to. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(LU, pivots, info)`.
+
+.. _LAPACK's getrf:
+    https://www.netlib.org/lapack/explore-html-3.6.1/dd/d9a/group__double_g_ecomputational_ga0019443faea08275ca60a734d0593e60.html
+""",
+)
+
+lu_solve = _add_docstr(
+    _linalg.linalg_lu_solve,
+    r"""
+linalg.lu_solve(LU, pivots, B, *, left=True, adjoint=False, out=None) -> Tensor
+
+Computes the solution of a square system of linear equations with a unique solution given an LU decomposition.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+this function computes the solution :math:`X \in \mathbb{K}^{n \times k}` of the **linear system** associated to
+:math:`A \in \mathbb{K}^{n \times n}, B \in \mathbb{K}^{n \times k}`, which is defined as
+
+.. math:: AX = B
+
+where :math:`A` is given factorized as returned by :func:`~lu_factor`.
+
+If :attr:`left`\ `= False`, this function returns the matrix :math:`X \in \mathbb{K}^{n \times k}` that solves the system
+
+.. math::
+
+    XA = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.}
+
+If  :attr:`adjoint`\ `= True` (and :attr:`left`\ `= True`), given an LU factorization of :math:`A`
+this function function returns the :math:`X \in \mathbb{K}^{n \times k}` that solves the system
+
+.. math::
+
+    A^{\text{H}}X = B\mathrlap{\qquad A \in \mathbb{K}^{k \times k}, B \in \mathbb{K}^{n \times k}.}
+
+where :math:`A^{\text{H}}` is the conjugate transpose when :math:`A` is complex, and the
+transpose when :math:`A` is real-valued. The :attr:`left`\ `= False` case is analogous.
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if the inputs are batches of matrices then
+the output has the same batch dimensions.
+
+Args:
+    LU (Tensor): tensor of shape `(*, n, n)` (or `(*, k, k)` if :attr:`left`\ `= True`)
+                 where `*` is zero or more batch dimensions as returned by :func:`~lu_factor`.
+    pivots (Tensor): tensor of shape `(*, n)` (or `(*, k)` if :attr:`left`\ `= True`)
+                     where `*` is zero or more batch dimensions as returned by :func:`~lu_factor`.
+    B (Tensor): right-hand side tensor of shape `(*, n, k)`.
+
+Keyword args:
+    left (bool, optional): whether to solve the system :math:`AX=B` or :math:`XA = B`. Default: `True`.
+    adjoint (bool, optional): whether to solve the system :math:`AX=B` or :math:`A^{\text{H}}X = B`. Default: `False`.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> A = torch.randn(3, 3)
+    >>> LU, pivots = torch.linalg.lu_factor(A)
+    >>> B = torch.randn(3, 2)
+    >>> X = torch.linalg.lu_solve(LU, pivots, B)
+    >>> torch.allclose(A @ X, B)
+    True
+
+    >>> B = torch.randn(3, 3, 2)   # Broadcasting rules apply: A is broadcasted
+    >>> X = torch.linalg.lu_solve(LU, pivots, B)
+    >>> torch.allclose(A @ X, B)
+    True
+
+    >>> B = torch.randn(3, 5, 3)
+    >>> X = torch.linalg.lu_solve(LU, pivots, B, left=False)
+    >>> torch.allclose(X @ A, B)
+    True
+
+    >>> B = torch.randn(3, 3, 4)   # Now solve for A^T
+    >>> X = torch.linalg.lu_solve(LU, pivots, B, adjoint=True)
+    >>> torch.allclose(A.mT @ X, B)
+    True
+
+.. _invertible:
+    https://en.wikipedia.org/wiki/Invertible_matrix#The_invertible_matrix_theorem
+""",
+)
+
+lu = _add_docstr(
+    _linalg.linalg_lu,
+    r"""
+lu(A, *, pivot=True, out=None) -> (Tensor, Tensor, Tensor)
+
+Computes the LU decomposition with partial pivoting of a matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **LU decomposition with partial pivoting** of a matrix
+:math:`A \in \mathbb{K}^{m \times n}` is defined as
+
+.. math::
+
+    A = PLU\mathrlap{\qquad P \in \mathbb{K}^{m \times m}, L \in \mathbb{K}^{m \times k}, U \in \mathbb{K}^{k \times n}}
+
+where `k = min(m,n)`, :math:`P` is a `permutation matrix`_, :math:`L` is lower triangular with ones on the diagonal
+and :math:`U` is upper triangular.
+
+If :attr:`pivot`\ `= False` and :attr:`A` is on GPU, then the **LU decomposition without pivoting** is computed
+
+.. math::
+
+    A = LU\mathrlap{\qquad L \in \mathbb{K}^{m \times k}, U \in \mathbb{K}^{k \times n}}
+
+When :attr:`pivot`\ `= False`, the returned matrix :attr:`P` will be empty.
+The LU decomposition without pivoting `may not exist`_ if any of the principal minors of :attr:`A` is singular.
+In this case, the output matrix may contain `inf` or `NaN`.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+.. seealso::
+
+        :func:`torch.linalg.solve` solves a system of linear equations using the LU decomposition
+        with partial pivoting.
+
+.. warning:: The LU decomposition is almost never unique, as often there are different permutation
+             matrices that can yield different LU decompositions.
+             As such, different platforms, like SciPy, or inputs on different devices,
+             may produce different valid decompositions.
+
+.. warning:: Gradient computations are only supported if the input matrix is full-rank.
+             If this condition is not met, no error will be thrown, but the gradient
+             may not be finite.
+             This is because the LU decomposition with pivoting is not differentiable at these points.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+    pivot (bool, optional): Controls whether to compute the LU decomposition with partial pivoting or
+        no pivoting. Default: `True`.
+
+Keyword args:
+    out (tuple, optional): output tuple of three tensors. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(P, L, U)`.
+
+Examples::
+
+    >>> A = torch.randn(3, 2)
+    >>> P, L, U = torch.linalg.lu(A)
+    >>> P
+    tensor([[0., 1., 0.],
+            [0., 0., 1.],
+            [1., 0., 0.]])
+    >>> L
+    tensor([[1.0000, 0.0000],
+            [0.5007, 1.0000],
+            [0.0633, 0.9755]])
+    >>> U
+    tensor([[0.3771, 0.0489],
+            [0.0000, 0.9644]])
+    >>> torch.dist(A, P @ L @ U)
+    tensor(5.9605e-08)
+
+    >>> A = torch.randn(2, 5, 7, device="cuda")
+    >>> P, L, U = torch.linalg.lu(A, pivot=False)
+    >>> P
+    tensor([], device='cuda:0')
+    >>> torch.dist(A, L @ U)
+    tensor(1.0376e-06, device='cuda:0')
+
+.. _permutation matrix:
+    https://en.wikipedia.org/wiki/Permutation_matrix
+.. _may not exist:
+    https://en.wikipedia.org/wiki/LU_decomposition#Definitions
+""",
+)
+
+tensorinv = _add_docstr(
+    _linalg.linalg_tensorinv,
+    r"""
+linalg.tensorinv(A, ind=2, *, out=None) -> Tensor
+
+Computes the multiplicative inverse of :func:`torch.tensordot`.
+
+If `m` is the product of the first :attr:`ind` dimensions of :attr:`A` and `n` is the product of
+the rest of the dimensions, this function expects `m` and `n` to be equal.
+If this is the case, it computes a tensor `X` such that
+`tensordot(\ `:attr:`A`\ `, X, \ `:attr:`ind`\ `)` is the identity matrix in dimension `m`.
+`X` will have the shape of :attr:`A` but with the first :attr:`ind` dimensions pushed back to the end
+
+.. code:: text
+
+    X.shape == A.shape[ind:] + A.shape[:ind]
+
+Supports input of float, double, cfloat and cdouble dtypes.
+
+.. note:: When :attr:`A` is a `2`-dimensional tensor and :attr:`ind`\ `= 1`,
+          this function computes the (multiplicative) inverse of :attr:`A`
+          (see :func:`torch.linalg.inv`).
+
+.. note::
+    Consider using :func:`torch.linalg.tensorsolve` if possible for multiplying a tensor on the left
+    by the tensor inverse, as::
+
+        linalg.tensorsolve(A, B) == torch.tensordot(linalg.tensorinv(A), B)  # When B is a tensor with shape A.shape[:B.ndim]
+
+    It is always preferred to use :func:`~tensorsolve` when possible, as it is faster and more
+    numerically stable than computing the pseudoinverse explicitly.
+
+.. seealso::
+
+        :func:`torch.linalg.tensorsolve` computes
+        `torch.tensordot(tensorinv(\ `:attr:`A`\ `), \ `:attr:`B`\ `)`.
+
+Args:
+    A (Tensor): tensor to invert. Its shape must satisfy
+                    `prod(\ `:attr:`A`\ `.shape[:\ `:attr:`ind`\ `]) ==
+                    prod(\ `:attr:`A`\ `.shape[\ `:attr:`ind`\ `:])`.
+    ind (int): index at which to compute the inverse of :func:`torch.tensordot`. Default: `2`.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Raises:
+    RuntimeError: if the reshaped :attr:`A` is not invertible or the product of the first
+                  :attr:`ind` dimensions is not equal to the product of the rest.
+
+Examples::
+
+    >>> A = torch.eye(4 * 6).reshape((4, 6, 8, 3))
+    >>> Ainv = torch.linalg.tensorinv(A, ind=2)
+    >>> Ainv.shape
+    torch.Size([8, 3, 4, 6])
+    >>> B = torch.randn(4, 6)
+    >>> torch.allclose(torch.tensordot(Ainv, B), torch.linalg.tensorsolve(A, B))
+    True
+
+    >>> A = torch.randn(4, 4)
+    >>> Atensorinv = torch.linalg.tensorinv(A, ind=1)
+    >>> Ainv = torch.linalg.inv(A)
+    >>> torch.allclose(Atensorinv, Ainv)
+    True
+""",
+)
+
+tensorsolve = _add_docstr(
+    _linalg.linalg_tensorsolve,
+    r"""
+linalg.tensorsolve(A, B, dims=None, *, out=None) -> Tensor
+
+Computes the solution `X` to the system `torch.tensordot(A, X) = B`.
+
+If `m` is the product of the first :attr:`B`\ `.ndim`  dimensions of :attr:`A` and
+`n` is the product of the rest of the dimensions, this function expects `m` and `n` to be equal.
+
+The returned tensor `x` satisfies
+`tensordot(\ `:attr:`A`\ `, x, dims=x.ndim) == \ `:attr:`B`.
+`x` has shape :attr:`A`\ `[B.ndim:]`.
+
+If :attr:`dims` is specified, :attr:`A` will be reshaped as
+
+.. code:: text
+
+    A = movedim(A, dims, range(len(dims) - A.ndim + 1, 0))
+
+Supports inputs of float, double, cfloat and cdouble dtypes.
+
+.. seealso::
+
+        :func:`torch.linalg.tensorinv` computes the multiplicative inverse of
+        :func:`torch.tensordot`.
+
+Args:
+    A (Tensor): tensor to solve for. Its shape must satisfy
+                    `prod(\ `:attr:`A`\ `.shape[:\ `:attr:`B`\ `.ndim]) ==
+                    prod(\ `:attr:`A`\ `.shape[\ `:attr:`B`\ `.ndim:])`.
+    B (Tensor): tensor of shape :attr:`A`\ `.shape[:\ `:attr:`B`\ `.ndim]`.
+    dims (Tuple[int], optional): dimensions of :attr:`A` to be moved.
+        If `None`, no dimensions are moved. Default: `None`.
+
+Keyword args:
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Raises:
+    RuntimeError: if the reshaped :attr:`A`\ `.view(m, m)` with `m` as above  is not
+                  invertible or the product of the first :attr:`ind` dimensions is not equal
+                  to the product of the rest of the dimensions.
+
+Examples::
+
+    >>> A = torch.eye(2 * 3 * 4).reshape((2 * 3, 4, 2, 3, 4))
+    >>> B = torch.randn(2 * 3, 4)
+    >>> X = torch.linalg.tensorsolve(A, B)
+    >>> X.shape
+    torch.Size([2, 3, 4])
+    >>> torch.allclose(torch.tensordot(A, X, dims=X.ndim), B)
+    True
+
+    >>> A = torch.randn(6, 4, 4, 3, 2)
+    >>> B = torch.randn(4, 3, 2)
+    >>> X = torch.linalg.tensorsolve(A, B, dims=(0, 2))
+    >>> X.shape
+    torch.Size([6, 4])
+    >>> A = A.permute(1, 3, 4, 0, 2)
+    >>> A.shape[B.ndim:]
+    torch.Size([6, 4])
+    >>> torch.allclose(torch.tensordot(A, X, dims=X.ndim), B, atol=1e-6)
+    True
+""",
+)
+
+qr = _add_docstr(
+    _linalg.linalg_qr,
+    r"""
+qr(A, mode='reduced', *, out=None) -> (Tensor, Tensor)
+
+Computes the QR decomposition of a matrix.
+
+Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`,
+the **full QR decomposition** of a matrix
+:math:`A \in \mathbb{K}^{m \times n}` is defined as
+
+.. math::
+
+    A = QR\mathrlap{\qquad Q \in \mathbb{K}^{m \times m}, R \in \mathbb{K}^{m \times n}}
+
+where :math:`Q` is orthogonal in the real case and unitary in the complex case,
+and :math:`R` is upper triangular with real diagonal (even in the complex case).
+
+When `m > n` (tall matrix), as `R` is upper triangular, its last `m - n` rows are zero.
+In this case, we can drop the last `m - n` columns of `Q` to form the
+**reduced QR decomposition**:
+
+.. math::
+
+    A = QR\mathrlap{\qquad Q \in \mathbb{K}^{m \times n}, R \in \mathbb{K}^{n \times n}}
+
+The reduced QR decomposition agrees with the full QR decomposition when `n >= m` (wide matrix).
+
+Supports input of float, double, cfloat and cdouble dtypes.
+Also supports batches of matrices, and if :attr:`A` is a batch of matrices then
+the output has the same batch dimensions.
+
+The parameter :attr:`mode` chooses between the full and reduced QR decomposition.
+If :attr:`A` has shape `(*, m, n)`, denoting `k = min(m, n)`
+
+- :attr:`mode`\ `= 'reduced'` (default): Returns `(Q, R)` of shapes `(*, m, k)`, `(*, k, n)` respectively.
+  It is always differentiable.
+- :attr:`mode`\ `= 'complete'`: Returns `(Q, R)` of shapes `(*, m, m)`, `(*, m, n)` respectively.
+  It is differentiable for `m <= n`.
+- :attr:`mode`\ `= 'r'`: Computes only the reduced `R`. Returns `(Q, R)` with `Q` empty and `R` of shape `(*, k, n)`.
+  It is never differentiable.
+
+Differences with `numpy.linalg.qr`:
+
+- :attr:`mode`\ `= 'raw'` is not implemented.
+- Unlike `numpy.linalg.qr`, this function always returns a tuple of two tensors.
+  When :attr:`mode`\ `= 'r'`, the `Q` tensor is an empty tensor.
+
+.. warning:: The elements in the diagonal of `R` are not necessarily positive.
+             As such, the returned QR decomposition is only unique up to the sign of the diagonal of `R`.
+             Therefore, different platforms, like NumPy, or inputs on different devices,
+             may produce different valid decompositions.
+
+.. warning:: The QR decomposition is only well-defined if the first `k = min(m, n)` columns
+             of every matrix in :attr:`A` are linearly independent.
+             If this condition is not met, no error will be thrown, but the QR produced
+             may be incorrect and its autodiff may fail or produce incorrect results.
+
+Args:
+    A (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions.
+    mode (str, optional): one of `'reduced'`, `'complete'`, `'r'`.
+                          Controls the shape of the returned tensors. Default: `'reduced'`.
+
+Keyword args:
+    out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`.
+
+Returns:
+    A named tuple `(Q, R)`.
+
+Examples::
+
+    >>> A = torch.tensor([[12., -51, 4], [6, 167, -68], [-4, 24, -41]])
+    >>> Q, R = torch.linalg.qr(A)
+    >>> Q
+    tensor([[-0.8571,  0.3943,  0.3314],
+            [-0.4286, -0.9029, -0.0343],
+            [ 0.2857, -0.1714,  0.9429]])
+    >>> R
+    tensor([[ -14.0000,  -21.0000,   14.0000],
+            [   0.0000, -175.0000,   70.0000],
+            [   0.0000,    0.0000,  -35.0000]])
+    >>> (Q @ R).round()
+    tensor([[  12.,  -51.,    4.],
+            [   6.,  167.,  -68.],
+            [  -4.,   24.,  -41.]])
+    >>> (Q.T @ Q).round()
+    tensor([[ 1.,  0.,  0.],
+            [ 0.,  1., -0.],
+            [ 0., -0.,  1.]])
+    >>> Q2, R2 = torch.linalg.qr(A, mode='r')
+    >>> Q2
+    tensor([])
+    >>> torch.equal(R, R2)
+    True
+    >>> A = torch.randn(3, 4, 5)
+    >>> Q, R = torch.linalg.qr(A, mode='complete')
+    >>> torch.dist(Q @ R, A)
+    tensor(1.6099e-06)
+    >>> torch.dist(Q.mT @ Q, torch.eye(4))
+    tensor(6.2158e-07)
+""",
+)
+
+vander = _add_docstr(
+    _linalg.linalg_vander,
+    r"""
+vander(x, N=None) -> Tensor
+
+Generates a Vandermonde matrix.
+
+Returns the Vandermonde matrix :math:`V`
+
+.. math::
+
+    V = \begin{pmatrix}
+            1 & x_1 & x_1^2 & \dots & x_1^{N-1}\\
+            1 & x_2 & x_2^2 & \dots & x_2^{N-1}\\
+            1 & x_3 & x_3^2 & \dots & x_3^{N-1}\\
+            \vdots & \vdots & \vdots & \ddots &\vdots \\
+            1 & x_n & x_n^2 & \dots & x_n^{N-1}
+        \end{pmatrix}.
+
+for `N > 1`.
+If :attr:`N`\ `= None`, then `N = x.size(-1)` so that the output is a square matrix.
+
+Supports inputs of float, double, cfloat, cdouble, and integral dtypes.
+Also supports batches of vectors, and if :attr:`x` is a batch of vectors then
+the output has the same batch dimensions.
+
+Differences with `numpy.vander`:
+
+- Unlike `numpy.vander`, this function returns the powers of :attr:`x` in ascending order.
+  To get them in the reverse order call ``linalg.vander(x, N).flip(-1)``.
+
+Args:
+    x (Tensor): tensor of shape `(*, n)` where `*` is zero or more batch dimensions
+                consisting of vectors.
+
+Keyword args:
+    N (int, optional): Number of columns in the output. Default: `x.size(-1)`
+
+Example::
+
+    >>> x = torch.tensor([1, 2, 3, 5])
+    >>> linalg.vander(x)
+    tensor([[  1,   1,   1,   1],
+            [  1,   2,   4,   8],
+            [  1,   3,   9,  27],
+            [  1,   5,  25, 125]])
+    >>> linalg.vander(x, N=3)
+    tensor([[ 1,  1,  1],
+            [ 1,  2,  4],
+            [ 1,  3,  9],
+            [ 1,  5, 25]])
+""",
+)
+
+vecdot = _add_docstr(
+    _linalg.linalg_vecdot,
+    r"""
+linalg.vecdot(x, y, *, dim=-1, out=None) -> Tensor
+
+Computes the dot product of two batches of vectors along a dimension.
+
+In symbols, this function computes
+
+.. math::
+
+    \sum_{i=1}^n \overline{x_i}y_i.
+
+over the dimension :attr:`dim` where :math:`\overline{x_i}` denotes the conjugate for complex
+vectors, and it is the identity for real vectors.
+
+Supports input of half, bfloat16, float, double, cfloat, cdouble and integral dtypes.
+It also supports broadcasting.
+
+Args:
+    x (Tensor): first batch of vectors of shape `(*, n)`.
+    y (Tensor): second batch of vectors of shape `(*, n)`.
+
+Keyword args:
+    dim (int): Dimension along which to compute the dot product. Default: `-1`.
+    out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`.
+
+Examples::
+
+    >>> v1 = torch.randn(3, 2)
+    >>> v2 = torch.randn(3, 2)
+    >>> linalg.vecdot(v1, v2)
+    tensor([ 0.3223,  0.2815, -0.1944])
+    >>> torch.vdot(v1[0], v2[0])
+    tensor(0.3223)
+""",
+)
+
+_powsum = _add_docstr(
+    _linalg.linalg__powsum,
+    r"""
+linalg._powsum(x, ord, dim=None, keepdim=False, *, dtype=None, out=None) -> Tensor
+
+Computes the sum of the absolute values raised to the power ``ord``.
+
+This function computes ``sum(abs(x)**ord)`` without applying the final root,
+which is useful for distributed computing where the root should only be applied
+once after reducing across all ranks.
+
+Supports input of float, double, cfloat and cdouble dtypes.
+
+Args:
+    x (Tensor): tensor, flattened by default, or optionally over dimension(s)
+                specified by :attr:`dim`.
+    ord (int, float): the exponent value. Can be any real number.
+
+Keyword args:
+    dim (int, Tuple[int], optional): dimension(s) to reduce over.
+                                     Default: ``None`` (all dimensions).
+    keepdim (bool, optional): whether the output has :attr:`dim` retained. Default: ``False``.
+    dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.
+          If specified, the input tensor is cast to :attr:`dtype` before the operation
+          is performed. Default: ``None``.
+    out (Tensor, optional): output tensor. Ignored if ``None``. Default: ``None``.
+
+Returns:
+    A real-valued tensor, even when :attr:`x` is complex.
+
+Example::
+
+    >>> x = torch.tensor([1., 2., 3.])
+    >>> torch.linalg._powsum(x, 2)
+    tensor(14.)
+    >>> torch.linalg.vector_norm(x, 2)
+    tensor(3.7417)
+    >>> torch.linalg._powsum(x, 2) ** 0.5  # equivalent to vector_norm
+    tensor(3.7417)
+""",
+)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d00ba1e8d5aff6a18490ac7b16a629ac36e3dcb5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/__init__.py
@@ -0,0 +1,57 @@
+from torch.masked._ops import (
+    _canonical_dim,
+    _combine_input_and_mask,
+    _generate_docstring,
+    _input_mask,
+    _output_mask,
+    _reduction_identity,
+    _where,
+    amax,
+    amin,
+    argmax,
+    argmin,
+    cumprod,
+    cumsum,
+    log_softmax,
+    logaddexp,
+    logsumexp,
+    mean,
+    median,
+    norm,
+    normalize,
+    prod,
+    softmax,
+    softmin,
+    std,
+    sum,
+    var,
+)
+from torch.masked.maskedtensor.core import is_masked_tensor, MaskedTensor
+from torch.masked.maskedtensor.creation import as_masked_tensor, masked_tensor
+
+
+__all__ = [
+    "amax",
+    "amin",
+    "argmax",
+    "argmin",
+    "as_masked_tensor",
+    "cumprod",
+    "cumsum",
+    "is_masked_tensor",
+    "log_softmax",
+    "logaddexp",
+    "logsumexp",
+    "masked_tensor",
+    "MaskedTensor",
+    "mean",
+    "median",
+    "norm",
+    "normalize",
+    "prod",
+    "softmax",
+    "softmin",
+    "std",
+    "sum",
+    "var",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/_docs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/_docs.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa130bbefbc5caa7373459ef2fc3dc5292239948
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/_docs.py
@@ -0,0 +1,1177 @@
+# This file is generated, do not modify it!
+#
+# To update this file, run the update masked docs script as follows:
+#
+#   python tools/update_masked_docs.py
+#
+# The script must be called from an environment where the development
+# version of torch package can be imported and is functional.
+#
+
+amax_docstring = """amax(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+
+Returns maximum of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+
+The identity value of maximum operation, which is used to start the
+reduction, depends on input dtype. For instance, for float32, uint8,
+and int32 dtypes, the identity values are ``-inf``, ``0``, and ``-2147483648``, respectively.
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in maximum computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of maximum operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.amax(input, 1, mask=mask)
+    tensor([                  -1, -9223372036854775808])
+"""
+
+amin_docstring = """amin(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+
+Returns minimum of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+
+The identity value of minimum operation, which is used to start the
+reduction, depends on input dtype. For instance, for float32, uint8,
+and int32 dtypes, the identity values are ``inf``, ``255``, and ``2147483647``, respectively.
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in minimum computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of minimum operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.amin(input, 1, mask=mask)
+    tensor([                 -3, 9223372036854775807])
+"""
+
+argmax_docstring = """argmax(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+Returns argmax of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+The identity value of argmax operation, which is used to start the
+reduction, depends on input dtype. For instance, for float32, uint8,
+and int32 dtypes, the identity values are ``-inf``, ``0``, and ``-2147483648``, respectively.
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in argmax computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of argmax operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which argmax is computed.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.argmax(input, 1, mask=mask)
+    tensor([2, 0])
+"""
+
+argmin_docstring = """argmin(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+Returns argmin of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+The identity value of argmin operation, which is used to start the
+reduction, depends on input dtype. For instance, for float32, uint8,
+and int32 dtypes, the identity values are ``inf``, ``255``, and ``2147483647``, respectively.
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in argmin computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of argmin operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which argmin is computed.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.argmin(input, 1, mask=mask)
+    tensor([0, 0])
+"""
+
+cumprod_docstring = """cumprod(input, dim, *, dtype=None, mask=None) -> Tensor
+
+Returns cumulative_prod of all the slices in the :attr:`input` tensor
+along :attr:`dim` while the :attr:`input` elements are masked out
+according to the boolean tensor :attr:`mask`.
+
+Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is
+defined as ``prod(x[:i])``.
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True then
+the corresponding element in :attr:`input` tensor will be included in
+cumulative_prod computation, otherwise the element is ignored.
+
+The values of masked-out elements of the output tensor have undefined
+value: it may or may not be set to zero or nan; the choice may correspond to
+the value that leads to the most efficient storage of :attr:`output`
+tensor.
+
+The mask of the cumulative_prod output tensor can be computed as
+``torch.broadcast_to(mask, input.shape)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which cumulative_prod is computed.
+
+Keyword args:
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.cumprod(input, 1, mask=mask)
+    tensor([[-3., -3.,  3.],
+            [ 1.,  1.,  1.]])
+"""
+
+cumsum_docstring = """cumsum(input, dim, *, dtype=None, mask=None) -> Tensor
+
+Returns cumulative_sum of all the slices in the :attr:`input` tensor
+along :attr:`dim` while the :attr:`input` elements are masked out
+according to the boolean tensor :attr:`mask`.
+
+Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is
+defined as ``sum(x[:i])``.
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True then
+the corresponding element in :attr:`input` tensor will be included in
+cumulative_sum computation, otherwise the element is ignored.
+
+The values of masked-out elements of the output tensor have undefined
+value: it may or may not be set to zero or nan; the choice may correspond to
+the value that leads to the most efficient storage of :attr:`output`
+tensor.
+
+The mask of the cumulative_sum output tensor can be computed as
+``torch.broadcast_to(mask, input.shape)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which cumulative_sum is computed.
+
+Keyword args:
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.cumsum(input, 1, mask=mask)
+    tensor([[-3., -3., -4.],
+            [ 0.,  0.,  0.]])
+"""
+
+log_softmax_docstring = """log_softmax(input, dim, *, dtype=None, mask=None) -> Tensor
+
+Returns log_softmax of all the slices in the :attr:`input` tensor
+along :attr:`dim` while the :attr:`input` elements are masked out
+according to the boolean tensor :attr:`mask`.
+
+Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+of the :attr:`input` tensor. LogSoftmax of i-th element in ``x`` is
+defined as ``log(exp(x[i])/sum(exp(x)))``.
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True then
+the corresponding element in :attr:`input` tensor will be included in
+log_softmax computation, otherwise the element is ignored.
+
+The values of masked-out elements of the output tensor have undefined
+value: it may or may not be set to zero or nan; the choice may correspond to
+the value that leads to the most efficient storage of :attr:`output`
+tensor.
+
+The mask of the log_softmax output tensor can be computed as
+``torch.broadcast_to(mask, input.shape)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which log_softmax is computed.
+
+Keyword args:
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.log_softmax(input, 1, mask=mask)
+    tensor([[-2.1269,    -inf, -0.1269],
+            [    nan,     nan,     nan]])
+"""
+
+logsumexp_docstring = """logsumexp(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+
+Returns logsumexp of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+
+The identity value of logsumexp operation, which is used to start the reduction, is ``-2147483648``.
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in logsumexp computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of logsumexp operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.logsumexp(input, 1, mask=mask)
+    tensor([                   0, -9223372036854775808])
+"""
+
+mean_docstring = """mean(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+
+Returns mean of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+
+By definition, the identity value of a mean operation is the mean
+value of the tensor. If all elements of the input tensor along given
+dimension(s) :attr:`dim` are masked-out, the identity value of the
+mean is undefined.  Due to this ambiguity, the elements of output
+tensor with strided layout, that correspond to fully masked-out
+elements, have ``nan`` values.
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in mean computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of mean operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.mean(input, 1, mask=mask)
+    tensor([-2., nan])
+"""
+
+median_docstring = """median(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+Returns median of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+By definition, the identity value of a median operation is the median
+value of the tensor. If all elements of the input tensor along given
+dimension(s) :attr:`dim` are masked-out, the identity value of the
+median is undefined.  Due to this ambiguity, the elements of output
+tensor with strided layout, that correspond to fully masked-out
+elements, have ``nan`` values.
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in median computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of median operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which median is computed.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.median(input, 1, mask=mask)
+    tensor([-3., nan])
+"""
+
+norm_docstring = """norm(input, ord, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+
+Returns norm of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+
+The identity value of norm operation, which is used to start the
+reduction, is ``0.0``, except for ``ord=-inf`` it is
+``inf``.
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in norm computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of norm operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    ord (int, float, optional): the order of vector norm. Default: 2.
+      See :func:`torch.linalg.vector_norm` for a list of supported norms.
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.norm(input, 2.0, 1, mask=mask)
+    tensor([3.1623, 0.0000])
+"""
+
+normalize_docstring = """normalize(input, ord, dim, *, eps=1e-12, dtype=None, mask=None) -> Tensor
+
+Returns normalize of all the slices in the :attr:`input` tensor
+along :attr:`dim` while the :attr:`input` elements are masked out
+according to the boolean tensor :attr:`mask`.
+
+Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+of the :attr:`input` tensor. Normalize of i-th element in ``x`` is
+defined as ``x[i]/max(norm(x, p), eps)``.
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True then
+the corresponding element in :attr:`input` tensor will be included in
+normalize computation, otherwise the element is ignored.
+
+The values of masked-out elements of the output tensor have undefined
+value: it may or may not be set to zero or nan; the choice may correspond to
+the value that leads to the most efficient storage of :attr:`output`
+tensor.
+
+The mask of the normalize output tensor can be computed as
+``torch.broadcast_to(mask, input.shape)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    ord (int, float): the order of vector norm. Default: 2.
+      See :func:`torch.linalg.vector_norm` for a list of supported norms.
+    dim (int): the dimension along which normalize is computed.
+
+Keyword args:
+    eps (float, optional): small value to avoid division by zero. Default: 1e-12.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.normalize(input, 2.0, 1, mask=mask)
+    tensor([[-0.9487,  0.0000, -0.3162],
+            [ 0.0000,  0.0000,  0.0000]])
+"""
+
+prod_docstring = """prod(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+
+Returns product of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+
+The identity value of product operation, which is used to start the reduction, is ``1``.
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in product computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of product operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.prod(input, 1, mask=mask)
+    tensor([3, 1])
+"""
+
+softmax_docstring = """softmax(input, dim, *, dtype=None, mask=None) -> Tensor
+
+Returns softmax of all the slices in the :attr:`input` tensor
+along :attr:`dim` while the :attr:`input` elements are masked out
+according to the boolean tensor :attr:`mask`.
+
+Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+of the :attr:`input` tensor. Softmax of i-th element in ``x`` is
+defined as ``exp(x[i])/sum(exp(x))``.
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True then
+the corresponding element in :attr:`input` tensor will be included in
+softmax computation, otherwise the element is ignored.
+
+The values of masked-out elements of the output tensor have undefined
+value: it may or may not be set to zero or nan; the choice may correspond to
+the value that leads to the most efficient storage of :attr:`output`
+tensor.
+
+The mask of the softmax output tensor can be computed as
+``torch.broadcast_to(mask, input.shape)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which softmax is computed.
+
+Keyword args:
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.softmax(input, 1, mask=mask)
+    tensor([[0.1192, 0.0000, 0.8808],
+            [   nan,    nan,    nan]])
+"""
+
+softmin_docstring = """softmin(input, dim, *, dtype=None, mask=None) -> Tensor
+
+Returns softmin of all the slices in the :attr:`input` tensor
+along :attr:`dim` while the :attr:`input` elements are masked out
+according to the boolean tensor :attr:`mask`.
+
+Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+of the :attr:`input` tensor. Softmin of i-th element in ``x`` is
+defined as ``exp(-x[i])/sum(exp(-x))``.
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True then
+the corresponding element in :attr:`input` tensor will be included in
+softmin computation, otherwise the element is ignored.
+
+The values of masked-out elements of the output tensor have undefined
+value: it may or may not be set to zero or nan; the choice may correspond to
+the value that leads to the most efficient storage of :attr:`output`
+tensor.
+
+The mask of the softmin output tensor can be computed as
+``torch.broadcast_to(mask, input.shape)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int): the dimension along which softmin is computed.
+
+Keyword args:
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3., -2., -1.], [ 0., 1., 2.]])
+    >>> input
+    tensor([[-3., -2., -1.],
+            [ 0.,  1.,  2.]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.softmin(input, 1, mask=mask)
+    tensor([[0.8808, 0.0000, 0.1192],
+            [   nan,    nan,    nan]])
+"""
+
+std_docstring = """std(input, dim, unbiased, *, keepdim=False, dtype=None, mask=None) -> Tensor
+Returns standard_deviation of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+The identity value of sample standard deviation operation is undefined. The
+elements of output tensor with strided layout, that correspond to
+fully masked-out elements, have ``nan`` values.
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in standard_deviation computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of standard_deviation operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+    unbiased (bool): when True, use Bessel's correction, otherwise, compute
+      the uncorrected sample variance.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.std(input, 1, False, mask=mask)
+    tensor([1., nan])
+"""
+
+sum_docstring = """sum(input, dim, *, keepdim=False, dtype=None, mask=None) -> Tensor
+
+Returns sum of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+
+The identity value of sum operation, which is used to start the reduction, is ``0``.
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in sum computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of sum operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.sum(input, 1, mask=mask)
+    tensor([-4,  0])
+"""
+
+var_docstring = """var(input, dim, unbiased, *, keepdim=False, dtype=None, mask=None) -> Tensor
+Returns variance of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.
+The identity value of sample variance operation is undefined. The
+elements of output tensor with strided layout, that correspond to
+fully masked-out elements, have ``nan`` values.
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in variance computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of variance operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+      Default: None that is equivalent to ``tuple(range(input.ndim))``.
+    unbiased (bool): when True, use Bessel's correction, otherwise, compute
+      the uncorrected sample variance.
+
+Keyword args:
+    keepdim (bool, optional): whether the output tensor has
+      :attr:`dim` retained or not. Default: False.
+    dtype (:class:`torch.dtype`, optional): the desired data type
+      of returned tensor.  If specified, the input tensor is
+      casted to :attr:`dtype` before the operation is
+      performed. Default: None.
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+      containing the binary mask of validity of input tensor
+      elements.
+      Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+Example::
+
+    >>> input = tensor([[-3, -2, -1], [ 0, 1, 2]])
+    >>> input
+    tensor([[-3, -2, -1],
+            [ 0,  1,  2]])
+    >>> mask = tensor([[ True, False, True], [False, False, False]])
+    >>> mask
+    tensor([[ True, False,  True],
+            [False, False, False]])
+    >>> torch.masked._ops.var(input, 1, False, mask=mask)
+    tensor([1., nan])
+"""
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..509cf6a056e234d244b36cfebb45c548c55af89c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/_ops.py
@@ -0,0 +1,1839 @@
+# mypy: allow-untyped-defs
+import warnings
+from collections.abc import Callable
+from typing import Any, TYPE_CHECKING, TypeAlias, TypeVar
+from typing_extensions import ParamSpec
+
+import torch
+from torch import sym_float, Tensor
+from torch._prims_common import corresponding_real_dtype
+from torch.masked import _docs
+from torch.masked.maskedtensor.core import is_masked_tensor, MaskedTensor
+from torch.masked.maskedtensor.creation import as_masked_tensor
+
+
+if TYPE_CHECKING:
+    from torch._prims_common import DimsType
+    from torch.types import _dtype as DType
+
+    DimOrDims: TypeAlias = DimsType | None
+else:
+    # The JIT doesn't understand Union, nor torch.dtype here
+    DType = int
+    DimOrDims = tuple[int, ...] | None
+
+
+__all__: list[str] = []
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+# All masked reduction/normalization operations have the same
+# signatures. Here we introduce docstring templates that are applied
+# to docstrings of reduction/normalization functions via
+# _apply_docstring_templates decorator.
+
+
+def _apply_docstring_templates(func: Callable[_P, _T]) -> Callable[_P, _T]:
+    """Decorator that applies docstring templates to function docstring
+    and returns the function instance.
+    """
+
+    doc_string = getattr(_docs, f"{func.__name__}_docstring", None)
+    if doc_string is None:
+        warnings.warn(
+            f"No documentation string available for {func.__name__}."
+            " PyTorch team should run `python tools/update_masked_docs.py`"
+            " to generate the missing docstrings.",
+            stacklevel=2,
+        )
+    else:
+        func.__doc__ = doc_string
+
+    # Expose function as public symbol
+    __all__.append(func.__name__)
+
+    return func
+
+
+def _generate_docstring(func):
+    """A utility function called from tools/update_masked_docs.py
+    script to update the module torch.masked._docs.py
+    """
+    docstring_templates = dict(
+        reduction_signature="""\
+{function_name}(input, {operation_args}, *, {operation_kwargs}) -> Tensor""",
+        reduction_descr="""\
+Returns {operation name} of all the elements in the :attr:`input`
+tensor along the given dimension(s) :attr:`dim` while the :attr:`input`
+elements are masked out according to the boolean tensor
+:attr:`mask`.""",
+        reduction_args="""\
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension(s) :attr:`dim` where it is of
+size 1. Otherwise, :attr:`dim` is squeezed (see
+:func:`torch.squeeze`), resulting in the output tensor having 1 (or
+``len(dim)``) fewer dimension(s).
+
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True
+then the corresponding element in :attr:`input` tensor will be
+included in {operation name} computation, otherwise the element is
+ignored.
+
+When all elements of :attr:`input` along the given dimension
+:attr:`dim` are ignored (fully masked-out), the corresponding element
+of the output tensor will have undefined value: it may or may not
+correspond to the identity value of {operation name} operation; the
+choice may correspond to the value that leads to the most efficient
+storage of :attr:`output` tensor.
+
+The mask of the output tensor can be computed as
+``torch.any(torch.broadcast_to(mask, input.shape), dim, keepdim=keepdim,
+dtype=torch.bool)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    {args_declarations}
+
+Keyword args:
+    {kwargs_declarations}""",
+        reduction_example="""\
+Example::
+
+    >>> input = {example_input}
+    >>> input
+    {indent_example_input}
+    >>> mask = {example_mask}
+    >>> mask
+    {indent_example_mask}
+    >>> {full_function_name}(input, {example_args}, mask=mask)
+    {indent_example_output}
+""",
+        reduction_identity="""\
+The identity value of {operation name} operation, which is used to start the reduction, is ``{identity_int32}``.""",
+        reduction_identity_dtype="""\
+The identity value of {operation name} operation, which is used to start the
+reduction, depends on input dtype. For instance, for float32, uint8,
+and int32 dtypes, the identity values are ``{identity_float32}``, ``{identity_uint8}``, and ``{identity_int32}``, respectively.""",
+        normalization_signature="""\
+{function_name}(input, {operation_args}, *, {operation_kwargs}) -> Tensor""",
+        normalization_descr="""\
+Returns {operation name} of all the slices in the :attr:`input` tensor
+along :attr:`dim` while the :attr:`input` elements are masked out
+according to the boolean tensor :attr:`mask`.
+
+{definition}""",
+        normalization_args="""\
+The boolean tensor :attr:`mask` defines the "validity" of
+:attr:`input` tensor elements: if :attr:`mask` element is True then
+the corresponding element in :attr:`input` tensor will be included in
+{operation name} computation, otherwise the element is ignored.
+
+The values of masked-out elements of the output tensor have undefined
+value: it may or may not be set to zero or nan; the choice may correspond to
+the value that leads to the most efficient storage of :attr:`output`
+tensor.
+
+The mask of the {operation name} output tensor can be computed as
+``torch.broadcast_to(mask, input.shape)``.
+
+The shapes of the :attr:`mask` tensor and the :attr:`input` tensor
+don't need to match, but they must be :ref:`broadcastable
+<broadcasting-semantics>` and the dimensionality of the :attr:`mask`
+tensor must not be greater than of the :attr:`input` tensor.
+
+Args:
+    input (Tensor): the input tensor
+    {args_declarations}
+
+Keyword args:
+    {kwargs_declarations}""",
+        normalization_example="""\
+Example::
+
+    >>> input = {example_input}
+    >>> input
+    {indent_example_input}
+    >>> mask = {example_mask}
+    >>> mask
+    {indent_example_mask}
+    >>> {full_function_name}(input, {example_args}, mask=mask)
+    {indent_example_output}
+""",
+    )
+
+    args_and_kwargs = {
+        # argument name sufficies separated by double underscore will
+        # be removed in the final documentation string.
+        "sum": (("dim",), ("keepdim=False", "dtype=None", "mask=None")),
+        "prod": (("dim",), ("keepdim=False", "dtype=None", "mask=None")),
+        "cumsum": (("dim__as_int",), ("dtype=None", "mask=None")),
+        "cumprod": (("dim__as_int",), ("dtype=None", "mask=None")),
+        "amin": (("dim",), ("keepdim=False", "dtype=None", "mask=None")),
+        "amax": (("dim",), ("keepdim=False", "dtype=None", "mask=None")),
+        "argmin": (("dim__as_int",), ("keepdim=False", "dtype=None", "mask=None")),
+        "argmax": (("dim__as_int",), ("keepdim=False", "dtype=None", "mask=None")),
+        "mean": (("dim",), ("keepdim=False", "dtype=None", "mask=None")),
+        "median": (("dim__as_int",), ("keepdim=False", "dtype=None", "mask=None")),
+        "norm": (
+            (
+                "ord",
+                "dim",
+            ),
+            ("keepdim=False", "dtype=None", "mask=None"),
+        ),
+        "var": (("dim", "unbiased"), ("keepdim=False", "dtype=None", "mask=None")),
+        "std": (("dim", "unbiased"), ("keepdim=False", "dtype=None", "mask=None")),
+        "logsumexp": (("dim",), ("keepdim=False", "dtype=None", "mask=None")),
+        "softmax": (("dim__as_int",), ("dtype=None", "mask=None")),
+        "log_softmax": (("dim__as_int",), ("dtype=None", "mask=None")),
+        "softmin": (("dim__as_int",), ("dtype=None", "mask=None")),
+        "normalize": (
+            (
+                "ord__required",
+                "dim__as_int",
+            ),
+            ("eps=1e-12", "dtype=None", "mask=None"),
+        ),
+    }
+
+    argument_declarations = {
+        "dim": """\
+    dim (int or tuple of ints, optional): the dimension or dimensions to reduce.
+    Default: None that is equivalent to ``tuple(range(input.ndim))``.""",
+        "dim__as_int": """\
+    dim (int): the dimension along which {operation name} is computed.""",
+        "ord": """\
+    ord (int, float, optional): the order of vector norm. Default: 2.
+    See :func:`torch.linalg.vector_norm` for a list of supported norms.""",
+        "ord__required": """\
+    ord (int, float): the order of vector norm. Default: 2.
+    See :func:`torch.linalg.vector_norm` for a list of supported norms.""",
+        "unbiased": """\
+    unbiased (bool): when True, use Bessel's correction, otherwise, compute
+    the uncorrected sample variance.""",
+        "eps": """\
+    eps (float, optional): small value to avoid division by zero. Default: {default}.""",
+        "keepdim": """\
+    keepdim (bool, optional): whether the output tensor has
+    :attr:`dim` retained or not. Default: {default}.""",
+        "dtype": """\
+    dtype (:class:`torch.dtype`, optional): the desired data type
+    of returned tensor.  If specified, the input tensor is
+    casted to :attr:`dtype` before the operation is
+    performed. Default: {default}.""",
+        "mask": """\
+    mask (:class:`torch.Tensor`, optional): the boolean tensor
+    containing the binary mask of validity of input tensor
+    elements.
+    Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.""",
+    }
+
+    definitions = {
+        "softmax": """\
+    Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+    of the :attr:`input` tensor. Softmax of i-th element in ``x`` is
+    defined as ``exp(x[i])/sum(exp(x))``.""",
+        "log_softmax": """\
+    Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+    of the :attr:`input` tensor. LogSoftmax of i-th element in ``x`` is
+    defined as ``log(exp(x[i])/sum(exp(x)))``.""",
+        "softmin": """\
+    Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+    of the :attr:`input` tensor. Softmin of i-th element in ``x`` is
+    defined as ``exp(-x[i])/sum(exp(-x))``.""",
+        "normalize": """\
+    Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+    of the :attr:`input` tensor. Normalize of i-th element in ``x`` is
+    defined as ``x[i]/max(norm(x, p), eps)``.""",
+        "cumsum": """\
+    Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+    of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is
+    defined as ``sum(x[:i])``.""",
+        "cumprod": """\
+    Let ``x`` be a sequence of unmasked elements of one-dimensional slice
+    of the :attr:`input` tensor. Cumsum of i-th element in ``x`` is
+    defined as ``prod(x[:i])``.""",
+    }
+
+    reduction_names = {
+        "sum": "sum",
+        "prod": "product",
+        "amax": "maximum",
+        "amin": "minimum",
+        "argmax": "argmax",
+        "argmin": "argmin",
+        "mean": "mean",
+        "median": "median",
+        "norm": "norm",
+        "var": "variance",
+        "std": "standard_deviation",
+        "logsumexp": "logsumexp",
+    }
+
+    normalization_names = {
+        "softmax": "softmax",
+        "log_softmax": "log_softmax",
+        "softmin": "softmin",
+        "normalize": "normalize",
+        "cumsum": "cumulative_sum",
+        "cumprod": "cumulative_prod",
+    }
+
+    operation_names = {}
+    operation_names.update(reduction_names)
+    operation_names.update(normalization_names)
+
+    # Default example data:
+    example_dim = 1
+    example_input = torch.tensor([[-3, -2, -1], [0, 1, 2]])
+    example_mask = torch.tensor([[True, False, True], [False, False, False]])
+    example_args: tuple[Any, ...]
+    if func.__name__ in {"norm", "normalize"}:
+        example_args = (2.0, example_dim)
+        example_input = example_input.to(dtype=torch.float32)
+    elif func.__name__ in {"var", "std"}:
+        example_args = (example_dim, False)
+    elif func.__name__ == "median":
+        example_args = (example_dim,)
+        example_input = example_input.to(dtype=torch.float32)
+    else:
+        example_args = (example_dim,)
+
+    operation_args: tuple[str, ...]
+    operation_kwargs: tuple[str, ...]
+    operation_args, operation_kwargs = args_and_kwargs[func.__name__]
+    arg_declarations = [
+        "\n    ".join(
+            argument_declarations.get(a, f"{a.split('__', 1)[0]}: TBD.").splitlines()
+        )
+        for a in operation_args
+    ]
+    kwarg_declarations = [
+        "\n    ".join(
+            argument_declarations.get(
+                a.split("=", 1)[0], f"{a.split('__', 1)[0]}: TBD."
+            )
+            .format(default=a.split("=", 1)[1])
+            .splitlines()
+        )
+        for a in operation_kwargs
+    ]
+
+    if func.__name__ in reduction_names:
+        op_kind = "reduction"
+        doc_sections = ["signature", "descr", "identity", "args", "example"]
+    elif func.__name__ in normalization_names:
+        op_kind = "normalization"
+        doc_sections = ["signature", "descr", "args", "example"]
+        example_input = example_input.to(dtype=torch.float32)
+    else:
+        # add function name to operation names dictionaries
+        raise AssertionError(f"unknown function {func.__name__}")
+    example_output = func(example_input, *example_args, mask=example_mask)
+
+    template_data = {
+        "function_name": func.__name__,
+        "full_function_name": func.__module__ + "." + func.__name__,
+        "operation name": operation_names[func.__name__],
+        "operation_args": ", ".join(a.split("__", 1)[0] for a in operation_args),
+        "operation_kwargs": ", ".join(a.split("__", 1)[0] for a in operation_kwargs),
+        # one-line representation of a tensor:
+        "example_input": " ".join(str(example_input).split()),
+        "example_args": ", ".join(map(str, example_args)),
+        "example_mask": " ".join(str(example_mask).split()),
+        # multi-line representation of a tensor with indent
+        "indent_example_input": ("\n    ").join(str(example_input).splitlines()),
+        "indent_example_mask": ("\n    ").join(str(example_mask).splitlines()),
+        "indent_example_output": ("\n    ").join(str(example_output).splitlines()),
+    }
+
+    if func.__name__ in reduction_names:
+        template_data.update(
+            identity_uint8=_reduction_identity(
+                func.__name__, torch.tensor(0, dtype=torch.uint8)
+            ),
+            identity_int32=_reduction_identity(
+                func.__name__, torch.tensor(0, dtype=torch.int32)
+            ),
+            identity_float32=_reduction_identity(
+                func.__name__, torch.tensor(0, dtype=torch.float32)
+            ),
+        )
+        if func.__name__ == "norm":
+            template_data.update(
+                identity_ord_ninf=_reduction_identity(
+                    func.__name__, torch.tensor(0, dtype=torch.float32), float("-inf")
+                )
+            )
+    elif func.__name__ in normalization_names:
+        template_data.update(definition=definitions[func.__name__])
+    else:
+        # add function name to operation names dictionaries
+        raise AssertionError(f"unknown function {func.__name__}")
+    template_data.update(
+        args_declarations=("\n    ".join(arg_declarations)).format_map(template_data)
+    )
+    template_data.update(
+        kwargs_declarations=("\n    ".join(kwarg_declarations)).format_map(
+            template_data
+        )
+    )
+
+    # Apply function name info to docstring templates:
+    templates = {
+        k: v.format_map(template_data)
+        for k, v in docstring_templates.items()
+        if k.startswith(op_kind)
+    }
+    templates.update(
+        (k, v.format_map(template_data) if isinstance(v, str) else v)
+        for k, v in template_data.items()
+    )
+
+    # Apply docstring templates to function doctring:
+    if func.__doc__ is None:
+        doc_template = "\n\n".join([f"{{{op_kind}_{sec}}}" for sec in doc_sections])
+    else:
+        doc_template = func.__doc__
+    return doc_template.format_map(templates)
+
+
+def _reduction_identity(op_name: str, input: Tensor, *args):
+    """Return identity value as scalar tensor of a reduction operation on
+    given input, or None, if the identity value cannot be uniquely
+    defined for the given input.
+
+    The identity value of the operation is defined as the initial
+    value to reduction operation that has a property ``op(op_identity,
+    value) == value`` for any value in the domain of the operation.
+    Or put it another way, including or excluding the identity value in
+    a list of operands will not change the reduction result.
+
+    See https://github.com/pytorch/rfcs/pull/27 for more information.
+
+    """
+    dtype: DType = input.dtype
+    device = input.device
+    op_name = op_name.rsplit(".", 1)[-1]  # lstrip module name when present
+    if op_name in {"sum", "cumsum"}:
+        return torch.tensor(0, dtype=dtype, device=device)
+    elif op_name in {"prod", "cumprod"}:
+        return torch.tensor(1, dtype=dtype, device=device)
+    elif op_name in {"amax", "argmax", "logaddexp"}:
+        if torch.is_floating_point(input):
+            return torch.tensor(-torch.inf, dtype=dtype, device=device)
+        elif torch.is_signed(input) or dtype == torch.uint8:
+            return torch.tensor(torch.iinfo(dtype).min, dtype=dtype, device=device)
+    elif op_name == "logsumexp":
+        if torch.is_floating_point(input):
+            return torch.tensor(-torch.inf, dtype=dtype, device=device)
+        elif torch.is_complex(input):
+            return torch.tensor(-torch.inf + 0j, dtype=dtype, device=device)
+        elif torch.is_signed(input) or dtype == torch.uint8:
+            return torch.tensor(torch.iinfo(dtype).min, dtype=dtype, device=device)
+    elif op_name in {"amin", "argmin"}:
+        if torch.is_floating_point(input):
+            return torch.tensor(torch.inf, dtype=dtype, device=device)
+        elif torch.is_signed(input) or dtype == torch.uint8:
+            return torch.tensor(torch.iinfo(dtype).max, dtype=dtype, device=device)
+    elif op_name == "mean":
+        # Strictly speaking, the identity value of the mean operation
+        # is the mean of the input. Since the mean value depends on
+        # the dim argument and it may be a non-scalar tensor, we
+        # consider the identity value of the mean operation ambiguous.
+        # Moreover, the mean value of empty input is undefined.
+        return None
+    elif op_name == "norm":
+        ord = args[0] if args else 2
+        if ord == float("-inf"):
+            if not torch.is_floating_point(input):
+                raise AssertionError(f"input must be floating point, got {input.dtype}")
+            return torch.tensor(torch.inf, dtype=dtype, device=device)
+        return torch.tensor(0, dtype=dtype, device=device)
+    elif op_name == "median":
+        # We use NaN for now because the implementation is currently using torch.nanmedian
+        # and NaN is the identity for that function since it gets ignored
+        dtype = input.dtype if torch.is_floating_point(input) else torch.float
+        return torch.tensor(torch.nan, dtype=dtype, device=device)
+    elif op_name in {"var", "std"}:
+        return None
+    raise NotImplementedError(f"identity of {op_name} on {dtype} input")
+
+
+def _canonical_dim(dim: DimOrDims, ndim: int) -> tuple[int, ...]:
+    """Return dim argument as a tuple of sorted dim values."""
+    dims: list[int] = []
+    if dim == ():
+        # Currently, `dim=()` in reductions operations means "reduce
+        # over all dimensions" while in future, it will read "no
+        # reduce". See https://github.com/pytorch/pytorch/issues/29137
+        # When gh-29137 is resolved, this if-block must be deleted.
+        dim = None
+    if dim is None:
+        return tuple(range(ndim))
+    ndim = max(ndim, 1)
+    dim_ = (dim,) if isinstance(dim, (int, torch.SymInt)) else dim
+    for d in dim_:
+        if d in dims:
+            raise RuntimeError(f"dim={d} appears multiple times in the list of dims")
+        if d >= ndim or d < -ndim:
+            raise IndexError(
+                f"Dimension out of range (expected to be in range of [{-ndim}, {ndim - 1}], but got {d})"
+            )
+        # pyrefly: ignore [bad-argument-type]
+        dims.append(d % ndim)
+    return tuple(sorted(dims))
+
+
+def _sparse_coo_flatten_indices(indices: Tensor, shape: tuple):
+    # Flatted N-D indices to 1-D indices
+    flat_indices = indices.new_zeros(indices.size(1))
+    for d, sz in enumerate(shape):
+        flat_indices.mul_(sz)
+        flat_indices.add_(indices[d])
+    return flat_indices
+
+
+def _any(input: Tensor, dim: tuple, keepdim: bool):
+    # Support torch.any with tuple dim argument.
+    # Workaround of https://github.com/pytorch/pytorch/issues/56586
+    r = input
+    for d in reversed(dim):
+        r = r.any(dim=d, keepdim=keepdim)
+    return r
+
+
+def _sparse_coo_where(mask: Tensor, input: Tensor, fill_value: Tensor) -> Tensor:
+    """Sparse variant of torch.where. Supports sparse COO and hybrid sparse COO tensors.
+
+    _sparse_coo_where implements the following invariant:
+
+      _sparse_coo_where(mask, input, fill_value).to_dense(fill_value) ==
+        torch.where(mask.to_dense(), input.to_dense(), torch.full(input.shape, fill_value))
+
+    where `a == b` means `assertEqual(a, b)`, mask is boolean sparse
+    tensor, and `to_dense(fill_value)` is like `to_dense()` except
+    that the unspecified elements are mapped to `fill_value` rather
+    than to `0`.
+
+    Returns a sparse COO tensor with the following features:
+
+    - all specified elements correspond to masked-in elements that
+      have the values of the input tensor. If there exists a masked-in
+      element (as specified by mask) that is not specified in the
+      input, in the result tensor, the corresponding element has value
+      0. In the dense part of the sparse tensor, the masked-out
+      elements are replaced with fill_value.
+
+    - all unspecified elements correspond to masked-out elements.
+    """
+
+    if input.layout != torch.sparse_coo:
+        raise AssertionError(f"input.layout must be sparse_coo, got {input.layout}")
+    if mask.layout != input.layout:
+        raise AssertionError(f"mask.layout must match input.layout, got {mask.layout}")
+    if mask.shape != input.shape:
+        raise AssertionError(
+            f"mask.shape must match input.shape: {mask.shape} vs {input.shape}"
+        )
+    if mask.dense_dim() != input.dense_dim():
+        # TODO: eliminate this restriction
+        raise AssertionError(
+            f"mask.dense_dim() must match input.dense_dim(): "
+            f"{mask.dense_dim()} vs {input.dense_dim()}"
+        )
+
+    input = input.coalesce()
+
+    # For set operations on sparse tensor indices, we'll convert
+    # multi-dimensional indices to 1-D indices for efficiency.
+    input_flat_indices = _sparse_coo_flatten_indices(
+        input.indices(), input.shape[: input.sparse_dim()]
+    )
+    mask_flat_indices = _sparse_coo_flatten_indices(
+        mask.indices(), mask.shape[: mask.sparse_dim()]
+    )
+
+    # the set of mask flat indices that define masked-in elements:
+    if mask.dense_dim() > 0:
+        mask_values = _any(
+            mask.values(), tuple(range(1, input.sparse_dim() + 1)), False
+        )
+    else:
+        mask_values = mask.values()
+    maskin_flat_indices = mask_flat_indices[mask_values.nonzero()[:, 0]]
+
+    def intersection(i1, i2):
+        union, counts = torch.cat([i1, i2]).unique(return_counts=True)
+        return union, torch.where(counts.gt(1))
+
+    def minus(i1, i2):
+        union, counts = torch.cat([i1, i2]).unique(return_counts=True)
+        return intersection(union[torch.where(counts.eq(1))], i1)
+
+    def _apply(a):
+        obj, w = a
+        return obj[w]
+
+    # the set of input flat indices of specified and masked-in elements:
+    maskin_input_flat_indices = _apply(
+        intersection(maskin_flat_indices, input_flat_indices)
+    )
+    _, w = intersection(input_flat_indices, maskin_input_flat_indices)
+
+    # the indices and values of masked-in elements
+    where_input_indices = input.indices()[(slice(None),) + w]
+    where_input_values = input.values()[w]
+
+    if mask.dense_dim() > 0:
+        # apply mask to the dense part of the input values:
+        _, w1 = intersection(mask_flat_indices, maskin_input_flat_indices)
+        where_mask_values = mask.values()[w1]
+        where_input_values = torch.where(
+            where_mask_values, where_input_values, fill_value
+        )
+
+    # the set of flat indices of unspecified input and masked-in elements:
+    maskin_zero_flat_indices = _apply(
+        minus(maskin_flat_indices, maskin_input_flat_indices)
+    )
+
+    # the indices of masked-in zero elements
+    _, w = intersection(mask_flat_indices, maskin_zero_flat_indices)
+    where_zero_indices = mask.indices()[(slice(None),) + w]
+
+    # construct result
+    n = where_zero_indices.size(1)
+    if n == 0:
+        # the input is coalesced, hence input_flat_indices are ordered
+        # and the result is guaranteed to be coalesced:
+        result = torch.sparse_coo_tensor(
+            where_input_indices, where_input_values, input.shape
+        )
+        return result._coalesced_(True)
+
+    where_indices = torch.cat([where_input_indices, where_zero_indices], dim=1)
+    where_values = torch.cat(
+        [
+            where_input_values,
+            where_input_values.new_zeros((n,) + where_input_values.shape[1:]),
+        ]
+    )
+    result = torch.sparse_coo_tensor(where_indices, where_values, input.shape)
+
+    # appending zero elements leads to uncoalesced sparse tensor
+    return result.coalesce()
+
+
+def _sparse_coo_scatter_reduction_helper(
+    op,
+    mask_input: Tensor,
+    dims: tuple[int, ...],
+    keepdim: bool,
+    dtype: DType | None = None,
+) -> Tensor:
+    reduce = op.__name__
+    valid_reductions = ["sum", "prod", "amax", "amin"]
+    if reduce not in valid_reductions:
+        raise ValueError(
+            f"op must be one of {' '.join(valid_reductions)}, but got {reduce} instead"
+        )
+
+    output_dtype = dtype
+    values, indices = mask_input._values(), mask_input._indices()
+    input_dims = mask_input.dim()
+    num_sparse_dims = mask_input.sparse_dim()
+    reduced_sparse_dims = []
+    retained_sparse_dims = []
+    reduced_dense_dims = []
+
+    # promote dtype if specified
+    if values.dtype != output_dtype:
+        values = values.to(output_dtype)
+
+    if keepdim:
+        output_shape = tuple(
+            1 if i in dims else si for (i, si) in enumerate(mask_input.shape)
+        )
+    else:
+        output_shape = tuple(
+            si for (i, si) in enumerate(mask_input.shape) if i not in dims
+        )
+
+    for d in dims:
+        if d >= input_dims:
+            continue
+
+        if d < num_sparse_dims:
+            reduced_sparse_dims.append(d)
+        else:
+            reduced_dense_dims.append(d + 1 - num_sparse_dims)
+
+    # Reduce dense dimensions
+    if len(reduced_dense_dims) > 0:
+        if reduce == "sum":
+            new_values = values
+            new_values = op(new_values, dim=reduced_dense_dims, keepdim=bool(keepdim))
+        else:
+            # FIXME: Implement reductions for dense dimensions for ops with non-zero reduction identities
+            return NotImplemented
+    else:
+        new_values = values.clone()
+
+    # Reduce sparse dimensions
+    if len(reduced_sparse_dims) == num_sparse_dims:
+        if reduce in {"amax", "amin"} and new_values.size(0) == 0:
+            # IndexError: amax(): Expected reduction dim 0 to have non-zero size.
+            # sum()/prod() return the reduction identity when dim has size 0 but amax()/amin() do not
+            # See https://github.com/pytorch/pytorch/issues/61901
+            new_values = _reduction_identity(reduce, new_values)
+        else:
+            new_values = op(new_values, dim=0)
+        if keepdim:
+            for _ in range(num_sparse_dims):
+                new_values = new_values.unsqueeze(0)
+        return new_values.to(dtype=output_dtype).to_sparse()
+    else:
+        new_indices = indices.clone()
+        if keepdim:
+            # zero out reduced sparse dimensions if keepdim = True
+            # ensures that the call to torch.unique folds duplicated indices together while preserving the dimension
+            new_indices[reduced_sparse_dims, :] = 0
+        else:
+            # remove reduced sparse dimensions if keepdim = False
+            if len(reduced_sparse_dims) > 0:
+                retained_sparse_dims = [
+                    i
+                    for i in range(num_sparse_dims)
+                    if i not in set(reduced_sparse_dims)
+                ]
+                new_indices = new_indices.index_select(
+                    0, torch.tensor(retained_sparse_dims).to(mask_input.device)
+                )
+
+    # Use scatter_reduce to reduce items in the new_values tensor that correspond to the same indices in new_indices
+    if new_indices.numel() > 0:
+        # lexsort indices and get index tensor for scatter reduction
+        new_indices, inverse_indices = torch.unique(
+            new_indices, return_inverse=True, dim=1
+        )
+        out_shape = list(new_values.shape)
+        out_shape[0] = new_indices.shape[1]
+        for _ in range(new_values.ndim - 1):
+            inverse_indices = inverse_indices.unsqueeze(-1)
+        scatter_indices = inverse_indices.expand(new_values.shape)
+        # FIXME: temporary workaround for issue with bfloat16/float16 remove when acctype is implemented for scatter_reduce
+        if output_dtype in {torch.bfloat16, torch.float16}:
+            new_values = new_values.to(torch.float)
+            out = new_values.new_empty(out_shape)
+            new_values = out.scatter_reduce_(
+                0, scatter_indices, new_values, reduce=reduce, include_self=False
+            )
+            new_values = new_values.to(dtype=output_dtype)
+        else:
+            out = new_values.new_empty(out_shape)
+            new_values = out.scatter_reduce_(
+                0, scatter_indices, new_values, reduce=reduce, include_self=False
+            )
+
+    return torch.sparse_coo_tensor(
+        new_indices,
+        new_values,
+        output_shape,
+        dtype=output_dtype,
+        device=mask_input.device,
+    )
+
+
+def _sparse_csr_segment_reduction_helper(
+    op,
+    mask_input: Tensor,
+    dims: tuple[int, ...],
+    keepdim: bool,
+    dtype: DType | None = None,
+) -> Tensor:
+    # Currently, while sparse CSR is always 2D with no dense dimensions keepdim must be True
+    # FIXME: when dense dimensions are implemented for CSR tensors
+    if not keepdim:
+        raise AssertionError(
+            "reduction operations on CSR tensors with keepdim=False is unsupported"
+        )
+    reduce = op.__name__
+    valid_reductions = ["sum", "prod", "mean", "amax", "amin"]
+    if reduce not in valid_reductions:
+        raise ValueError(
+            f"op must be one of {' '.join(valid_reductions)}, but got {reduce} instead"
+        )
+    device = mask_input.device
+    output_dtype = dtype
+    values, crow_indices, col_indices = (
+        mask_input.values(),
+        mask_input.crow_indices(),
+        mask_input.col_indices(),
+    )
+
+    # promote dtype if specified
+    if values.dtype != output_dtype:
+        values = values.to(output_dtype)
+
+    if len(dims) == 0:
+        return mask_input
+    if len(dims) == 1:
+        if dims[0] == 0:
+            new_col_indices, scatter_indices = torch.unique(
+                col_indices, return_inverse=True
+            )
+            new_nnz = new_col_indices.shape[0]
+            new_crow_indices = torch.tensor([0, new_nnz])
+            new_values = values.new_empty(new_col_indices.shape)
+            new_values.scatter_reduce_(
+                0, scatter_indices, values, reduce, include_self=False
+            )
+            new_shape = [1, mask_input.size(1)]
+        else:
+            if dims[0] != 1:
+                raise AssertionError(
+                    "Sparse CSR tensors are 2D and only support reduction along dim 0 or 1."
+                )
+            # all intervals new_crow_indices[i] - new_crow_indices[i-1] are 1
+            # except for where crow_indices[i] == crow_indices[i-1] where the interval remains as 0
+            new_crow_indices = torch.cat(
+                (
+                    crow_indices.new_zeros(1),
+                    torch.cumsum(torch.diff(crow_indices) != 0, 0),
+                ),
+                0,
+            )
+            new_nnz = new_crow_indices[-1]
+            new_col_indices = col_indices.new_zeros(new_nnz)  # type: ignore[call-overload]
+            new_values = torch._segment_reduce(values, reduce, offsets=crow_indices)  # type: ignore[attr-defined]
+            new_shape = [mask_input.size(0), 1]
+    else:
+        if len(dims) != 2:
+            raise AssertionError(f"expected len(dims) == 2, got {len(dims)}")
+        nnz = min(1, values.numel())
+        if nnz == 1:
+            op_kwargs = {"keepdim": True, "dtype": output_dtype}
+            # amax and amin do not support dtype kwarg
+            if reduce in ["amax", "amin"]:
+                del op_kwargs["dtype"]
+            new_values = op(values, 0, **op_kwargs)
+        else:
+            new_values = torch.empty(0, dtype=output_dtype)
+        new_col_indices = col_indices.new_zeros(nnz)
+        new_crow_indices = torch.tensor([0, nnz])
+        new_shape = [1, nnz]
+
+    return torch.sparse_csr_tensor(
+        new_crow_indices,
+        new_col_indices,
+        new_values,
+        new_shape,
+        dtype=output_dtype,
+        device=device,
+    )
+
+
+def _sparse_csr_where(mask: Tensor, input: Tensor, fill_value: Tensor) -> Tensor:
+    """Sparse variant of torch.where. Supports sparse CSR tensors."""
+    # TODO: implement sparse CSR specific where operator for efficiency
+    return _sparse_coo_where(
+        mask.to_sparse_coo(), input.to_sparse_coo(), fill_value
+    ).to_sparse_csr()
+
+
+def _where(mask: Tensor, input: Tensor, fill_value: Tensor) -> Tensor:
+    """torch.where with sparse inputs support.
+
+    _where implements the following invariant:
+
+      _where(mask, input, fill_value).to_dense(fill_value) ==
+        torch.where(mask.to_dense(), input.to_dense(), torch.full(input.shape, fill_value))
+
+    where `a == b` means `assertEqual(a, b)`, mask is boolean sparse
+    tensor, and `to_dense(fill_value)` is like `to_dense()` except
+    that the unspecified elements are mapped to `fill_value` rather
+    than to `0`.
+
+    Returns a sparse tensor with the following features:
+
+    - all specified elements correspond to masked-in elements that
+      have the values of the input tensor. If there exists a masked-in
+      element (as specified by mask) that is not specified in the
+      input, in the result tensor, the corresponding element has value
+      0. In the dense part of the sparse tensor, the masked-out
+      elements are replaced with fill_value.
+
+    - all unspecified elements correspond to masked-out elements.
+    """
+    if mask.layout == torch.strided:
+        return torch.where(mask, input, fill_value)
+    elif mask.layout == torch.sparse_coo:
+        return _sparse_coo_where(mask, input, fill_value)
+    elif mask.layout == torch.sparse_csr:
+        return _sparse_csr_where(mask, input, fill_value)
+    else:
+        raise ValueError(
+            f"_where expects strided or sparse COO or sparse CSR tensor but got {mask.layout}"
+        )
+
+
+def _input_mask(input: Tensor | MaskedTensor, *args, **kwargs) -> Tensor:
+    """Return canonical input mask.
+
+    A canonical input mask is defined as a boolean mask tensor that
+    shape and layout matches with the shape and the layout of the
+    input.
+
+    The canonical input mask is computed from the :attr:`mask` tensor
+    content to meet the following criteria:
+
+    1. The shape of the canonical input mask is the same as the shape
+       of :attr:`input` tensor. If the mask tensor has a smaller shape
+       than the shape of the :attr:`input`, broadcasting rules will be
+       applied. Downcasting of mask is not supported.
+
+    2. The layout of the canonical input mask is the same as the
+       layout of the :attr:`input` tensor. If the mask has different
+       layout, it will be converted to the expected layout.  In the
+       case of sparse COO layout, the canonical input mask will be
+       coalesced.
+
+    3. The dtype of the canonical input mask is torch.bool. If the
+       mask dtype is not bool then it will be converted to bool dtype
+       using `.to(dtype=bool)` method call.
+
+    4. The elements of the canonical input mask have boolean values
+       copied from the content of the :attr:`mask` tensor (after
+       possible broadcasting and dtype conversion transforms).  In
+       general, the sparsity pattern of the sparse canonical input
+       mask need not to be the same as the sparsity pattern of the
+       sparse :attr:`input` tensor.
+
+    """
+    if input.layout not in {torch.strided, torch.sparse_coo, torch.sparse_csr}:
+        raise ValueError(
+            f"_input_mask expects strided or sparse COO or sparse CSR tensor but got {input.layout}"
+        )
+
+    mask = kwargs.get("mask")
+
+    # default mask
+    if mask is None:
+        raise ValueError("_input_mask requires explicit mask")
+
+    # mask shape must match with input shape
+    if mask.shape != input.shape:
+        if mask.ndim > input.ndim:
+            raise IndexError(
+                "_input_mask expected broadcastable mask (got mask dimensionality higher than of the input)"
+            )
+        if mask.layout == torch.strided:
+            mask = torch.broadcast_to(mask.clone(), input.shape).to(dtype=torch.bool)
+        elif mask.layout == torch.sparse_coo:
+            mask = torch._sparse_broadcast_to(mask, input.shape)
+        else:
+            if mask.layout != torch.sparse_csr:
+                raise AssertionError(f"expected sparse_csr layout, got {mask.layout}")
+            # Broadcasting of CSR tensors is not implemented. Working
+            # around by using COO layout.
+            mask = torch._sparse_broadcast_to(
+                mask.to_sparse(), input.shape
+            ).to_sparse_csr()
+
+    # mask layout must match with input layout
+    if mask.layout != input.layout:
+        if input.layout == torch.strided:
+            mask = mask.to_dense()
+        elif input.layout == torch.sparse_coo:
+            if mask.layout == torch.strided:
+                mask = mask.to_sparse(input.sparse_dim())
+            else:
+                mask = mask.to_sparse()
+        else:
+            if input.layout != torch.sparse_csr:
+                raise AssertionError(f"expected sparse_csr layout, got {input.layout}")
+            mask = mask.to_sparse_csr()
+
+    # sparse mask must be coalesced
+    if mask.layout == torch.sparse_coo:
+        mask = mask.coalesce()
+
+    # mask is a boolean tensor
+    mask = mask.to(dtype=torch.bool)
+
+    return mask
+
+
+def _output_mask(op, input: Tensor, *args, **kwargs) -> Tensor:
+    """Return output mask of masked operation applied to given arguments."""
+    if callable(op):
+        is_reduction = op.__name__ in {
+            "sum",
+            "prod",
+            "amax",
+            "amin",
+            "argmax",
+            "argmin",
+            "mean",
+            "median",
+            "norm",
+            "var",
+            "std",
+            "logsumexp",
+        }
+        is_normalization = op.__name__ in {
+            "softmax",
+            "log_softmax",
+            "softmin",
+            "normalize",
+            "cumsum",
+            "cumprod",
+        }
+        if is_reduction:
+            if op.__name__ == "norm":
+                if args:
+                    args = args[1:]  # lstrip ord argument
+            dim = args[0] if args else kwargs.get("dim")
+            outmask = _input_mask(input, *args, **kwargs)
+            keepdim = kwargs.get("keepdim", False)
+            dim_ = _canonical_dim(dim, input.ndim)
+            return _any(outmask, dim_, bool(keepdim))
+        elif is_normalization:
+            return _input_mask(input, *args, **kwargs)
+        else:
+            raise ValueError(
+                f"_output_mask expected masked operation (got callable {op.__module__}.{op.__name__})"
+            )
+    else:
+        raise ValueError(
+            f"_output_mask expected masked operation (got {type(op).__name__} object)"
+        )
+
+
+def _combine_input_and_mask(op, input: MaskedTensor | Tensor, mask, *args) -> Tensor:
+    def helper(input, mask):
+        if mask is None:
+            return input
+        canonical_mask = _input_mask(input, mask=mask)
+        if callable(op):
+            fill_value = _reduction_identity(op.__name__, input, *args)
+            return _where(canonical_mask, input, fill_value)
+        else:
+            raise ValueError(
+                f"_combine_input_and_mask expected masked operation (got {type(op).__name__} object)"
+            )
+
+    class Combine(torch.autograd.Function):
+        @staticmethod
+        # pyrefly: ignore [bad-override]
+        def forward(ctx, input, mask):
+            """Return input with masked-out elements eliminated for the given operations."""
+            ctx.save_for_backward(mask)
+
+            if mask is not None:
+                ctx.mark_non_differentiable(mask)
+
+            return helper(input, mask)
+
+        @staticmethod
+        # pyrefly: ignore [bad-override]
+        def backward(ctx, grad_output):
+            (mask,) = ctx.saved_tensors
+            grad_data = (
+                grad_output.get_data() if is_masked_tensor(grad_output) else grad_output
+            )
+            result = as_masked_tensor(grad_data, mask)
+            return result, None
+
+    return (
+        Combine.apply(input.get_data(), input.get_mask())  # type: ignore[union-attr]
+        if is_masked_tensor(input)
+        else helper(input, mask)
+    )
+
+
+@_apply_docstring_templates
+def sum(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    # __doc__ is generated by _apply_docstring_templates decorator
+    if dtype is None:
+        # promote integer types to int64 when output dtype is not specified
+        if input.layout == torch.sparse_csr:
+            if input.dtype in {
+                torch.uint8,
+                torch.bool,
+                torch.int8,
+                torch.int16,
+                torch.int32,
+            }:
+                # csr.to(dtype=torch.int64) is not implemented, so
+                # using coo.to on input to ensure the promoted dtype
+                input = input.to_sparse_coo().to(dtype=torch.int64).to_sparse_csr()
+            else:
+                dtype = input.dtype
+        else:
+            dtype = input.dtype
+            if input.dtype in {
+                torch.uint8,
+                torch.bool,
+                torch.int8,
+                torch.int16,
+                torch.int32,
+            }:
+                dtype = torch.int64
+    dim_ = _canonical_dim(dim, input.ndim)
+    mask_input = _combine_input_and_mask(sum, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.sum(mask_input, dim_, bool(keepdim), dtype=dtype)
+    elif mask_input.layout == torch.sparse_coo:
+        return _sparse_coo_scatter_reduction_helper(
+            torch.sum, mask_input, dim_, bool(keepdim), dtype
+        )
+    elif mask_input.layout == torch.sparse_csr:
+        return torch._sparse_csr_sum(
+            mask_input, dim=list(dim_), keepdim=bool(keepdim), dtype=dtype
+        )
+    else:
+        raise ValueError(
+            f"masked sum expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def prod(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    # __doc__ is generated by _apply_docstring_templates decorator
+    if dtype is None:
+        # promote integer types to int64 when output dtype is not specified
+        if input.layout == torch.sparse_csr:
+            if input.dtype in {
+                torch.uint8,
+                torch.bool,
+                torch.int8,
+                torch.int16,
+                torch.int32,
+            }:
+                # csr.to(dtype=torch.int64) is not implemented, so
+                # using coo.to on input to ensure the promoted dtype
+                input = input.to_sparse_coo().to(dtype=torch.int64).to_sparse_csr()
+            else:
+                dtype = input.dtype
+        else:
+            dtype = input.dtype
+            if input.dtype in {
+                torch.uint8,
+                torch.bool,
+                torch.int8,
+                torch.int16,
+                torch.int32,
+            }:
+                dtype = torch.int64
+    dim_ = _canonical_dim(dim, input.ndim)
+    mask_input = _combine_input_and_mask(prod, input, mask)
+    if mask_input.layout == torch.strided:
+        # Workaround https://github.com/pytorch/pytorch/issues/56586
+        result = mask_input
+        result = result.to(dtype=dtype)
+        for d in reversed(dim_):
+            result = result.prod(dim=d, keepdim=bool(keepdim))
+        return result
+    elif mask_input.layout == torch.sparse_coo:
+        if mask is None:
+            # See comment in the sparse_csr branch, the same issue arises for sparse_coo tensors
+            raise ValueError(
+                "masked prod expects explicit mask for sparse_coo tensor input"
+            )
+        return _sparse_coo_scatter_reduction_helper(
+            torch.prod, mask_input, dim_, bool(keepdim), dtype
+        )
+    elif mask_input.layout == torch.sparse_csr:
+        if mask is None:
+            # mask is None corresponds to all-True mask. The
+            # unspecified elements in the CSR tensor correspond to
+            # zero values. Hence, the prod reduction result is
+            # automatically zero unless all elements are specified.
+            # A semi-optimal way to take this into account is to use:
+            #
+            #   masked_prod(csr, ..., mask=None) == torch._sparse_csr_prod(csr, ...) * all(csr.nonzero(), ...)
+            #
+            # but that requires implementing `all` and `nonzero`
+            # support for sparse csr tensors.
+            raise ValueError(
+                "masked prod expects explicit mask for sparse_csr tensor input"
+            )
+        return torch._sparse_csr_prod(
+            mask_input, dim=list(dim_), keepdim=bool(keepdim), dtype=dtype
+        )
+    else:
+        raise ValueError(
+            f"masked prod expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def cumsum(
+    input: Tensor,
+    dim: int,
+    *,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    if dtype is None:
+        dtype = input.dtype
+    dim_ = _canonical_dim(dim, input.ndim)[0]
+    mask_input = _combine_input_and_mask(sum, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.cumsum(mask_input, dim_, dtype=dtype).to(dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked cumsum expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def cumprod(
+    input: Tensor,
+    dim: int,
+    *,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    if dtype is None:
+        dtype = input.dtype
+    dim_ = _canonical_dim(dim, input.ndim)[0]
+    mask_input = _combine_input_and_mask(prod, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.cumprod(mask_input, dim_, dtype=dtype).to(dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked cumprod expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def amax(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+
+{reduction_descr}
+
+{reduction_identity_dtype}
+
+{reduction_args}
+
+{reduction_example}"""
+    if dtype is None:
+        dtype = input.dtype
+
+    mask_input = _combine_input_and_mask(amax, input, mask)
+    dim_ = _canonical_dim(dim, mask_input.ndim)
+    if mask_input.layout == torch.strided:
+        return torch.amax(mask_input, dim_, bool(keepdim)).to(dtype=dtype)
+    elif mask_input.layout == torch.sparse_coo:
+        if mask is None:
+            # See comment in the sparse_csr branch of prod, a similar issue arises here
+            # where unspecified elements along a dimension may need to be reduced with the result
+            raise ValueError(
+                "masked amax expects explicit mask for sparse_coo tensor input"
+            )
+        return _sparse_coo_scatter_reduction_helper(
+            torch.amax, mask_input, dim_, bool(keepdim), dtype
+        )
+    elif mask_input.layout == torch.sparse_csr:
+        if mask is None:
+            raise ValueError(
+                "masked amax expects explicit mask for sparse_csr tensor input"
+            )
+        return _sparse_csr_segment_reduction_helper(
+            torch.amax, mask_input, dim_, bool(keepdim), dtype
+        )
+    else:
+        raise ValueError(
+            f"masked amax expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def amin(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+
+{reduction_descr}
+
+{reduction_identity_dtype}
+
+{reduction_args}
+
+{reduction_example}"""
+    if dtype is None:
+        dtype = input.dtype
+
+    mask_input = _combine_input_and_mask(amin, input, mask)
+    dim_ = _canonical_dim(dim, mask_input.ndim)
+    if mask_input.layout == torch.strided:
+        return torch.amin(mask_input, dim_, bool(keepdim)).to(dtype=dtype)
+    elif mask_input.layout == torch.sparse_coo:
+        if mask is None:
+            # See comment in the sparse_csr branch of prod, a similar issue arises here
+            # where unspecified elements along a dimension may need to be reduced with the result
+            raise ValueError(
+                "masked amax expects explicit mask for sparse_coo tensor input"
+            )
+        return _sparse_coo_scatter_reduction_helper(
+            torch.amin, mask_input, dim_, bool(keepdim), dtype
+        )
+    elif mask_input.layout == torch.sparse_csr:
+        if mask is None:
+            raise ValueError(
+                "masked amin expects explicit mask for sparse_csr tensor input"
+            )
+        return _sparse_csr_segment_reduction_helper(
+            torch.amin, mask_input, dim_, bool(keepdim), dtype
+        )
+    else:
+        raise ValueError(
+            f"masked amin expects strided, sparse_coo or sparse_csr tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def argmax(
+    input: Tensor | MaskedTensor,
+    dim: int | None = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+{reduction_descr}
+{reduction_identity_dtype}
+{reduction_args}
+{reduction_example}"""
+    if dtype is None:
+        dtype = input.dtype
+    mask_input = _combine_input_and_mask(argmax, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.argmax(mask_input, dim, bool(keepdim)).to(dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked argmax expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def argmin(
+    input: Tensor | MaskedTensor,
+    dim: int | None = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+{reduction_descr}
+{reduction_identity_dtype}
+{reduction_args}
+{reduction_example}"""
+    if dtype is None:
+        dtype = input.dtype
+    mask_input = _combine_input_and_mask(argmin, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.argmin(mask_input, dim, bool(keepdim)).to(dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked argmin expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def mean(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+
+{reduction_descr}
+
+By definition, the identity value of a mean operation is the mean
+value of the tensor. If all elements of the input tensor along given
+dimension(s) :attr:`dim` are masked-out, the identity value of the
+mean is undefined.  Due to this ambiguity, the elements of output
+tensor with strided layout, that correspond to fully masked-out
+elements, have ``nan`` values.
+
+{reduction_args}
+
+{reduction_example}"""
+    dtype_source = "Optional"
+    if dtype is None:
+        dtype = input.dtype
+        dtype_source = "Input"
+
+    if not (dtype.is_floating_point or dtype.is_complex):
+        raise ValueError(
+            f"mean(): Could not infer output dtype. {dtype_source} dtype must be either "
+            f"a floating point or complex dtype. Got: {dtype}"
+        )
+    if input.layout == torch.strided:
+        if mask is None:
+            # TODO: compute count analytically
+            # pyrefly: ignore [no-matching-overload]
+            count = sum(
+                torch.ones(input.shape, dtype=torch.int64, device=input.device),
+                dim,
+                keepdim=keepdim,
+            )
+            # pyrefly: ignore [no-matching-overload]
+            total = sum(input, dim, keepdim=keepdim, dtype=dtype)
+        else:
+            inmask = _input_mask(input, mask=mask)
+            count = inmask.sum(dim=dim, keepdim=bool(keepdim))
+            # pyrefly: ignore [no-matching-overload]
+            total = sum(input, dim, keepdim=keepdim, dtype=dtype, mask=inmask)
+        return total / count
+    elif input.layout == torch.sparse_csr:
+        mask_input = _combine_input_and_mask(mean, input, mask)
+        dim_ = _canonical_dim(dim, mask_input.ndim)
+        if mask is None:
+            raise ValueError(
+                "masked mean expects explicit mask for sparse_csr tensor input"
+            )
+        return _sparse_csr_segment_reduction_helper(
+            torch.mean, mask_input, dim_, bool(keepdim), dtype
+        )
+    else:
+        raise ValueError(
+            f"masked mean expects strided or sparse_csr tensor (got {input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def median(
+    input: Tensor | MaskedTensor,
+    dim: int = -1,
+    *,
+    keepdim: bool = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+{reduction_descr}
+By definition, the identity value of a median operation is the median
+value of the tensor. If all elements of the input tensor along given
+dimension(s) :attr:`dim` are masked-out, the identity value of the
+median is undefined.  Due to this ambiguity, the elements of output
+tensor with strided layout, that correspond to fully masked-out
+elements, have ``nan`` values.
+{reduction_args}
+{reduction_example}"""
+    if dtype is None:
+        dtype = input.dtype
+    dim_ = _canonical_dim(dim, input.ndim)[0]
+    is_float = torch.is_floating_point(input)
+    if not is_float:
+        input = input.to(dtype=torch.float)
+    mask_input = _combine_input_and_mask(median, input, mask)
+    if mask_input.layout == torch.strided:
+        output = torch.nanmedian(mask_input, dim_, keepdim).values
+        if is_float:
+            return output
+        elif not is_float and not torch.isnan(output).any():
+            return output.to(dtype=dtype)
+        else:
+            raise ValueError(
+                "masked median expects no fully masked out rows if dtype is not floating point"
+            )
+    else:
+        raise ValueError(
+            f"masked median expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def logsumexp(
+    input: Tensor,
+    dim: DimOrDims = None,
+    *,
+    keepdim: bool = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    if dtype is None:
+        dtype = input.dtype
+    dim_ = _canonical_dim(dim, input.ndim)
+    mask_input = _combine_input_and_mask(logsumexp, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.logsumexp(mask_input, dim_, keepdim=keepdim).to(dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked logsumexp expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+# Cannot use _apply_docstring_templates as it is only set up for reductions and normalizations
+def logaddexp(
+    input: Tensor | MaskedTensor,
+    other: Tensor | MaskedTensor,
+    *,
+    dtype: DType | None = None,
+    input_mask: Tensor | None = None,
+    other_mask: Tensor | None = None,
+) -> Tensor:
+    """logaddexp(input, other, *, dtype=None, input_mask=None, other_mask=None) -> Tensor
+
+    Returns logaddexp of all the elements in the :attr:`input` and the :attr:`other`
+    tensor. The :attr:`input` elements are masked out according to the boolean tensor
+    :attr:`input_mask` and the attr:`other` elements are masked out according to the boolean tensor
+    :attr:`other_mask`.
+
+    The shapes of a mask tensor and the tensor to be masked
+    don't need to match, but they must be :ref:`broadcastable
+    <broadcasting-semantics>` and the dimensionality of the mask
+    tensor must not be greater than of the tensor to be masked.
+
+    Args:
+        input (Tensor): the input tensor
+        other (Tensor): the second input tensor
+
+    Keyword args:
+        dtype (:class:`torch.dtype`, optional): the desired data type
+          of returned tensor.  If specified, the output tensor is
+          casted to :attr:`dtype` after the operation is
+          performed. Default: None.
+        input_mask (:class:`torch.Tensor`, optional): the boolean tensor
+          containing the binary mask of validity of :attr:`input` tensor elements.
+          Default: None that is equivalent to ``torch.ones(input.shape, dtype=torch.bool)``.
+        other_mask (:class:`torch.Tensor`, optional): the boolean tensor
+          containing the binary mask of validity of :attr:`other` tensor elements.
+          Default: None that is equivalent to ``torch.ones(other.shape, dtype=torch.bool)``.
+
+    Example::
+
+        >>> input = torch.tensor([-100.0, -200, -300])
+        >>> input
+        tensor([-100., -200., -300.])
+        >>> other = torch.tensor([-1.0, -2, -3])
+        >>> other
+        tensor([-1., -2., -3.])
+        >>> mask = torch.tensor([True, False, True])
+        >>> mask
+        tensor([ True, False,  True])
+        >>> torch.masked._ops.logaddexp(input, other, input_mask=mask, other_mask=mask)
+        tensor([-1., -inf, -3.])"""
+    if dtype is None:
+        dtype = input.dtype
+    if input.layout == torch.strided and other.layout == torch.strided:
+        mask_input = _combine_input_and_mask(logaddexp, input, input_mask)
+        mask_other = _combine_input_and_mask(logaddexp, other, other_mask)
+        return torch.logaddexp(mask_input, mask_other).to(dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked logaddexp expects strided tensors (got {input.layout} tensor for input, {other.layout} for other)"
+        )
+
+
+@_apply_docstring_templates
+def norm(
+    input: Tensor | MaskedTensor,
+    ord: float | None = 2.0,
+    dim: DimOrDims = None,
+    *,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+
+{reduction_descr}
+
+The identity value of norm operation, which is used to start the
+reduction, is ``{identity_float32}``, except for ``ord=-inf`` it is
+``{identity_ord_ninf}``.
+
+{reduction_args}
+
+{reduction_example}"""
+    if dtype is None:
+        dtype = input.dtype
+    mask_input = _combine_input_and_mask(norm, input, mask, ord)
+    if mask_input.layout == torch.strided:
+        dim_ = _canonical_dim(dim, input.ndim)
+        return torch.linalg.vector_norm(
+            mask_input, ord, dim_, bool(keepdim), dtype=dtype
+        )
+    else:
+        raise ValueError(
+            f"masked norm expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+def _std_var(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims,
+    unbiased: bool | None,
+    *,
+    correction_opt: int | float | None,
+    keepdim: bool | None,
+    dtype: DType | None,
+    mask: Tensor | None,
+    take_sqrt: bool | None,
+) -> Tensor:
+    if unbiased is not None and correction_opt is not None:
+        raise AssertionError("Only one of unbiased and correction may be given")
+    correction = 1.0
+    if unbiased is not None:
+        correction = 1.0 if unbiased else 0.0
+    if correction_opt is not None:
+        correction = sym_float(correction_opt)
+
+    if dtype is None:
+        dtype = input.dtype
+        if not (dtype.is_floating_point or dtype.is_complex):
+            dtype = torch.float32
+    compute_dtype = dtype
+    if not (compute_dtype.is_floating_point or compute_dtype.is_complex):
+        compute_dtype = torch.float32
+    if input.layout == torch.strided:
+        if mask is None:
+            # TODO: compute count analytically
+            # pyrefly: ignore [no-matching-overload]
+            count = sum(
+                torch.ones(input.shape, dtype=torch.int64, device=input.device),
+                dim,
+                keepdim=True,
+            )
+            # pyrefly: ignore [no-matching-overload]
+            sample_total = sum(input, dim, keepdim=True, dtype=dtype)
+        else:
+            inmask = _input_mask(input, mask=mask)
+            count = inmask.sum(dim=dim, keepdim=True)
+            # pyrefly: ignore [no-matching-overload]
+            sample_total = sum(input, dim, keepdim=True, dtype=dtype, mask=inmask)
+        # TODO: replace torch.subtract/divide/square/maximum with
+        # masked subtract/divide/square/maximum when these will be
+        # available.
+        sample_mean = torch.divide(sample_total, count)
+        x = torch.subtract(input, sample_mean)
+        if mask is None:
+            # pyrefly: ignore [no-matching-overload]
+            total = sum(x * x.conj(), dim, keepdim=keepdim, dtype=compute_dtype)
+        else:
+            # pyrefly: ignore [no-matching-overload]
+            total = sum(
+                x * x.conj(),
+                dim,
+                keepdim=keepdim,
+                dtype=compute_dtype,
+                mask=inmask,  # type: ignore[possibly-undefined]
+            )
+        if not keepdim:
+            count = count.reshape(total.shape)
+        if correction != 0:
+            real_dtype = (
+                corresponding_real_dtype(compute_dtype)
+                if compute_dtype.is_complex
+                else compute_dtype
+            )
+            count = count.to(real_dtype)
+            count = torch.subtract(count, correction)
+            count = torch.maximum(count, count.new_zeros([]))
+        output = torch.divide(total, count).to(dtype=dtype)
+        if take_sqrt:
+            output = torch.sqrt(output)
+        return output
+    else:
+        raise ValueError(
+            f"masked std/var expects strided tensor (got {input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def var(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims = None,
+    unbiased: bool | None = None,
+    *,
+    correction: int | float | None = None,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+{reduction_descr}
+The identity value of sample variance operation is undefined. The
+elements of output tensor with strided layout, that correspond to
+fully masked-out elements, have ``nan`` values.
+{reduction_args}
+{reduction_example}"""
+    return _std_var(
+        input=input,
+        dim=dim,
+        unbiased=unbiased,
+        correction_opt=correction,
+        keepdim=keepdim,
+        dtype=dtype,
+        mask=mask,
+        take_sqrt=False,
+    )
+
+
+@_apply_docstring_templates
+def std(
+    input: Tensor | MaskedTensor,
+    dim: DimOrDims = None,
+    unbiased: bool | None = None,
+    *,
+    correction: int | None = None,
+    keepdim: bool | None = False,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    """\
+{reduction_signature}
+{reduction_descr}
+The identity value of sample standard deviation operation is undefined. The
+elements of output tensor with strided layout, that correspond to
+fully masked-out elements, have ``nan`` values.
+{reduction_args}
+{reduction_example}"""
+    return _std_var(
+        input=input,
+        dim=dim,
+        unbiased=unbiased,
+        correction_opt=correction,
+        keepdim=keepdim,
+        dtype=dtype,
+        mask=mask,
+        take_sqrt=True,
+    )
+
+
+@_apply_docstring_templates
+def softmax(
+    input: Tensor | MaskedTensor,
+    dim: int,
+    *,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    if dtype is None:
+        dtype = input.dtype
+    dim_ = _canonical_dim(dim, input.ndim)[0]
+    mask_input = _combine_input_and_mask(amax, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.nn.functional.softmax(mask_input, dim_, dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked softmax expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def log_softmax(
+    input: Tensor | MaskedTensor,
+    dim: int,
+    *,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    if dtype is None:
+        dtype = input.dtype
+    dim_ = _canonical_dim(dim, input.ndim)[0]
+    mask_input = _combine_input_and_mask(amax, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.nn.functional.log_softmax(mask_input, dim_, dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked log_softmax expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def softmin(
+    input: Tensor | MaskedTensor,
+    dim: int,
+    *,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    if dtype is None:
+        dtype = input.dtype
+    dim_ = _canonical_dim(dim, input.ndim)[0]
+    mask_input = _combine_input_and_mask(amin, input, mask)
+    if mask_input.layout == torch.strided:
+        return torch.nn.functional.softmin(mask_input, dim_, dtype=dtype)
+    else:
+        raise ValueError(
+            f"masked softmin expects strided tensor (got {mask_input.layout} tensor)"
+        )
+
+
+@_apply_docstring_templates
+def normalize(
+    input: Tensor | MaskedTensor,
+    ord: float,
+    dim: int,
+    *,
+    eps: float = 1e-12,
+    dtype: DType | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    if dtype is None:
+        dtype = input.dtype
+    # TODO: eliminate mask_input as unnecessary when using masked divide.
+    mask_input = _combine_input_and_mask(sum, input, mask)
+    if mask_input.layout == torch.strided:
+        nrm_ = norm(input, ord, dim, keepdim=True, dtype=dtype, mask=mask)
+        # TODO: replace torch.maximum with masked maximum when available.
+        denom = torch.maximum(nrm_, nrm_.new_full([], eps))
+        # TODO: replace torch.divide with masked divide when available.
+        return torch.divide(mask_input, denom)
+    else:
+        raise ValueError(
+            f"masked normalize expects strided tensor (got {mask_input.layout} tensor)"
+        )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9ef878d3c4b20ef38c7dfd6e14631e99b2fddcc1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+from .binary import _apply_native_binary, _is_native_binary
+from .core import is_masked_tensor, MaskedTensor
+from .passthrough import _apply_pass_through_fn, _is_pass_through_fn
+from .reductions import _apply_reduction, _is_reduction
+from .unary import _apply_native_unary, _is_native_unary
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/_ops_refs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/_ops_refs.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc6ee102a66b5640932c5d10d6f1f0e54768a3cd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/_ops_refs.py
@@ -0,0 +1,547 @@
+# mypy: allow-untyped-defs
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+from collections.abc import Callable
+from functools import partial
+from typing import Any, TYPE_CHECKING
+
+import torch
+
+from .binary import _apply_native_binary, NATIVE_BINARY_FNS, NATIVE_INPLACE_BINARY_FNS
+from .core import (
+    _get_data,
+    _masks_match,
+    _maybe_get_mask,
+    is_masked_tensor,
+    MaskedTensor,
+)
+from .passthrough import _apply_pass_through_fn, PASSTHROUGH_FNS
+from .reductions import (
+    _apply_reduction,
+    NATIVE_REDUCE_FNS,
+    TENSOR_REDUCE_FNS,
+    TORCH_REDUCE_FNS,
+)
+from .unary import _apply_native_unary, NATIVE_INPLACE_UNARY_FNS, NATIVE_UNARY_FNS
+
+
+if TYPE_CHECKING:
+    from torch._ops import OpOverload
+
+
+__all__ = []  # type: ignore[var-annotated]
+
+
+def _check_args_kwargs_length(
+    args, kwargs, error_prefix, len_args=None, len_kwargs=None
+):
+    if len_args is not None and len_args != len(args):
+        raise ValueError(
+            f"{error_prefix}: len(args) must be {len_args} but got {len(args)}"
+        )
+    if len_kwargs is not None and len_kwargs != len(kwargs):
+        raise ValueError(
+            f"{error_prefix}: len(kwargs) must be {len_kwargs} but got {len(kwargs)}"
+        )
+
+
+class _MaskedContiguous(torch.autograd.Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, input):
+        if not is_masked_tensor(input):
+            raise ValueError("MaskedContiguous forward: input must be a MaskedTensor.")
+
+        if input.is_contiguous():
+            return input
+
+        data = input.get_data()
+        mask = input.get_mask()
+
+        return MaskedTensor(data.contiguous(), mask.contiguous())
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class _MaskedToDense(torch.autograd.Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, input):
+        if not is_masked_tensor(input):
+            raise ValueError("MaskedToDense forward: input must be a MaskedTensor.")
+
+        if input.layout == torch.strided:
+            return input
+
+        ctx.layout = input.layout
+        data = input.get_data()
+        mask = input.get_mask()
+
+        return MaskedTensor(data.to_dense(), mask.to_dense())
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(ctx, grad_output):
+        layout = ctx.layout
+
+        if layout == torch.sparse_coo:
+            return grad_output.to_sparse_coo()
+        elif layout == torch.sparse_csr:
+            return grad_output.to_sparse_csr()
+        elif layout == torch.strided:
+            return grad_output.to_dense()
+        raise ValueError(f"to_dense: Unsupported input layout: {layout}")
+
+
+class _MaskedToSparse(torch.autograd.Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, input):
+        if not is_masked_tensor(input):
+            raise ValueError("MaskedToSparse forward: input must be a MaskedTensor.")
+
+        # Following the convention from sparse tensors that to_sparse always means that we convert to sparse_coo
+        if input.layout == torch.sparse_coo:
+            return input
+
+        data = input.get_data()
+        mask = input.get_mask()
+        sparse_mask = mask.to_sparse_coo().coalesce()
+        sparse_data = data.sparse_mask(sparse_mask)
+
+        return MaskedTensor(sparse_data, sparse_mask)
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(ctx, grad_output):
+        return grad_output.to_dense()
+
+
+class _MaskedToSparseCsr(torch.autograd.Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, input):
+        if not is_masked_tensor(input):
+            raise ValueError("MaskedToSparseCsr forward: input must be a MaskedTensor.")
+
+        if input._masked_data.ndim != 2:
+            raise ValueError(
+                f"Only 2D tensors can be converted to the SparseCsr layout but got shape: {input._masked_data.size()}"
+            )
+
+        if input.layout == torch.sparse_csr:
+            return input
+
+        data = input.get_data()
+        mask = input.get_mask()
+        sparse_mask = mask.to_sparse_csr()
+        sparse_data = data.sparse_mask(sparse_mask)
+
+        return MaskedTensor(sparse_data, sparse_mask)
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(ctx, grad_output):
+        return grad_output.to_dense()
+
+
+class _MaskedWhere(torch.autograd.Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, cond, self, other):
+        ctx.mark_non_differentiable(cond)
+        ctx.save_for_backward(cond)
+        return torch.ops.aten.where(cond, self, other)
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(ctx, grad_output):
+        (cond,) = ctx.saved_tensors
+
+        def masked_out_like(mt):
+            return MaskedTensor(mt.get_data(), torch.zeros_like(mt.get_mask()).bool())
+
+        return (
+            None,
+            torch.ops.aten.where(cond, grad_output, masked_out_like(grad_output)),
+            torch.ops.aten.where(cond, masked_out_like(grad_output), grad_output),
+        )
+
+
+_MASKEDTENSOR_FUNCTION_TABLE = {}
+
+_function_fn_apply_map = {
+    (
+        tuple(NATIVE_REDUCE_FNS),
+        tuple(TORCH_REDUCE_FNS),
+        tuple(TENSOR_REDUCE_FNS),
+    ): _apply_reduction,
+}
+
+for fn_map_list, apply_fn in _function_fn_apply_map.items():
+    for fn_map in fn_map_list:
+        for fn in fn_map:
+            _MASKEDTENSOR_FUNCTION_TABLE[fn] = partial(apply_fn, fn)
+
+
+def register_function_func(ops):
+    """
+    Used for registering a new __torch_function__ function to MaskedTensor
+    Called via _MASKEDTENSOR_FUNCTION_TABLE[func](*args, **kwargs)
+
+    The code to register a new function looks like:
+
+    @register_function_func(list_of_ops)
+    def foo(func, *args, **kwargs):
+        <implementation>
+    """
+
+    def wrapper(func):
+        for op in ops:
+            _MASKEDTENSOR_FUNCTION_TABLE[op] = partial(func, op)
+
+    return wrapper
+
+
+@register_function_func(NATIVE_REDUCE_FNS + TORCH_REDUCE_FNS + TENSOR_REDUCE_FNS)
+def _general_function_reductions(func, *args, **kwargs):
+    return _apply_reduction(func, *args, **kwargs)
+
+
+@register_function_func([torch.Tensor.where, torch.where])
+def _function_where(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, "__torch_function__, torch.where", len_args=3, len_kwargs=0
+    )
+    return _MaskedWhere.apply(*args)
+
+
+@register_function_func([torch.Tensor.contiguous])
+def _function_contiguous(func, *args, **kwargs):
+    return _MaskedContiguous.apply(args[0])
+
+
+@register_function_func([torch.Tensor.to_dense])
+def _function_to_dense(func, *args, **kwargs):
+    return _MaskedToDense.apply(args[0])
+
+
+@register_function_func([torch.Tensor.to_sparse])
+def _function_to_sparse(func, *args, **kwargs):
+    return _MaskedToSparse.apply(args[0])
+
+
+@register_function_func([torch.Tensor.to_sparse_csr])
+def _function_to_sparse_csr(func, *args, **kwargs):
+    return _MaskedToSparseCsr.apply(args[0])
+
+
+_MASKEDTENSOR_DISPATCH_TABLE: dict["OpOverload", Callable[..., Any]] = {}
+
+
+def register_dispatch_func(aten_ops):
+    """
+    Used for registering a new __torch_dispatch__ function to MaskedTensor
+    Called via _MASKEDTENSOR_DISPATCH_TABLE[func](*args, **kwargs)
+
+    The code to register a new function looks like:
+
+    @register_dispatch_func(list_of_ops)
+    def foo(func, *args, **kwargs):
+        <implementation>
+    """
+
+    def wrapper(func):
+        for aten_op in aten_ops:
+            _MASKEDTENSOR_DISPATCH_TABLE[aten_op] = partial(func, aten_op)
+
+    return wrapper
+
+
+@register_dispatch_func(NATIVE_REDUCE_FNS + TORCH_REDUCE_FNS + TENSOR_REDUCE_FNS)
+def _general_reduction(func, *args, **kwargs):
+    return _apply_reduction(func, *args, **kwargs)
+
+
+@register_dispatch_func(PASSTHROUGH_FNS)
+def _general_passthrough(func, *args, **kwargs):
+    return _apply_pass_through_fn(func, *args, **kwargs)
+
+
+@register_dispatch_func(NATIVE_UNARY_FNS + NATIVE_INPLACE_UNARY_FNS)
+def _general_unary(func, *args, **kwargs):
+    return _apply_native_unary(func, *args, **kwargs)
+
+
+@register_dispatch_func(NATIVE_BINARY_FNS + NATIVE_INPLACE_BINARY_FNS)
+def _general_binary(func, *args, **kwargs):
+    return _apply_native_binary(func, *args, **kwargs)
+
+
+@register_dispatch_func([torch.ops.aten.stride])
+def stride(func, *args, **kwargs):
+    return None
+
+
+@register_dispatch_func([torch.ops.aten.sym_stride])
+def sym_stride(func, *args, **kwargs):
+    return None
+
+
+@register_dispatch_func([torch.ops.prim.layout])
+def layout(func, *args, **kwargs):
+    return _get_data(args[0]).layout
+
+
+@register_dispatch_func(
+    [torch.ops.aten.is_contiguous, torch.ops.aten.sym_is_contiguous]
+)
+def is_contiguous(func, *args, **kwargs):
+    data = _get_data(args[0])
+    if data.is_sparse:
+        raise ValueError("MaskedTensors with sparse data do not have is_contiguous")
+    return func(data, *args[1:], **kwargs)
+
+
+@register_dispatch_func([torch.ops.aten.is_strides_like_format])
+def is_strides_like_format(func, *args, **kwargs):
+    data = _get_data(args[0])
+    if data.is_sparse:
+        raise ValueError(
+            "MaskedTensors with sparse data do not have is_strides_like_format"
+        )
+    return func(data, *args[1:], **kwargs)
+
+
+@register_dispatch_func([torch.ops.aten.is_non_overlapping_and_dense])
+def is_non_overlapping_and_dense(func, *args, **kwargs):
+    data = _get_data(args[0])
+    if data.is_sparse:
+        raise ValueError(
+            "MaskedTensors with sparse data do not have is_non_overlapping_and_dense"
+        )
+    return func(data, *args[1:], **kwargs)
+
+
+@register_dispatch_func([torch.ops.aten.contiguous])
+def contiguous(func, *args, **kwargs):
+    if _get_data(args[0]).is_sparse:
+        raise ValueError("MaskedTensors with sparse data do not have contiguous")
+    return _MaskedContiguous.apply(args[0])
+
+
+@register_dispatch_func([torch.ops.aten.new_empty_strided])
+def new_empty_strided(func, *args, **kwargs):
+    _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=3)
+    data = _get_data(args[0])
+    mask = _maybe_get_mask(args[0])
+    if tuple(args[1]) != tuple(data.size()):
+        raise ValueError(
+            f"__torch_dispatch__, {func}: args[1] expected to be the same as data.size()"
+        )
+    if tuple(args[2]) != tuple(data.stride()):
+        raise ValueError(
+            f"__torch_dispatch__, {func}: args[2] expected to be the same as data.stride()"
+        )
+    return MaskedTensor(func(data, args[1], args[2], **kwargs), mask)
+
+
+@register_dispatch_func([torch.ops.aten._local_scalar_dense])
+def _local_scalar_dense(func, *args, **kwargs):
+    if not _maybe_get_mask(args[0]):
+        raise ValueError(f"__torch_dispatch__, {func}: expected a mask tensor")
+    return torch.ops.aten._local_scalar_dense(_get_data(args[0]))
+
+
+@register_dispatch_func([torch.ops.aten.detach, torch.ops.aten.clone])
+def _apply_fn_on_data(func, *args, **kwargs):
+    return MaskedTensor(func(_get_data(args[0])), _maybe_get_mask(args[0]))
+
+
+@register_dispatch_func([torch.ops.aten._to_copy])
+def _to_copy(func, *args, **kwargs):
+    new_data = func(_get_data(args[0]), *args[1:], **kwargs)
+    cloned_kwargs = kwargs.copy()
+    cloned_kwargs["dtype"] = torch.bool
+    new_mask = func(_maybe_get_mask(args[0]), *args[1:], **cloned_kwargs)
+    return MaskedTensor(new_data, new_mask)
+
+
+@register_dispatch_func([torch.ops.aten._softmax])
+def _softmax(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=3, len_kwargs=0
+    )
+    data = _get_data(args[0])
+    mask = _maybe_get_mask(args[0])
+    result_data = torch.ops.aten._masked_softmax(data, ~mask, args[1], 2)
+    return MaskedTensor(result_data, mask)
+
+
+@register_dispatch_func([torch.ops.aten.ones_like])
+def ones_like(func, *args, **kwargs):
+    _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=1)
+    result_data = func(_get_data(args[0]), **kwargs)
+    return MaskedTensor(result_data, _maybe_get_mask(args[0]))
+
+
+@register_dispatch_func([torch.ops.aten._softmax_backward_data])
+def _softmax_backward_data(func, *args, **kwargs):
+    _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=4)
+    grad, output, dim, _input_dtype = args
+    if is_masked_tensor(grad) and is_masked_tensor(output):
+        if not _masks_match(grad, output):
+            raise ValueError(
+                f"__torch_dispatch__, {func}: expected the masks of grad and output to match"
+            )
+        grad_data = _get_data(grad)
+        new_grad_data = torch.ops.aten._masked_softmax_backward(
+            grad_data,
+            _get_data(output),
+            ~_maybe_get_mask(grad),
+            dim % grad_data.ndim,
+        )
+        res = MaskedTensor(new_grad_data, _maybe_get_mask(grad))
+        return res
+    else:
+        raise ValueError(
+            f"__torch_dispatch__, {func}: grad and output must both be MaskedTensors"
+        )
+
+
+@register_dispatch_func([torch.ops.aten.copy_])
+def copy_(func, *args, **kwargs):
+    _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=2)
+    if not _masks_match(_maybe_get_mask(args[0]), _maybe_get_mask(args[1])):
+        raise ValueError("args[0] mask and args[1] mask must match but do not")
+    func(_get_data(args[0]), _get_data(args[1]))
+    return args[0]
+
+
+@register_dispatch_func([torch.ops.aten.where])
+def where(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=3, len_kwargs=0
+    )
+    if not torch.is_tensor(args[0]):
+        raise ValueError(f"__torch_dispatch__, {func}: expected args[0] to be a tensor")
+    mx = args[1]
+    my = args[2]
+    if not is_masked_tensor(mx):
+        mx = MaskedTensor(mx, torch.ones_like(mx, dtype=torch.bool))
+    if not is_masked_tensor(my):
+        my = MaskedTensor(my, torch.ones_like(my, dtype=torch.bool))
+    new_data = func(args[0], mx.get_data(), my.get_data())
+    new_mask = func(args[0], mx.get_mask(), my.get_mask())
+    return MaskedTensor(new_data, new_mask)
+
+
+@register_dispatch_func([torch.ops.aten._to_sparse])
+def _to_sparse(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0
+    )
+    if not torch.is_tensor(args[0]):
+        raise TypeError(f"__torch_dispatch__, {func}: expected args[0] to be a tensor")
+    mt = args[0]
+    if not is_masked_tensor(mt):
+        mt = MaskedTensor(mt, torch.ones_like(mt, dtype=torch.bool))
+    if mt.is_sparse_coo():
+        return mt
+    new_mask = func(_maybe_get_mask(args[0])).coalesce()
+    new_data = _get_data(args[0]).sparse_mask(new_mask)
+    return MaskedTensor(new_data, new_mask)
+
+
+@register_dispatch_func([torch.ops.aten._to_sparse_csr])
+def _to_sparse_csr(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0
+    )
+    if not torch.is_tensor(args[0]):
+        raise ValueError(f"__torch_dispatch__, {func}: expected args[0] to be a tensor")
+    mt = args[0]
+    if not is_masked_tensor(mt):
+        mt = MaskedTensor(mt, torch.ones_like(mt).bool())
+    if mt.is_sparse_csr():
+        return mt
+    new_mask = func(_maybe_get_mask(args[0]))
+    new_data = _get_data(args[0]).sparse_mask(new_mask)
+    return MaskedTensor(new_data, new_mask)
+
+
+@register_dispatch_func([torch.ops.aten._to_dense])
+def _to_dense(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0
+    )
+    if not torch.is_tensor(args[0]):
+        raise ValueError(f"__torch_dispatch__, {func}: expected args[0] to be a tensor")
+    mt = args[0]
+    if not is_masked_tensor(mt):
+        mt = MaskedTensor(mt, torch.ones_like(mt).bool())
+    new_data = func(_get_data(args[0]))
+    new_mask = func(_maybe_get_mask(args[0]))
+    return MaskedTensor(new_data, new_mask)
+
+
+@register_dispatch_func([torch.ops.aten._indices])
+def _indices(func, *args, **kwargs):
+    # Assumes data is sparse
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0
+    )
+    data = _get_data(args[0]).indices()
+    return MaskedTensor(data, torch.ones_like(data).bool())
+
+
+@register_dispatch_func([torch.ops.aten._values])
+def _values(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0
+    )
+    data = _get_data(args[0]).values()
+    return MaskedTensor(data, torch.ones_like(data).bool())
+
+
+@register_dispatch_func([torch.ops.aten._sparse_coo_tensor_with_dims_and_tensors])
+def _sparse_coo_tensor_with_dims_and_tensors(func, *args, **kwargs):
+    new_args = list(args)
+    if is_masked_tensor(args[-1]):
+        new_args[-1] = args[-1].get_data()
+    if is_masked_tensor(args[-2]):
+        new_args[-2] = args[-2].get_data()
+
+    new_data = func(*new_args, **kwargs)
+    new_args[-1] = torch.ones_like(new_args[-1])
+    new_mask = func(*new_args, **kwargs).bool()
+
+    return MaskedTensor(new_data, new_mask)
+
+
+@register_dispatch_func([torch.ops.aten.is_same_size])
+def is_same_size(func, *args, **kwargs):
+    _check_args_kwargs_length(args, kwargs, f"__torch_dispatch__, {func}", len_args=2)
+    return _get_data(args[0]).is_same_size(_get_data(args[1]))
+
+
+@register_dispatch_func([torch.ops.aten._is_any_true])
+def _is_any_true(func, *args, **kwargs):
+    _check_args_kwargs_length(
+        args, kwargs, f"__torch_dispatch__, {func}", len_args=1, len_kwargs=0
+    )
+    data = _get_data(args[0])
+    mask = _maybe_get_mask(args[0])
+    if mask is None:
+        raise ValueError(
+            f"__torch_dispatch__, {func}: expected args[0] to be a MaskedTensor"
+        )
+    if data.dtype != torch.bool:
+        raise ValueError(f"__torch_dispatch__, {func}: expected a boolean tensor")
+    if data.is_sparse:
+        raise ValueError(f"MaskedTensors with sparse data do not have {func}")
+
+    return MaskedTensor(func(data & mask), torch.tensor(True))
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/binary.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/binary.py
new file mode 100644
index 0000000000000000000000000000000000000000..8315ae11be7175c2b5aaef178a4bc4785dcbcb29
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/binary.py
@@ -0,0 +1,200 @@
+# mypy: allow-untyped-defs
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import torch
+
+from .core import (
+    _map_mt_args_kwargs,
+    _masks_match,
+    _tensors_match,
+    _wrap_result,
+    is_masked_tensor,
+)
+
+
+__all__ = []  # type: ignore[var-annotated]
+
+BINARY_NAMES = [
+    "add",
+    "atan2",
+    "arctan2",
+    "bitwise_and",
+    "bitwise_or",
+    "bitwise_xor",
+    "bitwise_left_shift",
+    "bitwise_right_shift",
+    "div",
+    "divide",
+    "floor_divide",
+    "fmod",
+    "logaddexp",
+    "logaddexp2",
+    "mul",
+    "multiply",
+    "nextafter",
+    "remainder",
+    "sub",
+    "subtract",
+    "true_divide",
+    "eq",
+    "ne",
+    "le",
+    "ge",
+    "greater",
+    "greater_equal",
+    "gt",
+    "less_equal",
+    "lt",
+    "less",
+    "maximum",
+    "minimum",
+    "fmax",
+    "fmin",
+    "not_equal",
+]
+
+INPLACE_BINARY_NAMES = [
+    n + "_"
+    for n in (
+        list(
+            set(BINARY_NAMES)
+            - {
+                "logaddexp",
+                "logaddexp2",
+                "equal",
+                "fmin",
+                "minimum",
+                "maximum",
+                "fmax",
+            }
+        )
+    )
+]
+
+
+def _get_at_least_one_mask(a, b):
+    if not is_masked_tensor(a) and not is_masked_tensor(b):
+        raise TypeError("At least one of `a` and `b` must be a MaskedTensor")
+    if not _masks_match(a, b):
+        raise ValueError("a and b must have matching masks")
+    if is_masked_tensor(a):
+        return a.get_mask()
+    return b.get_mask()
+
+
+def _binary_helper(fn, args, kwargs, inplace):
+    if len(kwargs) != 0:
+        raise ValueError("len(kwargs) must equal 0")
+    for a in args[2:]:
+        if torch.is_tensor(a):
+            raise TypeError(
+                "MaskedTensor binary ops do not support Tensor arguments aside from the lhs and rhs"
+            )
+
+    if not _masks_match(*args[:2]):
+        raise ValueError(
+            "Input masks must match. If you need support for this, please open an issue on Github."
+        )
+
+    data_args, _data_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_data())
+    mask_args, _mask_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_mask())
+
+    args0_layout = data_args[0].layout
+    same_layout = (
+        torch.is_tensor(data_args[1]) or is_masked_tensor(data_args[1])
+    ) and (args0_layout == data_args[1].layout)
+
+    if args0_layout == torch.sparse_coo:
+        if same_layout:
+            if not _tensors_match(data_args[0].indices(), data_args[1].indices()):
+                raise ValueError(
+                    "sparse_coo indices must match. If you need support for this, please open an issue on Github."
+                )
+            if data_args[0].size() != data_args[1].size():
+                raise ValueError(
+                    "input1 and input2 must have the same size for binary functions."
+                )
+
+            data_args[1] = data_args[1].values()
+
+        i = data_args[0].indices()
+        size = data_args[0].size()
+        data_args[0] = data_args[0].values()
+        v = fn(*data_args)
+        result_data = torch.sparse_coo_tensor(i, v, size)
+
+    elif args0_layout == torch.sparse_csr:
+        if same_layout:
+            if not (
+                _tensors_match(data_args[0].crow_indices(), data_args[1].crow_indices())
+                and _tensors_match(
+                    data_args[0].col_indices(), data_args[1].col_indices()
+                )
+            ):
+                raise ValueError(
+                    "sparse_csr indices must match. If you need support for this, please open an issue on Github."
+                )
+
+            data_args[1] = data_args[1].values()
+
+        crow = data_args[0].crow_indices()
+        col = data_args[0].col_indices()
+        size = data_args[0].size()
+        data_args[0] = data_args[0].values()
+        v = fn(*data_args)
+        result_data = torch.sparse_csr_tensor(crow, col, v, size)
+
+    else:
+        result_data = fn(*data_args)
+
+    if inplace:
+        args[0]._set_data_mask(result_data, mask_args[0])
+        return args[0]
+    else:
+        result_mask = _get_at_least_one_mask(*args[:2])
+        # sparse tensors don't have strides so we can only expand if the layout is strided
+        if args0_layout == torch.strided:
+            result_mask = result_mask.expand_as(result_data)
+        return _wrap_result(result_data, result_mask)
+
+
+def _torch_binary(fn_name):
+    fn = getattr(torch.ops.aten, fn_name)
+
+    def binary_fn(*args, **kwargs):
+        return _binary_helper(fn, args, kwargs, inplace=False)
+
+    return binary_fn
+
+
+def _torch_inplace_binary(fn_name):
+    fn = getattr(torch.ops.aten, fn_name)
+
+    def binary_fn(*args, **kwargs):
+        return _binary_helper(fn, args, kwargs, inplace=True)
+
+    return binary_fn
+
+
+NATIVE_BINARY_MAP = {
+    getattr(torch.ops.aten, name): _torch_binary(name) for name in BINARY_NAMES
+}
+NATIVE_INPLACE_BINARY_MAP = {
+    getattr(torch.ops.aten, name): _torch_inplace_binary(name)
+    for name in INPLACE_BINARY_NAMES
+}
+
+NATIVE_BINARY_FNS = list(NATIVE_BINARY_MAP.keys())
+NATIVE_INPLACE_BINARY_FNS = list(NATIVE_INPLACE_BINARY_MAP.keys())
+
+
+def _is_native_binary(fn):
+    return fn in NATIVE_BINARY_FNS or fn in NATIVE_INPLACE_BINARY_FNS
+
+
+def _apply_native_binary(fn, *args, **kwargs):
+    if fn in NATIVE_BINARY_FNS:
+        return NATIVE_BINARY_MAP[fn](*args, **kwargs)
+    if fn in NATIVE_INPLACE_BINARY_FNS:
+        return NATIVE_INPLACE_BINARY_MAP[fn](*args, **kwargs)
+    return NotImplemented
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/core.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/core.py
new file mode 100644
index 0000000000000000000000000000000000000000..cad5621b29bd663ef4462f1be6c8f8f2c4762c2d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/core.py
@@ -0,0 +1,364 @@
+# mypy: allow-untyped-defs
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import warnings
+from typing import Any
+from typing_extensions import TypeIs
+
+import torch
+from torch.overrides import get_default_nowrap_functions
+
+
+__all__ = [
+    "MaskedTensor",
+    "is_masked_tensor",
+]
+
+
+def is_masked_tensor(obj: Any, /) -> TypeIs["MaskedTensor"]:
+    r"""Returns True if the input is a MaskedTensor, else False
+
+    Args:
+        a: any input
+
+    Examples:
+
+        >>> # xdoctest: +SKIP
+        >>> from torch.masked import MaskedTensor
+        >>> data = torch.arange(6).reshape(2, 3)
+        >>> mask = torch.tensor([[True, False, False], [True, True, False]])
+        >>> mt = MaskedTensor(data, mask)
+        >>> is_masked_tensor(mt)
+        True
+    """
+    return isinstance(obj, MaskedTensor)
+
+
+def _tensors_match(a, b, exact=True, rtol=1e-05, atol=1e-08):
+    if is_masked_tensor(a) or is_masked_tensor(b):
+        raise ValueError("Neither `a` nor `b` can be a MaskedTensor.")
+    if a.layout != b.layout:
+        raise ValueError(
+            f"`a` and `b` must have the same layout. Got {a.layout} and {b.layout}"
+        )
+
+    if a.dtype != b.dtype:
+        b = b.type(a.dtype)
+    if a.layout == b.layout == torch.sparse_coo:
+        return _tensors_match(a.values(), b.values(), exact) and _tensors_match(
+            a.indices(), b.indices(), exact
+        )
+    elif a.layout == b.layout == torch.sparse_csr:
+        return (
+            _tensors_match(a.crow_indices(), b.crow_indices(), exact)
+            and _tensors_match(a.col_indices(), b.col_indices(), exact)
+            and _tensors_match(a.values(), b.values(), exact)
+        )
+    if exact:
+        return (a.dim() == b.dim()) and torch.eq(a, b).all().item()
+    return (a.dim() == b.dim()) and torch.allclose(a, b, rtol=rtol, atol=atol)
+
+
+def _masks_match(a, b):
+    if is_masked_tensor(a) and is_masked_tensor(b):
+        mask_a = a.get_mask()
+        mask_b = b.get_mask()
+        return _tensors_match(mask_a, mask_b, exact=True)
+    return True
+
+
+def _map_mt_args_kwargs(args, kwargs, map_fn):
+    def _helper(a, map_fn):
+        if is_masked_tensor(a):
+            return map_fn(a)
+        elif torch.is_tensor(a):
+            return a
+        elif isinstance(a, list):
+            a_impl, _ = _map_mt_args_kwargs(a, {}, map_fn)
+            return a_impl
+        elif isinstance(a, tuple):
+            a_impl, _ = _map_mt_args_kwargs(a, {}, map_fn)
+            return tuple(a_impl)
+        else:
+            return a
+
+    if kwargs is None:
+        kwargs = {}
+    impl_args = []
+    for a in args:
+        impl_args.append(_helper(a, map_fn))
+    impl_kwargs = {}
+    for k in kwargs:
+        impl_kwargs[k] = _helper(a, map_fn)
+    return impl_args, impl_kwargs
+
+
+def _wrap_result(result_data, result_mask):
+    if isinstance(result_data, list):
+        return [_wrap_result(r, m) for (r, m) in zip(result_data, result_mask)]
+    if isinstance(result_data, tuple):
+        return tuple(_wrap_result(r, m) for (r, m) in zip(result_data, result_mask))
+    if torch.is_tensor(result_data):
+        return MaskedTensor(result_data, result_mask)
+    # Expect result_data and result_mask to be Tensors only
+    return NotImplemented
+
+
+def _masked_tensor_str(data, mask, formatter):
+    if data.layout in {torch.sparse_coo, torch.sparse_csr}:
+        data = data.to_dense()
+        mask = mask.to_dense()
+    if data.dim() == 1:
+        formatted_elements = [
+            formatter.format(d.item()) if isinstance(d.item(), float) else str(d.item())
+            for d in data
+        ]
+        max_len = max(8 if x[1] else len(x[0]) for x in zip(formatted_elements, ~mask))
+        return (
+            "["
+            + ", ".join(
+                [
+                    "--".rjust(max_len) if m else e
+                    for (e, m) in zip(formatted_elements, ~mask)
+                ]
+            )
+            + "]"
+        )
+    sub_strings = [_masked_tensor_str(d, m, formatter) for (d, m) in zip(data, mask)]
+    sub_strings = ["\n".join(["  " + si for si in s.split("\n")]) for s in sub_strings]
+    return "[\n" + ",\n".join(sub_strings) + "\n]"
+
+
+def _get_data(a):
+    if is_masked_tensor(a):
+        return a._masked_data
+    return a
+
+
+def _maybe_get_mask(a):
+    if is_masked_tensor(a):
+        return a.get_mask()
+    return None
+
+
+class MaskedTensor(torch.Tensor):
+    @staticmethod
+    def __new__(cls, data, mask, requires_grad=False):
+        if is_masked_tensor(data) or not torch.is_tensor(data):
+            raise TypeError("data must be a Tensor")
+        if is_masked_tensor(mask) or not torch.is_tensor(mask):
+            raise TypeError("mask must be a Tensor")
+        # Use a Tensor that of the give size for the wrapper.
+        kwargs = {
+            "device": data.device,
+            "dtype": data.dtype,
+            "layout": data.layout,
+            "requires_grad": requires_grad,
+            "dispatch_sizes_strides_policy": "strides",
+            "dispatch_layout": True,
+        }
+        warnings.warn(
+            (
+                "The PyTorch API of MaskedTensors is in prototype stage "
+                "and will change in the near future. Please open a Github issue "
+                "for features requests and see our documentation on the torch.masked "
+                "module for further information about the project."
+            ),
+            UserWarning,
+            stacklevel=2,
+        )
+        if data.requires_grad:
+            warnings.warn(
+                "It is not recommended to create a MaskedTensor with a tensor that requires_grad. "
+                "To avoid this, you can use data.detach().clone()",
+                UserWarning,
+                stacklevel=2,
+            )
+        # pyrefly: ignore [bad-argument-type]
+        return torch.Tensor._make_wrapper_subclass(cls, data.size(), **kwargs)
+
+    def _preprocess_data(self, data, mask):
+        from .._ops import _sparse_coo_where, _sparse_csr_where
+
+        if data.layout != mask.layout:
+            raise TypeError("data and mask must have the same layout.")
+        if data.layout == torch.sparse_coo:
+            data = data.coalesce()
+            mask = mask.coalesce()
+            if data._nnz() != mask._nnz():
+                data = _sparse_coo_where(mask, data, torch.tensor(0))
+        elif data.layout == torch.sparse_csr:
+            if data._nnz() != mask._nnz():
+                data = _sparse_csr_where(mask, data, torch.tensor(0))
+
+        # Have to pick awkward names to not conflict with existing fields such as data
+        self._masked_data = data.clone()
+        self._masked_mask = mask.clone()
+
+    def _validate_members(self):
+        data = self._masked_data
+        mask = self.get_mask()
+        if type(data) is not type(mask):
+            raise TypeError(
+                f"data and mask must have the same type. Got {type(data)} and {type(mask)}"
+            )
+        if data.layout not in {torch.strided, torch.sparse_coo, torch.sparse_csr}:
+            raise TypeError(f"data layout of {data.layout} is not supported.")
+        if data.layout == torch.sparse_coo:
+            if not _tensors_match(data.indices(), mask.indices(), exact=True):
+                raise ValueError(
+                    "data and mask are both sparse COO tensors but do not have the same indices."
+                )
+        elif data.layout == torch.sparse_csr:
+            if not _tensors_match(
+                data.crow_indices(), mask.crow_indices(), exact=True
+            ) or not _tensors_match(data.col_indices(), mask.col_indices(), exact=True):
+                raise ValueError(
+                    "data and mask are both sparse CSR tensors but do not share either crow or col indices."
+                )
+        if mask.dtype != torch.bool:
+            raise TypeError("mask must have dtype bool.")
+        if not (
+            data.dtype == torch.float16
+            or data.dtype == torch.float32
+            or data.dtype == torch.float64
+            or data.dtype == torch.bool
+            or data.dtype == torch.int8
+            or data.dtype == torch.int16
+            or data.dtype == torch.int32
+            or data.dtype == torch.int64
+        ):
+            raise TypeError(f"{data.dtype} is not supported in MaskedTensor.")
+        if data.dim() != mask.dim():
+            raise ValueError("data.dim() must equal mask.dim()")
+        if data.size() != mask.size():
+            raise ValueError("data.size() must equal mask.size()")
+
+    def __init__(self, data, mask, requires_grad=False):
+        self._preprocess_data(data, mask)
+        self._validate_members()
+
+    @staticmethod
+    def _from_values(data, mask):
+        """Differentiable constructor for MaskedTensor"""
+
+        class Constructor(torch.autograd.Function):
+            @staticmethod
+            # pyrefly: ignore [bad-override]
+            def forward(ctx, data, mask):
+                return MaskedTensor(data, mask)
+
+            @staticmethod
+            # pyrefly: ignore [bad-override]
+            def backward(ctx, grad_output):
+                return grad_output, None
+
+        result = Constructor.apply(data, mask)
+        return result
+
+    def _set_data_mask(self, data, mask):
+        self._masked_data = data
+        self._masked_mask = mask
+        self._validate_members()
+
+    def __repr__(self):  # type: ignore[override]
+        formatter = "{0:8.4f}"
+        if self.dim() == 0:
+            scalar_data = self.get_data().item()
+            data_formatted = (
+                formatter.format(scalar_data)
+                if isinstance(scalar_data, float)
+                else str(scalar_data)
+            )
+            if not self.get_mask().item():
+                data_formatted = "--"
+            return (
+                "MaskedTensor("
+                + data_formatted
+                + ", "
+                + str(self.get_mask().item())
+                + ")"
+            )
+        s = _masked_tensor_str(self.get_data(), self.get_mask(), formatter)
+        s = "\n".join("  " + si for si in s.split("\n"))
+        return "MaskedTensor(\n" + s + "\n)"
+
+    # Seems like this needs to be defined before torch_dispatch to work
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+
+        from ._ops_refs import _MASKEDTENSOR_FUNCTION_TABLE
+
+        if func in _MASKEDTENSOR_FUNCTION_TABLE:
+            return _MASKEDTENSOR_FUNCTION_TABLE[func](*args, **kwargs)
+
+        if not all(issubclass(cls, t) for t in types):
+            return NotImplemented
+        with torch._C.DisableTorchFunctionSubclass():
+            ret = func(*args, **kwargs)
+            if func in get_default_nowrap_functions():
+                return ret
+            else:
+                return torch._tensor._convert(ret, cls)
+
+    @classmethod
+    def unary(cls, fn, data, mask):
+        return MaskedTensor(fn(data), mask)
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs):  # type: ignore[override]
+        func = func.overloadpacket
+
+        from ._ops_refs import _MASKEDTENSOR_DISPATCH_TABLE
+
+        if func in _MASKEDTENSOR_DISPATCH_TABLE:
+            return _MASKEDTENSOR_DISPATCH_TABLE[func](*args, **kwargs)
+
+        msg = (
+            f"{func.__name__} is not implemented in __torch_dispatch__ for MaskedTensor.\n"
+            "If you would like this operator to be supported, please file an issue for a feature request at "
+            "https://github.com/pytorch/maskedtensor/issues with a minimal reproducible code snippet.\n"
+            "In the case that the semantics for the operator are not trivial, it would be appreciated "
+            "to also include a proposal for the semantics."
+        )
+        warnings.warn(msg, stacklevel=2)
+        return NotImplemented
+
+    def __lt__(self, other):
+        if is_masked_tensor(other):
+            return MaskedTensor(self.get_data() < _get_data(other), self.get_mask())
+        return MaskedTensor(self.get_data() < other, self.get_mask())
+
+    def to_tensor(self, value):
+        return self.get_data().masked_fill(~self.get_mask(), value)
+
+    def get_data(self):
+        class GetData(torch.autograd.Function):
+            @staticmethod
+            # pyrefly: ignore [bad-override]
+            def forward(ctx, self):
+                return self._masked_data.detach()
+
+            @staticmethod
+            # pyrefly: ignore [bad-override]
+            def backward(ctx, grad_output):
+                if is_masked_tensor(grad_output):
+                    return grad_output
+                return MaskedTensor(grad_output, self.get_mask())
+
+        return GetData.apply(self)
+
+    def get_mask(self):
+        return self._masked_mask
+
+    def is_sparse_coo(self):
+        return self.layout == torch.sparse_coo
+
+    def is_sparse_csr(self):  # type: ignore[override]
+        return self.layout == torch.sparse_csr
+
+    # Update later to support more sparse layouts
+    @property
+    def is_sparse(self):  # type: ignore[override]
+        return self.is_sparse_coo() or self.is_sparse_csr()
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/creation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/creation.py
new file mode 100644
index 0000000000000000000000000000000000000000..35c8e3d2aa9438dbcfc7995a1cdcd3c5cc8dc1fc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/creation.py
@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+from .core import MaskedTensor
+
+
+__all__ = [
+    "as_masked_tensor",
+    "masked_tensor",
+]
+
+
+# These two factory functions are intended to mirror
+#     torch.tensor - guaranteed to be a leaf node
+#     torch.as_tensor - differentiable constructor that preserves the autograd history
+
+
+def masked_tensor(
+    data: object, mask: object, requires_grad: bool = False
+) -> MaskedTensor:
+    return MaskedTensor(data, mask, requires_grad)
+
+
+def as_masked_tensor(data: object, mask: object) -> MaskedTensor:
+    return MaskedTensor._from_values(data, mask)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/passthrough.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/passthrough.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba13f50c1fee9c9fc10563ffc9f4ff3211c0dca6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/passthrough.py
@@ -0,0 +1,50 @@
+# mypy: allow-untyped-defs
+# Copyright (c) Meta Platforms, Inc. and affiliates
+"""
+These are functions that should simply be applied to both mask and data.
+Take select or stack as an example. This operation can be applied to
+both the mask and data of a MaskedTensor and the result wrapped into
+a new MaskedTensor as a result.
+"""
+
+import torch
+
+from .core import _map_mt_args_kwargs, _wrap_result
+
+
+__all__ = []  # type: ignore[var-annotated]
+
+
+PASSTHROUGH_FNS = [
+    torch.ops.aten.select,
+    torch.ops.aten.transpose,
+    torch.ops.aten.split,
+    torch.ops.aten.t,
+    torch.ops.aten.slice,
+    torch.ops.aten.slice_backward,
+    torch.ops.aten.select_backward,
+    torch.ops.aten.index,
+    torch.ops.aten.expand,
+    torch.ops.aten.view,
+    torch.ops.aten._unsafe_view,
+    torch.ops.aten._reshape_alias,
+    torch.ops.aten.cat,
+    torch.ops.aten.unsqueeze,
+    torch.ops.aten.unfold,
+    torch.ops.aten.unfold_backward,
+    torch.ops.aten.im2col,
+    torch.ops.aten.col2im,
+    torch.ops.aten.stack,
+]
+
+
+def _is_pass_through_fn(fn):
+    return fn in PASSTHROUGH_FNS
+
+
+def _apply_pass_through_fn(fn, *args, **kwargs):
+    data_args, data_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_data())
+    result_data = fn(*data_args, **data_kwargs)
+    mask_args, mask_kwargs = _map_mt_args_kwargs(args, kwargs, lambda x: x.get_mask())
+    result_mask = fn(*mask_args, **mask_kwargs)
+    return _wrap_result(result_data, result_mask)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/reductions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/reductions.py
new file mode 100644
index 0000000000000000000000000000000000000000..6acc8415267bb9fdd7fe6af707cfbbaa74869184
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/reductions.py
@@ -0,0 +1,176 @@
+# mypy: allow-untyped-defs
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import warnings
+
+import torch
+
+from .core import is_masked_tensor
+from .creation import as_masked_tensor, masked_tensor
+
+
+__all__ = []  # type: ignore[var-annotated]
+
+
+def _masked_all_all(data, mask=None):
+    if mask is None:
+        return data.all()
+    return data.masked_fill(~mask, True).all()
+
+
+def _masked_all_dim(data, dim, keepdim=False, mask=None):
+    if mask is None:
+        return torch.all(data, dim=dim, keepdim=keepdim)
+    return torch.all(data.masked_fill(~mask, True), dim=dim, keepdim=keepdim)
+
+
+def _masked_all(*args, **kwargs):
+    if len(args) == 1 and len(kwargs) == 1:
+        return _masked_all_all(args[0], mask=kwargs["mask"])
+    return _masked_all_dim(*args, **kwargs)
+
+
+def _multidim_any(mask, dim, keepdim):
+    if isinstance(dim, int):
+        return _multidim_any(mask, [dim], keepdim)
+    for d in sorted(dim, reverse=True):
+        mask = torch.any(mask, dim=d, keepdim=keepdim)
+    return mask
+
+
+def _get_masked_fn(fn):
+    if fn == "all":
+        return _masked_all
+    return getattr(torch.masked, fn)
+
+
+def _torch_reduce_all(fn):
+    def reduce_all(self):
+        masked_fn = _get_masked_fn(fn)
+        data = self.get_data()
+        mask = self.get_mask().values() if self.is_sparse else self.get_mask()
+        # When reduction is "all", then torch.argmin/torch.argmax needs to return the index of the
+        # element corresponding to the min/max, but this operation isn't supported correctly for sparse layouts.
+        # Therefore, this implementation calculates it using the strides.
+        if fn == "all":
+            result_data = masked_fn(data, mask=mask)
+
+        elif fn in {"argmin", "argmax"} and self.is_sparse_coo():
+            sparse_idx = masked_fn(data.values(), mask=mask).to(dtype=torch.int)
+            indices = (
+                data.to_sparse_coo().indices()
+                if not self.is_sparse_coo()
+                else data.indices()
+            )
+            idx = indices.unbind(1)[sparse_idx]
+            stride = data.size().numel() / torch.tensor(
+                data.size(), device=data.device
+            ).cumprod(0)
+            result_data = torch.sum(idx * stride)
+
+        # we simply pass in the values for sparse COO/CSR tensors
+        elif self.is_sparse:
+            result_data = masked_fn(masked_tensor(data.values(), mask))
+
+        else:
+            result_data = masked_fn(self, mask=mask)
+
+        return as_masked_tensor(result_data, torch.any(mask))
+
+    return reduce_all
+
+
+def _torch_reduce_dim(fn):
+    def reduce_dim(self, dim, keepdim=False, dtype=None):
+        if self.is_sparse:
+            msg = (
+                f"The sparse version of {fn} is not implemented in reductions.\n"
+                "If you would like this operator to be supported, please file an issue for a feature request at "
+                "https://github.com/pytorch/maskedtensor/issues with a minimal reproducible code snippet.\n"
+                "In the case that the semantics for the operator are not trivial, it would be appreciated "
+                "to also include a proposal for the semantics."
+            )
+            warnings.warn(msg, stacklevel=2)
+            return NotImplemented
+        if not is_masked_tensor(self):
+            raise TypeError("Input to reduce_dim must be a MaskedTensor")
+
+        masked_fn = _get_masked_fn(fn)
+        data = self.get_data()
+        mask = self.get_mask()
+        if fn == "all":
+            result_data = masked_fn(data, dim=dim, keepdim=keepdim, mask=mask)
+        else:
+            result_data = masked_fn(
+                self, dim=dim, keepdim=keepdim, dtype=dtype, mask=self.get_mask()
+            )
+        return as_masked_tensor(result_data, _multidim_any(mask, dim, keepdim))
+
+    return reduce_dim
+
+
+def _torch_reduce(fn):
+    def reduce_fn(*args, **kwargs):
+        if len(args) == 1 and len(kwargs) == 0:
+            return _torch_reduce_all(fn)(args[0])
+        return _torch_reduce_dim(fn)(*args, **kwargs)
+
+    return reduce_fn
+
+
+def _reduce_dim_args(input, dim, keepdim=False, dtype=None):
+    return input, dim, keepdim, dtype
+
+
+def _torch_grad_reduce(fn):
+    def grad_reduce(*args, **kwargs):
+        if len(args) == 1 and len(kwargs) == 0:
+            return _torch_reduce_all(fn)(args[0])
+        # TODO: autograd.Function doesn't support kwarg
+        input, dim, keepdim, dtype = _reduce_dim_args(*args, **kwargs)
+        return _torch_reduce_dim(fn)(input, dim, keepdim, dtype)
+
+    return grad_reduce
+
+
+REDUCE_NAMES = [
+    "sum",
+    "mean",
+    "amin",
+    "amax",
+    "argmin",
+    "argmax",
+    "prod",
+    "all",
+    "norm",
+    "var",
+    "std",
+]
+
+NATIVE_REDUCE_MAP = {
+    getattr(torch.ops.aten, name): _torch_reduce(name) for name in REDUCE_NAMES
+}
+TORCH_REDUCE_MAP = {
+    getattr(torch, name): _torch_grad_reduce(name) for name in REDUCE_NAMES
+}
+TENSOR_REDUCE_MAP = {
+    getattr(torch.Tensor, name): _torch_grad_reduce(name) for name in REDUCE_NAMES
+}
+
+NATIVE_REDUCE_FNS = list(NATIVE_REDUCE_MAP.keys())
+TORCH_REDUCE_FNS = list(TORCH_REDUCE_MAP.keys())
+TENSOR_REDUCE_FNS = list(TENSOR_REDUCE_MAP.keys())
+
+
+def _is_reduction(fn):
+    return fn in NATIVE_REDUCE_MAP or fn in TORCH_REDUCE_MAP or fn in TENSOR_REDUCE_MAP
+
+
+def _apply_reduction(fn, *args, **kwargs):
+    if fn in NATIVE_REDUCE_MAP:
+        return NATIVE_REDUCE_MAP[fn](*args, **kwargs)
+    if fn in TORCH_REDUCE_MAP:
+        return TORCH_REDUCE_MAP[fn](*args, **kwargs)
+    if fn in TENSOR_REDUCE_MAP:
+        return TENSOR_REDUCE_MAP[fn](*args, **kwargs)
+    return NotImplemented
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/unary.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/unary.py
new file mode 100644
index 0000000000000000000000000000000000000000..e04ee6e810a7418829b68323097612391017b14e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/masked/maskedtensor/unary.py
@@ -0,0 +1,194 @@
+# mypy: allow-untyped-defs
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import torch
+
+from .core import _map_mt_args_kwargs, _wrap_result
+
+
+__all__ = []  # type: ignore[var-annotated]
+
+
+UNARY_NAMES = [
+    "abs",
+    "absolute",
+    "acos",
+    "arccos",
+    "acosh",
+    "arccosh",
+    "angle",
+    "asin",
+    "arcsin",
+    "asinh",
+    "arcsinh",
+    "atan",
+    "arctan",
+    "atanh",
+    "arctanh",
+    "bitwise_not",
+    "ceil",
+    "clamp",
+    "clip",
+    "conj_physical",
+    "cos",
+    "cosh",
+    "deg2rad",
+    "digamma",
+    "erf",
+    "erfc",
+    "erfinv",
+    "exp",
+    "exp2",
+    "expm1",
+    "fix",
+    "floor",
+    "frac",
+    "lgamma",
+    "log",
+    "log10",
+    "log1p",
+    "log2",
+    "logit",
+    "i0",
+    "isnan",
+    "nan_to_num",
+    "neg",
+    "negative",
+    "positive",
+    "pow",
+    "rad2deg",
+    "reciprocal",
+    "round",
+    "rsqrt",
+    "sigmoid",
+    "sign",
+    "sgn",
+    "signbit",
+    "sin",
+    "sinc",
+    "sinh",
+    "sqrt",
+    "square",
+    "tan",
+    "tanh",
+    "trunc",
+]
+
+INPLACE_UNARY_NAMES = [
+    n + "_"
+    for n in (list(set(UNARY_NAMES) - {"angle", "positive", "signbit", "isnan"}))
+]
+
+# Explicitly tracking functions we know are currently not supported
+# This might be due to missing code gen or because of complex semantics
+UNARY_NAMES_UNSUPPORTED = [
+    "atan2",
+    "arctan2",
+    "bitwise_left_shift",
+    "bitwise_right_shift",
+    "copysign",
+    "float_power",
+    "fmod",
+    "frexp",
+    "gradient",
+    "imag",
+    "ldexp",
+    "lerp",
+    "logical_not",
+    "hypot",
+    "igamma",
+    "igammac",
+    "mvlgamma",
+    "nextafter",
+    "polygamma",
+    "real",
+    "remainder",
+    "true_divide",
+    "xlogy",
+]
+
+
+def _unary_helper(fn, args, kwargs, inplace):
+    if len(kwargs) != 0:
+        raise ValueError(
+            "MaskedTensor unary ops require that len(kwargs) == 0. "
+            "If you need support for this, please open an issue on Github."
+        )
+    for a in args[1:]:
+        if torch.is_tensor(a):
+            raise TypeError(
+                "MaskedTensor unary ops do not support additional Tensor arguments"
+            )
+
+    mask_args, _mask_kwargs = _map_mt_args_kwargs(
+        args, kwargs, lambda x: x._masked_mask
+    )
+    data_args, _data_kwargs = _map_mt_args_kwargs(
+        args, kwargs, lambda x: x._masked_data
+    )
+
+    if args[0].layout == torch.sparse_coo:
+        data_args[0] = data_args[0].coalesce()
+        s = data_args[0].size()
+        i = data_args[0].indices()
+        data_args[0] = data_args[0].coalesce().values()
+        v = fn(*data_args)
+        result_data = torch.sparse_coo_tensor(i, v, size=s)
+
+    elif args[0].layout == torch.sparse_csr:
+        crow = data_args[0].crow_indices()
+        col = data_args[0].col_indices()
+        data_args[0] = data_args[0].values()
+        v = fn(*data_args)
+        result_data = torch.sparse_csr_tensor(crow, col, v)
+
+    else:
+        result_data = fn(*data_args)
+
+    if inplace:
+        args[0]._set_data_mask(result_data, mask_args[0])
+        return args[0]
+    else:
+        return _wrap_result(result_data, mask_args[0])
+
+
+def _torch_unary(fn_name):
+    fn = getattr(torch.ops.aten, fn_name)
+
+    def unary_fn(*args, **kwargs):
+        return _unary_helper(fn, args, kwargs, inplace=False)
+
+    return unary_fn
+
+
+def _torch_inplace_unary(fn_name):
+    fn = getattr(torch.ops.aten, fn_name)
+
+    def unary_fn(*args, **kwargs):
+        return _unary_helper(fn, args, kwargs, inplace=True)
+
+    return unary_fn
+
+
+NATIVE_UNARY_MAP = {
+    getattr(torch.ops.aten, name): _torch_unary(name) for name in UNARY_NAMES
+}
+NATIVE_INPLACE_UNARY_MAP = {
+    getattr(torch.ops.aten, name): _torch_inplace_unary(name)
+    for name in INPLACE_UNARY_NAMES
+}
+
+NATIVE_UNARY_FNS = list(NATIVE_UNARY_MAP.keys())
+NATIVE_INPLACE_UNARY_FNS = list(NATIVE_INPLACE_UNARY_MAP.keys())
+
+
+def _is_native_unary(fn):
+    return fn in NATIVE_UNARY_FNS or fn in NATIVE_INPLACE_UNARY_FNS
+
+
+def _apply_native_unary(fn, *args, **kwargs):
+    if fn in NATIVE_UNARY_FNS:
+        return NATIVE_UNARY_MAP[fn](*args, **kwargs)
+    if fn in NATIVE_INPLACE_UNARY_FNS:
+        return NATIVE_INPLACE_UNARY_MAP[fn](*args, **kwargs)
+    return NotImplemented
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/monitor/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/monitor/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..89af9a227ecda7e71a43cc00db2551e7ae10732e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/monitor/__init__.py
@@ -0,0 +1,39 @@
+from typing import TYPE_CHECKING
+
+from torch._C._monitor import *  # noqa: F403
+from torch._C._monitor import _WaitCounter, _WaitCounterTracker
+
+
+if TYPE_CHECKING:
+    from torch.utils.tensorboard import SummaryWriter
+
+STAT_EVENT = "torch.monitor.Stat"
+
+
+class TensorboardEventHandler:
+    """
+    TensorboardEventHandler is an event handler that will write known events to
+    the provided SummaryWriter.
+
+    This currently only supports ``torch.monitor.Stat`` events which are logged
+    as scalars.
+
+    Example:
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_MONITOR)
+        >>> # xdoctest: +REQUIRES(module:tensorboard)
+        >>> from torch.utils.tensorboard import SummaryWriter
+        >>> from torch.monitor import TensorboardEventHandler, register_event_handler
+        >>> writer = SummaryWriter("log_dir")
+        >>> register_event_handler(TensorboardEventHandler(writer))
+    """
+
+    def __init__(self, writer: "SummaryWriter") -> None:
+        """
+        Constructs the ``TensorboardEventHandler``.
+        """
+        self._writer = writer
+
+    def __call__(self, event: Event) -> None:
+        if event.name == STAT_EVENT:
+            for k, v in event.data.items():
+                self._writer.add_scalar(k, v, walltime=event.timestamp.timestamp())
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7125f1ef2eec0f5cffa0f3fa2fd010e02cb743a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/__init__.py
@@ -0,0 +1,190 @@
+# mypy: allow-untyped-defs
+r"""
+This package enables an interface for accessing MPS (Metal Performance Shaders) backend in Python.
+Metal is Apple's API for programming metal GPU (graphics processor unit). Using MPS means that increased
+performance can be achieved, by running work on the metal GPU(s).
+See https://developer.apple.com/documentation/metalperformanceshaders for more details.
+"""
+
+import torch
+from torch import Tensor
+
+
+_is_in_bad_fork = getattr(torch._C, "_mps_is_in_bad_fork", lambda: False)
+_default_mps_generator: torch._C.Generator = None  # type: ignore[assignment]
+
+
+# local helper function (not public or exported)
+def _get_default_mps_generator() -> torch._C.Generator:
+    global _default_mps_generator
+    if _default_mps_generator is None:
+        _default_mps_generator = torch._C._mps_get_default_generator()
+    return _default_mps_generator
+
+
+def device_count() -> int:
+    r"""Returns the number of available MPS devices."""
+    return int(torch._C._has_mps and torch._C._mps_is_available())
+
+
+def synchronize() -> None:
+    r"""Waits for all kernels in all streams on a MPS device to complete."""
+    return torch._C._mps_deviceSynchronize()
+
+
+def get_rng_state(device: int | str | torch.device = "mps") -> Tensor:
+    r"""Returns the random number generator state as a ByteTensor.
+
+    Args:
+        device (torch.device or int, optional): The device to return the RNG state of.
+            Default: ``'mps'`` (i.e., ``torch.device('mps')``, the current MPS device).
+    """
+    return _get_default_mps_generator().get_state()
+
+
+def set_rng_state(new_state: Tensor, device: int | str | torch.device = "mps") -> None:
+    r"""Sets the random number generator state.
+
+    Args:
+        new_state (torch.ByteTensor): The desired state
+        device (torch.device or int, optional): The device to set the RNG state.
+            Default: ``'mps'`` (i.e., ``torch.device('mps')``, the current MPS device).
+    """
+    new_state_copy = new_state.clone(memory_format=torch.contiguous_format)
+    _get_default_mps_generator().set_state(new_state_copy)
+
+
+def manual_seed(seed: int) -> None:
+    r"""Sets the seed for generating random numbers.
+
+    Args:
+        seed (int): The desired seed.
+    """
+    # the torch.mps.manual_seed() can be called from the global
+    # torch.manual_seed() in torch/random.py. So we need to make
+    # sure mps is available (otherwise we just return without
+    # erroring out)
+    if not torch._C._has_mps:
+        return
+    seed = int(seed)
+    _get_default_mps_generator().manual_seed(seed)
+
+
+def seed() -> None:
+    r"""Sets the seed for generating random numbers to a random number."""
+    _get_default_mps_generator().seed()
+
+
+def empty_cache() -> None:
+    r"""Releases all unoccupied cached memory currently held by the caching
+    allocator so that those can be used in other GPU applications.
+    """
+    torch._C._mps_emptyCache()
+
+
+def set_per_process_memory_fraction(fraction) -> None:
+    r"""Set memory fraction for limiting process's memory allocation on MPS device.
+    The allowed value equals the fraction multiplied by recommended maximum device memory
+    (obtained from Metal API device.recommendedMaxWorkingSetSize).
+    If trying to allocate more than the allowed value in a process, it will raise an out of
+    memory error in allocator.
+
+    Args:
+        fraction(float): Range: 0~2. Allowed memory equals total_memory * fraction.
+
+    .. note::
+       Passing 0 to fraction means unlimited allocations
+       (may cause system failure if out of memory).
+       Passing fraction greater than 1.0 allows limits beyond the value
+       returned from device.recommendedMaxWorkingSetSize.
+    """
+
+    if not isinstance(fraction, float):
+        raise TypeError("Invalid type for fraction argument, must be `float`")
+    if fraction < 0 or fraction > 2:
+        raise ValueError(f"Invalid fraction value: {fraction}. Allowed range: 0~2")
+
+    torch._C._mps_setMemoryFraction(fraction)
+
+
+def current_allocated_memory() -> int:
+    r"""Returns the current GPU memory occupied by tensors in bytes.
+
+    .. note::
+       The returned size does not include cached allocations in
+       memory pools of MPSAllocator.
+    """
+    return torch._C._mps_currentAllocatedMemory()
+
+
+def driver_allocated_memory() -> int:
+    r"""Returns total GPU memory allocated by Metal driver for the process in bytes.
+
+    .. note::
+       The returned size includes cached allocations in MPSAllocator pools
+       as well as allocations from MPS/MPSGraph frameworks.
+    """
+    return torch._C._mps_driverAllocatedMemory()
+
+
+def recommended_max_memory() -> int:
+    r"""Returns recommended max Working set size for GPU memory in bytes.
+
+    .. note::
+       Recommended max working set size for Metal.
+       returned from device.recommendedMaxWorkingSetSize.
+    """
+    return torch._C._mps_recommendedMaxMemory()
+
+
+def compile_shader(source: str):
+    r"""Compiles compute shader from source and allows one to invoke kernels
+    defined there from the comfort of Python runtime
+    Example::
+
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_MPS)
+        >>> lib = torch.mps.compile_shader(
+        ... "kernel void full(device float* out, constant float& val, uint idx [[thread_position_in_grid]]) { out[idx] = val; }"
+        ...  )
+        >>> x = torch.zeros(16, device="mps")
+        >>> lib.full(x, 3.14)
+    """
+    from pathlib import Path
+
+    from torch.utils._cpp_embed_headers import _embed_headers
+
+    if not hasattr(torch._C, "_mps_compileShader"):
+        raise RuntimeError("MPS is not available")
+    source = _embed_headers(
+        [l + "\n" for l in source.split("\n")],
+        [Path(__file__).parent.parent / "include"],
+        set(),
+    )
+    return torch._C._mps_compileShader(source)
+
+
+def is_available() -> bool:
+    return device_count() > 0
+
+
+from . import profiler
+from .event import Event
+
+
+__all__ = [
+    "compile_shader",
+    "device_count",
+    "get_rng_state",
+    "manual_seed",
+    "seed",
+    "set_rng_state",
+    "synchronize",
+    "empty_cache",
+    "set_per_process_memory_fraction",
+    "current_allocated_memory",
+    "driver_allocated_memory",
+    "Event",
+    "profiler",
+    "recommended_max_memory",
+    "is_available",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/event.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/event.py
new file mode 100644
index 0000000000000000000000000000000000000000..3f597c66a41f835fb4a7c58438c788b5025b8383
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/event.py
@@ -0,0 +1,45 @@
+import torch
+
+
+class Event:
+    r"""Wrapper around an MPS event.
+
+    MPS events are synchronization markers that can be used to monitor the
+    device's progress, to accurately measure timing, and to synchronize MPS streams.
+
+    Args:
+        enable_timing (bool, optional): indicates if the event should measure time
+            (default: ``False``)
+    """
+
+    def __init__(self, enable_timing: bool = False) -> None:
+        self.__eventId = torch._C._mps_acquireEvent(enable_timing)
+
+    def __del__(self) -> None:
+        # checks if torch._C is already destroyed
+        if hasattr(torch._C, "_mps_releaseEvent") and self.__eventId > 0:
+            torch._C._mps_releaseEvent(self.__eventId)
+
+    def record(self) -> None:
+        r"""Records the event in the default stream."""
+        torch._C._mps_recordEvent(self.__eventId)
+
+    def wait(self) -> None:
+        r"""Makes all future work submitted to the default stream wait for this event."""
+        torch._C._mps_waitForEvent(self.__eventId)
+
+    def query(self) -> bool:
+        r"""Returns True if all work currently captured by event has completed."""
+        return torch._C._mps_queryEvent(self.__eventId)
+
+    def synchronize(self) -> None:
+        r"""Waits until the completion of all work currently captured in this event.
+        This prevents the CPU thread from proceeding until the event completes.
+        """
+        torch._C._mps_synchronizeEvent(self.__eventId)
+
+    def elapsed_time(self, end_event: "Event") -> float:
+        r"""Returns the time elapsed in milliseconds after the event was
+        recorded and before the end_event was recorded.
+        """
+        return torch._C._mps_elapsedTimeOfEvents(self.__eventId, end_event.__eventId)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/profiler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/profiler.py
new file mode 100644
index 0000000000000000000000000000000000000000..daaaf729ff2e699a9d93514d182116d366daa38d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mps/profiler.py
@@ -0,0 +1,100 @@
+import contextlib
+from collections.abc import Iterator
+from typing import Literal
+
+import torch
+
+
+__all__ = [
+    "start",
+    "stop",
+    "profile",
+    "metal_capture",
+    "is_metal_capture_enabled",
+    "is_capturing_metal",
+]
+
+
+ProfilerMode = Literal["interval", "event", "interval,event"]
+
+
+def start(mode: ProfilerMode = "interval", wait_until_completed: bool = False) -> None:
+    r"""Start OS Signpost tracing from MPS backend.
+
+    The generated OS Signposts could be recorded and viewed in
+    XCode Instruments Logging tool.
+
+    Args:
+        mode(str): OS Signpost tracing mode could be "interval", "event",
+            or both "interval,event".
+            The interval mode traces the duration of execution of the operations,
+            whereas event mode marks the completion of executions.
+            See document `Recording Performance Data`_ for more info.
+        wait_until_completed(bool): Waits until the MPS Stream complete
+            executing each encoded GPU operation. This helps generating single
+            dispatches on the trace's timeline.
+            Note that enabling this option would affect the performance negatively.
+
+    .. _Recording Performance Data:
+       https://developer.apple.com/documentation/os/logging/recording_performance_data
+    """
+    mode_normalized = mode.lower().replace(" ", "")
+    torch._C._mps_profilerStartTrace(  # type: ignore[attr-defined]
+        mode_normalized, wait_until_completed
+    )
+
+
+def stop() -> None:
+    r"""Stops generating OS Signpost tracing from MPS backend."""
+    torch._C._mps_profilerStopTrace()  # type: ignore[attr-defined]
+
+
+@contextlib.contextmanager
+def profile(
+    mode: ProfilerMode = "interval", wait_until_completed: bool = False
+) -> Iterator[None]:
+    r"""Context Manager to enabling generating OS Signpost tracing from MPS backend.
+
+    Args:
+        mode(str): OS Signpost tracing mode could be "interval", "event",
+            or both "interval,event".
+            The interval mode traces the duration of execution of the operations,
+            whereas event mode marks the completion of executions.
+            See document `Recording Performance Data`_ for more info.
+        wait_until_completed(bool): Waits until the MPS Stream complete
+            executing each encoded GPU operation. This helps generating single
+            dispatches on the trace's timeline.
+            Note that enabling this option would affect the performance negatively.
+
+    .. _Recording Performance Data:
+       https://developer.apple.com/documentation/os/logging/recording_performance_data
+    """
+    try:
+        start(mode, wait_until_completed)
+        yield
+    finally:
+        stop()
+
+
+def is_metal_capture_enabled() -> bool:
+    """Checks if `metal_capture` context manager is usable
+    To enable metal capture, set MTL_CAPTURE_ENABLED envvar
+    """
+    return torch._C._mps_isCaptureEnabled()  # type: ignore[attr-defined, no-any-return]
+
+
+def is_capturing_metal() -> bool:
+    """Checks if metal capture is in progress"""
+    return torch._C._mps_isCapturing()  # type: ignore[attr-defined, no-any-return]
+
+
+@contextlib.contextmanager
+def metal_capture(fname: str) -> Iterator[None]:
+    """Context manager that enables capturing of Metal calls into gputrace"""
+    try:
+        torch._C._mps_startCapture(fname)  # type: ignore[attr-defined]
+        yield
+        # Drain all the work that were enqueued during the context call
+        torch.mps.synchronize()
+    finally:
+        torch._C._mps_stopCapture()  # type: ignore[attr-defined]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..23fb8e4f79d7f459c6b7be494117e696c560b851
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/__init__.py
@@ -0,0 +1,584 @@
+# mypy: allow-untyped-defs
+r"""
+This package enables an interface for accessing MTIA backend in python
+"""
+
+import threading
+import traceback
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch import Tensor
+from torch._environment import is_fbcode, is_prod
+from torch._utils import _LazySeedTracker
+from torch.types import Device
+
+from ._utils import _get_device_index
+
+
+# torch.mtia.Event/Stream is alias of torch.Event/Stream
+Event = torch.Event
+Stream = torch.Stream
+
+# Default generators are initialized inside _mtia_init
+default_generators: tuple[torch._C.Generator, ...] = ()  # type: ignore[assignment]
+
+
+_initialized = False
+_queued_calls: list[
+    tuple[Callable[[], None], list[str]]
+] = []  # don't invoke these until initialization occurs
+_tls = threading.local()
+_initialization_lock = threading.Lock()
+_lazy_seed_tracker = _LazySeedTracker()
+
+
+if hasattr(torch._C, "_mtia_exchangeDevice"):
+    _exchange_device = torch._C._mtia_exchangeDevice
+else:
+
+    def _exchange_device(device: int) -> int:
+        if device < 0:
+            return -1
+        raise RuntimeError("PyTorch was compiled without MTIA support")
+
+
+if hasattr(torch._C, "_mtia_maybeExchangeDevice"):
+    _maybe_exchange_device = torch._C._mtia_maybeExchangeDevice
+else:
+
+    def _maybe_exchange_device(device: int) -> int:
+        if device < 0:
+            return -1
+        raise RuntimeError("PyTorch was compiled without MTIA support")
+
+
+def init():
+    _lazy_init()
+
+
+def is_initialized():
+    r"""Return whether PyTorch's MTIA state has been initialized."""
+    return _initialized and not _is_in_bad_fork()
+
+
+def _lazy_call(callable, **kwargs):
+    with _initialization_lock:
+        if is_initialized():
+            return callable()
+        else:
+            global _lazy_seed_tracker
+            if kwargs.get("seed_all", False):
+                _lazy_seed_tracker.queue_seed_all(callable, traceback.format_stack())
+            elif kwargs.get("seed", False):
+                _lazy_seed_tracker.queue_seed(callable, traceback.format_stack())
+            else:
+                _queued_calls.append((callable, traceback.format_stack()))
+
+
+def _is_in_bad_fork() -> bool:
+    return torch._C._mtia_isInBadFork()
+
+
+def _lazy_init() -> None:
+    global _initialized, _queued_calls
+    if is_initialized() or hasattr(_tls, "is_initializing"):
+        return
+    with _initialization_lock:
+        # We be double-checking locking, boys! This is OK because
+        # the above test was GIL protected anyway. The inner test
+        # is for when a thread blocked on some other thread which was
+        # doing the initialization; when they get the lock, they will
+        # find there is nothing left to do.
+        if is_initialized():
+            return
+        # It is important to prevent other threads from entering _lazy_init
+        # immediately, while we are still guaranteed to have the GIL, because some
+        # of the C calls we make below will release the GIL
+        if _is_in_bad_fork():
+            raise RuntimeError(
+                "Cannot re-initialize MTIA in forked subprocess. To use MTIA with "
+                "multiprocessing, you must use the 'spawn' start method"
+            )
+        if not _is_compiled():
+            raise AssertionError(
+                "Torch not compiled with MTIA enabled. "
+                "Ensure you have `import mtia.host_runtime.torch_mtia.dynamic_library` in your python "
+                "src file and include `//mtia/host_runtime/torch_mtia:torch_mtia` as "
+                "your target dependency!"
+            )
+
+        # Install the C++ resource manager to enable Buck resource lookup from Python.
+        # This must be called before _mtia_init() which may access Buck resources.
+        if is_fbcode() and is_prod():
+            try:
+                from libfb.py.cxx_resources import cxx_resource_manager
+
+                cxx_resource_manager.install()
+            except ModuleNotFoundError:
+                # cxx_resource_manager is not available in all build configurations
+                pass
+
+        torch._C._mtia_init()
+        # Some of the queued calls may reentrantly call _lazy_init();
+        # we need to just return without initializing in that case.
+        # However, we must not let any *other* threads in!
+        _tls.is_initializing = True
+
+        _queued_calls.extend(calls for calls in _lazy_seed_tracker.get_calls() if calls)
+
+        try:
+            for queued_call, orig_traceback in _queued_calls:
+                try:
+                    queued_call()
+                except Exception as e:
+                    msg = (
+                        f"MTIA call failed lazily at initialization with error: {str(e)}\n\n"
+                        f"MTIA call was originally invoked at:\n\n{''.join(orig_traceback)}"
+                    )
+                    raise DeferredMtiaCallError(msg) from e
+        finally:
+            delattr(_tls, "is_initializing")
+        _initialized = True
+
+
+class DeferredMtiaCallError(Exception):
+    pass
+
+
+def _is_compiled() -> bool:
+    r"""Return true if compiled with MTIA support."""
+    return torch._C._mtia_isBuilt()
+
+
+def is_available() -> bool:
+    r"""Return true if MTIA device is available"""
+    if not _is_compiled():
+        return False
+    # MTIA has to init devices first to know if there is any devices available.
+    return device_count() > 0
+
+
+def synchronize(device: Device = None) -> None:
+    r"""Waits for all jobs in all streams on a MTIA device to complete."""
+    with torch.mtia.device(device):
+        return torch._C._mtia_deviceSynchronize()
+
+
+def device_count() -> int:
+    r"""Return the number of MTIA devices available."""
+    return torch._C._mtia_getDeviceCount()
+
+
+def current_device() -> int:
+    r"""Return the index of a currently selected device."""
+    # pyrefly: ignore [missing-attribute]
+    return torch._C._mtia_getDevice()
+
+
+def current_stream(device: Device = None) -> Stream:
+    r"""Return the currently selected :class:`Stream` for a given device.
+
+    Args:
+        device (torch.device or int, optional): selected device. Returns
+            the currently selected :class:`Stream` for the current device, given
+            by :func:`~torch.mtia.current_device`, if :attr:`device` is ``None``
+            (default).
+    """
+    device_idx = _get_device_index(device, optional=True)
+    return torch._C._mtia_getCurrentStream(device_idx)
+
+
+def default_stream(device: Device = None) -> Stream:
+    r"""Return the default :class:`Stream` for a given device.
+
+    Args:
+        device (torch.device or int, optional): selected device. Returns
+            the default :class:`Stream` for the current device, given by
+            :func:`~torch.mtia.current_device`, if :attr:`device` is ``None``
+            (default).
+    """
+    return torch._C._mtia_getDefaultStream(_get_device_index(device, optional=True))
+
+
+def record_memory_history(
+    enabled: str | None = "all", stacks: str = "python", max_entries: int = 0
+) -> None:
+    r"""Enable/Disable the memory profiler on MTIA allocator
+
+    Args:
+        enabled (all or state, optional) selected device. Returns
+            statistics for the current device, given by current_device(),
+            if device is None (default).
+
+        stacks ("python" or "cpp", optional). Select the stack trace to record.
+
+        max_entries (int, optional). Maximum number of entries to record.
+    """
+    if not is_initialized():
+        return
+    torch._C._mtia_recordMemoryHistory(enabled, stacks, max_entries)
+
+
+def snapshot() -> dict[str, Any]:
+    r"""Return a dictionary of MTIA memory allocator history"""
+
+    return torch._C._mtia_memorySnapshot()
+
+
+def attach_out_of_memory_observer(
+    observer: Callable[[int, int, int, int], None],
+) -> None:
+    r"""Attach an out-of-memory observer to MTIA memory allocator"""
+    torch._C._mtia_attachOutOfMemoryObserver(observer)
+
+
+def is_bf16_supported(including_emulation: bool = True):
+    r"""Return a bool indicating if the current MTIA device supports dtype bfloat16."""
+    return True
+
+
+def get_device_capability(device: Device = None) -> tuple[int, int]:
+    r"""Return capability of a given device as a tuple of (major version, minor version).
+
+    Args:
+        device (torch.device or int, optional) selected device. Returns
+            statistics for the current device, given by current_device(),
+            if device is None (default).
+    """
+    return torch._C._mtia_getDeviceCapability(_get_device_index(device, optional=True))
+
+
+def empty_cache() -> None:
+    r"""Empty the MTIA device cache."""
+    return torch._C._mtia_emptyCache()
+
+
+def set_stream(stream: Stream):
+    r"""Set the current stream. This is a wrapper API to set the stream.
+        Usage of this function is discouraged in favor of the ``stream``
+        context manager.
+
+    Args:
+        stream (Stream): selected stream. This function is a no-op
+            if this argument is ``None``.
+    """
+    if stream is None:
+        return
+    torch._C._mtia_setCurrentStream(stream)
+
+
+def set_device(device: Device) -> None:
+    r"""Set the current device.
+
+    Args:
+        device (torch.device or int): selected device. This function is a no-op
+            if this argument is negative.
+    """
+    device = _get_device_index(device)
+    if device >= 0:
+        # pyrefly: ignore [missing-attribute]
+        torch._C._mtia_setDevice(device)
+
+
+def get_device_properties(device: Device = None) -> dict[str, Any]:
+    r"""Return a dictionary of MTIA device properties
+
+    Args:
+        device (torch.device or int, optional) selected device. Returns
+            statistics for the current device, given by current_device(),
+            if device is None (default).
+    """
+    return torch._C._mtia_getDeviceProperties(_get_device_index(device, optional=True))
+
+
+class device:
+    r"""Context-manager that changes the selected device.
+
+    Args:
+        device (torch.device or int): device index to select. It's a no-op if
+            this argument is a negative integer or ``None``.
+    """
+
+    def __init__(self, device: Any):
+        self.idx = _get_device_index(device, optional=True)
+        self.prev_idx = -1
+
+    def __enter__(self):
+        # pyrefly: ignore [missing-attribute]
+        self.prev_idx = torch._C._mtia_maybeExchangeDevice(self.idx)
+
+    def __exit__(self, type: Any, value: Any, traceback: Any):
+        # pyrefly: ignore [missing-attribute]
+        self.idx = torch._C._mtia_maybeExchangeDevice(self.prev_idx)
+        return False
+
+
+class StreamContext:
+    r"""Context-manager that selects a given stream.
+
+    All MTIA kernels queued within its context will be enqueued on a selected
+    stream.
+
+    Args:
+        Stream (Stream): selected stream. This manager is a no-op if it's
+            ``None``.
+    .. note:: Streams are per-device.
+    """
+
+    cur_stream: Stream | None
+
+    def __init__(self, stream: Stream | None):
+        self.cur_stream = None
+        self.stream = stream
+        self.idx = _get_device_index(None, True)
+        if not torch.jit.is_scripting():
+            if self.idx is None:
+                self.idx = -1  # pyrefly: ignore [bad-assignment]
+
+        self.src_prev_stream = (
+            None if not torch.jit.is_scripting() else torch.mtia.default_stream(None)
+        )
+        self.dst_prev_stream = (
+            None if not torch.jit.is_scripting() else torch.mtia.default_stream(None)
+        )
+
+    def __enter__(self):
+        # Local cur_stream variable for type refinement
+        cur_stream = self.stream
+        # Return if stream is None or MTIA device not available
+        if cur_stream is None or self.idx == -1:
+            return
+        self.src_prev_stream = torch.mtia.current_stream(None)
+
+        # If the stream is not on the current device, then
+        # set the current stream on the device
+        if self.src_prev_stream.device != cur_stream.device:
+            with device(cur_stream.device):
+                self.dst_prev_stream = torch.mtia.current_stream(cur_stream.device)
+        torch.mtia.set_stream(cur_stream)
+
+    def __exit__(self, type: Any, value: Any, traceback: Any):
+        # Local cur_stream variable for type refinement
+        cur_stream = self.stream
+        # If stream is None or no MTIA device available, return
+        if cur_stream is None or self.idx == -1:
+            return
+
+        # Reset the stream on the original device
+        # and destination device
+        if self.src_prev_stream.device != cur_stream.device:  # type: ignore[union-attr]
+            torch.mtia.set_stream(self.dst_prev_stream)  # type: ignore[arg-type]
+        torch.mtia.set_stream(self.src_prev_stream)  # type: ignore[arg-type]
+
+
+def _set_stream_by_id(stream_id, device_index, device_type):
+    r"""set stream specified by the stream id, device index and
+        device type
+
+    Args: stream_id (int): stream id in stream pool
+          device_index (int): device index in topo
+          device_type (int): enum device type
+    """
+    torch._C._mtia_setStream(stream_id, device_index, device_type)
+
+
+def stream(stream: Stream | None) -> StreamContext:
+    r"""Wrap around the Context-manager StreamContext that selects a given stream.
+
+    Arguments:
+        stream (Stream): selected stream. This manager is a no-op if it's
+            ``None``.
+    .. note:: In eager mode stream is of type Stream class while in JIT it doesn't support torch.mtia.stream
+    """
+    return StreamContext(stream)
+
+
+def get_rng_state(device: Device = "mtia") -> Tensor:
+    r"""Returns the random number generator state of the specified MTIA device as a ByteTensor.
+
+    Args:
+        device (torch.device or int, optional): The device to return the RNG state of.
+            Default: ``'mtia'`` (i.e., ``torch.device('mtia')``, the current mtia device).
+
+    .. warning::
+        This function eagerly initializes MTIA.
+    """
+    _lazy_init()
+    idx = _get_device_index(device, optional=True)
+    if idx is None:
+        idx = current_device()
+    default_generator = default_generators[idx]
+    return default_generator.get_state()
+
+
+def get_rng_state_all() -> list[Tensor]:
+    r"""Returns a list of ByteTensor representing the random number states of all devices."""
+    results = [get_rng_state(i) for i in range(device_count())]
+    return results
+
+
+def set_rng_state(new_state: Tensor, device: Device = "mtia") -> None:
+    r"""Sets the random number generator state of the specified MTIA device.
+
+    Args:
+        new_state (torch.ByteTensor): The desired state
+        device (torch.device or int, optional): The device to set the RNG state.
+            Default: ``'mtia'`` (i.e., ``torch.device('mtia')``, the current mtia device).
+    """
+    if not is_initialized():
+        with torch._C._DisableFuncTorch():
+            # Clone the state because the callback will be triggered
+            # later when MTIA is lazy initialized.
+            new_state = new_state.clone(memory_format=torch.contiguous_format)
+
+    idx = _get_device_index(device, optional=True)
+    if idx is None:
+        idx = current_device()
+
+    def cb():
+        default_generator = default_generators[idx]
+        default_generator.set_state(new_state)
+
+    _lazy_call(cb)
+
+
+def set_rng_state_all(new_states: list[Tensor]) -> None:
+    r"""Sets the random number generator state of all devices.
+
+    Args:
+        new_states (Iterable of torch.ByteTensor): The desired state for each device.
+    """
+    for i, state in enumerate(new_states):
+        set_rng_state(state, i)
+
+
+def manual_seed(seed: int) -> None:
+    r"""Sets the seed for generating random numbers for the current MTIA device.
+    It's safe to call this function if MTIA is not available; in that case, it is silently ignored.
+
+    Args:
+        seed (int): The desired seed.
+
+    .. warning::
+        If you are working with a multi-GPU model, this function is insufficient
+        to get determinism.  To seed all GPUs, use :func:`manual_seed_all`.
+    """
+    seed = int(seed)
+
+    def cb():
+        idx = current_device()
+        default_generator = default_generators[idx]
+        default_generator.manual_seed(seed)
+
+    _lazy_call(cb, seed=True)
+
+
+def manual_seed_all(seed: int) -> None:
+    r"""Sets the seed for generating random numbers on all MTIA devices.
+    It's safe to call this function if MTIA is not available; in that case, it is silently ignored.
+
+    Args:
+        seed (int): The desired seed.
+    """
+    seed = int(seed)
+
+    def cb():
+        for i in range(device_count()):
+            default_generator = default_generators[i]
+            default_generator.manual_seed(seed)
+
+    _lazy_call(cb, seed_all=True)
+
+
+def seed() -> int:
+    r"""Sets the seed for generating random numbers to a random number for the current MTIA device.
+    It's safe to call this function if MTIA is not available; in that case, it is silently ignored.
+
+    .. warning::
+        If you are working with a multi-GPU model, this function will only initialize
+        the seed on one GPU.  To initialize all GPUs, use :func:`seed_all`.
+    """
+
+    def cb():
+        idx = current_device()
+        default_generator = default_generators[idx]
+        return default_generator.seed()
+
+    return _lazy_call(cb, seed=True)
+
+
+def seed_all() -> int:
+    r"""Sets the seed for generating random numbers to a random number on all MTIA devices.
+
+    It's safe to call this function if MTIA is not available; in that case, it is silently ignored.
+    """
+
+    def cb():
+        random_seed = 0
+        seeded = False
+        for i in range(device_count()):
+            default_generator = default_generators[i]
+            if not seeded:
+                random_seed = default_generator.seed()
+                seeded = True
+            else:
+                default_generator.manual_seed(random_seed)
+        return random_seed
+
+    return _lazy_call(cb, seed_all=True)
+
+
+def initial_seed() -> int:
+    r"""Returns the current random seed of the current MTIA device.
+
+    .. warning::
+        This function eagerly initializes MTIA.
+    """
+    _lazy_init()
+    idx = current_device()
+    return default_generators[idx].initial_seed()
+
+
+from .memory import *  # noqa: F403
+from .mtia_graph import *  # noqa: F403
+
+
+__all__ = [
+    "init",
+    "is_available",
+    "is_initialized",
+    "synchronize",
+    "device_count",
+    "current_device",
+    "current_stream",
+    "default_stream",
+    "memory_stats",
+    "max_memory_allocated",
+    "memory_allocated",
+    "reset_peak_memory_stats",
+    "get_device_capability",
+    "get_device_properties",
+    "record_memory_history",
+    "snapshot",
+    "attach_out_of_memory_observer",
+    "empty_cache",
+    "set_device",
+    "set_stream",
+    "stream",
+    "device",
+    "set_rng_state",
+    "get_rng_state",
+    "set_rng_state_all",
+    "get_rng_state_all",
+    "manual_seed",
+    "manual_seed_all",
+    "seed",
+    "seed_all",
+    "initial_seed",
+    "is_bf16_supported",
+    "MTIAGraph",
+    "graph",
+    "graph_pool_handle",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..16710c244c61ce696ec96818340f95cb4d10af32
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/_utils.py
@@ -0,0 +1,47 @@
+from typing import Any
+
+import torch
+
+
+def _get_device_index(
+    device: Any, optional: bool = False, allow_cpu: bool = False
+) -> int:
+    r"""Get the device index from :attr:`device`, which can be a torch.device object, a Python integer, or ``None``.
+
+    If :attr:`device` is a torch.device object, returns the device index if it
+    is a MTIA device. Note that for a MTIA device without a specified index,
+    i.e., ``torch.device('mtia')``, this will return the current default MTIA
+    device if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``,
+    CPU devices will be accepted and ``-1`` will be returned in this case.
+
+    If :attr:`device` is a Python integer, it is returned as is.
+
+    If :attr:`device` is ``None``, this will return the current default MTIA
+    device if :attr:`optional` is ``True``.
+    """
+
+    if device is None and optional:
+        # If device is None (frequent), then we can can short-circuit the logic
+        return torch._C._mtia_getDevice()
+    if isinstance(device, int):
+        return device
+    if not torch.jit.is_scripting():
+        if isinstance(device, torch.mtia.device):
+            return device.idx
+    if isinstance(device, str):
+        device = torch.device(device)
+    device_idx: int | None = None
+    if isinstance(device, torch.device):
+        if not allow_cpu and device.type == "cpu":
+            raise ValueError(f"Expected a non cpu device, but got: {device}")
+        if device.type not in ["mtia", "cpu"]:
+            raise ValueError(f"Expected a mtia or cpu device, but got: {device}")
+        device_idx = -1 if device.type == "cpu" else device.index
+    if device_idx is None:
+        if optional:
+            device_idx = torch._C._mtia_getDevice()
+        else:
+            raise ValueError(
+                f"Expected a torch.device with a specified index or an integer, but got: {device}"
+            )
+    return device_idx
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/memory.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/memory.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0fbffc9489115da72023de417ee619d80a8b011
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/memory.py
@@ -0,0 +1,71 @@
+# pyre-strict
+
+r"""This package adds support for device memory management implemented in MTIA."""
+
+from typing import Any
+
+import torch
+
+from . import Device, is_initialized
+from ._utils import _get_device_index
+
+
+def memory_stats(device: Device = None) -> dict[str, Any]:
+    r"""Return a dictionary of MTIA memory allocator statistics for a given device.
+
+    Args:
+        device (torch.device, str, or int, optional) selected device. Returns
+            statistics for the current device, given by current_device(),
+            if device is None (default).
+    """
+    if not is_initialized():
+        return {}
+    return torch._C._mtia_memoryStats(_get_device_index(device, optional=True))
+
+
+def max_memory_allocated(device: Device = None) -> int:
+    r"""Return the maximum memory allocated in bytes for a given device.
+
+    Args:
+        device (torch.device, str, or int, optional) selected device. Returns
+            statistics for the current device, given by current_device(),
+            if device is None (default).
+    """
+    if not is_initialized():
+        return 0
+    return memory_stats(device).get("dram", 0).get("peak_bytes", 0)
+
+
+def memory_allocated(device: Device = None) -> int:
+    r"""Return the current MTIA memory occupied by tensors in bytes for a given device.
+
+    Args:
+        device (torch.device or int or str, optional): selected device. Returns
+            statistic for the current device, given by :func:`~torch.mtia.current_device`,
+            if :attr:`device` is ``None`` (default).
+    """
+    if not is_initialized():
+        return 0
+    return memory_stats(device).get("dram", 0).get("allocated_bytes", 0)
+
+
+def reset_peak_memory_stats(device: Device = None) -> None:
+    r"""Reset the peak memory stats for a given device.
+
+
+    Args:
+        device (torch.device, str, or int, optional) selected device. Returns
+            statistics for the current device, given by current_device(),
+            if device is None (default).
+    """
+    if not is_initialized():
+        return
+    torch._C._mtia_resetPeakMemoryStats(_get_device_index(device, optional=True))
+
+
+__all__ = [
+    "memory_stats",
+    "max_memory_allocated",
+    "memory_allocated",
+    "reset_peak_memory_stats",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/mtia_graph.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/mtia_graph.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fc61b67f0903158209dd7f9d404d935ce96410f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/mtia/mtia_graph.py
@@ -0,0 +1,104 @@
+# pylint: disable=useless-parent-delegation
+from __future__ import annotations
+
+from typing_extensions import Self
+
+import torch
+
+
+_POOL_HANDLE = tuple[int, int]
+
+
+def graph_pool_handle() -> _POOL_HANDLE:
+    """
+    Return an opaque token representing the id of a graph memory pool.
+    """
+    # pyrefly: ignore [missing-attribute]
+    return torch._C._mtia_graphPoolHandle()
+
+
+class MTIAGraph(torch._C._MTIAGraph):
+    """
+    Wrapper around a MTIA graph.
+    """
+
+    def __new__(cls, keep_graph: bool = False) -> Self:
+        return super().__new__(cls, keep_graph)
+
+    def capture_begin(self, pool: _POOL_HANDLE) -> None:
+        """
+        Begin capturing a MTIA graph.
+        """
+        super().capture_begin(pool)
+
+    def capture_end(self) -> None:
+        """
+        End the capture of a MTIA graph.
+        """
+        super().capture_end()
+
+    def instantiate(self) -> None:
+        """
+        Instantiate the captured MTIA graph.
+        """
+        super().instantiate()
+
+    def replay(self) -> None:
+        """
+        Replay the captured MTIA graph.
+        """
+        super().replay()
+
+    def reset(self) -> None:
+        """
+        Destroy the captured graph and reset the states.
+        """
+        super().reset()
+
+    def pool(self) -> _POOL_HANDLE:
+        """
+        Return an opaque token representing the id of this graph's memory pool
+        """
+        return super().pool()
+
+
+class graph:
+    default_capture_stream: torch.mtia.Stream | None = None
+
+    def __init__(
+        self,
+        mtia_graph: MTIAGraph,
+        pool: _POOL_HANDLE | None = None,
+        stream: torch.mtia.Stream | None = None,
+    ):
+        if self.__class__.default_capture_stream is None:
+            self.__class__.default_capture_stream = torch.mtia.current_stream()
+
+        self.pool: tuple[()] | tuple[_POOL_HANDLE] = () if pool is None else (pool,)
+        self.capture_stream = (
+            stream if stream is not None else self.__class__.default_capture_stream
+        )
+        if self.capture_stream is None:
+            raise AssertionError("capture_stream must not be None")
+        self.stream_ctx = torch.mtia.stream(self.capture_stream)
+        self.mtia_graph = mtia_graph
+
+    def __enter__(self) -> None:
+        torch.mtia.synchronize()
+        torch.mtia.empty_cache()
+
+        self.stream_ctx.__enter__()
+
+        pool_arg = self.pool[0] if self.pool else (0, 0)
+        self.mtia_graph.capture_begin(pool_arg)
+
+    def __exit__(self, *args: object) -> None:
+        self.mtia_graph.capture_end()
+        self.stream_ctx.__exit__(*args)
+
+
+__all__ = [
+    "MTIAGraph",
+    "graph",
+    "graph_pool_handle",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d73f3d028820b91d683459c77fe7a7269205f074
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/__init__.py
@@ -0,0 +1,126 @@
+# mypy: allow-untyped-defs
+"""torch.multiprocessing is a wrapper around the native :mod:`multiprocessing` module.
+
+It registers custom reducers, that use shared memory to provide shared
+views on the same data in different processes. Once the tensor/storage is moved
+to shared_memory (see :func:`~torch.Tensor.share_memory_`), it will be possible
+to send it to other processes without making any copies.
+
+The API is 100% compatible with the original module - it's enough to change
+``import multiprocessing`` to ``import torch.multiprocessing`` to have all the
+tensors sent through the queues or shared via other mechanisms, moved to shared
+memory.
+
+Because of the similarity of APIs we do not document most of this package
+contents, and we recommend referring to very good docs of the original module.
+"""
+
+import multiprocessing
+import sys
+
+import torch
+
+from .reductions import init_reductions
+
+
+__all__ = ["set_sharing_strategy", "get_sharing_strategy", "get_all_sharing_strategies"]
+
+
+from multiprocessing import *  # noqa: F403
+
+
+__all__ += multiprocessing.__all__  # noqa: PLE0605 type: ignore[attr-defined]
+
+
+# This call adds a Linux specific prctl(2) wrapper function to this module.
+# See https://github.com/pytorch/pytorch/pull/14391 for more information.
+torch._C._multiprocessing_init()
+
+
+"""Add helper function to spawn N processes and wait for completion of any of
+them."""
+from .spawn import (
+    ENV_VAR_PARALLEL_START,
+    ProcessContext,
+    ProcessExitedException,
+    ProcessRaisedException,
+    spawn,
+    SpawnContext,
+    start_processes,
+)
+
+
+if sys.platform == "darwin" or sys.platform == "win32":
+    _sharing_strategy = "file_system"
+    _all_sharing_strategies = {"file_system"}
+else:
+    _sharing_strategy = "file_descriptor"
+    _all_sharing_strategies = {"file_descriptor", "file_system"}
+
+
+def set_sharing_strategy(new_strategy):
+    """Set the strategy for sharing CPU tensors.
+
+    Args:
+        new_strategy (str): Name of the selected strategy. Should be one of
+            the values returned by :func:`get_all_sharing_strategies()`.
+    """
+    global _sharing_strategy
+    if new_strategy not in _all_sharing_strategies:
+        raise AssertionError(
+            f"invalid sharing strategy {new_strategy!r}, "
+            f"expected one of {_all_sharing_strategies}"
+        )
+    _sharing_strategy = new_strategy
+
+
+def get_sharing_strategy():
+    """Return the current strategy for sharing CPU tensors."""
+    return _sharing_strategy
+
+
+def get_all_sharing_strategies():
+    """Return a set of sharing strategies supported on a current system."""
+    return _all_sharing_strategies
+
+
+def _set_thread_name(name: str) -> None:
+    """Set the name of the current thread.
+
+    Args:
+        name (str): Name of the current thread.
+    """
+    torch._C._set_thread_name(name)
+
+
+def _get_thread_name() -> str:
+    """Get the name of the current thread.
+
+    Returns:
+        str: Name of the current thread.
+    """
+    return torch._C._get_thread_name()
+
+
+init_reductions()
+
+# Leak ResourceTracker at exit for Python-3.12 on MacOS
+# See https://github.com/pytorch/pytorch/issues/153050 and
+# https://github.com/python/cpython/issues/88887 for more details
+from multiprocessing.resource_tracker import ResourceTracker as _RT
+
+
+if (
+    sys.platform == "darwin"
+    and sys.version_info >= (3, 12, 2)
+    and hasattr(_RT, "__del__")
+):
+    import atexit
+
+    def _leak_RT_at_exit():
+        def _noop(x):
+            pass
+
+        _RT.__del__ = _noop  # type: ignore[attr-defined]
+
+    atexit.register(_leak_RT_at_exit)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/_atfork.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/_atfork.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac4a97c10c07ae680765b0f362ef33c4bfb2308b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/_atfork.py
@@ -0,0 +1,35 @@
+# mypy: allow-untyped-defs
+import sys
+
+
+__all__ = ["register_after_fork"]
+
+if sys.platform == "win32":
+    import multiprocessing.util as _util
+
+    def _register(func):
+        def wrapper(arg):
+            func()
+
+        _util.register_after_fork(_register, wrapper)
+
+else:
+    import os
+
+    def _register(func):
+        os.register_at_fork(after_in_child=func)
+
+
+def register_after_fork(func):
+    """Register a callable to be executed in the child process after a fork.
+
+    Note:
+        In python < 3.7 this will only work with processes created using the
+        ``multiprocessing`` module. In python >= 3.7 it also works with
+        ``os.fork()``.
+
+    Args:
+        func (function): Function taking no arguments to be called in the child after fork
+
+    """
+    _register(func)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/pool.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/pool.py
new file mode 100644
index 0000000000000000000000000000000000000000..32a47efac0d6ea6f0c8700924aafa67564813f5b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/pool.py
@@ -0,0 +1,52 @@
+import multiprocessing.pool
+import multiprocessing.util as util
+
+from .queue import SimpleQueue
+
+
+def clean_worker(*args, **kwargs):
+    import gc
+
+    multiprocessing.pool.worker(*args, **kwargs)
+    # Regular multiprocessing workers don't fully clean up after themselves,
+    # so we have to explicitly trigger garbage collection to make sure that all
+    # destructors are called...
+    gc.collect()
+
+
+class Pool(multiprocessing.pool.Pool):
+    """Pool implementation which uses our version of SimpleQueue.
+
+    This lets us pass tensors in shared memory across processes instead of
+    serializing the underlying data.
+    """
+
+    def _setup_queues(self):
+        self._inqueue = SimpleQueue()
+        self._outqueue = SimpleQueue()
+        self._quick_put = self._inqueue._writer.send
+        self._quick_get = self._outqueue._reader.recv
+
+    def _repopulate_pool(self):
+        """Increase the number of pool processes to the specified number.
+
+        Bring the number of pool processes up to the specified number, for use after
+        reaping workers which have exited.
+        """
+        for _ in range(self._processes - len(self._pool)):
+            # changed worker -> clean_worker
+            args = (
+                self._inqueue,
+                self._outqueue,
+                self._initializer,
+                self._initargs,
+                self._maxtasksperchild,
+            )
+            if hasattr(self, "_wrap_exception"):
+                args += (self._wrap_exception,)
+            w = self.Process(target=clean_worker, args=args)
+            self._pool.append(w)
+            w.name = w.name.replace("Process", "PoolWorker")
+            w.daemon = True
+            w.start()
+            util.debug("added worker")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/queue.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/queue.py
new file mode 100644
index 0000000000000000000000000000000000000000..876bf8d0e7459b60a41b59b0a093608e515ba455
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/queue.py
@@ -0,0 +1,43 @@
+# mypy: allow-untyped-defs
+import io
+import multiprocessing.queues
+import pickle
+from multiprocessing.reduction import ForkingPickler
+
+
+class ConnectionWrapper:
+    """Proxy class for _multiprocessing.Connection which uses ForkingPickler for object serialization."""
+
+    def __init__(self, conn):
+        self.conn = conn
+
+    def send(self, obj):
+        buf = io.BytesIO()
+        ForkingPickler(buf, pickle.HIGHEST_PROTOCOL).dump(obj)
+        self.send_bytes(buf.getvalue())
+
+    def recv(self):
+        buf = self.recv_bytes()
+        return pickle.loads(buf)
+
+    def __getattr__(self, name):
+        if "conn" in self.__dict__:
+            return getattr(self.conn, name)
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute 'conn'")
+
+
+class Queue(multiprocessing.queues.Queue):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._reader: ConnectionWrapper = ConnectionWrapper(self._reader)
+        self._writer: ConnectionWrapper = ConnectionWrapper(self._writer)
+        self._send = self._writer.send
+        self._recv = self._reader.recv
+
+
+class SimpleQueue(multiprocessing.queues.SimpleQueue):
+    def _make_methods(self):
+        if not isinstance(self._reader, ConnectionWrapper):
+            self._reader: ConnectionWrapper = ConnectionWrapper(self._reader)
+            self._writer: ConnectionWrapper = ConnectionWrapper(self._writer)
+        super()._make_methods()  # type: ignore[misc]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/reductions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/reductions.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0b6d9468bcb51f82a2e5e146fc285f8ac242fe5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/reductions.py
@@ -0,0 +1,644 @@
+# mypy: allow-untyped-defs
+import multiprocessing
+import os
+import threading
+from multiprocessing import reduction
+from multiprocessing.util import register_after_fork
+
+import torch
+from torch._namedtensor_internals import check_serializing_named_tensor
+
+
+try:
+    # Early load resource_sharer to prevent a partially initialized instance
+    # from being inherited in a forked child process. The reduce_storage method
+    # requires this module indirectly through DupFd(). The built-in mp.Queue
+    # class pickles arguments in a background thread which may overlap with the
+    # fork.
+    import multiprocessing.resource_sharer
+except ImportError:
+    pass
+
+
+class StorageWeakRef:
+    r"""A weak reference to a Storage.
+
+    The cdata member is a Python number containing the integer representation of
+    the Storage pointer.
+    """
+
+    __slots__ = ["cdata", "_free_weak_ref"]
+
+    def __init__(self, storage):
+        self.cdata = storage._weak_ref()
+        # Save a direct reference to _free_weak_ref because the `torch` module
+        # might be cleared during Python shutdown before this module is cleared.
+        self._free_weak_ref = torch.Storage._free_weak_ref  # type: ignore[attr-defined]
+
+    @classmethod
+    def from_weakref(cls, cdata):
+        instance = cls.__new__(cls)
+        instance.cdata = cdata
+        instance._free_weak_ref = torch.Storage._free_weak_ref  # type: ignore[attr-defined]
+        return instance
+
+    def expired(self):
+        return torch.Storage._expired(self.cdata)  # type: ignore[attr-defined]
+
+    def __del__(self):
+        self._free_weak_ref(self.cdata)
+
+    def __hash__(self):
+        return self.cdata
+
+    def __eq__(self, other):
+        if id(self) == id(other):
+            return True
+        return self.cdata == other.cdata
+
+
+class SharedCache(dict):
+    """Dictionary from multiprocessing handles to StorageWeakRef."""
+
+    def __init__(self) -> None:
+        # free_dead_references() is called if the len exceeds the current
+        # limit. The limit scales with the number of remaining live objects.
+        self.limit = 128
+        # `fork` inherits lock state, so in case we fork when the lock is held,
+        # we register a function to reset the lock to a new object to avoid
+        # possible deadlocks, following python multiprocessing library design.
+        self._after_fork()
+        register_after_fork(self, SharedCache._after_fork)
+
+    def _after_fork(self):
+        self.lock = threading.Lock()
+
+    def get(self, key):  # type: ignore[override]
+        with self.lock:
+            return dict.get(self, key)
+
+    def __setitem__(self, key, storage_ref):
+        with self.lock:
+            dict.__setitem__(self, key, storage_ref)
+            if len(self) > self.limit:
+                self.free_dead_references()
+
+    def free_dead_references(self):
+        live = 0
+        for key, storage_ref in list(self.items()):
+            if storage_ref.expired():
+                del self[key]
+            else:
+                live += 1
+        self.limit = max(128, live * 2)
+
+
+# mapping from handles to StorageWeakRef objects
+shared_cache = SharedCache()
+
+
+def rebuild_event(device, handle):
+    return torch.cuda.Event.from_ipc_handle(device, handle)
+
+
+def reduce_event(event):
+    handle = event.ipc_handle()
+    return (rebuild_event, (event.device, handle))
+
+
+def rebuild_tensor(cls, storage, metadata):
+    storage_offset, size, stride, requires_grad = metadata
+    t = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)
+    if cls == torch.nn.parameter.Parameter:
+        # we have to pass requires_grad into constructor, rather than set it as an
+        # attribute later, because it's an important check for Integer Tensors to
+        # have requires_grad=False (or else they raise an error)
+        t = torch.nn.parameter.Parameter(t, requires_grad=requires_grad)
+    else:
+        t.requires_grad = requires_grad
+    return t
+
+
+def rebuild_meta_tensor(
+    tensor_cls,
+    tensor_size,
+    tensor_stride,
+    tensor_offset,
+    dtype,
+    storage_size_bytes,
+    requires_grad,
+):
+    untyped_storage = torch.UntypedStorage(storage_size_bytes, device="meta")
+
+    typed_storage = torch.TypedStorage(
+        wrap_storage=untyped_storage, dtype=dtype, _internal=True
+    )
+
+    t = torch._utils._rebuild_tensor(
+        typed_storage,
+        tensor_offset,
+        tensor_size,
+        tensor_stride,
+    )
+
+    if tensor_cls == torch.nn.parameter.Parameter:
+        # It is crucial for integer tensors to receive
+        # the requires_grad=False as an argument in the constructor
+        t = torch.nn.parameter.Parameter(t, requires_grad=requires_grad)
+    else:
+        t.requires_grad = requires_grad
+
+    return t
+
+
+def rebuild_cuda_tensor(
+    tensor_cls,
+    tensor_size,
+    tensor_stride,
+    tensor_offset,
+    storage_cls,
+    dtype,
+    storage_device,
+    storage_handle,
+    storage_size_bytes,
+    storage_offset_bytes,
+    requires_grad,
+    ref_counter_handle,
+    ref_counter_offset,
+    event_handle,
+    event_sync_required,
+):
+    # If storage_handle is None, storage points to nullptr.
+    if storage_handle is None or storage_size_bytes == 0:
+        storage = storage_cls(0, dtype=dtype, device=storage_device, _internal=True)
+    else:
+        storage = storage_from_cache(
+            storage_cls, (storage_handle, storage_offset_bytes)
+        )
+        if storage is None:
+            torch.cuda._lazy_init()
+            storage = storage_cls._new_shared_cuda(
+                storage_device,
+                storage_handle,
+                storage_size_bytes,
+                storage_offset_bytes,
+                ref_counter_handle,
+                ref_counter_offset,
+                event_handle,
+                event_sync_required,
+            )
+            shared_cache[(storage_handle, storage_offset_bytes)] = StorageWeakRef(
+                storage
+            )
+        else:
+            # We already ref counting this Storage, but producer needs new ref-counters to be released.
+            storage_cls._release_ipc_counter(
+                ref_counter_handle, ref_counter_offset, device=storage_device
+            )
+
+    _storage = (
+        storage
+        if isinstance(storage, torch.UntypedStorage)
+        else storage._untyped_storage
+    )
+
+    t = torch._utils._rebuild_tensor(
+        torch.storage.TypedStorage(wrap_storage=_storage, dtype=dtype, _internal=True),
+        tensor_offset,
+        tensor_size,
+        tensor_stride,
+    )
+
+    if tensor_cls == torch.nn.parameter.Parameter:
+        # It is crucial for integer tensors to receive
+        # the requires_grad=False as an argument in the constructor
+        t = torch.nn.parameter.Parameter(t, requires_grad=requires_grad)
+    else:
+        t.requires_grad = requires_grad
+
+    return t
+
+
+def reduce_tensor(tensor):
+    if tensor.requires_grad and not tensor.is_leaf:
+        raise RuntimeError(
+            "Cowardly refusing to serialize non-leaf tensor which requires_grad, "
+            "since autograd does not support crossing process boundaries.  "
+            "If you just want to transfer the data, call detach() on the tensor "
+            "before serializing (e.g., putting it on the queue)."
+        )
+
+    check_serializing_named_tensor(tensor)
+    torch.utils.hooks.warn_if_has_hooks(tensor)
+
+    # Note [CUDA IPC and the caching allocator]
+    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    # When you send a CUDA tensor over IPC, you might expect that you will
+    # get out the same storage from the other end.  However, the CUDA caching
+    # allocator makes it difficult to preserve this invariant.  Consider
+    # the following situation: a tensor of size 0x100 points to offset 0x20 of
+    # a storage at 0xA100 of size 0x100.  (For simplicity, all of these
+    # sizes are given in bytes).  HOWEVER, with the caching allocator, this storage
+    # might be part of a larger cudaMalloc allocation 0xA000 of size 0x4000.
+    #
+    # When we want to send this CUDA tensor over IPC, we must send the
+    # *entire* cudaMalloc allocation, i.e., the 0xA000 region, not just
+    # the storage 0xA100 (because that is what CUDA supports).  So, on the
+    # other end, there simply isn't any way to say, "Wait, you gave me
+    # a bigger region (0xA000) than the one I wanted (0xA100)".
+    #
+    # OK, so if you sent the cudaMalloc allocation, can you just wrap that up as
+    # one storage itself? No, because this cudaMalloc allocation might contain
+    # storages of mixed types: float, bytes, double... If you make the entire
+    # allocation a single storage of a type A, we'll hit an error when constructing
+    # a tensor of type B on the storage.
+    #
+    # cudaIpcMemHandle is an identifier to access the sender cudaMalloc allocation on the
+    # receiver side. However, cudaIpcMemHandles from each device in a given process may
+    # only be opened by one context per device per other process.
+    # If we open and close a memory handle multiples times in a process, CUDA is allowed
+    # to give it a different address; similarly, once we close the memory, we're not
+    # allowed to access it(and the storage/tensor built on top of it), even if it is
+    # still live in the original process. As we cannot make a cudaMalloc allocation
+    # to a single storage in one go, this requires us to cache the device pointer for
+    # each cudaIpcMemHandle on C++ side to reconstruct types of storages, while keep
+    # the old ones alives.
+    # See [https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html]
+    #
+    # This is fine, because all we need to do is to save our position in the allocation,
+    # and reconstruct storage and tensor from it.
+    # 0xA000 ->  -------CUDA Allocation------
+    #           |                            |
+    #           |                            |
+    #           |                            |
+    #           |                            |
+    # 0xA100 ->  --------storage1 begin------
+    #           |                            |
+    # 0xA120 ->  --------tensor1 begin ------
+    #           |                            |
+    #           |                            |
+    #           |                            |
+    #           |                            |
+    #           |                            |
+    # 0xA160 ->  --------tensor1 end---------
+    #           |                            |
+    #           |                            |
+    #           |                            |
+    # 0xA200 ->  --------storage1 end--------
+    #           |                            |
+    # 0xE000 ->  --------CUDA allocation-----
+    #
+    # To send tensor1, the following info are required from sender to receiver for
+    # storage reconstruction.
+    #   1. cudaIpcMemHandle of 0xA000(which can be mapped to a basePtr in receiver process).
+    #      basePtr may not be exactly 0xA000 since it's a different process.
+    #   2. offset(0xA100) of storage1 in the CUDA allocation.
+    #   3. size of storage1(0x100).
+    #
+    # On receiver side:
+    #   1. Get the devPtr of the MemHandle to access the memory, reconstruct a storage
+    #      of the same type using (basePtr, offset, size).
+    #   2. we can reconstruct the tensor on top of the reconstructed storage
+    #   Tensor(size=0x040, offset=0x020, storage=Storage(data=basePtr+0xA100, size=0x0100))
+    #
+    # This strategy has a few implications:
+    #
+    # 1. When we serialize a CUDA tensor for IPC, we cannot do it all in one
+    #    go (non-compositionally), and this requires to have a global map
+    #    memHandle -> devPtr for each process.
+    #
+    # 2. We MUST NOT let the new IPC tensor be resizable.  Originally, a resize
+    #    of the storage beyond 0x100 would merely have caused us to do a
+    #    reallocation.  You don't really want to do this, but if you did,
+    #    all that would happen is that you would lose IPC sharing.  But if
+    #    you do this in the new world, we will happily let you write out of
+    #    bounds of your "allocation", clobbering unrelated data in the cached
+    #    allocator block.  BAD!
+    #
+    # By the way, in old versions of PyTorch, we supported this situation
+    # natively using a "storage view", which permitted multiple storages to be
+    # views on each other.  But this was the *only* use of storage views, so we
+    # eliminated it so that we could just use tensor views to implement the same
+    # thing.
+    #
+
+    # TODO: Handle distinguishing between subclass and non-subclass versions of NT better
+    # https://github.com/pytorch/pytorch/issues/110543
+    from torch.nested._internal.nested_tensor import NestedTensor
+
+    if tensor.is_nested and not isinstance(tensor, NestedTensor):
+        return reduce_nested_tensor(tensor)
+
+    if tensor.layout in {
+        torch.sparse_coo,
+        torch.sparse_csr,
+        torch.sparse_bsr,
+        torch.sparse_csc,
+        torch.sparse_bsc,
+    }:
+        return reduce_sparse_tensor(tensor)
+
+    storage = tensor._typed_storage()
+
+    if storage._untyped_storage.device.type == "cuda":
+        (
+            device,
+            handle,
+            storage_size_bytes,
+            storage_offset_bytes,
+            ref_counter_handle,
+            ref_counter_offset,
+            event_handle,
+            event_sync_required,
+        ) = storage._share_cuda_()
+        tensor_offset = tensor.storage_offset()
+        shared_cache[handle] = StorageWeakRef(storage)
+        # _backward_hooks purposely omitted here, see
+        # Note [Don't serialize hooks]
+        return (
+            rebuild_cuda_tensor,
+            (
+                type(tensor),
+                tensor.size(),
+                tensor.stride(),
+                tensor_offset,  # tensor offset in its storage
+                type(storage),
+                tensor.dtype,
+                device,
+                handle,  # identifier which CUDA allocation is the storage in.
+                storage_size_bytes,  # size(in bytes) of the storage
+                storage_offset_bytes,  # offset(in bytes) of the storage in the CUDA allocation
+                tensor.requires_grad,
+                ref_counter_handle,
+                ref_counter_offset,
+                event_handle,
+                event_sync_required,
+            ),
+        )
+    elif storage._untyped_storage.device.type == "meta":
+        return (
+            rebuild_meta_tensor,
+            (
+                type(tensor),
+                tensor.size(),
+                tensor.stride(),
+                tensor.storage_offset(),
+                tensor.dtype,
+                tensor.untyped_storage().size(),
+                tensor.requires_grad,
+            ),
+        )
+
+    # _backward_hooks purposely omitted here, see Note [Don't serialize hooks]
+    metadata = (
+        tensor.storage_offset(),
+        tensor.size(),
+        tensor.stride(),
+        tensor.requires_grad,
+    )
+    return (rebuild_tensor, (type(tensor), storage, metadata))
+
+
+def rebuild_nested_tensor(
+    rebuild_buffer_func,
+    rebuild_buffer_args,
+    rebuild_sizes_func,
+    rebuild_sizes_args,
+    rebuild_strides_func,
+    rebuild_strides_args,
+    rebuild_offsets_func,
+    rebuild_offsets_args,
+):
+    buffer = rebuild_buffer_func(*rebuild_buffer_args)
+    sizes = rebuild_sizes_func(*rebuild_sizes_args)
+    strides = rebuild_strides_func(*rebuild_strides_args)
+    offsets = rebuild_offsets_func(*rebuild_offsets_args)
+    return torch._nested_view_from_buffer_copy(buffer, sizes, strides, offsets)
+
+
+def reduce_nested_tensor(nt):
+    rebuild_buffer_func, rebuild_buffer_args = reduce_tensor(nt.values())
+    rebuild_sizes_func, rebuild_sizes_args = reduce_tensor(nt._nested_tensor_size())
+    rebuild_strides_func, rebuild_strides_args = reduce_tensor(
+        nt._nested_tensor_strides()
+    )
+    rebuild_offsets_func, rebuild_offsets_args = reduce_tensor(
+        nt._nested_tensor_storage_offsets()
+    )
+
+    return (
+        rebuild_nested_tensor,
+        (
+            rebuild_buffer_func,
+            rebuild_buffer_args,
+            rebuild_sizes_func,
+            rebuild_sizes_args,
+            rebuild_strides_func,
+            rebuild_strides_args,
+            rebuild_offsets_func,
+            rebuild_offsets_args,
+        ),
+    )
+
+
+def rebuild_sparse_coo_tensor(
+    rebuild_indices_func,
+    rebuild_indices_args,
+    rebuild_values_func,
+    rebuild_values_args,
+    shape,
+    is_coalesced,
+):
+    indices = rebuild_indices_func(*rebuild_indices_args)
+    values = rebuild_values_func(*rebuild_values_args)
+    return torch.sparse_coo_tensor(indices, values, shape, is_coalesced=is_coalesced)
+
+
+def rebuild_sparse_compressed_tensor(
+    rebuild_compressed_indices_func,
+    rebuild_compressed_indices_args,
+    rebuild_plain_indices_func,
+    rebuild_plain_indices_args,
+    rebuild_values_func,
+    rebuild_values_args,
+    shape,
+    layout,
+):
+    compressed_indices = rebuild_compressed_indices_func(
+        *rebuild_compressed_indices_args
+    )
+    plain_indices = rebuild_plain_indices_func(*rebuild_plain_indices_args)
+    values = rebuild_values_func(*rebuild_values_args)
+    return torch.sparse_compressed_tensor(
+        compressed_indices, plain_indices, values, shape, layout=layout
+    )
+
+
+def reduce_sparse_tensor(sparse):
+    if sparse.layout is torch.sparse_coo:
+        rebuild_indices_func, rebuild_indices_args = reduce_tensor(sparse._indices())
+        rebuild_values_func, rebuild_values_args = reduce_tensor(sparse._values())
+        return (
+            rebuild_sparse_coo_tensor,
+            (
+                rebuild_indices_func,
+                rebuild_indices_args,
+                rebuild_values_func,
+                rebuild_values_args,
+                sparse.shape,
+                sparse.is_coalesced(),
+            ),
+        )
+    else:
+        if sparse.layout in {torch.sparse_csr, torch.sparse_bsr}:
+            compressed_indices = sparse.crow_indices()
+            plain_indices = sparse.col_indices()
+        elif sparse.layout in {torch.sparse_csc, torch.sparse_bsc}:
+            compressed_indices = sparse.ccol_indices()
+            plain_indices = sparse.row_indices()
+        else:
+            raise NotImplementedError(sparse.layout)
+        (
+            rebuild_compressed_indices_func,
+            rebuild_compressed_indices_args,
+        ) = reduce_tensor(compressed_indices)
+        rebuild_plain_indices_func, rebuild_plain_indices_args = reduce_tensor(
+            plain_indices
+        )
+        rebuild_values_func, rebuild_values_args = reduce_tensor(sparse.values())
+        return (
+            rebuild_sparse_compressed_tensor,
+            (
+                rebuild_compressed_indices_func,
+                rebuild_compressed_indices_args,
+                rebuild_plain_indices_func,
+                rebuild_plain_indices_args,
+                rebuild_values_func,
+                rebuild_values_args,
+                sparse.shape,
+                sparse.layout,
+            ),
+        )
+
+
+def fd_id(fd):
+    # Returns a tuple which uniquely identifies a file descriptor. In Mac OS,
+    # this doesn't work with shared memory handles, which is why we don't
+    # support the "file_descriptor" sharing method on that platform.
+    stat = os.fstat(fd)
+    return (stat.st_ino, stat.st_dev)
+
+
+def storage_from_cache(cls, key):
+    storage_ref = shared_cache.get(key)
+    if storage_ref is None:
+        return None
+    return torch.UntypedStorage._new_with_weak_ptr(storage_ref.cdata)
+
+
+def rebuild_storage_fd(cls, df, size):
+    fd = df.detach()
+    try:
+        storage = storage_from_cache(cls, fd_id(fd))
+        if storage is not None:
+            return storage
+        storage = cls._new_shared_fd_cpu(fd, size)
+        shared_cache[fd_id(fd)] = StorageWeakRef(storage)
+        return storage
+    finally:
+        os.close(fd)
+
+
+def rebuild_storage_filename(cls, manager, handle, size, dtype=None):
+    storage: torch.TypedStorage | torch.UntypedStorage = storage_from_cache(cls, handle)
+    if storage is not None:
+        return storage._shared_decref()
+    if dtype is None:
+        storage = torch.UntypedStorage._new_shared_filename_cpu(manager, handle, size)
+    else:
+        byte_size = size * torch._utils._element_size(dtype)
+        untyped_storage: torch.UntypedStorage = (
+            torch.UntypedStorage._new_shared_filename_cpu(manager, handle, byte_size)
+        )
+        storage = torch.TypedStorage(
+            wrap_storage=untyped_storage, dtype=dtype, _internal=True
+        )
+    shared_cache[handle] = StorageWeakRef(storage)
+    return storage._shared_decref()
+
+
+def rebuild_storage_empty(cls):
+    return cls()
+
+
+def rebuild_typed_storage(storage, dtype):
+    return torch.storage.TypedStorage(wrap_storage=storage, dtype=dtype, _internal=True)
+
+
+# Use for torch.storage.TypedStorage
+def reduce_typed_storage(storage):
+    return (rebuild_typed_storage, (storage._untyped_storage, storage.dtype))
+
+
+def rebuild_typed_storage_child(storage, storage_type):
+    return storage_type(wrap_storage=storage, _internal=True)
+
+
+# Use for child classes of torch.storage.TypedStorage, like torch.FloatStorage
+def reduce_typed_storage_child(storage):
+    return (rebuild_typed_storage_child, (storage._untyped_storage, type(storage)))
+
+
+def reduce_storage(storage):
+    from . import get_sharing_strategy
+
+    if storage.is_cuda:
+        raise RuntimeError(
+            "Cannot pickle CUDA storage; try pickling a CUDA tensor instead"
+        )
+    elif storage.device.type == "meta":
+        raise RuntimeError(
+            "Cannot pickle meta storage; try pickling a meta tensor instead"
+        )
+    elif get_sharing_strategy() == "file_system":
+        metadata = storage._share_filename_cpu_()
+        cache_key = metadata[1]
+        rebuild = rebuild_storage_filename
+        if isinstance(storage, torch.TypedStorage):
+            metadata += (storage.dtype,)
+        storage._shared_incref()
+    elif storage.size() == 0:
+        # This is special cased because Empty tensors
+        # (with size 0) cannot be mmapped.
+        return (rebuild_storage_empty, (type(storage),))
+    else:
+        fd, size = storage._share_fd_cpu_()
+        df = multiprocessing.reduction.DupFd(fd)
+        cache_key = fd_id(fd)
+        metadata = (df, size)
+        rebuild = rebuild_storage_fd  # type: ignore[assignment]
+
+    shared_cache[cache_key] = StorageWeakRef(storage)
+    return (rebuild, (type(storage),) + metadata)
+
+
+def init_reductions():
+    reduction.register(torch.cuda.Event, reduce_event)
+
+    for t in torch._storage_classes:
+        if t.__name__ == "UntypedStorage":
+            reduction.register(t, reduce_storage)
+        else:
+            reduction.register(t, reduce_typed_storage_child)
+
+    reduction.register(torch.storage.TypedStorage, reduce_typed_storage)
+
+    for t in torch._tensor_classes:
+        reduction.register(t, reduce_tensor)
+
+    # TODO: Maybe this should be in tensor_classes? :)
+    reduction.register(torch.Tensor, reduce_tensor)
+
+    from torch.nn.parameter import Parameter
+
+    reduction.register(Parameter, reduce_tensor)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/spawn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/spawn.py
new file mode 100644
index 0000000000000000000000000000000000000000..ecf8a43ad502683eb08306f5e5a95a641cee0595
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/multiprocessing/spawn.py
@@ -0,0 +1,340 @@
+# mypy: allow-untyped-defs
+import logging
+import multiprocessing
+import multiprocessing.connection
+import os
+import pickle
+import signal
+import sys
+import tempfile
+import time
+import warnings
+from concurrent.futures import as_completed, ThreadPoolExecutor
+
+from . import _prctl_pr_set_pdeathsig  # type: ignore[attr-defined]
+
+
+ENV_VAR_PARALLEL_START = "TORCH_MP_PARALLEL_START"
+
+log = logging.getLogger(__name__)
+
+__all__ = [
+    "ProcessContext",
+    "ProcessException",
+    "ProcessExitedException",
+    "ProcessRaisedException",
+    "spawn",
+    "SpawnContext",
+    "start_processes",
+]
+
+
+class ProcessException(Exception):
+    __slots__ = ["error_index", "error_pid"]
+
+    def __init__(self, msg: str, error_index: int, error_pid: int):
+        super().__init__(msg)
+        self.msg = msg
+        self.error_index = error_index
+        self.error_pid = error_pid
+
+    def __reduce__(self):
+        return type(self), (self.msg, self.error_index, self.error_pid)
+
+
+class ProcessRaisedException(ProcessException):
+    """Exception raised when a process failed due to an exception raised by the code."""
+
+
+class ProcessExitedException(ProcessException):
+    """Exception raised when a process failed due to signal or exited with a specific code."""
+
+    __slots__ = ["exit_code"]
+
+    def __init__(
+        self,
+        msg: str,
+        error_index: int,
+        error_pid: int,
+        exit_code: int,
+        signal_name: str | None = None,
+    ):
+        super().__init__(msg, error_index, error_pid)
+        self.exit_code = exit_code
+        self.signal_name = signal_name
+
+    def __reduce__(self):
+        return (
+            type(self),
+            (
+                self.msg,
+                self.error_index,
+                self.error_pid,
+                self.exit_code,
+                self.signal_name,
+            ),
+        )
+
+
+def _wrap(fn, i, args, error_file):
+    # prctl(2) is a Linux specific system call.
+    # On other systems the following function call has no effect.
+    # This is set to ensure that non-daemonic child processes can
+    # terminate if their parent terminates before they do.
+    _prctl_pr_set_pdeathsig(signal.SIGINT)
+
+    try:
+        fn(i, *args)
+    except KeyboardInterrupt:
+        pass  # SIGINT; Killed by parent, do nothing
+    except Exception:
+        # Propagate exception to parent process, keeping original traceback
+        import traceback
+
+        with open(error_file, "wb") as fh:
+            pickle.dump(traceback.format_exc(), fh)
+        sys.exit(1)
+
+
+class ProcessContext:
+    def __init__(self, processes, error_files):
+        self.error_files = error_files
+        self.processes = processes
+        self.sentinels = {
+            process.sentinel: index for index, process in enumerate(processes)
+        }
+
+    def pids(self):
+        return [int(process.pid) for process in self.processes]
+
+    def _join_procs_with_timeout(self, timeout: float):
+        """Attempt to join all processes with a shared timeout."""
+        end = time.monotonic() + timeout
+        for process in self.processes:
+            # pyrefly: ignore [no-matching-overload]
+            time_to_wait = max(0, end - time.monotonic())
+            process.join(time_to_wait)
+
+    def join(self, timeout: float | None = None, grace_period: float | None = None):
+        r"""Join one or more processes within spawn context.
+
+        Attempt to join one or more processes in this spawn context.
+        If one of them exited with a non-zero exit status, this function
+        kills the remaining processes (optionally with a grace period)
+        and raises an exception with the cause of the first process exiting.
+
+        Returns ``True`` if all processes have been joined successfully,
+        ``False`` if there are more processes that need to be joined.
+
+        Args:
+            timeout (float): Wait this long (in seconds) before giving up on waiting.
+            grace_period (float): When any processes fail, wait this long (in seconds)
+                for others to shutdown gracefully before terminating them. If they
+                still don't exit, wait another grace period before killing them.
+        """
+        # Ensure this function can be called even when we're done.
+        if len(self.sentinels) == 0:
+            return True
+
+        # Wait for any process to fail or all of them to succeed.
+        ready = multiprocessing.connection.wait(
+            self.sentinels.keys(),
+            timeout=timeout,
+        )
+
+        error_index = None
+        for sentinel in ready:
+            index = self.sentinels.pop(sentinel)
+            process = self.processes[index]
+            process.join()
+            if process.exitcode != 0:
+                error_index = index
+                break
+
+        # Return if there was no error.
+        if error_index is None:
+            # Return whether or not all processes have been joined.
+            return len(self.sentinels) == 0
+        # An error occurred. Clean-up all processes before returning.
+        # First, allow a grace period for processes to shutdown themselves.
+        if grace_period is not None:
+            self._join_procs_with_timeout(grace_period)
+        # Then, terminate processes that are still alive. Try SIGTERM first.
+        for process in self.processes:
+            if process.is_alive():
+                log.warning("Terminating process %s via signal SIGTERM", process.pid)
+                process.terminate()
+
+        # Try SIGKILL if the process isn't going down after another grace_period.
+        # The reason is related to python signal handling is limited
+        # to main thread and if that is in c/c++ land and stuck it won't
+        # to handle it. We have seen processes getting stuck not handling
+        # SIGTERM for the above reason.
+        self._join_procs_with_timeout(30 if grace_period is None else grace_period)
+        for process in self.processes:
+            if process.is_alive():
+                log.warning(
+                    "Unable to shutdown process %s via SIGTERM , forcefully exiting via SIGKILL",
+                    process.pid,
+                )
+                process.kill()
+            process.join()
+
+        # The file will only be created if the process crashed.
+        failed_process = self.processes[error_index]
+        if not os.access(self.error_files[error_index], os.R_OK):
+            exitcode = self.processes[error_index].exitcode
+            if exitcode < 0:
+                try:
+                    name = signal.Signals(-exitcode).name
+                except ValueError:
+                    name = f"<Unknown signal {-exitcode}>"
+                raise ProcessExitedException(
+                    f"process {error_index:d} terminated with signal {name}",
+                    error_index=error_index,
+                    error_pid=failed_process.pid,
+                    exit_code=exitcode,
+                    signal_name=name,
+                )
+            else:
+                raise ProcessExitedException(
+                    f"process {error_index:d} terminated with exit code {exitcode:d}",
+                    error_index=error_index,
+                    error_pid=failed_process.pid,
+                    exit_code=exitcode,
+                )
+
+        with open(self.error_files[error_index], "rb") as fh:
+            original_trace = pickle.load(fh)
+        msg = f"\n\n-- Process {error_index:d} terminated with the following error:\n"
+        msg += original_trace
+        raise ProcessRaisedException(msg, error_index, failed_process.pid)
+
+
+class SpawnContext(ProcessContext):
+    def __init__(self, processes, error_files):
+        warnings.warn(
+            "SpawnContext is renamed to ProcessContext since 1.4 release.", stacklevel=2
+        )
+        super().__init__(processes, error_files)
+
+
+# Note: [start_processes]
+# mp.start_processes handles both start_method='spawn' and 'fork'. It's supposed to be a
+# more generalized API than mp.spawn. Currently we only document mp.spawn as it's the
+# CUDA compatible start_method. However, in environments like Ipython notebooks, 'fork'
+# works better than 'spawn'. Every helper function we created for mp.spawn is indeed
+# general enough, and backends like XLA can reuse them in Colab notebooks as well.
+# Currently we only add this API first, we can consider adding it to documentation as
+# needed in the future.
+def start_processes(
+    fn,
+    args=(),
+    nprocs=1,
+    join=True,
+    daemon=False,
+    start_method="spawn",
+):
+    # To speed up performance in certain cases (see https://github.com/pytorch/pytorch/issues/133010),
+    # this func will start processes in parallel if start_method is 'forkserver'.
+    # Please opt in to this perf optimization by setting env var (TORCH_MP_PARALLEL_START) to 1.
+    # todo: investigate why spawn does not work with threadpool and raises SIGINT
+    if (
+        start_method == "forkserver"
+        and os.environ.get(ENV_VAR_PARALLEL_START, "0") == "1"
+    ):
+        log.info("Starting processes in parallel.")
+        start_parallel = True
+    else:
+        # Set env var TORCH_MP_PARALLEL_START to 0 to disable parallel start
+        start_parallel = False
+
+    mp = multiprocessing.get_context(start_method)
+    error_files = [None] * nprocs
+    processes = [None] * nprocs
+
+    def start_process(i):
+        # Each process is assigned a file to write tracebacks to.  We
+        # use the file being non-empty to indicate an exception
+        # occurred (vs an expected shutdown).  Note: this previously
+        # used a multiprocessing.Queue but that can be prone to
+        # deadlocks, so we went with a simpler solution for a one-shot
+        # message between processes.
+        tf = tempfile.NamedTemporaryFile(  # noqa: SIM115
+            prefix="pytorch-errorfile-", suffix=".pickle", delete=False
+        )
+        tf.close()
+        os.unlink(tf.name)
+
+        process = mp.Process(  # pyrefly: ignore  # missing-attribute
+            target=_wrap,
+            args=(fn, i, args, tf.name),
+            daemon=daemon,
+        )
+
+        process.start()
+        return i, process, tf.name
+
+    if not start_parallel:
+        for i in range(nprocs):
+            idx, process, tf_name = start_process(i)
+            error_files[idx] = tf_name
+            processes[idx] = process
+    else:
+        with ThreadPoolExecutor(max_workers=nprocs) as executor:
+            futures = [executor.submit(start_process, i) for i in range(nprocs)]
+            for fut in as_completed(futures):
+                idx, process, tf_name = fut.result()
+                # idx and process rank needs to be the same.
+                error_files[idx] = tf_name
+                processes[idx] = process
+    context = ProcessContext(processes, error_files)
+    if not join:
+        return context
+
+    # Loop on join until it returns True or raises an exception.
+    while not context.join():
+        pass
+
+
+def spawn(fn, args=(), nprocs=1, join=True, daemon=False, start_method="spawn"):
+    r"""Spawns ``nprocs`` processes that run ``fn`` with ``args``.
+
+    If one of the processes exits with a non-zero exit status, the
+    remaining processes are killed and an exception is raised with the
+    cause of termination. In the case an exception was caught in the
+    child process, it is forwarded and its traceback is included in
+    the exception raised in the parent process.
+
+    Args:
+        fn (function): Function is called as the entrypoint of the
+            spawned process. This function must be defined at the top
+            level of a module so it can be pickled and spawned. This
+            is a requirement imposed by multiprocessing.
+
+            The function is called as ``fn(i, *args)``, where ``i`` is
+            the process index and ``args`` is the passed through tuple
+            of arguments.
+
+        args (tuple): Arguments passed to ``fn``.
+        nprocs (int): Number of processes to spawn.
+        join (bool): Perform a blocking join on all processes.
+        daemon (bool): The spawned processes' daemon flag. If set to True,
+                       daemonic processes will be created.
+        start_method (str): (deprecated) this method will always use ``spawn``
+                               as the start method. To use a different start method
+                               use ``start_processes()``.
+
+    Returns:
+        None if ``join`` is ``True``,
+        :class:`~ProcessContext` if ``join`` is ``False``
+
+    """
+    if start_method != "spawn":
+        msg = (
+            f"This method only supports start_method=spawn (got: {start_method}).\n"
+            "To use a different start_method use:\n\t\t"
+            " torch.multiprocessing.start_processes(...)"
+        )
+        warnings.warn(msg, FutureWarning, stacklevel=2)
+    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/_lower_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/_lower_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..548329757b346ba25d6e0ab9f22e5b45e17f79b9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/_lower_utils.py
@@ -0,0 +1,104 @@
+import types
+
+import torch
+import torch.utils._pytree as pytree
+from torch.export import ExportedProgram
+from torch.export.pt2_archive._package import AOTI_FILES, package_pt2
+from torch.types import FileLike
+
+from ._lowered_aoti_module import LoweredBackendModule
+
+
+def get_new_ep_with_flat_inputs_outputs(ep: ExportedProgram) -> ExportedProgram:
+    class FlattenedModule(torch.nn.Module):
+        def __init__(
+            self,
+            original_module: torch.fx.GraphModule,
+            in_spec: pytree.TreeSpec,
+            out_spec: pytree.TreeSpec,
+        ) -> None:
+            super().__init__()
+            self.original_module = original_module
+            self.in_spec = in_spec
+            self.out_spec = out_spec
+
+        def forward(self, *flat_inputs):  # type: ignore[no-untyped-def]
+            # Unflatten inputs to original structure
+            inputs = pytree.tree_unflatten(flat_inputs, self.in_spec)
+            args, kwargs = inputs
+            outputs = self.original_module(*args, **kwargs)
+            # Flatten outputs
+            flat_outputs, _ = pytree.tree_flatten(outputs)
+            return tuple(flat_outputs)
+
+    flattened_module = FlattenedModule(
+        ep.module(), ep.call_spec.in_spec, ep.call_spec.out_spec
+    )
+    args, kwargs = ep.example_inputs
+    flat_inputs, _ = pytree.tree_flatten((args, kwargs))
+    flat_ep = torch.export.export(flattened_module, tuple(flat_inputs))
+
+    return flat_ep
+
+
+def lower_exported_program(
+    exported_program: ExportedProgram, model_name: str, backend_id: str
+) -> tuple[ExportedProgram, AOTI_FILES]:
+    """
+    Lower an exported program to AOTInductor and return a delegate ExportedProgram
+    with the `executorch_call_delegate` HOP
+    """
+    args, kwargs = exported_program.example_inputs
+    out_spec = exported_program.call_spec.out_spec
+    flat_ep = get_new_ep_with_flat_inputs_outputs(exported_program)
+    flat_inputs, _ = pytree.tree_flatten((args, kwargs))
+
+    aoti_files = torch._inductor.aot_compile(
+        flat_ep.module(), tuple(flat_inputs), options={"aot_inductor.package": True}
+    )
+    if not isinstance(aoti_files, list):
+        raise AssertionError(
+            f"aoti_files must be a list, got {type(aoti_files).__name__}"
+        )
+
+    lowered_aoti_module = LoweredBackendModule(
+        flat_ep, backend_id, module_name=model_name
+    )
+
+    def patched_forward(self, *args, **kwargs):  # type: ignore[no-untyped-def]
+        flat_inputs, _ = pytree.tree_flatten((args, kwargs))
+        flat_outputs = torch._higher_order_ops.executorch_call_delegate(
+            self, *flat_inputs
+        )
+        if out_spec is not None and flat_outputs is not None:
+            return pytree.tree_unflatten(flat_outputs, out_spec)
+        else:
+            return flat_outputs
+
+    lowered_aoti_module.forward = types.MethodType(patched_forward, lowered_aoti_module)  # type: ignore[method-assign]
+
+    aoti_delegate_ep = torch.export.export(lowered_aoti_module, args, kwargs)
+
+    return aoti_delegate_ep, aoti_files
+
+
+def package_nativert_with_aoti_delegate(
+    f: FileLike,
+    model_name: str,
+    backend_id: str,
+    original_ep: ExportedProgram,
+    delegate_ep: ExportedProgram,
+    delegate_files: AOTI_FILES,
+) -> None:
+    """
+    Package a pt2 archive file that can be consumed by NativeRT with AOTI Delegate
+    """
+    package_pt2(
+        f,
+        exported_programs={
+            model_name: original_ep,
+            f"{model_name}-{backend_id}": delegate_ep,
+        },
+        aoti_files={f"{model_name}-{backend_id}": delegate_files},  # type: ignore[dict-item]
+    )
+    return
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/_lowered_aoti_module.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/_lowered_aoti_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..c08e83211ef330be11788b5ca82a1dcc9a0c9f9d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nativert/backends/_lowered_aoti_module.py
@@ -0,0 +1,31 @@
+import torch
+from torch.export import ExportedProgram
+
+
+class LoweredBackendModule(torch.nn.Module):
+    def __init__(
+        self,
+        original_exported_program: ExportedProgram,
+        backend_id: str,
+        *,
+        module_name: str | None = None,
+    ) -> None:
+        super().__init__()
+        self._backend_id = backend_id
+        self._module_name = module_name
+        self._original_exported_program = original_exported_program
+
+    @property
+    def backend_id(self) -> str:
+        return self._backend_id
+
+    @property
+    def module_name(self) -> str | None:
+        return self._module_name
+
+    @property
+    def original_module(self) -> ExportedProgram:
+        return self._original_exported_program
+
+    def forward(self, *args, **kwargs):  # type: ignore[no-untyped-def]
+        return torch._higher_order_ops.executorch_call_delegate(self, *args, **kwargs)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2782c13769449bb1f99f7d317a26c42701b11833
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/__init__.py
@@ -0,0 +1,524 @@
+# mypy: allow-untyped-defs
+from typing import Optional, Union
+
+import torch
+import torch.nn.functional as F
+from torch import SymInt, Tensor
+from torch._C import _add_docstr, _nested  # type: ignore[attr-defined]
+from torch.types import _device as Device, _dtype as DType
+
+
+__all__ = [
+    "to_padded_tensor",
+    "as_nested_tensor",
+    "nested_tensor",
+    "nested_tensor_from_jagged",
+    "narrow",
+    "masked_select",
+]
+
+# Allowlist these for weights_only load of NJT
+from ._internal.nested_tensor import _rebuild_njt, NestedTensor as _NestedTensor
+
+
+torch.serialization.add_safe_globals([_NestedTensor, _rebuild_njt])
+
+
+def as_nested_tensor(
+    ts: Tensor | list[Tensor] | tuple[Tensor, ...],
+    dtype: DType | None = None,
+    device: Device | None = None,
+    layout=None,
+) -> Tensor:
+    r"""
+    Constructs a nested tensor preserving autograd history from a tensor or a list / tuple of
+    tensors.
+
+    If a nested tensor is passed, it will be returned directly unless the device / dtype / layout
+    differ. Note that converting device / dtype will result in a copy, while converting layout
+    is not currently supported by this function.
+
+    If a non-nested tensor is passed, it is treated as a batch of constituents of consistent size.
+    A copy will be incurred if the passed device / dtype differ from those of the input OR if
+    the input is non-contiguous. Otherwise, the input's storage will be used directly.
+
+    If a tensor list is provided, tensors in the list are always copied during construction of
+    the nested tensor.
+
+    Args:
+        ts (Tensor or List[Tensor] or Tuple[Tensor]): a tensor to treat as a nested tensor OR a
+            list / tuple of tensors with the same ndim
+
+    Keyword arguments:
+        dtype (:class:`torch.dtype`, optional): the desired type of returned nested tensor.
+            Default: if None, same :class:`torch.dtype` as leftmost tensor in the list.
+        device (:class:`torch.device`, optional): the desired device of returned nested tensor.
+            Default: if None, same :class:`torch.device` as leftmost tensor in the list
+        layout (:class:`torch.layout`, optional): the desired layout of returned nested tensor.
+            Only strided and jagged layouts are supported. Default: if None, the strided layout.
+
+    Example::
+
+        >>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
+        >>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
+        >>> nt = torch.nested.as_nested_tensor([a, b])
+        >>> nt.is_leaf
+        False
+        >>> fake_grad = torch.nested.nested_tensor([torch.ones_like(a), torch.zeros_like(b)])
+        >>> nt.backward(fake_grad)
+        >>> a.grad
+        tensor([1., 1., 1.])
+        >>> b.grad
+        tensor([0., 0., 0., 0., 0.])
+        >>> c = torch.randn(3, 5, requires_grad=True)
+        >>> nt2 = torch.nested.as_nested_tensor(c)
+    """
+    is_tensor_list = isinstance(ts, (list, tuple)) and all(
+        isinstance(t, Tensor) for t in ts
+    )
+    if not isinstance(ts, Tensor) and not is_tensor_list:
+        raise TypeError(
+            "as_nested_tensor(): Expected first argument to be a tensor or a list / tuple of tensors "
+        )
+    # convert tuple -> list if needed
+    if is_tensor_list and not isinstance(ts, list):
+        ts = list(ts)
+
+    if isinstance(ts, Tensor) and ts.dim() < 2:
+        raise RuntimeError(
+            "as_nested_tensor(): Expected tensor argument to have dim() > 1"
+        )
+
+    if isinstance(ts, Tensor) and ts.is_nested:
+        if layout == ts.layout:
+            # return input directly or input copied to device / dtype
+            return ts.to(device=device, dtype=dtype)
+        else:
+            # TODO: Just use nt.to(layout=layout) when it exists.
+            raise RuntimeError(
+                "as_nested_tensor(): Converting between nested tensor layouts is not supported"
+            )
+
+    if layout is None:
+        layout = torch.strided
+    if layout == torch.strided:
+        if isinstance(ts, Tensor):
+            # contiguous() might be necessary to get flattened view.
+            # we could probably be more precise about when to do this as an optimization
+            buffer = ts.contiguous().view(-1).to(device=device, dtype=dtype)
+            nested_sizes = torch.tensor([t.shape for t in ts])
+            return torch._nested_view_from_buffer(
+                buffer,
+                nested_sizes,
+                *torch._nested_compute_contiguous_strides_offsets(nested_sizes),
+            )
+        else:
+            if not isinstance(ts, list):
+                raise AssertionError(
+                    f"Expected ts to be a list, but got {type(ts).__name__}"
+                )
+            return torch._nested_tensor_from_tensor_list(ts, dtype, None, device, None)
+    elif layout == torch.jagged:
+        if isinstance(ts, Tensor):
+            if device is None:
+                device = ts.device
+
+            # contiguous() might be necessary to get flattened view.
+            # we could probably be more precise about when to do this as an optimization
+            values = ts.contiguous().flatten(0, 1).to(device=device, dtype=dtype)
+            batch_size = ts.shape[0]
+            seq_len = ts.shape[1]
+            offsets = torch.arange(
+                0, batch_size * seq_len + 1, seq_len, device=device, dtype=torch.int64
+            )
+
+            from torch.nested._internal.nested_tensor import (
+                nested_view_from_values_offsets,
+            )
+
+            return nested_view_from_values_offsets(
+                values, offsets, min_seqlen=seq_len, max_seqlen=seq_len
+            )
+        else:
+            from torch.nested._internal.nested_tensor import jagged_from_list
+
+            if not isinstance(ts, list):
+                raise AssertionError(
+                    f"Expected ts to be a list, but got {type(ts).__name__}"
+                )
+            nt, _ = jagged_from_list(ts, offsets=None, device=device, dtype=dtype)
+            return nt
+    else:
+        raise RuntimeError(
+            f"Specified layout is unsupported for nested tensors: {layout}"
+        )
+
+
+# Note: This not only adds doc strings for the nested ops, but
+# also connects the torch.nested Python namespace to the torch._C._nested builtins.
+
+to_padded_tensor = _add_docstr(
+    _nested.nested_to_padded_tensor,
+    r"""
+to_padded_tensor(input, padding, output_size=None, out=None) -> Tensor
+
+Returns a new (non-nested) Tensor by padding the :attr:`input` nested tensor.
+The leading entries will be filled with the nested data,
+while the trailing entries will be padded.
+
+.. warning::
+
+    :func:`to_padded_tensor` always copies the underlying data,
+    since the nested and the non-nested tensors differ in memory layout.
+
+Args:
+    padding (float): The padding value for the trailing entries.
+
+Keyword args:
+    output_size (Tuple[int]): The size of the output tensor.
+                              If given, it must be large enough to contain all nested data;
+                              else, will infer by taking the max size of each nested sub-tensor along each dimension.
+    out (Tensor, optional): the output tensor.
+
+Example::
+
+    >>> nt = torch.nested.nested_tensor([torch.randn((2, 5)), torch.randn((3, 4))])
+    nested_tensor([
+      tensor([[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
+              [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995]]),
+      tensor([[-1.8546, -0.7194, -0.2918, -0.1846],
+              [ 0.2773,  0.8793, -0.5183, -0.6447],
+              [ 1.8009,  1.8468, -0.9832, -1.5272]])
+    ])
+    >>> pt_infer = torch.nested.to_padded_tensor(nt, 0.0)
+    tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
+             [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995],
+             [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],
+            [[-1.8546, -0.7194, -0.2918, -0.1846,  0.0000],
+             [ 0.2773,  0.8793, -0.5183, -0.6447,  0.0000],
+             [ 1.8009,  1.8468, -0.9832, -1.5272,  0.0000]]])
+    >>> pt_large = torch.nested.to_padded_tensor(nt, 1.0, (2, 4, 6))
+    tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276,  1.0000],
+             [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995,  1.0000],
+             [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000],
+             [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]],
+            [[-1.8546, -0.7194, -0.2918, -0.1846,  1.0000,  1.0000],
+             [ 0.2773,  0.8793, -0.5183, -0.6447,  1.0000,  1.0000],
+             [ 1.8009,  1.8468, -0.9832, -1.5272,  1.0000,  1.0000],
+             [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]]])
+    >>> pt_small = torch.nested.to_padded_tensor(nt, 2.0, (2, 2, 2))
+    RuntimeError: Value in output_size is less than NestedTensor padded size. Truncation is not supported.
+
+""",
+)
+
+
+def nested_tensor(
+    tensor_list,
+    *,
+    dtype=None,
+    layout=None,
+    device=None,
+    requires_grad=False,
+    pin_memory=False,
+) -> Tensor:
+    r"""
+    Constructs a nested tensor with no autograd history (also known as a "leaf tensor", see
+    :ref:`Autograd mechanics <autograd-mechanics>`) from :attr:`tensor_list` a list of tensors.
+
+    Args:
+        tensor_list (List[array_like]): a list of tensors, or anything that can be passed to torch.tensor,
+        where each element of the list has the same dimensionality.
+
+    Keyword arguments:
+        dtype (:class:`torch.dtype`, optional): the desired type of returned nested tensor.
+            Default: if None, same :class:`torch.dtype` as leftmost tensor in the list.
+        layout (:class:`torch.layout`, optional): the desired layout of returned nested tensor.
+            Only strided and jagged layouts are supported. Default: if None, the strided layout.
+        device (:class:`torch.device`, optional): the desired device of returned nested tensor.
+            Default: if None, same :class:`torch.device` as leftmost tensor in the list
+        requires_grad (bool, optional): If autograd should record operations on the
+            returned nested tensor. Default: ``False``.
+        pin_memory (bool, optional): If set, returned nested tensor would be allocated in
+            the pinned memory. Works only for CPU tensors. Default: ``False``.
+
+    Example::
+
+        >>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
+        >>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
+        >>> nt = torch.nested.nested_tensor([a, b], requires_grad=True)
+        >>> nt.is_leaf
+        True
+    """
+    if layout is None:
+        layout = torch.strided
+    if layout == torch.strided:
+        return _nested.nested_tensor(
+            tensor_list,
+            dtype=dtype,
+            device=device,
+            requires_grad=requires_grad,
+            pin_memory=pin_memory,
+        )
+    elif layout == torch.jagged:
+        # Need to wrap lists of scalars as tensors
+        list_of_tensors = [
+            t if isinstance(t, Tensor) else torch.as_tensor(t) for t in tensor_list
+        ]
+
+        from torch.nested._internal.nested_tensor import jagged_from_list
+
+        with torch.no_grad():
+            nt, _ = jagged_from_list(
+                list_of_tensors, offsets=None, device=device, dtype=dtype
+            )
+
+        nt.requires_grad_(requires_grad)
+        if pin_memory:
+            nt = nt.pin_memory()  # type: ignore[assignment]
+
+        return nt
+    else:
+        raise RuntimeError(
+            f"Specified layout is unsupported for nested tensors: {layout}"
+        )
+
+
+def narrow(
+    tensor: Tensor,
+    dim: int,
+    start: int | Tensor,
+    length: int | Tensor,
+    layout=torch.strided,
+) -> Tensor:
+    r"""
+    Constructs a nested tensor (which might be a view) from :attr:`tensor`, a strided tensor. This follows
+    similar semantics to torch.Tensor.narrow, where in the :attr:`dim`-th dimension the new nested tensor
+    shows only the elements in the interval `[start, start+length)`. As nested representations
+    allow for a different `start` and `length` at each 'row' of that dimension, :attr:`start` and :attr:`length`
+    can also be tensors of shape `tensor.shape[0]`.
+
+    There's some differences depending on the layout you use for the nested tensor. If using strided layout,
+    torch.narrow will do a copy of the narrowed data into a contiguous NT with strided layout, while
+    jagged layout narrow() will create a non-contiguous view of your original strided tensor. This particular
+    representation is really useful for representing kv-caches in Transformer models, as specialized
+    SDPA kernels can deal with format easily, resulting in performance improvements.
+
+
+    Args:
+        tensor (:class:`torch.Tensor`): a strided tensor, which will be used as the underlying data
+            for the nested tensor if using the jagged layout or will be copied for the strided layout.
+        dim (int): the dimension where narrow will be applied. Only `dim=1` is supported for the
+            jagged layout, while strided supports all dim
+        start (Union[int, :class:`torch.Tensor`]): starting element for the narrow operation
+        length (Union[int, :class:`torch.Tensor`]): number of elements taken during the narrow op
+
+    Keyword arguments:
+        layout (:class:`torch.layout`, optional): the desired layout of returned nested tensor.
+            Only strided and jagged layouts are supported. Default: if None, the strided layout.
+
+    Example::
+
+        >>> starts = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64)
+        >>> lengths = torch.tensor([3, 2, 2, 1, 5], dtype=torch.int64)
+        >>> narrow_base = torch.randn(5, 10, 20)
+        >>> nt_narrowed = torch.nested.narrow(narrow_base, 1, starts, lengths, layout=torch.jagged)
+        >>> nt_narrowed.is_contiguous()
+        False
+    """
+    if not isinstance(start, (int, SymInt, Tensor)):
+        raise RuntimeError("start must be an integer or a tensor")
+
+    if not isinstance(length, (int, SymInt, Tensor)):
+        raise RuntimeError("length must be an integer or a tensor")
+
+    if layout == torch.strided:
+        if isinstance(start, Tensor) or isinstance(length, Tensor):
+            raise RuntimeError(
+                "start and length must be integers for the strided layout NT impl"
+            )
+        # TODO: switch to as_nested_tensor(tensor) when it is available
+        nt = as_nested_tensor(torch.unbind(tensor), layout=torch.strided).narrow(
+            dim, start, length
+        )
+    elif layout == torch.jagged:
+        if dim != 1:
+            raise RuntimeError("jagged layout only supports dim=1")
+
+        from torch.nested._internal.nested_tensor import jagged_from_tensor_and_lengths
+
+        if isinstance(start, (int, SymInt)):
+            start = torch.tensor([start], device=tensor.device, dtype=torch.int64)
+
+        if isinstance(length, (int, SymInt)):
+            length = torch.tensor([length], device=tensor.device, dtype=torch.int64)
+
+        nt, _, _ = jagged_from_tensor_and_lengths(tensor, start, length)
+    else:
+        raise RuntimeError(
+            f"Specified layout is unsupported for nested narrow: {layout}"
+        )
+
+    return nt
+
+
+def nested_tensor_from_jagged(
+    values: Tensor,
+    offsets: Tensor | None = None,
+    lengths: Tensor | None = None,
+    jagged_dim: int | None = None,
+    min_seqlen: int | None = None,
+    max_seqlen: int | None = None,
+) -> Tensor:
+    r"""
+    Constructs a jagged layout nested tensor from the given jagged components. The jagged layout
+    consists of a required values buffer with the jagged dimension packed into a single dimension.
+    The offsets / lengths metadata determines how this dimension is split into batch elements
+    and are expected to be allocated on the same device as the values buffer.
+
+    Expected metadata formats:
+        * offsets: Indices within the packed dimension splitting it into heterogeneously-sized
+          batch elements. Example: [0, 2, 3, 6] indicates that a packed jagged dim of size 6
+          should be conceptually split into batch elements of length [2, 1, 3]. Note that both the
+          beginning and ending offsets are required for kernel convenience (i.e. shape batch_size + 1).
+        * lengths: Lengths of the individual batch elements; shape == batch_size. Example: [2, 1, 3]
+          indicates that a packed jagged dim of size 6 should be conceptually split into batch
+          elements of length [2, 1, 3].
+
+    Note that it can be useful to provide both offsets and lengths. This describes a nested tensor
+    with "holes", where the offsets indicate the start position of each batch item and the length
+    specifies the total number of elements (see example below).
+
+    The returned jagged layout nested tensor will be a view of the input values tensor.
+
+    Args:
+        values (:class:`torch.Tensor`): The underlying buffer in the shape of
+            (sum_B(*), D_1, ..., D_N). The jagged dimension is packed into a single dimension,
+            with the offsets / lengths metadata used to distinguish batch elements.
+        offsets (optional :class:`torch.Tensor`): Offsets into the jagged dimension of shape B + 1.
+        lengths (optional :class:`torch.Tensor`): Lengths of the batch elements of shape B.
+        jagged_dim (optional int): Indicates which dimension in values is the packed jagged
+            dimension. Must be >= 1 as the batch dimension (dim=0) cannot be ragged.
+            If None, this is set to dim=1 (i.e. the dimension immediately following the batch dimension). Default: None
+        min_seqlen (optional int): If set, uses the specified value as the cached minimum sequence
+            length for the returned nested tensor. This can be a useful alternative to computing
+            this value on-demand, possibly avoiding a GPU -> CPU sync. Default: None
+        max_seqlen (optional int): If set, uses the specified value as the cached maximum sequence
+            length for the returned nested tensor. This can be a useful alternative to computing
+            this value on-demand, possibly avoiding a GPU -> CPU sync. Default: None
+
+    Example::
+
+        >>> values = torch.randn(12, 5)
+        >>> offsets = torch.tensor([0, 3, 5, 6, 10, 12])
+        >>> nt = nested_tensor_from_jagged(values, offsets)
+        >>> # 3D shape with the middle dimension jagged
+        >>> nt.shape
+        torch.Size([5, j2, 5])
+        >>> # Length of each item in the batch:
+        >>> offsets.diff()
+        tensor([3, 2, 1, 4, 2])
+
+        >>> values = torch.randn(6, 5)
+        >>> offsets = torch.tensor([0, 2, 3, 6])
+        >>> lengths = torch.tensor([1, 1, 2])
+        >>> # NT with holes
+        >>> nt = nested_tensor_from_jagged(values, offsets, lengths)
+        >>> a, b, c = nt.unbind()
+        >>> # Batch item 1 consists of indices [0, 1)
+        >>> torch.equal(a, values[0:1, :])
+        True
+        >>> # Batch item 2 consists of indices [2, 3)
+        >>> torch.equal(b, values[2:3, :])
+        True
+        >>> # Batch item 3 consists of indices [3, 5)
+        >>> torch.equal(c, values[3:5, :])
+        True
+    """
+    from torch.fx._symbolic_trace import is_fx_tracing
+
+    if is_fx_tracing():
+        raise RuntimeError(
+            "torch.nested.nested_tensor_from_jagged does not support tracing with fx.symbolic_trace. "
+            "Use fx.wrap to wrap the function that calls nested_tensor_from_jagged."
+        )
+
+    if offsets is None:
+        if lengths is None:
+            raise RuntimeError(
+                "nested_tensor_from_jagged(): At least one of offsets or lengths is required."
+            )
+        else:
+            # TODO: Truly support offsets=None at some point?
+            # For now, just convert lengths -> offsets for kernel convenience
+            offsets = F.pad(lengths.cumsum(0), (1, 0))
+            lengths = None
+
+    if jagged_dim is None:
+        jagged_dim = 1
+    elif jagged_dim < 1:
+        raise ValueError(f"Expected jagged_dim >=1, but got {jagged_dim}.")
+
+    from torch.nested._internal.nested_tensor import (
+        nested_view_from_values_offsets_lengths,
+    )
+
+    return nested_view_from_values_offsets_lengths(
+        values,
+        offsets,
+        lengths,
+        ragged_idx=jagged_dim,
+        min_seqlen=min_seqlen,
+        max_seqlen=max_seqlen,
+    )
+
+
+def masked_select(tensor: Tensor, mask: Tensor) -> Tensor:
+    r"""
+    Constructs a nested tensor given a strided tensor input and a strided mask, the resulting jagged layout nested tensor
+    will have values retain values where the mask is equal to True. The dimensionality of the mask is preserved and is
+    represented with the offsets, this is unlike :func:`masked_select` where the output is collapsed to a 1D tensor.
+
+    Args:
+    tensor (:class:`torch.Tensor`): a strided tensor from which the jagged layout nested tensor is constructed from.
+    mask (:class:`torch.Tensor`): a strided mask tensor which is applied to the tensor input
+
+    Example::
+
+        >>> tensor = torch.randn(3, 3)
+        >>> mask = torch.tensor([[False, False, True], [True, False, True], [False, False, True]])
+        >>> nt = torch.nested.masked_select(tensor, mask)
+        >>> nt.shape
+        torch.Size([3, j4])
+        >>> # Length of each item in the batch:
+        >>> nt.offsets().diff()
+        tensor([1, 2, 1])
+
+        >>> tensor = torch.randn(6, 5)
+        >>> mask = torch.tensor([False])
+        >>> nt = torch.nested.masked_select(tensor, mask)
+        >>> nt.shape
+        torch.Size([6, j5])
+        >>> # Length of each item in the batch:
+        >>> nt.offsets().diff()
+        tensor([0, 0, 0, 0, 0, 0])
+    """
+    if tensor.layout != torch.strided:
+        raise RuntimeError(
+            f"torch.nested.masked_select requires a strided tensor, given {tensor.layout}"
+        )
+
+    if mask.layout != torch.strided:
+        raise RuntimeError(
+            f"torch.nested.masked_select requires a strided mask, given: {mask.layout}"
+        )
+    res_values = tensor.masked_select(mask)
+    expanded_mask = mask.expand(tensor.shape)
+    res_lengths = expanded_mask.sum(dim=tensor.ndim - 1).view(-1)
+
+    from torch.nested._internal.nested_tensor import nested_view_from_values_offsets
+
+    return nested_view_from_values_offsets(
+        values=res_values,
+        offsets=F.pad(res_lengths.cumsum(dim=0), (1, 0)),
+    )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/nested_int.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/nested_int.py
new file mode 100644
index 0000000000000000000000000000000000000000..851fe2e11546d02e86cd63efdedf71525af4ce39
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/nested_int.py
@@ -0,0 +1,119 @@
+from typing import *  # noqa: F403
+
+import torch
+from torch.fx.experimental._constant_symnode import ConstantIntNode
+
+
+__all__ = ["NestedIntNode"]
+
+
+# Python version of aten/src/ATen/core/NestedIntSymNodeImpl.cpp
+def _eq(lhs: Any, rhs: Any) -> bool:
+    return (
+        isinstance(lhs, NestedIntNode)
+        and isinstance(rhs, NestedIntNode)
+        and lhs.t_id == rhs.t_id
+        and lhs.coeff == rhs.coeff
+    )
+
+
+def _ge(lhs: Any, rhs: Any) -> bool:
+    if isinstance(rhs, NestedIntNode) and isinstance(lhs, NestedIntNode):
+        if lhs.t_id == rhs.t_id:
+            return lhs.coeff >= rhs.coeff
+        raise ValueError("ge: relation is indeterminate")
+    elif isinstance(lhs, NestedIntNode):
+        if rhs.is_constant() and rhs.constant_int() <= 2:
+            return True
+        raise ValueError("ge: relation is indeterminate")
+    elif isinstance(rhs, NestedIntNode):
+        if lhs.is_constant() and lhs.constant_int() < 2:
+            return False
+        raise ValueError("ge: relation is indeterminate")
+    else:
+        raise ValueError("inputs unsupported")
+
+
+class NestedIntNode:
+    def __init__(self, t_id: int, coeff: int) -> None:
+        self.t_id = t_id
+        self.coeff = coeff
+
+    def nested_int_coeff(self) -> int:
+        return self.coeff
+
+    def maybe_as_int(self) -> Optional[int]:
+        return None
+
+    def is_int(self) -> bool:
+        return True
+
+    def is_float(self) -> bool:
+        return False
+
+    def is_bool(self) -> bool:
+        return False
+
+    def is_nested_int(self) -> bool:
+        return True
+
+    def clone(self) -> "NestedIntNode":
+        return self
+
+    def _str(self) -> Any:
+        if self.coeff == 1:
+            return f"j{self.t_id}"
+        return f"{self.coeff}*j{self.t_id}"
+
+    def str(self) -> Any:
+        return self._str()
+
+    def __str__(self) -> Any:
+        return self._str()
+
+    def __repr__(self) -> Any:
+        return self._str()
+
+    def _graph_repr(self) -> Any:
+        return self._str()
+
+    def mul(self, other: Any) -> "NestedIntNode":
+        if other.is_constant():
+            other = other.constant_int()
+        else:
+            raise ValueError(f"unsupported: {type(other)}")
+        return NestedIntNode(self.t_id, self.coeff * other)
+
+    def eq(self, other: Any) -> Any:
+        return torch._C._get_constant_bool_symnode(_eq(self, other))
+
+    def ne(self, other: Any) -> Any:
+        return torch._C._get_constant_bool_symnode(not _eq(self, other))
+
+    def gt(self, other: Any) -> Any:
+        return torch._C._get_constant_bool_symnode(not _ge(other, self))
+
+    def lt(self, other: Any) -> Any:
+        return torch._C._get_constant_bool_symnode(not _ge(self, other))
+
+    def le(self, other: Any) -> Any:
+        return torch._C._get_constant_bool_symnode(_ge(other, self))
+
+    def ge(self, other: Any) -> Any:
+        return torch._C._get_constant_bool_symnode(_ge(self, other))
+
+    def is_symbolic(self) -> bool:
+        return False
+
+    def nested_int(self) -> int:
+        return self.t_id
+
+    def is_constant(self) -> bool:
+        return False
+
+    def wrap_int(self, num: int) -> ConstantIntNode:
+        if type(num) is not int:
+            raise AssertionError(
+                f"Expected num to be int, but got {type(num).__name__}"
+            )
+        return ConstantIntNode(num)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/nested_tensor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/nested_tensor.py
new file mode 100644
index 0000000000000000000000000000000000000000..718c5c53d44116ef5637ecdc2598b0620ba7b33e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/nested_tensor.py
@@ -0,0 +1,692 @@
+# mypy: allow-untyped-defs
+from typing import *  # noqa: F403
+
+import torch
+from torch._C import DispatchKey, DispatchKeySet
+from torch._prims_common import is_expandable_to
+from torch.nested._internal.nested_int import NestedIntNode
+from torch.utils.weak import WeakTensorKeyDictionary
+
+
+_tensor_id_counter = 0
+_tensor_symint_registry = WeakTensorKeyDictionary()
+
+
+def get_tensor_symint(tensor, *, coeff=1):
+    from torch._subclasses.fake_tensor import FakeTensor
+    from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor
+
+    # NB: Only FakeTensor is associated with a memo
+    tensor = mb_unwrap_functional_tensor(tensor)
+    if isinstance(tensor, FakeTensor):
+        return tensor.get_nested_int(coeff=coeff)
+
+    global _tensor_id_counter
+
+    tensor_symint = _tensor_symint_registry.get(tensor)
+    if tensor_symint is None:
+        tensor_symint = torch.SymInt(NestedIntNode(_tensor_id_counter, coeff))
+        _tensor_id_counter += 1
+        _tensor_symint_registry[tensor] = tensor_symint
+    return tensor_symint
+
+
+# SDPA metadata; max / min seqlens are needed for e.g. flash
+def _get_sdpa_extreme_seqlen(func, tensor):
+    return int(func(tensor).item())
+
+
+def _store_val_in_tensor(val) -> torch.Tensor:
+    # hack to get dynamic shapes support: store in a (val, 0) shaped tensor
+    return torch.zeros(val, 0)
+
+
+def _load_val_from_tensor(t: torch.Tensor):
+    return t.shape[0]
+
+
+# serialization function must be defined at top level
+def _rebuild_njt(constructor_kwargs):
+    return NestedTensor(**constructor_kwargs)
+
+
+class NestedTensor(torch.Tensor):
+    _values: torch.Tensor  # type: ignore[assignment]
+    _offsets: torch.Tensor
+    _lengths: Optional[torch.Tensor]
+    # NOTE [ Nested ints for ragged sizes and strides ]
+    #
+    # Jagged layout tensors are tensors that represent a n-dim tensor with a
+    # ragged dimension, but are backed by an (n-1)-dim tensor underneath, e.g.,
+    # a jagged tensor with outer shape [B, x, D] is represented internally by a
+    # tensor with shape [sum(x), D] where we introduce what we call a nested int
+    # denoted as "x" here (but sometimes denoted with "*" to
+    # represent the ragged dimension, and sum(x) represents the dim of the inner
+    # tensor or equivalently the sum of all the sizes of the constituent
+    # tensors' varying lengths.
+    #
+    # We also use nested ints to represent the strides of this tensor.
+    # For example, a jagged tensor with shape [B, x, D] can be strided in two
+    # ways: [xD, D, 1] and [x, 1, sum(x)], where xD represents x multiplied by D
+    _size: tuple[int, ...]
+    _strides: tuple[int, ...]
+    # Indicates that the nth dimension is ragged
+    _ragged_idx: int
+    _metadata_cache: Dict[str, Any]
+
+    @staticmethod
+    def __new__(
+        cls,
+        values,
+        offsets,
+        *,
+        lengths=None,
+        **kwargs,
+    ):
+        ks = DispatchKeySet(DispatchKey.NestedTensor)
+        ks = ks.add(DispatchKey.AutogradNestedTensor)
+
+        # Only support jagged for now.
+        if offsets is None:
+            raise AssertionError("offsets must not be None")
+        if offsets.ndim != 1:
+            raise AssertionError(f"offsets must be 1D, but got {offsets.ndim}D")
+        if isinstance(values, NestedTensor):
+            raise AssertionError("values must not be a NestedTensor")
+        if values.device != offsets.device:
+            raise AssertionError(
+                f"values and offsets must be on the same device, but got "
+                f"values.device={values.device} and offsets.device={offsets.device}"
+            )
+
+        # Query cache for the symint associated with offsets or lengths
+        # (create a new one if needed).
+        ragged_source = offsets if lengths is None else lengths
+        ragged_size = get_tensor_symint(ragged_source, coeff=1)
+        _ragged_idx = kwargs.get("_ragged_idx", 1)
+        B = offsets.shape[0] - 1
+        if lengths is not None:
+            if B != lengths.shape[0]:
+                raise AssertionError(
+                    f"offsets and lengths batch sizes must match: "
+                    f"offsets.shape[0] - 1 = {B}, lengths.shape[0] = {lengths.shape[0]}"
+                )
+
+        # subtract 1 to convert to values dim space
+        r = _ragged_idx - 1
+        _size = (B, *values.shape[:r], ragged_size, *values.shape[r + 1 :])
+        stride = values.stride()
+        _strides = (ragged_size * stride[r], *stride)
+
+        r = torch.Tensor._make_wrapper_subclass(
+            cls,
+            _size,
+            _strides,
+            0,
+            torch.contiguous_format,
+            values.dtype,
+            torch.jagged,
+            values.device,
+            False,
+            kwargs.get("requires_grad", False),
+            "sizes",
+            False,
+            True,  # dispatch_layout
+            ks,
+            # don't try to calculate storage based on non-zero size
+            storage_size=values.untyped_storage().size(),
+        )
+        r._ragged_idx = _ragged_idx
+        r._size = _size
+        r._strides = _strides
+
+        return r
+
+    def __init__(self, values, offsets, *, lengths=None, **kwargs) -> None:
+        super().__init__()
+
+        self._values = values
+        self._offsets = offsets
+        self._lengths = lengths
+
+        # holds properties that are computed lazily
+        self._metadata_cache = kwargs.get("_metadata_cache") or {}
+
+        # collapsed ragged dim must always be dynamic
+        torch._dynamo.maybe_mark_dynamic(self, self._ragged_idx)
+        torch._dynamo.maybe_mark_dynamic(self._values, self._ragged_idx - 1)
+
+        # min / max sequence length should be dynamic if present
+        max_seqlen_tensor = self._metadata_cache.get("max_seqlen", None)
+        if max_seqlen_tensor is not None:
+            torch._dynamo.mark_dynamic(max_seqlen_tensor, 0)
+        min_seqlen_tensor = self._metadata_cache.get("min_seqlen", None)
+        if min_seqlen_tensor is not None:
+            torch._dynamo.mark_dynamic(min_seqlen_tensor, 0)
+
+    def values(self):
+        # dispatch to get proper view relationship
+        return torch._nested_get_values(self)  # type: ignore[attr-defined]
+
+    def offsets(self):
+        return self._offsets
+
+    def lengths(self):
+        return self._lengths
+
+    # Private accessor functions for min / max sequence length. They're
+    # purposefully not @properties because those don't work with PT2 (yet).
+    # These compute / cache if not present.
+    # TODO: Revisit this when @properties are better supported by PT2. I think the ideal
+    # state would be to have public @properties for min / max sequence length that compile
+    # (including setters).
+    def _get_max_seqlen(self):
+        max_seqlen_tensor = self._max_seqlen_tensor
+        if max_seqlen_tensor is None:
+            # compute & cache
+            max_val = _get_sdpa_extreme_seqlen(
+                torch.max,
+                self._offsets.diff() if self._lengths is None else self._lengths,
+            )
+            max_seqlen_tensor = _store_val_in_tensor(max_val)
+            self._metadata_cache["max_seqlen"] = max_seqlen_tensor
+        return _load_val_from_tensor(max_seqlen_tensor)
+
+    def _get_min_seqlen(self):
+        min_seqlen_tensor = self._min_seqlen_tensor
+        if min_seqlen_tensor is None:
+            # compute & cache
+            min_val = _get_sdpa_extreme_seqlen(
+                torch.min,
+                self._offsets.diff() if self._lengths is None else self._lengths,
+            )
+            min_seqlen_tensor = _store_val_in_tensor(min_val)
+            self._metadata_cache["min_seqlen"] = min_seqlen_tensor
+        return _load_val_from_tensor(min_seqlen_tensor)
+
+    # Private accessors used for treating min / max seqlen as inner tensors for
+    # flatten / unflatten. These must be properties to work with the traceable wrapper
+    # subclass logic. These do not compute / cache if not present.
+    @property
+    def _max_seqlen_tensor(self) -> Optional[torch.Tensor]:
+        return self._metadata_cache.get("max_seqlen", None)
+
+    @_max_seqlen_tensor.setter
+    def _max_seqlen_tensor(self, val: Optional[torch.Tensor]) -> None:
+        self._metadata_cache["max_seqlen"] = val
+
+    @property
+    def _min_seqlen_tensor(self) -> Optional[torch.Tensor]:
+        return self._metadata_cache.get("min_seqlen", None)
+
+    @_min_seqlen_tensor.setter
+    def _min_seqlen_tensor(self, val: Optional[torch.Tensor]) -> None:
+        self._metadata_cache["min_seqlen"] = val
+
+    # These are old private @property accessors that are kept around for internal BC
+    # reasons. TODO: Remove these!
+    @property
+    def _max_seqlen(self):
+        return self._get_max_seqlen()
+
+    @property
+    def _min_seqlen(self):
+        return self._get_min_seqlen()
+
+    # Convenience accessors that return a min / max seqlen if one is present and do NOT
+    # compute / cache them if they're not.
+    @property
+    def _maybe_max_seqlen(self) -> Optional[int]:
+        mt = self._max_seqlen_tensor
+        return None if mt is None else _load_val_from_tensor(mt)
+
+    @property
+    def _maybe_min_seqlen(self) -> Optional[int]:
+        mt = self._min_seqlen_tensor
+        return None if mt is None else _load_val_from_tensor(mt)
+
+    def _is_contiguous_or_false(self):
+        if self.lengths() is not None:
+            return False
+        from torch._prims_common import is_contiguous_for_memory_format_or_false
+
+        return is_contiguous_for_memory_format_or_false(
+            self._values, memory_format=torch.contiguous_format
+        )
+
+    def __repr__(self) -> str:  # type: ignore[override]
+        # We should implement this in torch/_tensor_str.py instead
+        grad_fn_str = (
+            f", requires_grad={self.requires_grad}" if self.requires_grad else ""
+        )
+
+        if self.grad_fn:
+            grad_fn_str = f", grad_fn={self.grad_fn}"
+
+        return f"NestedTensor(size={self._size}, offsets={self._offsets}{grad_fn_str}, contiguous={self._is_contiguous_or_false()})"
+
+    # TODO: Remove this in favor of the default tensor subclass serialization logic.
+    # We don't do this today because of https://github.com/pytorch/pytorch/issues/125622.
+    def __reduce_ex__(self, proto):
+        state = torch._utils._get_obj_state(self)
+
+        # Cached PyCapsules for sizes / strides are not serializable.
+        # See Note [Tensor Subclass custom size/stride caching strategy]
+        self._clear_non_serializable_cached_data()
+        # SymNodes are not serializable
+        if "_size" not in state or "_strides" not in state:
+            raise AssertionError("state must contain '_size' and '_strides'")
+        state = dict(state)
+        del state["_size"]
+        del state["_strides"]
+
+        func = _rebuild_njt
+        constructor_kwargs = {
+            "values": self._values,
+            "offsets": self._offsets,
+            "lengths": self._lengths,
+            "_ragged_idx": self._ragged_idx,
+            "_metadata_cache": self._metadata_cache,
+            "requires_grad": self.requires_grad,
+        }
+        args = (constructor_kwargs,)
+        return (torch._tensor._rebuild_from_type_v2, (func, type(self), args, state))
+
+    def __tensor_flatten__(self):
+        ctx = {
+            "requires_grad": self.requires_grad,
+            "ragged_idx": self._ragged_idx,
+        }
+        inner_tensors = ["_values", "_offsets"]
+        if self._lengths is not None:
+            inner_tensors.append("_lengths")
+        if self._min_seqlen_tensor is not None:
+            inner_tensors.append("_min_seqlen_tensor")
+        if self._max_seqlen_tensor is not None:
+            inner_tensors.append("_max_seqlen_tensor")
+        return inner_tensors, ctx
+
+    @staticmethod
+    def __tensor_unflatten__(inner_tensors: Dict, meta, outer_size, outer_stride):
+        from torch._subclasses.fake_tensor import FakeTensor
+
+        # inner tensors: _values, _offsets, [_lengths], [_min_seqlen], [_max_seqlen]
+        if not (len(inner_tensors) >= 2 and len(inner_tensors) <= 5):
+            raise AssertionError(
+                f"Expected 2-5 inner tensors, but got {len(inner_tensors)}"
+            )
+        values = inner_tensors["_values"]
+        offsets = inner_tensors["_offsets"]
+        lengths = inner_tensors.get("_lengths", None)
+        min_seqlen_tensor = inner_tensors.get("_min_seqlen_tensor", None)
+        max_seqlen_tensor = inner_tensors.get("_max_seqlen_tensor", None)
+
+        metadata_cache = {}
+        if min_seqlen_tensor is not None:
+            metadata_cache["min_seqlen"] = min_seqlen_tensor
+        if max_seqlen_tensor is not None:
+            metadata_cache["max_seqlen"] = max_seqlen_tensor
+        ragged_idx = meta["ragged_idx"]
+
+        # Alternatively, we could make it the caller's responsibility to
+        # cache it. But this heuristic seems simple enough.
+        ragged_source = offsets if lengths is None else lengths
+        if isinstance(ragged_source, FakeTensor):
+            ragged_size = outer_size[ragged_idx]
+            ragged_source.nested_int_memo = ragged_size
+
+        return NestedTensor(
+            values,
+            offsets=offsets,
+            lengths=lengths,
+            requires_grad=meta["requires_grad"],
+            _ragged_idx=ragged_idx,
+            _metadata_cache=metadata_cache,
+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):  # type: ignore[override]
+        # If you're wondering why there's a nested tensor with one of its
+        # size = -1, see note: [NJT outer_size in AOTDispatcher]
+        kwargs = {} if kwargs is None else kwargs
+
+        # Lazy import to avoid circular dependency
+        from .ops import lookup_jagged
+
+        fn = lookup_jagged(func, *args, **kwargs)
+        if fn is not None:
+            return fn(*args, **kwargs)
+
+        # Poor man's redispatch for composite ops. This becomes relevant under inference
+        # mode, where disabling autograd key dispatch prevents decomposition.
+        all_dks = (
+            # We want to handle both the cases where NestedTensor overrides the
+            # composite implicit autograd kernel, and the case where it doesn't.
+            # Prioritize calling into NestedTensor's kernel if it exists.
+            torch._C.DispatchKey.CompositeImplicitAutogradNestedTensor,
+            torch._C.DispatchKey.CompositeImplicitAutograd,
+        )
+        for dk in all_dks:
+            if torch._C._dispatch_has_kernel_for_dispatch_key(func.name(), dk):
+                with torch.overrides.enable_reentrant_dispatch():
+                    return func._op_dk(dk, *args, **kwargs)
+
+        raise NotImplementedError(func)
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+
+        from torch.fx.experimental.proxy_tensor import maybe_enable_thunkify
+
+        from .ops import jagged_torch_function
+
+        # This should be removed after
+        # https://github.com/pytorch/pytorch/pull/125941/ lands
+        with maybe_enable_thunkify():
+            try:
+                return jagged_torch_function(func, *args, **kwargs)
+            except NotImplementedError:
+                pass
+            with torch._C.DisableTorchFunctionSubclass():
+                return func(*args, **kwargs)
+
+
+# NB: These fake view autograd.Functions are superseded by real view ops. Don't use them!
+# TODO: Remove ViewBufferFromNested, ViewNestedFromBuffer, and buffer_from_jagged once the
+# internal BC period has passed.
+
+
+# Not actually a view!
+class ViewBufferFromNested(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: NestedTensor):  # type: ignore[override]
+        ctx.save_for_backward(x.offsets())
+        ctx.metadata_cache = x._metadata_cache
+        ctx.ragged_idx = x._ragged_idx
+        return x._values
+
+    @staticmethod
+    def backward(ctx, gO: torch.Tensor):  # type: ignore[override]
+        (offsets,) = ctx.saved_tensors
+        return NestedTensor(
+            gO,
+            offsets=offsets,
+            _metadata_cache=ctx.metadata_cache,
+            _ragged_idx=ctx.ragged_idx,
+        )
+
+
+# Not actually a view!
+class ViewNestedFromBuffer(torch.autograd.Function):
+    @staticmethod
+    def forward(  # pyrefly: ignore  # bad-override
+        ctx,
+        values: torch.Tensor,
+        offsets: torch.Tensor,
+        metadata_cache: Optional[Dict[str, Any]] = None,
+    ):  # type: ignore[override]
+        # maintain BC with this usages of this where the seqlens are stuffed
+        # directly into the metadata cache as non-Tensors / ints
+        if metadata_cache is not None:
+            min_seqlen = metadata_cache.get("min_seqlen", None)
+            max_seqlen = metadata_cache.get("max_seqlen", None)
+            if min_seqlen is not None and not isinstance(min_seqlen, torch.Tensor):
+                metadata_cache["min_seqlen"] = _store_val_in_tensor(min_seqlen)
+            if max_seqlen is not None and not isinstance(max_seqlen, torch.Tensor):
+                metadata_cache["max_seqlen"] = _store_val_in_tensor(max_seqlen)
+        return NestedTensor(
+            values.detach(),
+            offsets=offsets,
+            _metadata_cache=metadata_cache,
+        )
+
+    @staticmethod
+    def backward(ctx, gO: NestedTensor):  # type: ignore[override]
+        return gO._values, None, None
+
+
+def buffer_from_jagged(jagged):
+    return ViewBufferFromNested.apply(jagged)
+
+
+# Need to make it obvious that users should be passing in offsets
+def jagged_from_list(
+    tensors: List[torch.Tensor],
+    offsets: Optional[torch.Tensor],
+    dtype=None,
+    device=None,
+) -> tuple[NestedTensor, torch.Tensor]:
+    """Constructs a NestedTensor backed by jagged layout from a list of tensors"""
+
+    if len(tensors) == 0:
+        raise RuntimeError("Cannot construct a nested tensor from an empty tensor list")
+    if not len(set(t.dtype for t in tensors)) == 1:  # noqa: C401
+        raise RuntimeError(
+            "When constructing a nested tensor, all tensors in list must have the same dtype"
+        )
+    if not len(set(t.device for t in tensors)) == 1:  # noqa: C401
+        raise RuntimeError(
+            "When constructing a nested tensor, all tensors in list must be on the same device"
+        )
+    if not len(set(t.dim() for t in tensors)) == 1:  # noqa: C401
+        raise RuntimeError(
+            "When constructing a nested tensor, all tensors in list must have the same dim"
+        )
+    component_dim = tensors[0].dim()
+    if component_dim == 0:
+        raise RuntimeError(
+            "Cannot construct a nested tensor from a list of zero-dim tensors"
+        )
+
+    # Check that the NT is representable by the jagged layout, which
+    # allows for a single ragged dimension after the batch dim.
+    # e.g. (B, *, D_0, ..., D_N), (B, D_0, *, ..., D_N), etc.
+    sizes = [t.shape for t in tensors]
+    ragged_idx = None
+    for d in range(component_dim):
+        dim_is_ragged = any(size[d] != sizes[0][d] for size in sizes)
+        if dim_is_ragged:
+            if ragged_idx is None:
+                # add 1 to convert to outer NJT dim space
+                ragged_idx = d + 1
+            else:
+                raise RuntimeError(
+                    "Cannot represent given tensor list as a nested tensor with the jagged layout. "
+                    "Note that the jagged layout only allows for a single ragged dimension. "
+                    "For example: (B, *, D_0, D_1, ..., D_N), with ragged * dim."
+                )
+
+    # allow for a rectangular NJT and default the ragged dim next to the batch dim
+    if ragged_idx is None:
+        ragged_idx = 1
+
+    # Set properties appropriately.
+    values = torch.cat(tensors, dim=(ragged_idx - 1))
+    to_kwargs = {}
+    if device is not None:
+        to_kwargs["device"] = device
+    if dtype is not None:
+        to_kwargs["dtype"] = dtype
+    values = values.to(**to_kwargs)
+
+    # Calculate jagged offsets if not provided.
+    if offsets is None:
+        # Jagged layout specifies that offsets are stored as int64 on the same device as values.
+        # TODO: An alternative way to construct offsets is to use F.pad. This avoids creating
+        # an extra leaf tensor during the forward, potentially resolving compatibility issues.
+        offsets = torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int64, device=values.device),
+                torch.tensor(
+                    [s[ragged_idx - 1] for s in sizes], device=values.device
+                ).cumsum(dim=0),
+            ]
+        )
+
+    # compute this now since it's easy
+    min_seqlen = min(t.shape[ragged_idx - 1] for t in tensors)
+    max_seqlen = max(t.shape[ragged_idx - 1] for t in tensors)
+    ret_nt = nested_view_from_values_offsets(
+        values,
+        offsets,
+        min_seqlen=min_seqlen,
+        max_seqlen=max_seqlen,
+        ragged_idx=ragged_idx,
+    )
+    return (ret_nt, offsets)  # type: ignore[return-value]
+
+
+def jagged_from_tensor_and_lengths(
+    tensor: torch.Tensor, starts: torch.Tensor, lengths: torch.Tensor
+) -> tuple[NestedTensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Constructs a NestedTensor backed by jagged layout from a tensor, starts of sequences, and sequence lengths"""
+    batch_size = tensor.shape[0]
+    if is_expandable_to(starts.shape, (batch_size,)) and is_expandable_to(
+        lengths.shape, (batch_size,)
+    ):
+        start_list = starts.expand(batch_size)
+        length_list = lengths.expand(batch_size)
+    else:
+        raise RuntimeError(
+            "When constructing a jagged nested tensor using narrow(), "
+            "your start and length must be Tensors that broadcast to input.shape[0]"
+        )
+
+    # Calculate jagged offsets
+    if len(tensor.shape) < 2:
+        raise AssertionError(
+            "tensor must at least be 2D for the nested narrow op to work"
+        )
+    max_seq_len = tensor.shape[1]
+    offset_lengths = max_seq_len * torch.arange(
+        0, batch_size, dtype=torch.int64, device=tensor.device
+    )
+    # Jagged layout specifies that offsets are stored as int64 on the same device as values.
+    offsets = torch.cat(
+        [
+            start_list + offset_lengths,
+            (start_list[-1] + offset_lengths[-1] + length_list[-1]).unsqueeze(0),
+        ]
+    )
+
+    # Reshape buffer to flatten the 1st and 2nd dimension (view used to enforce non-copy)
+    if len(tensor.shape) > 2:
+        values = tensor.view(-1, *tensor.shape[2:])
+    else:
+        values = tensor.view(-1)
+
+    # Check if offsets and lengths make it possibly contiguous and return a regular NT
+    is_contiguous = True
+    orig_dim = tensor.shape[1]
+    if torch.any(length_list[1:-1].ne(orig_dim)):
+        is_contiguous = False
+    if torch.any(offsets[1:-2].diff().ne(orig_dim)):
+        is_contiguous = False
+    if offsets[0] + length_list[0] != orig_dim:
+        is_contiguous = False
+
+    actual_max_seqlen = int(torch.max(lengths).item())
+    min_seqlen = int(torch.min(lengths).item())
+
+    if is_contiguous:
+        ret_nt = nested_view_from_values_offsets(
+            values[offsets[0] : offsets[-1]],
+            offsets - offsets[0],
+            min_seqlen=min_seqlen,
+            max_seqlen=actual_max_seqlen,
+        )
+    else:
+        ret_nt = nested_view_from_values_offsets_lengths(
+            values,
+            offsets,
+            length_list,
+            min_seqlen=min_seqlen,
+            max_seqlen=actual_max_seqlen,
+        )
+
+    return (ret_nt, offsets, None if is_contiguous else length_list)
+
+
+# NB: A dummy arg is required so that NestedTensor.__torch_dispatch__() is invoked
+# for _nested_view_from_values_offsets(). Sizes don't matter much, but they shouldn't be
+# 0/1 because the dummy can be fake-ified and we want to avoid specializing.
+# This arg is otherwise unused.
+_dummy_instance: Optional[torch.Tensor] = None
+
+
+def _nt_view_dummy() -> torch.Tensor:
+    global _dummy_instance
+    if _dummy_instance is None:
+        _dummy_instance = NestedTensor(
+            values=torch.zeros(3, 3, device="meta"),
+            offsets=torch.zeros(3, device="meta", dtype=torch.int64),
+        ).detach()
+    return _dummy_instance
+
+
+def nested_view_from_values_offsets(
+    values, offsets, ragged_idx=1, min_seqlen=None, max_seqlen=None
+):
+    min_seqlen_tensor = None
+    if min_seqlen is not None:
+        min_seqlen_tensor = _store_val_in_tensor(min_seqlen)
+
+    max_seqlen_tensor = None
+    if max_seqlen is not None:
+        max_seqlen_tensor = _store_val_in_tensor(max_seqlen)
+
+    return torch._nested_view_from_jagged(  # type: ignore[attr-defined]
+        values,
+        offsets,
+        _nt_view_dummy(),
+        None,
+        ragged_idx,
+        min_seqlen_tensor,
+        max_seqlen_tensor,
+    )  # type: ignore[return-value]
+
+
+def nested_view_from_values_offsets_lengths(
+    values, offsets, lengths, ragged_idx=1, min_seqlen=None, max_seqlen=None
+):
+    min_seqlen_tensor = None
+    if min_seqlen is not None:
+        min_seqlen_tensor = _store_val_in_tensor(min_seqlen)
+
+    max_seqlen_tensor = None
+    if max_seqlen is not None:
+        max_seqlen_tensor = _store_val_in_tensor(max_seqlen)
+
+    return torch._nested_view_from_jagged(  # type: ignore[attr-defined]
+        values,
+        offsets,
+        _nt_view_dummy(),
+        lengths,
+        ragged_idx,
+        min_seqlen_tensor,
+        max_seqlen_tensor,
+    )  # type: ignore[return-value]
+
+
+def nested_from_padded(
+    padded, offsets, ragged_idx=1, min_seqlen=None, max_seqlen=None, sum_S=None
+):
+    min_seqlen_tensor = None
+    if min_seqlen is not None:
+        min_seqlen_tensor = _store_val_in_tensor(min_seqlen)
+
+    max_seqlen_tensor = None
+    if max_seqlen is not None:
+        max_seqlen_tensor = _store_val_in_tensor(max_seqlen)
+
+    return torch._nested_from_padded_tensor(
+        padded,
+        offsets,
+        _nt_view_dummy(),
+        ragged_idx,
+        min_seqlen_tensor,
+        max_seqlen_tensor,
+        sum_S,
+    )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a1b8e473069fc1e6017993d8b717455466d7467
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/ops.py
@@ -0,0 +1,2756 @@
+# mypy: allow-untyped-defs
+import functools
+import math
+import operator
+from typing import *  # noqa: F403
+
+import torch
+import torch.nn.functional as F
+from torch.fx.operator_schemas import normalize_function
+from torch.nested._internal.sdpa import jagged_scaled_dot_product_attention
+
+from .nested_tensor import NestedTensor
+
+
+__all__: list[Any] = []
+
+JAGGED_OPS_TABLE: Dict[Any, Any] = {}
+
+
+def _get_padding_value(dtype, padding_type):
+    if dtype.is_floating_point:
+        return (
+            torch.finfo(dtype).max if padding_type == "max" else torch.finfo(dtype).min
+        )
+    else:
+        # For integer dtypes, use infinity sentinels which the C++ implementation
+        # clamps to dtype min/max, avoiding precision loss through double.
+        return float("inf") if padding_type == "max" else float("-inf")
+
+
+def _outer_to_inner_dim(ndim, dim, ragged_dim, canonicalize=False):
+    from torch._prims_common import canonicalize_dims
+
+    if isinstance(dim, (tuple, list)):
+        output = type(dim)(_outer_to_inner_dim(ndim, d, ragged_dim) for d in dim)
+        # ensure no duplicates, which can result from both batch and ragged mapping to 0
+        return type(output)(dict.fromkeys(output))
+
+    if canonicalize:
+        dim = canonicalize_dims(ndim, dim)
+
+    if not (dim >= 0 and dim < ndim):  # pyrefly: ignore [unsupported-operation]
+        raise AssertionError(f"dim {dim} out of range for ndim {ndim}")
+
+    # Map dim=0 (AKA batch dim) -> packed dim i.e. outer ragged dim - 1.
+    # For other dims, subtract 1 to convert to inner space.
+    return (
+        # pyrefly: ignore [unsupported-operation]
+        ragged_dim - 1 if dim == 0 else dim - 1
+    )
+
+
+def _wrap_jagged_dim(
+    ndim,
+    dim,
+    ragged_dim,
+    op_name,
+    convert_to_inner_dim=True,
+    allow_ragged_dim=False,
+    allow_batch_dim=False,
+):
+    from torch._prims_common import canonicalize_dims
+
+    wrapped = canonicalize_dims(ndim, dim)
+    if wrapped == ragged_dim and not allow_ragged_dim:
+        raise RuntimeError(f"{op_name}(): not supported for NestedTensor on ragged dim")
+    elif wrapped == 0 and not allow_batch_dim:
+        raise RuntimeError(f"{op_name}(): not supported for NestedTensor on dim=0")
+    ret = (
+        _outer_to_inner_dim(ndim, wrapped, ragged_dim)
+        if convert_to_inner_dim
+        else wrapped
+    )
+    if allow_batch_dim:
+        # Need to disambiguate whether we're operating on the batch dim or not.
+        # Operating on dim=1 -> dim=0 after the inner dim conversion.
+        operating_on_batch = wrapped == 0
+        return (ret, operating_on_batch)
+    return ret
+
+
+def _wrap_jagged_dims(ndim, dims, op_name, ragged_idx=1):
+    """
+    For NestedTensor operators,
+    wraps dimensions to non-negative values,
+    and returns metadata related to reduction dimension(s).
+    """
+    from torch._prims_common import canonicalize_dims
+
+    if not isinstance(dims, (tuple, list)):
+        raise AssertionError(
+            f"_wrap_jagged_dims(): cannot iterate over dimensions of type {type(dims)}"
+        )
+
+    wrapped_dims = [
+        canonicalize_dims(ndim, d) for d in dims
+    ]  # convert all indices to non-negative values
+
+    operate_on_batch = 0 in wrapped_dims
+    operate_on_ragged = ragged_idx in wrapped_dims
+    operate_on_non_batch = any(d != 0 and d != ragged_idx for d in wrapped_dims)
+
+    # ensure no duplicates, which can result from both batch and ragged mapping to 0
+    outer_to_inner_dim = tuple(
+        dict.fromkeys(_outer_to_inner_dim(ndim, d, ragged_idx) for d in wrapped_dims)
+    )
+
+    return outer_to_inner_dim, operate_on_batch, operate_on_ragged, operate_on_non_batch
+
+
+def check_schema(schema_str: str, func, *args, **kwargs) -> None:
+    named_arg_types = schema_str.split(", ")
+    num_optional_args = [x.endswith("?") for x in named_arg_types].count(True)
+    min_args = len(named_arg_types) - num_optional_args
+
+    # special case: ellipses allows for any number of unchecked args at the end
+    if named_arg_types[-1] == "...":
+        named_arg_types = named_arg_types[:-1]
+    else:
+        if not (len(args) >= min_args and len(args) <= len(named_arg_types)):
+            raise ValueError(
+                f"NestedTensor {func.__name__}({schema_str}): expected at least {min_args} "
+                f"arguments and at most {len(named_arg_types)} arguments, but got: "
+                f"{len(args)} arguments"
+            )
+
+    arg_type_check_fns = {
+        "t": lambda x: isinstance(x, torch.Tensor) and not isinstance(x, NestedTensor),
+        "jt": lambda x: isinstance(x, NestedTensor)
+        and x._lengths is None
+        and x._ragged_idx == 1,  # ops with "jt" require contiguous JT only
+        "jt_all": lambda x: isinstance(
+            x, NestedTensor
+        ),  # ops with "jt_all" can accept all kinds of JT
+        "any": lambda x: True,
+    }
+    for i, named_arg_type in enumerate(named_arg_types):
+        name, arg_type = named_arg_type.split(": ")
+        is_optional = arg_type.endswith("?")
+        normalized_arg_type = arg_type[:-1] if is_optional else arg_type
+        if normalized_arg_type not in arg_type_check_fns:
+            raise AssertionError(f"Unknown arg type: {normalized_arg_type}")
+
+        if i >= len(args):
+            if not is_optional:
+                raise ValueError(
+                    f"NestedTensor {func.__name__}({schema_str}) "
+                    f"missing required argument: {name}"
+                )
+            continue
+
+        _check_fn = arg_type_check_fns[normalized_arg_type]
+
+        def check_fn(x, is_optional=is_optional):
+            if is_optional:
+                return x is None or _check_fn(x)
+            else:
+                return _check_fn(x)
+
+        if not check_fn(args[i]):
+            type_to_desc = {
+                "t": "tensor",
+                "t?": "optional tensor",
+                "jt": "contiguous jagged layout NestedTensor",
+                "jt_all": "jagged layout NestedTensor",
+                "any": "<any type>",
+            }
+
+            raise ValueError(
+                f"NestedTensor {func.__name__}({schema_str}): expected {name} to be a "
+                f"{type_to_desc[arg_type]}"
+            )
+
+
+def check_ragged_dim_same(
+    func, a: NestedTensor, a_name: str, b: NestedTensor, b_name: str
+) -> None:
+    # Calling into .shape here
+    if a._size[a._ragged_idx] != b._size[b._ragged_idx]:
+        raise RuntimeError(
+            f"NestedTensor {func.__name__}: expected {a_name} and {b_name} to have the "
+            "same exact offsets tensor."
+        )
+
+
+# returns True if the raggedness-relevant portions of the NT shape
+# match those of the specified size
+def raggedness_matches(nt, size):
+    end = nt._ragged_idx + 1
+    nt_ragged = nt._size[:end]
+    size_ragged = size[:end]
+    return len(nt_ragged) == len(size_ragged) and (
+        all(ns == s or s == -1 for ns, s in zip(nt_ragged, size_ragged))
+    )
+
+
+def squeeze_leading_ones(t):
+    # Note: [ Squeezing leading ones ]
+    #
+    # Squeeze leading ones from t.
+    #
+    # We want:
+    #   (B, j0, ?, ?) + (1, 1, ?, ?) -> (B, j0, ?, ?)
+    #   (B, j0, ?, ?) + (1, 1, 1, ?, ?) -> (1, B, j0, ?, ?)  (not yet supported)
+    #
+    # 1) Squeeze extra ones and grab values from NT
+    #   (1, 1, ?, ?) -> (?, ?)   and   (sum(*), ?, ?) -> (B, j0, ?, ?)
+    # 2) Do dense broadcasting:
+    #   (sum(*), ?, ?) + (?, ?) -> (sum(*), ?, ?)
+    # 3) Construct nested tensor
+    #   (sum(*), ?, ?) -> (B, j0, ?, ?)
+    #
+    # If unsqueezing on the 0th dim becomes supported, we would unsqueeze
+    # at step (4) and we would need to update this function to record how
+    # many ones we unsqueezed.
+    while t.dim() > 0 and t.shape[0] == 1:
+        t = t.squeeze(0)
+    return t
+
+
+def register_func(tables, aten_ops, schema_str):
+    if not isinstance(aten_ops, list):
+        aten_ops = [aten_ops]
+    if not isinstance(tables, list):
+        tables = [tables]
+
+    def wrapper(func):
+        for aten_op in aten_ops:
+
+            def get_inner(aten_op):
+                def inner(*args, **kwargs):
+                    check_schema(schema_str, func, *args, **kwargs)
+                    return func(aten_op, *args, **kwargs)
+
+                return inner
+
+            for table in tables:
+                table[aten_op] = get_inner(aten_op)
+        return func
+
+    return wrapper
+
+
+register_jagged_func = functools.partial(register_func, JAGGED_OPS_TABLE)
+
+
+def lookup_jagged(func, *args, **kwargs) -> Callable | None:
+    dispatch_func = JAGGED_OPS_TABLE.get(func, None)
+    if dispatch_func is not None:
+        return dispatch_func
+
+    # Handle pointwise fallbacks
+    if torch.Tag.pointwise in func.tags:
+        from torch.fx.experimental.symbolic_shapes import is_nested_int
+
+        # No pointwise ops legitimately accept nested int inputs. Without this check,
+        # they will be incorrectly interpreted as tensors.
+        # See https://github.com/pytorch/pytorch/issues/138496
+        for arg in args:
+            if is_nested_int(arg):
+                raise RuntimeError(
+                    f"NestedTensor {func.__name__}: invalid argument {arg}"
+                )
+
+        # Assume there aren't additional tensors that aren't the "unary/binary" args
+        num_tensor_args = sum(isinstance(x, torch.Tensor) for x in args)
+        if num_tensor_args == 1:
+            # Build up the check schema string. The first tensor arg is assumed to be
+            # an NJT and other args are sent through as-is.
+            schema_parts = []
+            for arg in func._schema.arguments:
+                if isinstance(arg.type, torch.TensorType):
+                    schema_parts.append(f"{arg.name}: jt_all")
+                    break
+                else:
+                    schema_parts.append(f"{arg.name}: any")
+            schema_parts.append("...")
+            check_schema_str = ", ".join(schema_parts)
+            check_schema(check_schema_str, func, *args, **kwargs)
+            return functools.partial(jagged_unary_pointwise, func)
+        elif num_tensor_args == 2:
+            check_schema("lhs: any, rhs: any, ...", func, *args, **kwargs)
+            return functools.partial(jagged_binary_pointwise, func)
+
+    return None
+
+
+def extract_kwargs(arg):
+    kwargs = {
+        "offsets": arg.offsets(),
+        "lengths": arg.lengths(),
+        "_metadata_cache": arg._metadata_cache,
+        "_ragged_idx": arg._ragged_idx,
+    }
+    return kwargs
+
+
+def jagged_unary_pointwise(func, *args, **kwargs):
+    # assume if we get here that there is a single NJT input in the args
+    njt = next(arg for arg in args if isinstance(arg, NestedTensor))
+    return NestedTensor(
+        func(*(arg._values if arg is njt else arg for arg in args), **kwargs),
+        **extract_kwargs(njt),
+    )
+
+
+def jagged_binary_pointwise(func, *args, **kwargs):
+    a, b = args[0], args[1]
+    if not (isinstance(a, NestedTensor) or isinstance(b, NestedTensor)):
+        raise AssertionError("At least one of the arguments must be a NestedTensor")
+
+    mismatch_error_msg = (
+        "cannot call binary pointwise function {} with inputs of shapes {} and {}"
+    )
+    # a is NT, b is NT
+    if isinstance(a, NestedTensor) and isinstance(b, NestedTensor):
+        # ex: (B, j0, D) + (B, j0, D)
+        # ex: (B, j0, D) + (B, j0, 1)
+        if raggedness_matches(a, b._size):
+            return NestedTensor(
+                func(a._values, b._values, *args[2:], **kwargs), **extract_kwargs(a)
+            )
+        raise RuntimeError(mismatch_error_msg.format(func.__name__, a._size, b._size))
+    # either a is NT or b is NT at this point
+    a_is_nt = isinstance(a, NestedTensor)
+    extracted_kwargs = extract_kwargs(a) if a_is_nt else extract_kwargs(b)
+
+    # === Handle broadcasting across the batch / ragged dims ===
+
+    # Easy case: take advantage of pre-existing broadcasting logic
+    # ex: (B, j0, ?, ?) + (?) -> (B, j0, ?, ?)
+    # ex: (B, j0, ?, ?) + (?, ?) -> (B, j0, ?, ?)
+    # ex: (B, j0, ?, ?) + (1, 1, ?, ?) -> (B, j0, ?, ?)
+    nt, t = (a, b) if a_is_nt else (b, a)
+    # See Note: [ Squeezing leading ones ]
+    if t.dim() > nt.dim():
+        raise NotImplementedError("NYI: broadcasting NT with T with larger dim")
+    t_squeezed = squeeze_leading_ones(t)
+    if nt.dim() >= t_squeezed.dim() + 2:
+        lhs, rhs = (nt._values, t_squeezed) if a_is_nt else (t_squeezed, nt._values)
+        return NestedTensor(func(lhs, rhs, *args[2:], **kwargs), **extracted_kwargs)
+
+    # Harder case: do manual broadcasting when NT dim == non-NT dim
+    # ex: (B, j0, D_0, D_1) + (B, 1, D_0, D_1) -> (B, j0, D_0, D_1)
+    if a.dim() == b.dim():
+        # ex: (B, j0, D_0, D_1) + (1, 1, D_0, D_1) -> should
+        # be (B, j0, D_0, D_1) but not yet supported
+        if a.shape[0] != b.shape[0]:
+            raise RuntimeError(
+                mismatch_error_msg.format(func.__name__, a.shape, b.shape)
+            )
+
+        from .nested_tensor import nested_from_padded
+
+        # handle broadcasting via padded dense -> jagged conversion
+        min_seqlen = nt._maybe_min_seqlen
+        max_seqlen = nt._maybe_max_seqlen
+        padded_max_S = max_seqlen
+        total_L = nt._values.shape[nt._ragged_idx - 1]
+        if padded_max_S is None:
+            # use upper bound on max seqlen if it's not present
+            padded_max_S = total_L
+
+        # convert dense tensor -> jagged
+        t = t.expand(
+            [x if i != nt._ragged_idx else padded_max_S for i, x in enumerate(t.shape)]
+        )
+        t_as_nt = nested_from_padded(
+            t,
+            offsets=nt._offsets,
+            ragged_idx=nt._ragged_idx,
+            sum_S=total_L,
+            min_seqlen=min_seqlen,
+            max_seqlen=max_seqlen,
+        )
+
+        # function call with two NJTs
+        lhs, rhs = (nt, t_as_nt) if a_is_nt else (t_as_nt, nt)
+        return func(lhs, rhs, *args[2:], **kwargs)
+
+    # ex: (B, j0, D_0, D_1) + (A, B, 1, D_0, D_1) -> error because this breaks the invariant
+    # that ragged dim is wrt left-most batch dim
+    raise RuntimeError(mismatch_error_msg.format(func.__name__, a.shape, b.shape))
+
+
+def jagged_torch_function(func, *args, **kwargs):
+    # SDPA has special kernels that handle nested tensors.
+    # Dispatch to the correct implementation here
+    if func is torch._C._nn.scaled_dot_product_attention:
+        return jagged_scaled_dot_product_attention(*args, **kwargs)
+
+    if func.__name__ == "apply_":
+        func(args[0]._values, *args[1:], **kwargs)
+        return args[0]
+
+    # Handle flatten() here because it's CompositeImplicit.
+    if func.__name__ == "flatten":
+
+        def _flatten_sig(input, start_dim=0, end_dim=-1) -> None:
+            pass
+
+        _, new_kwargs = normalize_function(  # type: ignore[misc]
+            _flatten_sig, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+        )
+
+        inp = new_kwargs.pop("input")
+
+        # NB: stay in outer dim space because we're going to redispatch on a NT input
+        start_dim = _wrap_jagged_dim(
+            inp.dim(),
+            new_kwargs["start_dim"],
+            inp._ragged_idx,
+            "flatten",
+            convert_to_inner_dim=False,
+        )
+        end_dim = _wrap_jagged_dim(
+            inp.dim(),
+            new_kwargs["end_dim"],
+            inp._ragged_idx,
+            "flatten",
+            convert_to_inner_dim=False,
+        )
+
+        if start_dim == end_dim:
+            return inp
+
+        product = functools.reduce(operator.mul, inp.shape[start_dim : end_dim + 1])
+        new_shape = (*inp.shape[:start_dim], product, *inp.shape[end_dim + 1 :])
+
+        return inp.reshape(*new_shape)
+
+    # Handle NestedTensor share_memory_.
+    if func.__name__ == "share_memory_":
+        nt = args[0]
+
+        if nt.is_cuda:
+            return nt
+
+        names, _ = nt.__tensor_flatten__()
+        with torch._C.DisableTorchFunctionSubclass():
+            for name in names:
+                component = getattr(nt, name, None)
+                if component is not None:
+                    component.share_memory_()
+        return nt
+
+    # Handle NestedTensor is_shared.
+    if func.__name__ == "is_shared":
+        nt = args[0]
+
+        if nt.is_cuda:
+            return False
+
+        names, _ = nt.__tensor_flatten__()
+        if not names:
+            return False
+        return all(
+            getattr(nt, name) is not None and getattr(nt, name).is_shared()
+            for name in names
+        )
+
+    # Handle nested-specific input validation for CompositeImplicit rms_norm
+    if func.__name__ == "rms_norm":
+
+        def _rms_norm_sig(input, normalized_shape, weight=None, eps=None) -> None:
+            pass
+
+        _, new_kwargs = normalize_function(  # type: ignore[misc]
+            _rms_norm_sig, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+        )
+
+        inp = new_kwargs.pop("input")
+        normalized_shape = new_kwargs.pop("normalized_shape")
+
+        # can't normalize over the ragged dim (yet)
+        max_normalizable = inp.dim() - inp._ragged_idx - 1
+        if len(normalized_shape) > max_normalizable:
+            raise ValueError(
+                "rms_norm(): Normalization over the ragged dim not supported for nested tensors"
+            )
+
+        with torch._C.DisableTorchFunctionSubclass():
+            return func(*args, **kwargs)
+
+    raise NotImplementedError(func)
+
+
+@register_jagged_func(
+    [
+        torch.ops.aten.is_non_overlapping_and_dense.default,
+        torch.ops.aten.sym_size.default,
+        torch.ops.aten.dim.default,
+        torch.ops.aten.numel.default,
+        torch.ops.aten.sym_numel.default,
+        torch.ops.aten.sym_stride.default,
+        torch.ops.aten.sym_storage_offset.default,
+    ],
+    "self: jt_all",
+)
+def tensor_attr_supported_getter(func, *args, **kwargs):
+    if func is torch.ops.aten.is_non_overlapping_and_dense.default:
+        return False
+
+    if func is torch.ops.aten.sym_size.default:
+        return args[0]._size
+
+    if func is torch.ops.aten.dim.default:
+        return len(args[0]._size)
+
+    if func in (torch.ops.aten.sym_numel.default, torch.ops.aten.numel.default):
+        if args[0]._lengths is not None:
+            return int(sum(args[0]._lengths) * math.prod(args[0]._size[2:]))
+        return args[0]._values.numel()
+
+    if func is torch.ops.aten.sym_stride.default:
+        return args[0]._strides
+
+    if func is torch.ops.aten.sym_storage_offset.default:
+        return args[0]._values.storage_offset()
+
+
+@register_jagged_func(torch.ops.prim.layout.default, "self: jt_all")
+def prim_layout_default(func, *args, **kwargs):
+    return torch.jagged
+
+
+@register_jagged_func(
+    [torch.ops.aten.size.default],
+    "self: jt_all",
+)
+def tensor_attr_unsupported_getter(func, *args, **kwargs) -> None:
+    if func is torch.ops.aten.size.default:
+        raise RuntimeError(
+            "NestedTensor does not support directly calling torch.ops.aten.size; "
+            "please use `nested_tensor.size()` instead."
+        )
+
+
+@register_jagged_func(torch.ops.aten.is_contiguous.default, "self: jt_all")
+def is_contiguous_general(func, *args, **kwargs):
+    from torch._prims_common import is_contiguous_for_memory_format
+
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    inp = new_kwargs.pop("input")
+
+    # If created from narrow() check for lengths
+    if inp.lengths() is not None:
+        return False
+
+    new_kwargs["memory_format"] = new_kwargs.get(
+        "memory_format", torch.contiguous_format
+    )
+    if new_kwargs["memory_format"] == torch.preserve_format:
+        return True
+    return is_contiguous_for_memory_format(inp._values, **new_kwargs)
+
+
+register_jagged_func(
+    torch.ops.aten.is_contiguous.memory_format, "self: jt_all, memory_format: any?"
+)(is_contiguous_general)
+
+
+@register_jagged_func(
+    torch.ops.aten.sym_is_contiguous.default, "self: jt_all, memory_format: any?"
+)
+def sym_is_contiguous_general(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    inp = new_kwargs.pop("input")
+
+    # If created from narrow() check for lengths
+    if inp.lengths() is not None:
+        return False
+
+    new_kwargs["memory_format"] = new_kwargs.get(
+        "memory_format", torch.contiguous_format
+    )
+
+    if new_kwargs["memory_format"] == torch.preserve_format:
+        return True
+
+    return torch.ops.aten.sym_is_contiguous.default(inp._values, **new_kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.clone.default, "input: jt_all, memory_format: any?"
+)
+def clone_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    new_meta = extract_kwargs(inp)
+
+    if inp._lengths is not None:
+        if new_kwargs["memory_format"] == torch.contiguous_format:
+            # need to copy to remove "holes" non-contiguity / lengths metadata
+            # TODO: write a kernel for this
+            from .nested_tensor import jagged_from_list
+
+            # TODO: We probably want the output to have the same ragged structure / nested int.
+            if inp._ragged_idx != 1:
+                raise AssertionError(
+                    "NJT with ragged_idx != 1 not supported for contiguous clone"
+                )
+            contig, _ = jagged_from_list(inp.unbind(), offsets=None)
+            return contig
+
+    return NestedTensor(func(inp._values, **new_kwargs), **new_meta)
+
+
+@register_jagged_func(torch.ops.aten.linear.default, "input: jt, weight: t, bias: t?")
+def linear_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(
+    torch.ops.aten.linear_backward.default,
+    "self: jt, grad_output: jt, weight: t, output_mask: any",
+)
+def linear_backward_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    grad_output = new_kwargs.pop("grad_output")
+    weight = new_kwargs.pop("weight")
+    output_mask = new_kwargs.pop("output_mask")
+
+    ds, dw, db = None, None, None
+    check_ragged_dim_same(func, inp, "self", grad_output, "grad_output")
+    if output_mask[0]:
+        ds = NestedTensor(
+            torch.matmul(grad_output._values, weight), **extract_kwargs(grad_output)
+        )
+    if output_mask[1]:
+        # NB: Fold dims of values for input and grad_output to treat them as 2D. This
+        # trick avoids materializing large intermediates and immediately reducing over
+        # them via sum(). This is equivalent to computing:
+        #     torch.matmul(grad_output._values.transpose(-2, -1), inp._values)
+        # and then summing over the leading dimensions to get a 2D weight grad.
+        grad_2d = grad_output._values.reshape(-1, weight.size(0))
+        input_2d = inp._values.reshape(-1, weight.size(1))
+        dw = torch.matmul(grad_2d.t(), input_2d)
+    if output_mask[2]:
+        # Sum over all but the last dim to get a 1D bias grad. We cannot
+        # rely on the autograd engine to reduce for us, because returning a
+        # tensor aliasing the input would violate the aten signature annotation
+        reduce_dims = tuple(range(grad_output._values.ndim - 1))
+        if reduce_dims == ():
+            db = grad_output._values.clone()
+        else:
+            db = torch.sum(grad_output._values, reduce_dims, keepdim=False)
+    return (ds, dw, db)
+
+
+@register_jagged_func(torch.ops.aten.to.dtype, "input: jt_all, dtype: any")
+def to_dtype(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(torch.ops.aten._to_copy.default, "self: jt_all")
+def to_copy_default(func, *args, **kwargs):
+    from .nested_tensor import _tensor_symint_registry
+
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    # don't change layout
+    new_kwargs.pop("layout")
+
+    new_values = func(inp._values, **new_kwargs)
+    new_offsets = inp._offsets.to(device=new_values.device)
+    new_lengths = None
+    if inp._lengths is not None:
+        new_lengths = inp._lengths.to(device=new_values.device)
+
+    from torch._subclasses.fake_tensor import FakeTensor
+    from torch._subclasses.functional_tensor import (
+        FunctionalTensor,
+        mb_unwrap_functional_tensor,
+    )
+
+    ragged_source = inp._offsets if inp._lengths is None else inp._lengths
+    new_thing = new_offsets if new_lengths is None else new_lengths
+    if isinstance(new_thing, (FakeTensor, FunctionalTensor)):
+        # Temporary hack until we have the union find
+        tgt = mb_unwrap_functional_tensor(new_thing)
+        src = mb_unwrap_functional_tensor(ragged_source)
+        # pyrefly: ignore[missing-attribute]
+        tgt.nested_int_memo = src.nested_int_memo
+    else:
+        _tensor_symint_registry[new_thing] = _tensor_symint_registry[ragged_source]
+    inp_kwargs = extract_kwargs(inp)
+    inp_kwargs["offsets"] = new_offsets
+    inp_kwargs["lengths"] = new_lengths
+
+    output = NestedTensor(new_values, **inp_kwargs)
+    return output
+
+
+@register_jagged_func(
+    torch.ops.aten.copy_.default, "self: jt_all, src: jt_all, non_blocking: any?"
+)
+def copy_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    inp = new_kwargs.pop("input")
+    src = new_kwargs.pop("src")
+    if inp._size != src._size:
+        # try to recursively copy_ on unbound components to get around nested int mismatch
+        # TODO: eventually do a direct copy when this is possible
+        inp_comps = inp.unbind()
+        inp_comp_shapes = [c.shape for c in inp_comps]
+        src_comps = src.unbind()
+        src_comp_shapes = [c.shape for c in src_comps]
+        if inp_comp_shapes != src_comp_shapes:
+            raise RuntimeError(
+                "copy_(): expected compatible input and src shapes, but got: "
+                f"{inp.shape} and {src.shape}"
+            )
+        for inp_comp, src_comp in zip(inp_comps, src_comps):
+            inp_comp.copy_(src_comp)
+
+    # AOTD allows mutations of inputs only, (not views of the inputs).
+    # NJT.values() returns _values.detach() to workaround some issues.
+    # To keep mutation in the graph, AOTD manually calls copy_ on the input (NJT).
+    # Here we directly mutate self._values to not emit .detach() in the graph, which would make it non-compilable.
+    inp._values.copy_(src._values)
+    return inp
+
+
+register_jagged_func(torch.ops.aten.detach.default, "self: jt_all")(
+    jagged_unary_pointwise
+)
+
+
+@register_jagged_func(
+    [
+        torch.ops.aten.empty_like.default,
+        torch.ops.aten.ones_like.default,
+        torch.ops.aten.zeros_like.default,
+        torch.ops.aten.rand_like.default,
+        torch.ops.aten.randn_like.default,
+    ],
+    "self: jt_all",
+)
+def like_factory_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    # Default layout is technically torch.strided but only jagged is supported here.
+    # Rather than force users to specify the layout, assume jagged.
+    # This should be set to strided for redispatching on values.
+    new_kwargs["layout"] = torch.strided
+
+    new_values = func(inp._values, **new_kwargs)
+    new_offsets = inp._offsets.to(device=new_values.device)
+    new_lengths = None
+    if inp._lengths is not None:
+        new_lengths = inp._lengths.to(device=new_values.device)
+    output_kwargs = extract_kwargs(inp)
+    if "offsets" in output_kwargs:
+        output_kwargs["offsets"] = new_offsets
+    if "lengths" in output_kwargs:
+        output_kwargs["lengths"] = new_lengths
+
+    if inp.device != new_values.device:
+        # Update the nested int registry to indicate that the ragged structure is the same
+        # between the two offsets / lengths on different devices.
+        from torch._subclasses.fake_tensor import FakeTensor
+        from torch._subclasses.functional_tensor import (
+            FunctionalTensor,
+            mb_unwrap_functional_tensor,
+        )
+
+        from .nested_tensor import _tensor_symint_registry
+
+        ragged_source = inp._offsets if inp._lengths is None else inp._lengths
+        new_thing = new_offsets if new_lengths is None else new_lengths
+        if isinstance(new_thing, (FakeTensor, FunctionalTensor)):
+            # Temporary hack until we have the union find
+            tgt = mb_unwrap_functional_tensor(new_thing)
+            src = mb_unwrap_functional_tensor(ragged_source)
+            # pyrefly: ignore[missing-attribute]
+            tgt.nested_int_memo = src.nested_int_memo
+        else:
+            _tensor_symint_registry[new_thing] = _tensor_symint_registry[ragged_source]
+
+    return NestedTensor(new_values, **output_kwargs)
+
+
+register_jagged_func(torch.ops.aten.full_like.default, "self: jt_all, fill_value: any")(
+    like_factory_default
+)
+
+register_jagged_func(torch.ops.aten.randint_like.default, "self: jt_all, high: any")(
+    like_factory_default
+)
+
+register_jagged_func(
+    torch.ops.aten.randint_like.low_dtype, "self: jt_all, low: any, high: any"
+)(like_factory_default)
+
+
+@register_jagged_func(torch.ops.aten.zero_.default, "self: jt_all")
+def zero__default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    func(inp._values)
+    return inp
+
+
+@register_jagged_func(
+    torch.ops.aten._softmax.default, "self: jt_all, dim: any, half_to_float: any"
+)
+def _softmax_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    if isinstance(new_kwargs["dim"], tuple):
+        raise RuntimeError(
+            "softmax(): not supported for dimensions of type 'tuple' for NestedTensor"
+        )
+
+    inp = new_kwargs.pop("input")
+
+    (
+        new_kwargs["dim"],
+        reduce_on_batch,
+        reduce_on_ragged,
+        _reduce_on_non_batch,
+    ) = _wrap_jagged_dims(
+        inp.dim(),
+        (new_kwargs["dim"],),
+        "softmax",
+        inp._ragged_idx,
+    )
+
+    if reduce_on_batch:
+        raise RuntimeError(
+            "softmax(): not supported when reducing across the batch dimension for NestedTensor"
+        )
+
+    if reduce_on_ragged and inp._ragged_idx > 1:
+        raise RuntimeError(
+            "softmax(): not supported when reducing along the ragged dimension for ragged_idx > 1 for NestedTensor"
+        )
+
+    if reduce_on_ragged and inp._lengths is not None:
+        raise RuntimeError(
+            "softmax(): not supported where lengths is not None "
+            + "if reducing across the ragged dimension for NestedTensor"
+        )
+
+    new_kwargs["dim"] = new_kwargs["dim"][
+        0
+    ]  # torch.softmax takes in the reduction dimension as an integer
+
+    if reduce_on_ragged:
+        padded_softmax_values = torch.nn.functional.softmax(
+            torch.ops.aten._jagged_to_padded_dense_forward(
+                inp._values.reshape(
+                    inp._values.shape[0], -1
+                ),  # values are required to be 2D tensors for j2pd
+                [inp._offsets],
+                max_lengths=[inp._max_seqlen],  # max length of ragged dimension
+                padding_value=float("-inf"),  # e^-inf = 0
+            ),
+            dim=inp._ragged_idx,
+        )
+
+        softmax_values = torch.ops.aten._padded_dense_to_jagged_forward(
+            padded_softmax_values,
+            [inp._offsets],
+            total_L=inp._values.shape[
+                0
+            ],  # providing this parameter helps avoid a GPU/CPU sync
+        ).reshape(
+            -1, *inp._values.shape[1:]
+        )  # expand softmax_values back to original shape (inp._values.shape)
+
+        return NestedTensor(softmax_values, **extract_kwargs(inp))
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(
+    torch.ops.aten._log_softmax.default, "self: jt_all, dim: any, half_to_float: any"
+)
+def _log_softmax_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    if isinstance(new_kwargs["dim"], tuple):
+        raise RuntimeError(
+            "log_softmax(): not supported for dimensions of type 'tuple' for NestedTensor"
+        )
+
+    inp = new_kwargs.pop("input")
+
+    (
+        new_kwargs["dim"],
+        reduce_on_batch,
+        reduce_on_ragged,
+        _reduce_on_non_batch,
+    ) = _wrap_jagged_dims(
+        inp.dim(), (new_kwargs["dim"],), "log_softmax", inp._ragged_idx
+    )
+
+    if reduce_on_batch:
+        raise RuntimeError(
+            "log_softmax(): not supported when reducing across the batch dimension for NestedTensor"
+        )
+
+    if reduce_on_ragged:
+        raise RuntimeError(
+            "log_softmax(): not supported when reducing along the ragged dimension for NestedTensor"
+        )
+
+    # torch.log_softmax takes in the reduction dimension as an integer
+    new_kwargs["dim"] = new_kwargs["dim"][0]
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(
+    torch.ops.aten._softmax_backward_data.default,
+    "grad_output: jt, output: jt, dim: any, input_dtype: any",
+)
+def _softmax_backward(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    grad_out = new_kwargs.pop("grad_output")
+    output = new_kwargs.pop("output")
+    return NestedTensor(
+        func(grad_out._values, output._values, **new_kwargs), **extract_kwargs(grad_out)
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.native_dropout.default, "self: jt, float: any, train: any?"
+)
+def native_dropout_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    out1, out2 = func(inp._values, **new_kwargs)
+    return (
+        NestedTensor(out1, **extract_kwargs(inp)),
+        NestedTensor(out2, **extract_kwargs(inp)),
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.native_dropout_backward.default,
+    "grad_output: jt, mask: jt, scale: any",
+)
+def native_dropout_backward_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    grad_output = new_kwargs.pop("grad_output")
+    mask = new_kwargs.pop("mask")
+    return NestedTensor(
+        func(grad_output._values, mask._values, **new_kwargs),
+        **extract_kwargs(grad_output),
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.prod.dim_int,
+    "self: jt_all, dim: any, keepdim: any?, dtype: any?",
+)
+def prod_dim_int(func, *args, **kwargs):
+    return _apply_reduction(func, "prod", 1, *args, **kwargs)
+
+
+@register_jagged_func(torch.ops.aten.prod.default, "self: jt_all, dtype: any?")
+def prod_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return func(inp._values, **new_kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.split.Tensor, "self: jt, split_size: any, dim: any?"
+)
+def split_tensor(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        inp.dim(), new_kwargs["dim"], inp._ragged_idx, "split"
+    )
+
+    return tuple(
+        NestedTensor(values=x, **extract_kwargs(inp))
+        for x in func(inp._values, **new_kwargs)
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.split_with_sizes.default, "self: jt, split_sizes: any, dim: any?"
+)
+def split_with_sizes_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        inp.dim(), new_kwargs["dim"], inp._ragged_idx, "split_with_sizes"
+    )
+
+    return [
+        NestedTensor(values=x, **extract_kwargs(inp))
+        for x in func(inp._values, **new_kwargs)
+    ]
+
+
+@register_jagged_func(
+    torch.ops.aten.narrow.default, "self: jt, dim: any, start: any, length: any"
+)
+def narrow(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    inp = new_kwargs.pop("input")
+
+    dim = _wrap_jagged_dim(inp.dim(), new_kwargs["dim"], inp._ragged_idx, "narrow")
+    values = func(
+        inp._values,
+        dim=dim,
+        start=new_kwargs["start"],
+        length=new_kwargs["length"],
+    )
+    return NestedTensor(values, **extract_kwargs(inp))
+
+
+@register_jagged_func(torch.ops.aten.chunk.default, "self: jt, chunks: any, dim: any?")
+def chunk_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    new_kwargs["dim"], operating_on_batch = _wrap_jagged_dim(
+        inp.dim(), new_kwargs["dim"], inp._ragged_idx, "chunk", allow_batch_dim=True
+    )
+
+    if operating_on_batch:
+        chunks = new_kwargs["chunks"]
+
+        # get _offsets of the chunks
+        lengths = inp._offsets.diff()
+        chunked_lengths = lengths.chunk(chunks)
+        chunked_offsets = [torch.cumsum(x, dim=0) for x in chunked_lengths]
+        chunked_offsets = [F.pad(x, (1, 0), value=0) for x in chunked_offsets]  # type: ignore[arg-type]
+        nested_kwargs = [
+            {"offsets": per_offsets, "_ragged_idx": inp._ragged_idx}
+            for per_offsets in chunked_offsets
+        ]
+
+        # get _values of the chunks
+        split_sizes = [x.sum().item() for x in chunked_lengths]
+        chunk_values = inp._values.split(split_sizes)
+
+        # Note that the actual number of chunks returned is not necessarily the same as
+        # the input number; it can be counter-intuitive, but it matches dense behavior.
+        return [
+            NestedTensor(values=chunk_values[i], **(nested_kwargs[i]))
+            for i in range(len(chunk_values))
+        ]
+    else:
+        return [
+            NestedTensor(values=x, **extract_kwargs(inp))
+            for x in func(inp._values, **new_kwargs)
+        ]
+
+
+@register_jagged_func(torch.ops.aten.unbind.int, "self: jt_all, dim: any?")
+def unbind_int(func, *args, **kwargs):
+    # Note that this specializes on the length of the offsets
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    dim = new_kwargs["dim"]
+    if dim != 0:
+        raise RuntimeError("unbind(): only supported for NestedTensor on dim=0")
+
+    inp = new_kwargs.pop("input")
+    values = inp.values()
+    offsets = inp.offsets()
+    lengths = inp.lengths()
+    ragged_idx = inp._ragged_idx
+
+    def _torch_check(_lengths: list[int], _offsets: list[int] | None = None) -> None:
+        # This torch._check are needed for torch.compile
+        # symbolic shapes processing.
+        # offsets and lengths are symbolic variables during compilation,
+        # we guarantee the correct offsets/lengths correspondence:
+        # sum of lengths <= total ragged_dim_size
+        # every length and offset are size-like variable (allows sym shapes to reason it as [2, inf))
+        # offset[i] + length[i] <= ragged_dim_size, for unbind and split dim correctness
+        # offsets[i] <= ragged_dim_size
+
+        lengths_sum = 0
+        ragged_dim_size = values.shape[ragged_idx - 1]
+        for i in range(len(_lengths)):
+            torch._check(_lengths[i] >= 0)
+            torch._check(_lengths[i] <= ragged_dim_size)
+
+            lengths_sum += _lengths[i]
+            if _offsets is not None:
+                torch._check(
+                    _offsets[i] + _lengths[i] <= ragged_dim_size,
+                    lambda: "unbind(): nested tensor offsets and lengths do not match ragged_idx dimension",
+                )
+        torch._check(lengths_sum <= ragged_dim_size)
+
+        if _offsets is not None:
+            for i in range(len(_offsets)):
+                torch._check(_offsets[i] >= 0)
+                torch._check(_offsets[i] <= ragged_dim_size)
+
+    if lengths is None:
+        lengths_scalars = offsets.diff().tolist()
+        _torch_check(lengths_scalars)
+
+        return torch.split(values, lengths_scalars, dim=(ragged_idx - 1))
+
+    if ragged_idx <= 0:
+        raise RuntimeError(
+            "unbind(): nested tensor ragged_idx out of bounds (should be >= 1)"
+        )
+
+    lengths_scalars = lengths.tolist()
+    offsets_scalars = offsets.tolist()
+
+    _torch_check(lengths_scalars, offsets_scalars)
+
+    return [
+        torch.narrow(
+            values,
+            dim=(ragged_idx - 1),
+            start=offsets_scalars[i],
+            length=lengths_scalars[i],
+        )
+        for i in range(lengths.shape[0])
+    ]
+
+
+@register_jagged_func(torch.ops.aten.squeeze.dim, "self: jt, dim: any")
+def squeeze_dim(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    values = inp._values
+
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        len(inp._size), new_kwargs["dim"], inp._ragged_idx, "squeeze"
+    )
+    return NestedTensor(func(values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(torch.ops.aten.unsqueeze.default, "self: jt_all, dim: any")
+def unsqueeze_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    values = inp._values
+
+    # Account for collapsed jagged dim
+    dim = new_kwargs["dim"]
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        len(inp._size) + 1, dim, inp._ragged_idx, "unsqueeze", allow_ragged_dim=True
+    )
+
+    # ragged_idx changes if a dimension is added before it
+    output_kwargs = extract_kwargs(inp)
+    if new_kwargs["dim"] <= inp._ragged_idx - 1:
+        output_kwargs["_ragged_idx"] += 1
+
+    return NestedTensor(func(values, **new_kwargs), **output_kwargs)
+
+
+@register_jagged_func(torch.ops.aten.cat.default, "tensors: any, dim: any?")
+def cat_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    tensors = new_kwargs.pop("tensors")
+
+    # Convert any non-nested to nested
+    nested = [t for t in tensors if t.is_nested]
+    if len(nested) == 0:
+        raise AssertionError("At least one tensor must be nested")
+    first = nested[0]
+    tensors = [t if t.is_nested else t.expand_as(first) for t in tensors]
+
+    # Account for collapsed jagged dim
+    dim = new_kwargs["dim"]
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        len(first.shape), dim, first._ragged_idx, "cat"
+    )
+
+    return NestedTensor(
+        func([t._values for t in tensors], **new_kwargs), **extract_kwargs(tensors[0])
+    )
+
+
+@register_jagged_func(torch.ops.aten.matmul.default, "self: any, other: any")
+def matmul_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    other = new_kwargs.pop("other")
+
+    def _unbind_impl(a, b):
+        return [
+            func(a_comp, b_comp) for (a_comp, b_comp) in zip(a.unbind(), b.unbind())
+        ]
+
+    def _padded_impl(a, b):
+        if a.is_nested:
+            nt = a
+        else:
+            nt = b
+
+        from .nested_tensor import nested_from_padded
+
+        min_seqlen = nt._maybe_min_seqlen
+        max_seqlen = nt._maybe_max_seqlen
+        padded_max_S = max_seqlen
+        total_L = nt._values.shape[nt._ragged_idx - 1]
+        if padded_max_S is None:
+            # use upper bound on max seqlen if it's not present
+            padded_max_S = total_L
+
+        padded_shape = (
+            *nt.shape[: nt._ragged_idx],
+            padded_max_S,
+            *nt.shape[nt._ragged_idx + 1 :],
+        )
+        padded_nt = nt.to_padded_tensor(0.0, output_size=padded_shape)
+        if a.is_nested:
+            padded_t = func(padded_nt, b)
+        else:
+            padded_t = func(a, padded_nt)
+        return nested_from_padded(
+            padded_t,
+            offsets=nt._offsets,
+            ragged_idx=nt._ragged_idx,
+            sum_S=total_L,
+            min_seqlen=min_seqlen,
+            max_seqlen=max_seqlen,
+        )
+
+    # TODO: Back these with proper kernels (e.g. grouped GEMM)
+    # NJT x dense
+    if inp.is_nested and not other.is_nested:
+        # (B, j1, D) x (B, D, E) => (B, j1, E)
+        if (
+            inp.dim() >= 3
+            and inp.dim() == other.dim()
+            and inp._ragged_idx < inp.dim() - 1
+        ):
+            # convert to padded for this
+            return _padded_impl(inp, other)
+        # Support broadcasting the dense:
+        # (B, j1, D) x (D, E) => (B, j1, E)
+        # (B, j1, D, E) x (E, F) => (B, j1, D, F)
+        # etc.
+        elif (
+            other.dim() == 2
+            and inp.dim() > other.dim()
+            and inp._ragged_idx < inp.dim() - 1
+        ):
+            return NestedTensor(
+                func(inp._values, other, **new_kwargs), **extract_kwargs(inp)
+            )
+    # Dense x NJT
+    elif not inp.is_nested and other.is_nested:
+        # (B, D, E) x (B, E, j1) => (B, E, j1)
+        if other.dim() >= 3 and other.dim() == inp.dim() and other._ragged_idx >= 2:
+            # convert to padded for this
+            return _padded_impl(inp, other)
+        # Support broadcasting the dense:
+        # (D, E) x (B, E, j1) => (B, D, j1)
+        # (D, E) x (B, E, j1, F) => (B, D, j1, F)
+        # etc.
+        elif inp.dim() == 2 and other.dim() > inp.dim() and other._ragged_idx >= 2:
+            return NestedTensor(
+                func(inp, other._values, **new_kwargs), **extract_kwargs(other)
+            )
+
+    # NJT x NJT
+    elif inp.is_nested and other.is_nested:
+        # Support ragged batch dim:
+        # (B, j1, D, E) x (B, j1, E, F) => (B, j1, D, F), etc.
+        if inp.dim() > 3 and other.dim() > 3 and raggedness_matches(inp, other._size):
+            return NestedTensor(func(inp._values, other._values), **extract_kwargs(inp))
+        # Support reducing over ragged with dense output:
+        # (B, D, j1) x (B, j1, E) => (B, D, E)
+        elif (
+            inp.dim() == 3
+            and other.dim() == 3
+            and inp._ragged_idx == 2
+            and other._ragged_idx == 1
+            and inp.size(inp._ragged_idx) == other.size(other._ragged_idx)
+        ):
+            # do unbind for this; can't use padded conversion due to j1 in last dim
+            return torch.stack(_unbind_impl(inp, other))
+
+    raise RuntimeError(
+        f"matmul(): not supported between inputs of shapes {inp._size} and {other.shape}"
+    )
+
+
+@register_jagged_func(torch.ops.aten.bmm.default, "self: jt_all, mat2: any")
+def bmm_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    other = new_kwargs.pop("mat2")
+
+    if inp.dim() != 3:
+        raise ValueError("bmm(): input must be 3D")
+    if other.dim() != 3:
+        raise ValueError("bmm(): mat2 must be 3D")
+
+    return matmul_default(torch.ops.aten.matmul.default, inp, other)
+
+
+@register_jagged_func(
+    torch.ops.aten.expand.default, "self: jt_all, size: any, implicit: any?"
+)
+def expand_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    size = new_kwargs["size"]
+
+    if "implicit" in new_kwargs and new_kwargs.pop("implicit"):
+        raise AssertionError("implicit expand is not supported")
+    if not raggedness_matches(inp, size):
+        raise RuntimeError(f"expand(): cannot expand shape {inp._size} -> {size}")
+
+    expand_arg = [-1 if d == inp._ragged_idx else size[d] for d in range(1, inp.dim())]
+    return NestedTensor(func(inp._values, expand_arg), **extract_kwargs(inp))
+
+
+@register_jagged_func(torch.ops.aten.expand_as.default, "self: t, other: jt")
+def expand_as_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    other = new_kwargs.pop("other")
+
+    return NestedTensor(func(inp, other._values), **extract_kwargs(other))
+
+
+@register_jagged_func(torch.ops.aten.broadcast_to.default, "self: jt_all, size: any")
+def broadcast_to(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    size = new_kwargs.pop("size")
+
+    if len(size) <= inp.dim():
+        return inp.expand([*(1 for _ in range(inp.dim() - len(size))), *size])
+
+    raise ValueError(
+        "broadcast_to(): broadcasting to a higher-dim shape is currently not supported "
+        "for nested tensors with the jagged layout"
+    )
+
+
+@register_jagged_func(torch.ops.aten.broadcast_tensors.default, "tensors: any")
+def broadcast_tensors(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    tensors = new_kwargs.pop("tensors")
+    if len(tensors) == 0:
+        raise ValueError("broadcast_tensors(): expected at least one tensor input")
+    if len(tensors) == 1:
+        return tensors[0]
+
+    outs = []
+    broadcast_shape = torch.broadcast_shapes(*(t.shape for t in tensors))
+    # Pull out the first NJT. If broadcast_shapes() worked, the nested ints are compatible.
+    njt = next(t for t in tensors if isinstance(t, NestedTensor))
+    for t in tensors:
+        if t.is_nested:
+            outs.append(t.broadcast_to(broadcast_shape))
+        elif t.dim() < len(broadcast_shape):
+            outs.append(
+                NestedTensor(t.broadcast_to(njt._values.shape), **extract_kwargs(njt))
+            )
+        else:
+            raise ValueError(
+                "broadcast_tensors(): broadcasting nested tensors with dense tensors of equal "
+                "or higher dim is not currently supported"
+            )
+
+    return tuple(outs)
+
+
+@register_jagged_func(
+    torch.ops.aten.where.self, "condition: jt_all, self: any, other: any"
+)
+def where_self(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    condition = new_kwargs.pop("condition")
+    inp = new_kwargs.pop("input")
+    other = new_kwargs.pop("other")
+
+    # if the tensors aren't compatible, broadcast_tensors() will let us know
+    condition, inp, other = torch.broadcast_tensors(condition, inp, other)
+
+    return NestedTensor(
+        func(condition._values, inp._values, other._values, **new_kwargs),
+        **extract_kwargs(condition),
+    )
+
+
+@register_jagged_func(torch.ops.aten._pin_memory.default, "self: jt, device: any?")
+def _pin_memory_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(torch.ops.aten.is_pinned.default, "self: jt, device: any?")
+def is_pinned_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return func(inp._values, **new_kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.is_same_size.default, "self: jt_all, other: jt_all"
+)
+def is_same_size_default(func, *args, **kwargs):
+    return args[0]._size == args[1]._size
+
+
+def _apply_reduction(func, func_name, identity_element, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    # some ops use dim=None to indicate a full reduction; some use an empty dim list
+    full_reduction = new_kwargs["dim"] is None or (
+        isinstance(new_kwargs["dim"], (tuple, list)) and len(new_kwargs["dim"]) == 0
+    )
+    if full_reduction:
+        out = func(inp._values, **new_kwargs)
+        if new_kwargs.get("keepdim", False):
+            if isinstance(out, (tuple, list)):
+                # some ops return multiple things; unsqueeze all of them
+                out = type(out)(o.unsqueeze(inp._ragged_idx) for o in out)
+            else:
+                out = out.unsqueeze(inp._ragged_idx)
+        return out
+
+    # some ops support lists of dims; some don't
+    dim_to_convert = new_kwargs["dim"]
+    is_dimlist = isinstance(new_kwargs["dim"], (tuple, list))
+    if not is_dimlist:
+        dim_to_convert = [dim_to_convert]
+
+    (
+        converted_dim,
+        reduce_on_batch,
+        reduce_on_ragged,
+        reduce_on_non_batch,
+    ) = _wrap_jagged_dims(
+        inp.dim(),
+        dim_to_convert,
+        f"{func_name}",
+        inp._ragged_idx,
+    )
+
+    if not is_dimlist:
+        # convert back from list
+        converted_dim = converted_dim[0]
+    new_kwargs["dim"] = converted_dim
+
+    if reduce_on_ragged and inp._lengths is not None:
+        raise RuntimeError(
+            f"{func_name}(): reducing across the ragged dimension is not supported "
+            "for non-contiguous nested tensors with holes"
+        )
+
+    from torch.utils._pytree import tree_map
+
+    # raggedness reduced away --> return dense tensor
+    if reduce_on_ragged:
+        # reduction cases: (batch, ragged), (batch, ragged, non-batch), etc.
+        if reduce_on_batch:
+            # no need to read offsets --> apply sum directly on values
+            out = func(inp._values, **new_kwargs)
+            if new_kwargs.get("keepdim", False):
+                # some ops return multiple things; unsqueeze all of them
+                out = tree_map(lambda o: o.unsqueeze(0), out)
+            return out
+        else:
+            # invalid reduction cases: (ragged, non-batch), etc.
+            if reduce_on_non_batch:
+                raise RuntimeError(
+                    f"{func_name}(): reducing along a ragged and non-batch dimension "
+                    "is not supported for nested tensors"
+                )
+
+            # reduction cases: (ragged)
+            # convert to padded dense and reduce
+            new_kwargs.pop("dim")
+            dim_to_pass = [inp._ragged_idx] if is_dimlist else inp._ragged_idx
+            return func(
+                inp.to_padded_tensor(identity_element), dim=dim_to_pass, **new_kwargs
+            )
+    # raggedness preserved --> return nested tensor
+    else:
+        # invalid reduction cases: (batch), (batch, non-batch), etc.
+        if reduce_on_batch:
+            raise RuntimeError(
+                f"{func_name}(): reducing along the batch dimension but not "
+                "the ragged dimension is not supported for nested tensors"
+            )
+
+        # reduction cases: (non-batch), (non-batch, non-batch), etc.
+        # apply sum directly on values
+        out = func(inp._values, **new_kwargs)
+        out_kwargs = extract_kwargs(inp)
+        if not new_kwargs.get("keepdim", False):
+            # dims are reduced away -> ragged_idx of output needs to be reevaluated
+            dimlist = (
+                new_kwargs["dim"]
+                if isinstance(new_kwargs["dim"], (tuple, list))
+                else [new_kwargs["dim"]]
+            )
+            for d in dimlist:
+                # adjust for all dims reduced before the ragged dim
+                if d < inp._ragged_idx - 1:
+                    out_kwargs["_ragged_idx"] -= 1
+
+        # some ops return multiple things; wrap each of them as an NJT
+        return tree_map(lambda o: NestedTensor(o, **out_kwargs), out)
+
+
+@register_jagged_func(torch.ops.aten.sum.default, "self: jt_all, dtype: any?")
+def sum_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return func(inp._values, **new_kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.sum.dim_IntList,
+    "self: jt_all, dim: any?, keepdim: any?, dtype: any?",
+)
+def sum_dim_IntList(func, *args, **kwargs):
+    return _apply_reduction(func, "sum", 0, *args, **kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.transpose.int, "self: jt_all, dim0: any, dim1: any"
+)
+def transpose_int(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    from torch._prims_common import canonicalize_dims
+
+    inp = new_kwargs.pop("input")
+    dim0, dim1 = canonicalize_dims(inp.dim(), (new_kwargs["dim0"], new_kwargs["dim1"]))
+
+    # To support the SDPA API, inputs need to have the ragged idx transposed to dim 2
+    # instead of 1, although the internal Flash and mem-effn implementations will
+    # use the inputs with raggedness in dim 1.
+    if dim0 == inp._ragged_idx or dim1 == inp._ragged_idx:
+        if dim0 == 0 or dim1 == 0:
+            raise ValueError(
+                "Transpose is not supported on the batch dimension for jagged NT"
+            )
+        if dim0 == inp._ragged_idx:
+            to_dim = dim1
+        else:
+            to_dim = dim0
+        inp_kwargs = extract_kwargs(inp)
+        inp_kwargs["_ragged_idx"] = to_dim
+        return NestedTensor(
+            inp.values().transpose(
+                _outer_to_inner_dim(len(inp._size), dim0, inp._ragged_idx),
+                _outer_to_inner_dim(len(inp._size), dim1, inp._ragged_idx),
+            ),
+            **inp_kwargs,
+        )
+
+    new_kwargs["dim0"] = _wrap_jagged_dim(
+        inp.dim(), new_kwargs["dim0"], inp._ragged_idx, "transpose"
+    )
+    new_kwargs["dim1"] = _wrap_jagged_dim(
+        inp.dim(), new_kwargs["dim1"], inp._ragged_idx, "transpose"
+    )
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(torch.ops.aten.permute.default, "self: jt_all, dims: any")
+def permute_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    inp = new_kwargs.pop("input")
+    dims = new_kwargs.pop("dims")
+    inp_kwargs = extract_kwargs(inp)
+    inp_dim = len(inp._size)
+
+    # The first two checks are the same as the checks in the normal permute implementation
+    if inp_dim != len(dims):
+        raise ValueError(
+            f"permute(): number of dimensions in the tensor input ({inp_dim}) "
+            + f"does not match the length of the desired ordering of dimensions ({len(dims)}).",
+        )
+
+    from torch._prims_common import canonicalize_dims
+
+    canonicalized_dims = canonicalize_dims(inp_dim, dims)
+
+    if len(canonicalized_dims) != len(set(canonicalized_dims)):
+        raise ValueError("permute(): duplicate dims are not allowed.")
+
+    if inp._lengths is not None:
+        raise ValueError(
+            "permute(): not supported on jagged layout nested tensor with holes"
+        )
+    if canonicalized_dims[0] != 0:
+        raise ValueError(
+            "Permute is not supported on the batch dimension for jagged NT"
+        )
+    inp_kwargs["_ragged_idx"] = canonicalized_dims.index(inp._ragged_idx)
+    inner_dims = [
+        _outer_to_inner_dim(inp_dim, dim, inp._ragged_idx)
+        for dim in canonicalized_dims[1:]
+    ]
+    new_kwargs["dims"] = inner_dims
+    return NestedTensor(func(inp._values, **new_kwargs), **inp_kwargs)
+
+
+@register_jagged_func(
+    [torch.ops.aten.view.default, torch.ops.aten._unsafe_view.default],
+    "self: jt_all, size: any",
+)
+def view_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    size = new_kwargs.pop("size")
+
+    if inp._ragged_idx != 1 and tuple(inp._size) != tuple(size):
+        raise RuntimeError(
+            f"view(): does not support ragged_idx != 1 except when inp._size == size. "
+            f"inp._size is ({inp._size}) and size is ({size})."
+        )
+
+    # Ensure specified size still includes batch and ragged dims
+    if len(size) < 3 or not raggedness_matches(inp, size):
+        raise RuntimeError(f"view(): cannot view shape {inp._size} as {size}")
+
+    # outer size: the size of the NT, e.g. [3, j0, 10]
+    # inner size: the size of the values, e.g. [8, 10] (e.g. for offsets = [0, 3, 5, 8])
+    # this function gets inner_size[inner_idx] for a given inner_idx.
+    #
+    # example: for outer size [a, b, c, j0, d, e, f]
+    #                         assume that j0 is ragged, other are concrete integers
+    #                         and ragged_idx=3
+    # inner size will be      [b, c, inp._values.size(ragged_idx), d, e, f]
+    # therefore:
+    #    inner_size[0] = outer_size[1]
+    #    inner_size[1] = outer_size[2]
+    #    inner_size[0] = inp._values.size(ragged_idx - 1)
+    #    inner_size[3] = outer_size[4]
+    #    inner_size[4] = outer_size[5]
+    def get_inner_size(inner_idx):
+        nonlocal inp, size
+        if inner_idx == inp._ragged_idx - 1:
+            return inp._values.size(inner_idx)
+        else:
+            return size[inner_idx + 1]
+
+    inner_size = [get_inner_size(i) for i in range(len(size) - 1)]
+
+    # Preserve inference-mode-ness of input.
+    # TODO: Do this for all other views!
+    with torch.inference_mode(inp.is_inference()):
+        return NestedTensor(func(inp._values, inner_size), **extract_kwargs(inp))
+
+
+@register_jagged_func(
+    torch.ops.aten.native_layer_norm.default,
+    "input: jt_all, normalized_shape: any, weight: any?, bias: any?, eps: any",
+)
+def native_layer_norm_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    if inp.dim() <= 2:
+        raise RuntimeError(
+            "layer_norm(): not supported for NestedTensor objects with 2 or fewer dimensions"
+        )
+
+    normalized_shape = new_kwargs["normalized_shape"]
+    ragged_size = inp.shape[inp._ragged_idx]
+
+    num_dims_not_normalized = inp.dim() - len(normalized_shape)
+
+    if (
+        num_dims_not_normalized == 0
+    ):  # error if trying to normalize over the batch dimension
+        raise RuntimeError(
+            "layer_norm(): not supported when normalizing over the batch dimension for NestedTensor"
+        )
+
+    if ragged_size in normalized_shape and inp._lengths is not None:
+        raise RuntimeError(
+            "layer_norm(): not supported where lengths is not None if operating on the ragged dimension for NestedTensor"
+        )
+
+    if (
+        ragged_size in normalized_shape
+    ):  # special handling for normalizing over the ragged dimension
+        padded_input = torch.ops.aten._jagged_to_padded_dense_forward(
+            inp._values.flatten(
+                start_dim=inp._ragged_idx
+            ),  # _jagged_to_padded_dense_forward requires values to be a 2D tensor
+            [inp._offsets],
+            max_lengths=[inp._max_seqlen],  # max length of ragged dimension
+        )
+
+        padded_mask = torch.ops.aten._jagged_to_padded_dense_forward(
+            torch.ones((inp._values.shape[0], 1), device=inp.device, dtype=inp.dtype),
+            [inp._offsets],
+            max_lengths=[inp._max_seqlen],  # max length of ragged dimension
+        ).expand(
+            padded_input.shape
+        )  # mask elements outside of the ragged dimension and expand to the same shape as padded input (3D dense tensor)
+
+        ragged_lengths = (
+            inp._offsets.diff().unsqueeze(1).unsqueeze(1) * padded_input.shape[2]
+        )  # ragged dim * inner dim, since we sum over dims (1, 2) (the layer on which we normalize)
+
+        mean = (
+            torch.sum(
+                padded_input,
+                dim=(1, 2),
+                keepdim=True,
+            )
+            / ragged_lengths
+        )  # a sum over (1, 2) ensures layer norm, whereas a sum over (1) would be an instance norm
+
+        padded_normalized = (
+            (padded_input - mean) * padded_mask
+        )  # mask elements outside of the ragged dimension size for correct variance calculation
+
+        variance = (
+            torch.sum(
+                torch.square(padded_normalized),
+                dim=(1, 2),
+                keepdim=True,
+            )
+            / ragged_lengths
+        )  # a sum over (1, 2) ensures layer norm, whereas a sum over (1) would be an instance norm
+
+        std = torch.sqrt(variance + new_kwargs["eps"])
+        padded_layer_norm = padded_normalized / std
+
+        jagged_layer_norm_values = torch.ops.aten._padded_dense_to_jagged_forward(
+            padded_layer_norm,
+            [inp._offsets],
+            total_L=inp._values.shape[
+                0
+            ],  # providing this parameter helps avoid a GPU/CPU sync
+        ).unflatten(
+            -1, inp.shape[inp._ragged_idx + 1 :]
+        )  # unflatten last dimension back into original nested tensor shape, e.g. (B, *, WH) --> (B, *, W, H)
+
+        return (
+            NestedTensor(jagged_layer_norm_values, **extract_kwargs(inp)),
+            mean,
+            std,
+        )
+
+    output, mean, std = func(inp._values, **new_kwargs)
+    return (NestedTensor(output, **extract_kwargs(inp)), mean, std)
+
+
+@register_jagged_func(
+    torch.ops.aten.native_layer_norm_backward.default,
+    "grad_out: jt, input: jt, normalized_shape: any, mean: any, rstd: any, weight: any?, bias: any?, output_mask: any",
+)
+def native_layer_norm_backward_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    grad_out = new_kwargs.pop("grad_out")
+    inp = new_kwargs.pop("input")
+    d_input, d_gamma, d_beta = func(grad_out._values, inp._values, **new_kwargs)
+    if d_input is None:
+        return (None, d_gamma, d_beta)
+
+    return (NestedTensor(d_input, **extract_kwargs(inp)), d_gamma, d_beta)
+
+
+@register_jagged_func(torch.ops.aten.select.int, "self: jt_all, dim: any, index: any")
+def select_int(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    new_kwargs["dim"], operating_on_batch = _wrap_jagged_dim(
+        inp.dim(), new_kwargs["dim"], inp._ragged_idx, "select", allow_batch_dim=True
+    )
+
+    # handle batch dim slicing via unbind() for now
+    # TODO: make this more efficient
+    if operating_on_batch:
+        return inp.unbind()[new_kwargs["index"]]
+
+    if inp._lengths is not None:
+        raise ValueError(
+            "select(): not yet supported on dim != 0 for non-contiguous nested tensor with holes"
+        )
+
+    # if selecting before the ragged dim, adjust output ragged_idx
+    out_kwargs = extract_kwargs(inp)
+    if new_kwargs["dim"] < inp._ragged_idx - 1:
+        out_kwargs["_ragged_idx"] -= 1
+
+    return NestedTensor(func(inp._values, **new_kwargs), **out_kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.slice.Tensor,
+    "self: jt, dim: any?, start: any?, end: any?, step: any?",
+)
+def slice_tensor(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        inp.dim(), new_kwargs["dim"], inp._ragged_idx, "slice"
+    )
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(
+    torch.ops.aten.index_put.default,
+    "input: jt_all, indices: any, values: t, accumulate: any?",
+)
+@register_jagged_func(
+    torch.ops.aten.index_put_.default,
+    "input: jt_all, indices: any, values: t, accumulate: any?",
+)
+def index_put_(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp: NestedTensor = new_kwargs.pop("input")
+
+    # For index_put_ to work, we add together the indices of the ragged dimension
+    # and the batch dimension, adding the offsets of each ragged dimension to its
+    # indices
+
+    indices = new_kwargs.pop("indices")
+
+    if len(indices) > inp.dim():
+        raise AssertionError(
+            f"Too many indices: got {len(indices)} but tensor has {inp.dim()} dimensions"
+        )
+
+    if len(indices) < inp._ragged_idx + 1:
+        if not inp.is_contiguous():
+            raise RuntimeError(
+                "index_put(): If ragged dimension is not part of indices, this only works on contiguous NJTs"
+            )
+        # Ragged dim is NOT part of indices, we need to pad the nested tensor to apply func
+        from .nested_tensor import nested_from_padded
+
+        min_seqlen = inp._maybe_min_seqlen
+        max_seqlen = inp._maybe_max_seqlen
+        padded_max_S = max_seqlen
+        total_L = inp._values.shape[inp._ragged_idx - 1]
+        if padded_max_S is None:
+            # use upper bound on max seqlen if it's not present
+            padded_max_S = total_L
+
+        padded_shape = (
+            *inp.shape[: inp._ragged_idx],
+            padded_max_S,
+            *inp.shape[inp._ragged_idx + 1 :],
+        )
+        padded_inp = inp.to_padded_tensor(0.0, output_size=padded_shape)
+        new_njt = nested_from_padded(
+            func(padded_inp, indices, **new_kwargs),
+            offsets=inp._offsets,
+            ragged_idx=inp._ragged_idx,
+            sum_S=total_L,
+            min_seqlen=min_seqlen,
+            max_seqlen=max_seqlen,
+        )
+
+        if func is torch.ops.aten.index_put_.default:
+            inp._values.copy_(new_njt.values())
+            return inp
+        return new_njt
+
+    # We can run on the underlying values directly
+
+    # Validate indices
+    if inp.lengths() is None:
+        lengths = inp.offsets().diff()
+    else:
+        lengths = inp.lengths()
+    torch._assert_async(
+        # pyrefly: ignore [no-matching-overload]
+        torch.all(indices[inp._ragged_idx] < lengths),
+        "Some indices in the ragged dimension are out of bounds!",
+    )
+
+    # Recompute indices for _values
+    ragged_indices = inp.offsets()[indices[0]] + indices[inp._ragged_idx]
+    func_indices = (
+        # before ragged dim
+        indices[1 : inp._ragged_idx]
+        # ragged dim (combined with batch)
+        + [ragged_indices]
+        # after ragged dim
+        + indices[inp._ragged_idx + 1 :]
+    )
+
+    if func is torch.ops.aten.index_put_.default:
+        inp._values = func(inp._values, func_indices, **new_kwargs)
+        return inp
+
+    return NestedTensor(
+        func(inp._values, func_indices, **new_kwargs),
+        **extract_kwargs(inp),
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.convolution.default,
+    "input: jt, weight: t, bias: t?, stride: any, padding: any, "
+    "dilation: any, transposed: any, output_padding: any, groups: any",
+)
+def convolution_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(
+    torch.ops.aten.mean.dim, "self: jt_all, dim: any?, keepdim: any?, dtype: any?"
+)
+def mean_dim(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs["input"]
+    (_, reduce_on_batch, reduce_on_ragged, reduce_on_non_batch) = _wrap_jagged_dims(
+        inp.dim(),
+        new_kwargs["dim"],
+        "mean",
+        inp._ragged_idx,
+    )
+
+    if reduce_on_ragged and not reduce_on_batch:
+        if reduce_on_non_batch:
+            raise AssertionError(
+                "Cannot reduce on both ragged and non-batch dimensions without also reducing on batch"
+            )
+        # calculate an intermediate sum and leave the dim in for normalization purposes
+        keepdim = new_kwargs["keepdim"]
+        new_kwargs["keepdim"] = True
+        intermediate_sum = _apply_reduction(
+            torch.ops.aten.sum.dim_IntList, "mean", 0, **new_kwargs
+        )
+
+        # normalize by sequence lengths
+        lengths = inp._lengths if inp._lengths is not None else inp._offsets.diff()
+        for _ in range(intermediate_sum.dim() - 1):
+            lengths = lengths.unsqueeze(-1)
+        out = intermediate_sum / lengths
+        if not keepdim:
+            out = out.squeeze(inp._ragged_idx)
+        return out
+
+    # at this point, we're just redispatching on the values buffer
+    # since we expect it to be unused, specify a weird intermediate value to
+    # hopefully make errors obvious
+    intermediate_value = 0.42
+    return _apply_reduction(func, "mean", intermediate_value, **new_kwargs)
+
+
+@register_jagged_func(torch.ops.aten.mean.default, "self: jt_all, dtype: any?")
+def mean_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return func(inp._values, **new_kwargs)
+
+
+@register_jagged_func(torch.ops.aten.any.dims, "self: jt_all, dim: any?, keepdim: any?")
+def any_dims(func, *args, **kwargs):
+    return _apply_reduction(func, "any", False, *args, **kwargs)
+
+
+@register_jagged_func(torch.ops.aten.any.dim, "self: jt_all, dim: any, keepdim: any?")
+def any_dim(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    # wrap dim in list to redispatch to dims overload
+    new_kwargs["dim"] = [new_kwargs["dim"]]
+    return any_dims(torch.ops.aten.any.dims, **new_kwargs)
+
+
+@register_jagged_func(torch.ops.aten.all.dims, "self: jt_all, dim: any?, keepdim: any?")
+def all_dims(func, *args, **kwargs):
+    return _apply_reduction(func, "all", True, *args, **kwargs)
+
+
+@register_jagged_func(torch.ops.aten.all.dim, "self: jt_all, dim: any, keepdim: any?")
+def all_dim(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    # wrap dim in list to redispatch to dims overload
+    new_kwargs["dim"] = [new_kwargs["dim"]]
+    return all_dims(torch.ops.aten.all.dims, **new_kwargs)
+
+
+@register_jagged_func(
+    [
+        torch.ops.aten.all.default,
+        torch.ops.aten.any.default,
+        torch.ops.aten.max.default,
+        torch.ops.aten.min.default,
+    ],
+    "self: jt_all",
+)
+def all_any_max_min_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return func(inp._values, **new_kwargs)
+
+
+@register_jagged_func(
+    [torch.ops.aten._is_all_true.default, torch.ops.aten._is_any_true.default],
+    "self: jt_all",
+)
+def _is_true_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    return func(inp._values)
+
+
+@register_jagged_func(torch.ops.aten.min.dim, "self: jt_all, dim: any, keepdim: any?")
+def min_dim(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    dtype = new_kwargs["input"].dtype
+    dtype_max = _get_padding_value(dtype, "max")
+    return _apply_reduction(func, "min", dtype_max, *args, **kwargs)
+
+
+@register_jagged_func(torch.ops.aten.max.dim, "self: jt_all, dim: any, keepdim: any?")
+def max_dim(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    dtype = new_kwargs["input"].dtype
+    dtype_min = _get_padding_value(dtype, "min")
+    return _apply_reduction(func, "max", dtype_min, *args, **kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.amin.default, "self: jt_all, dim: any?, keepdim: any?"
+)
+def amin_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    dtype = new_kwargs["input"].dtype
+    dtype_max = _get_padding_value(dtype, "max")
+    return _apply_reduction(func, "amin", dtype_max, *args, **kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.amax.default, "self: jt_all, dim: any?, keepdim: any?"
+)
+def amax_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    dtype = new_kwargs["input"].dtype
+    dtype_min = _get_padding_value(dtype, "min")
+    return _apply_reduction(func, "amax", dtype_min, *args, **kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.argmin.default, "self: jt_all, dim: any?, keepdim: any?"
+)
+def argmin_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    dtype = new_kwargs["input"].dtype
+    dtype_max = _get_padding_value(dtype, "max")
+    return _apply_reduction(func, "argmin", dtype_max, *args, **kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.argmax.default, "self: jt_all, dim: any?, keepdim: any?"
+)
+def argmax_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    dtype = new_kwargs["input"].dtype
+    dtype_min = _get_padding_value(dtype, "min")
+    return _apply_reduction(func, "argmax", dtype_min, *args, **kwargs)
+
+
+@register_jagged_func(
+    torch.ops.aten.value_selecting_reduction_backward.default,
+    "grad: jt_all, dim: any, indices: jt_all, sizes: any, keepdim: any",
+)
+def value_selecting_reduction_backward_default(func, *args, **kwargs):
+    from torch.fx.experimental.symbolic_shapes import is_nested_int
+
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    grad = new_kwargs.pop("grad")
+    new_kwargs["grad"] = grad._values
+    indices = new_kwargs.pop("indices")
+    new_kwargs["indices"] = indices._values
+    # should always succeed; sizes should contain a nested int
+    ragged_idx = next(i for i, s in enumerate(new_kwargs["sizes"]) if is_nested_int(s))
+    # convert dim -> values-space dim
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        len(new_kwargs["sizes"]),
+        new_kwargs["dim"],
+        ragged_idx,
+        "value_selecting_reduction_backward",
+    )
+    # convert saved NJT sizes -> values-space sizes
+    sizes = new_kwargs.pop("sizes")
+    sizes[ragged_idx] = indices._values.size(indices._ragged_idx - 1)
+    sizes = sizes[1:]
+    new_kwargs["sizes"] = sizes
+
+    output_kwargs = extract_kwargs(indices)
+    output_kwargs["_ragged_idx"] = ragged_idx
+
+    return NestedTensor(func(**new_kwargs), **output_kwargs)
+
+
+@register_jagged_func(torch.ops.aten.stack.default, "tensors: any, dim: any?")
+def stack_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    # guaranteed this is non-empty if we got here
+    tensors = new_kwargs.pop("tensors")
+    for t in tensors:
+        if not isinstance(t, NestedTensor):
+            raise RuntimeError("stack(): expected all nested tensors inputs")
+
+        if t.dim() != tensors[0].dim():
+            raise RuntimeError(
+                "stack(): expected all nested tensors to have the same dim"
+            )
+
+        if not raggedness_matches(t, tensors[0].shape):
+            raise RuntimeError(
+                "stack(): expected all nested tensors to have the same nested structure"
+            )
+
+    new_kwargs["dim"] = _wrap_jagged_dim(
+        tensors[0].dim() + 1, new_kwargs["dim"], tensors[0]._ragged_idx, "stack"
+    )
+
+    return NestedTensor(
+        func([t._values for t in tensors], **new_kwargs), **extract_kwargs(tensors[0])
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.embedding.default,
+    "weight: t, indices: jt, padding_idx: any?, scale_grad_by_freq: any?, sparse: any?",
+)
+def embedding_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    # guaranteed this is non-empty if we got here
+    indices = new_kwargs.pop("indices")
+    weight = new_kwargs.pop("weight")
+
+    return NestedTensor(
+        func(weight, indices._values, **new_kwargs), **extract_kwargs(indices)
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.embedding_dense_backward.default,
+    "grad_output: jt, indices: jt, num_weights: any, padding_idx: any, scale_grad_by_freq: any",
+)
+def embedding_dense_backward_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    indices = new_kwargs.pop("indices")
+    grad_output = new_kwargs.pop("grad_output")
+    return func(grad_output._values, indices._values, **new_kwargs)
+
+
+@register_jagged_func(
+    [
+        torch.ops.aten.values.default,
+        torch.ops.aten._nested_get_values.default,
+    ],
+    "self: jt_all",
+)
+def values_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    # TODO: Handle inference mode properly.
+    # See https://github.com/pytorch/pytorch/issues/112024#issuecomment-1779554292
+    return inp._values.detach()
+
+
+@register_jagged_func(torch.ops.aten.all.default, "self: jt_all")
+def all_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return func(inp._values)
+
+
+@register_jagged_func(
+    torch.ops.aten.to_padded_tensor.default,
+    "self: jt_all, padding: any, output_size: any?",
+)
+def to_padded_tensor_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    if inp._lengths is not None:
+        raise RuntimeError(
+            "to_padded_tensor(): not supported for nested tensors with holes"
+        )
+
+    # TODO: Handle the rest of output_size
+    output_size = new_kwargs["output_size"]
+    if output_size is not None:
+        max_seq_len = output_size[inp._ragged_idx]
+    else:
+        max_seq_len = (
+            inp._max_seqlen
+            if inp._max_seqlen_tensor is not None
+            else inp._values.size(0)
+        )
+
+    # only 2D values with ragged packed dim=0 is supported by the underlying FBGEMM
+    # kernel so do shape gymnastics if needed
+    values = inp.values()
+    if inp._ragged_idx > 1:
+        values = values.transpose(inp._ragged_idx - 1, 0)
+    values_shape = values.shape
+    if values.dim() > 2:
+        values = values.flatten(start_dim=1)
+    elif values.dim() == 1:
+        values = values.unsqueeze(-1)
+
+    # NB: The CUDA kernel for jagged -> padded dense conversion does not support
+    # integer / bool types; work around this by casting to half.
+    is_bool = values.dtype is torch.bool
+    if is_bool and values.is_cuda:
+        values = values.to(torch.half)
+    padded_out = torch.ops.aten._jagged_to_padded_dense_forward(
+        values,
+        [inp._offsets],
+        [max_seq_len],
+        new_kwargs["padding"],
+    )
+    if is_bool and padded_out.is_cuda:
+        padded_out = padded_out.to(torch.bool)
+
+    # shape gymnastics part 2
+    if len(values_shape) > 2:
+        padded_out = padded_out.unflatten(-1, values_shape[1:])
+    elif len(values_shape) == 1:
+        padded_out = padded_out.squeeze(-1)
+    if inp._ragged_idx > 1:
+        padded_out = padded_out.transpose(inp._ragged_idx, 1)
+
+    return padded_out
+
+
+@register_jagged_func(
+    torch.ops.aten._nested_from_padded_tensor.default,
+    "padded: t, offsets: t, dummy: jt, ragged_idx: any?, min_seqlen: any?, max_seqlen: any?, sum_S: any?",
+)
+def _nested_from_padded_tensor_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    padded, offsets = new_kwargs["padded"], new_kwargs["offsets"]
+    ragged_idx = new_kwargs.get("ragged_idx", 1)
+
+    # only 3D padded with ragged packed dim=0 is supported by the underlying FBGEMM
+    # kernel so do shape gymnastics
+    if ragged_idx > 1:
+        padded = padded.transpose(ragged_idx, 1)
+    padded_ragged_dim1_shape = padded.shape
+    if padded.dim() > 3:
+        padded = padded.flatten(start_dim=2)
+    elif padded.dim() < 3:
+        padded = padded.unsqueeze(-1)
+
+    # NB: The CUDA kernel for padded dense -> jagged conversion does not support
+    # integer / bool types; work around this by casting to half.
+    is_bool = padded.dtype is torch.bool
+    if is_bool and padded.is_cuda:
+        padded = padded.to(torch.half)
+    values = torch.ops.aten._padded_dense_to_jagged_forward(
+        padded, [offsets], new_kwargs["sum_S"]
+    )
+    if is_bool and values.is_cuda:
+        values = values.to(torch.bool)
+
+    # shape gymnastics part 2
+    if len(padded_ragged_dim1_shape) > 3:
+        values = values.unflatten(-1, padded_ragged_dim1_shape[2:])
+    elif len(padded_ragged_dim1_shape) < 3:
+        values = values.squeeze(-1)
+    if ragged_idx > 1:
+        values = values.transpose(ragged_idx - 1, 0)
+
+    min_seqlen = new_kwargs["min_seqlen"]
+    max_seqlen = new_kwargs["max_seqlen"]
+    metadata_cache = {}
+    if min_seqlen is not None:
+        metadata_cache["min_seqlen"] = min_seqlen
+    if max_seqlen is not None:
+        metadata_cache["max_seqlen"] = max_seqlen
+
+    return NestedTensor(
+        values,
+        offsets,
+        _ragged_idx=ragged_idx,
+        _metadata_cache=metadata_cache,
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten._nested_view_from_jagged.default,
+    "values: t, offsets: t, dummy: jt_all, lengths: t?, ragged_idx: any?, min_seqlen: t?, max_seqlen: t?",
+)
+def _nested_view_from_jagged_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    values, offsets, lengths = (
+        new_kwargs["input"],
+        new_kwargs["offsets"],
+        new_kwargs["lengths"],
+    )
+    ragged_idx = new_kwargs["ragged_idx"]
+    min_seqlen = new_kwargs["min_seqlen"]
+    max_seqlen = new_kwargs["max_seqlen"]
+    metadata_cache = {}
+    if min_seqlen is not None:
+        metadata_cache["min_seqlen"] = min_seqlen
+    if max_seqlen is not None:
+        metadata_cache["max_seqlen"] = max_seqlen
+
+    return NestedTensor(
+        values,
+        offsets,
+        lengths=lengths,
+        _ragged_idx=ragged_idx,
+        _metadata_cache=metadata_cache,
+    )
+
+
+@register_jagged_func(torch.ops.aten._nested_get_offsets.default, "self: jt_all")
+def _nested_get_offsets(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    return inp._offsets
+
+
+@register_jagged_func(torch.ops.aten._nested_get_lengths.default, "self: jt_all")
+def _nested_get_lengths(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    return inp._lengths
+
+
+@register_jagged_func(torch.ops.aten._nested_get_ragged_idx.default, "self: jt_all")
+def _nested_get_ragged_idx(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    return inp._ragged_idx
+
+
+@register_jagged_func(torch.ops.aten._nested_get_min_seqlen.default, "self: jt_all")
+def _nested_get_min_seqlen(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    return inp._metadata_cache.get("min_seqlen", None)
+
+
+@register_jagged_func(torch.ops.aten._nested_get_max_seqlen.default, "self: jt_all")
+def _nested_get_max_seqlen(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    return inp._metadata_cache.get("max_seqlen", None)
+
+
+# If a section of the Nested Tensor is fully masked out we still retain the section with a length of 0
+@register_jagged_func(torch.ops.aten.masked_select.default, "self: jt, mask: any")
+def masked_select_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+    inp = new_kwargs.pop("input")
+    mask = new_kwargs.pop("mask")
+
+    if inp.ndim > 2:
+        raise RuntimeError("masked_select only support 2-D selections currently")
+    elif inp.shape != mask.shape:
+        raise RuntimeError(
+            f"Mask with shape {mask.shape} is not compatible with input's shape {inp.shape}"
+        )
+    res_values = inp._values.masked_select(mask.values())
+    mask_cumsum = F.pad(mask.values().cumsum(dim=0), (1, 0))  # type: ignore[arg-type]
+
+    args = extract_kwargs(inp)
+    args["offsets"] = mask_cumsum[inp._offsets]
+    return NestedTensor(
+        values=res_values,
+        **args,
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten._nested_select_backward.default,
+    "grad_output: t, self: jt_all, dim: any, index: any",
+)
+def _nested_select_backward_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    grad_output = new_kwargs.pop("grad_output")
+
+    grad_input = torch.zeros_like(inp, dtype=grad_output.dtype)
+    grad_input.select(new_kwargs["dim"], new_kwargs["index"]).copy_(grad_output)
+
+    return grad_input
+
+
+@register_jagged_func(torch.ops.aten.record_stream.default, "self: jt_all, s: any")
+def record_stream_default(func, *args, **kwargs) -> None:
+    inp = args[0]
+    stream = args[1]
+    # ensure all components live until stream computation completes
+    func(inp._values, stream)
+    func(inp._offsets, stream)
+    if inp._lengths is not None:
+        func(inp._lengths, stream)
+
+
+@register_jagged_func(
+    [
+        torch.ops.aten.new_empty.default,
+        torch.ops.aten.new_zeros.default,
+        torch.ops.aten.new_ones.default,
+    ],
+    "self: jt_all, size: any, dtype: any?, layout: any?, device: any?, pin_memory: any?",
+)
+def new_empty_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    if len(new_kwargs["size"]) == 0:
+        return func(inp._values, **new_kwargs)
+
+    raise RuntimeError("new_empty() not supported for NJT with shape != ()")
+
+
+@register_jagged_func(
+    [
+        torch.ops.aten.elu_backward.default,
+        torch.ops.aten.hardshrink_backward.default,
+        torch.ops.aten.hardsigmoid_backward.default,
+        torch.ops.aten.hardtanh_backward.default,
+        torch.ops.aten.softplus_backward.default,
+        torch.ops.aten.softshrink_backward.default,
+    ],
+    "self: jt_all, ...",
+)
+def activation_backward(func, *args, **kwargs):
+    # first NJT arg is expected to be grad_output
+    grad_output = next(arg for arg in args if isinstance(arg, NestedTensor))
+    return NestedTensor(
+        func(
+            *(arg._values if isinstance(arg, NestedTensor) else arg for arg in args),
+            **kwargs,
+        ),
+        **extract_kwargs(grad_output),
+    )
+
+
+@register_jagged_func(torch.ops.aten.fill.Scalar, "self: jt_all, value: any")
+def fill_Scalar(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
+
+
+@register_jagged_func(torch.ops.aten.fill_.Scalar, "self: jt_all, value: any")
+def fill__Scalar(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+
+    func(inp._values, **new_kwargs)
+    return inp
+
+
+@register_jagged_func(torch.ops.aten.frexp.Tensor, "self: jt_all")
+def frexp_Tensor(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    output_kwargs = extract_kwargs(inp)
+
+    mantissa, exponent = func(inp._values)
+    return NestedTensor(mantissa, **output_kwargs), NestedTensor(
+        exponent, **output_kwargs
+    )
+
+
+@register_jagged_func(
+    torch.ops.aten.matmul_backward.default,
+    "grad: any, self: any, other: any, mask: any",
+)
+def matmul_backward_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    grad = new_kwargs.pop("grad")
+    inp = new_kwargs.pop("input")
+    other = new_kwargs.pop("other")
+    grad_input_mask = new_kwargs.pop("mask")
+
+    if grad is None:
+        return (None, None)
+
+    grad_self = None
+    if grad_input_mask[0]:
+        grad_self = torch.matmul(grad, other.transpose(-1, -2))
+
+    grad_other = None
+    if grad_input_mask[1]:
+        grad_other = torch.matmul(inp.transpose(-1, -2), grad)
+
+    return (grad_self, grad_other)
+
+
+# Make the dummy available on the C++ side.
+@register_jagged_func(torch.ops.aten._nested_get_jagged_dummy.default, "self: any")
+def _nested_get_jagged_dummy(func, *args, **kwargs):
+    from torch.nested._internal.nested_tensor import _nt_view_dummy
+
+    return _nt_view_dummy()
+
+
+with torch.library._scoped_library("aten", "IMPL") as aten:
+    aten.impl("_nested_get_jagged_dummy", _nested_get_jagged_dummy, "CPU")
+    aten.impl("_nested_get_jagged_dummy", _nested_get_jagged_dummy, "CUDA")
+    aten.impl("_nested_get_jagged_dummy", _nested_get_jagged_dummy, "Meta")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/sdpa.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/sdpa.py
new file mode 100644
index 0000000000000000000000000000000000000000..8854fafe1ff93d6a2491676057993fb06163e39a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nested/_internal/sdpa.py
@@ -0,0 +1,936 @@
+# mypy: allow-untyped-defs
+import logging
+
+import torch
+import torch.nn
+import torch.nn.functional as F
+from torch.backends.cuda import (
+    can_use_cudnn_attention,
+    can_use_efficient_attention,
+    can_use_flash_attention,
+    cudnn_sdp_enabled,
+    flash_sdp_enabled,
+    math_sdp_enabled,
+    mem_efficient_sdp_enabled,
+    SDPAParams,
+)
+from torch.nn.attention import SDPBackend
+
+from .nested_tensor import NestedTensor
+
+
+log = logging.getLogger(__name__)
+
+
+def _validate_sdpa_input(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None = None,
+    dropout_p=0.0,
+    is_causal=False,
+    scale=None,
+) -> None:
+    if (
+        not isinstance(query, NestedTensor)
+        or not isinstance(key, NestedTensor)
+        or not isinstance(value, NestedTensor)
+    ):
+        raise ValueError(
+            f"Expected query, key, and value to be nested tensors, "
+            f"but got query.is_nested: {query.is_nested}, key.is_nested: {key.is_nested}, "
+            f"and value.is_nested: {value.is_nested} instead."
+        )
+    if query.dtype != key.dtype or query.dtype != value.dtype:
+        raise ValueError(
+            f"Expected query, key, and value to have the same dtype, "
+            f"but got query.dtype: {query.dtype}, key.dtype: {key.dtype}, "
+            f"and value.dtype: {value.dtype} instead."
+        )
+    if query.device != key.device or query.device != value.device:
+        raise ValueError(
+            f"Expected query, key, and value to have the same device type, "
+            f"but got query.device: {query.device}, key.device: {key.device}, "
+            f"and value.device: {value.device} instead."
+        )
+    if query.dim() < 3 or key.dim() < 3 or value.dim() < 3:
+        raise ValueError(
+            f"Expected query, key, and value to all be  at least 3 dimensional, but got query.dim: "
+            f"{query.dim()}, key.dim: {key.dim()} and value.dim: {value.dim()} instead."
+        )
+    if query._ragged_idx != key._ragged_idx or query._ragged_idx != value._ragged_idx:
+        raise ValueError(
+            f"Expected query, key, and value to all be ragged on the same dimension, but got ragged "
+            f"dims {query._ragged_idx}, {key._ragged_idx}, and {value._ragged_idx}, respectively."
+        )
+    if attn_mask is not None:
+        # TODO: Figure out whether masks are actually supported for this layout or not
+        raise ValueError("Masks are not yet supported!")
+        if attn_mask.dtype != torch.bool and attn_mask.dtype != query.dtype:
+            raise ValueError(
+                f"Expected attn_mask dtype to be bool or to match query dtype, but got attn_mask.dtype: "
+                f"{attn_mask.dtype}, and query.dtype: {query.dtype} instead."
+            )
+
+
+def _check_batch_size_nested(params: SDPAParams, debug=False) -> bool:
+    # This is expected to be called after check_tensor_shapes ensuring that the
+    # size() calls won't error since the inputs are all 4 dimensional
+    q_batch_size = params.query.size(0)
+    k_batch_size = params.key.size(0)
+    v_batch_size = params.value.size(0)
+
+    # num_heads logic for nested input is checked in
+    # check_for_seq_len_0_nested_tensor as there is handling there to make sure
+    # num_heads is not ragged
+    return q_batch_size == k_batch_size and q_batch_size == v_batch_size
+
+
+def _check_head_dim_size_flash_nested(params: SDPAParams, debug=False) -> bool:
+    max_size = 256
+    query_size_last = params.query.size(-1)
+    key_size_last = params.key.size(-1)
+    value_size_last = params.value.size(-1)
+    same_head_dim_size = (
+        query_size_last == key_size_last and query_size_last == value_size_last
+    )
+    if not (
+        same_head_dim_size
+        and (query_size_last % 8 == 0)
+        and (query_size_last <= max_size)
+    ):
+        if debug:
+            log.warning(
+                "For NestedTensor inputs, Flash attention requires q,k,v to have the same "
+                "last dimension and to be a multiple of 8 and less than or equal to 256. "
+                "Got Query.size(-1): %d, Key.size(-1): %d, Value.size(-1): %d instead.",
+                query_size_last,
+                key_size_last,
+                value_size_last,
+            )
+        return False
+    return True
+
+
+def _check_head_dim_size_cudnn_nested(params: SDPAParams, debug=False) -> bool:
+    max_size = 128
+    query_size_last = params.query.size(-1)
+    key_size_last = params.key.size(-1)
+    value_size_last = params.value.size(-1)
+    same_head_dim_size = (
+        query_size_last == key_size_last and query_size_last == value_size_last
+    )
+    if not (
+        same_head_dim_size
+        and (query_size_last % 8 == 0)
+        and (query_size_last <= max_size)
+    ):
+        if debug:
+            log.warning(
+                "For NestedTensor inputs, cuDNN attention requires q,k,v to have the same "
+                "last dimension and to be a multiple of 8 and less than or equal to 128. "
+                "Got Query.size(-1): %d, Key.size(-1): %d, Value.size(-1): %d instead.",
+                query_size_last,
+                key_size_last,
+                value_size_last,
+            )
+        return False
+    return True
+
+
+def _check_for_seq_len_0_and_consistent_head_dim_nested_helper(
+    param: torch.Tensor, param_name: str, debug=False
+) -> bool:
+    if not isinstance(param, NestedTensor):
+        raise AssertionError("param should be a jagged NT")
+
+    if param._ragged_idx == 1:
+        # num_head_dims is ragged
+        if debug:
+            log.warning(
+                "Fused kernels do not support ragged num_head_dims, %s has a ragged num_heads.",
+                param_name,
+            )
+        return False
+
+    # This is being called inside sdp with shape [batch, heads, {seq_len}, dim]
+    if param._get_min_seqlen() == 0:
+        if debug:
+            log.warning(
+                "Fused kernels do not support seq_len == 0, %s has a seq len of 0.",
+                param_name,
+            )
+        return False
+
+    return True
+
+
+def _try_broadcast_param_size(q_size, k_size, v_size, param_name, debug=False) -> bool:
+    max_size = max(q_size, k_size, v_size)
+    if (
+        (q_size != max_size and q_size != 1)
+        or (k_size != max_size and k_size != 1)
+        or (v_size != max_size and v_size != 1)
+    ):
+        if debug:
+            log.warning(
+                "Both fused kernels require query, key and value to have broadcastable %s, "
+                "got Query %s %d, Key %s %d, Value %s %d instead.",
+                param_name,
+                param_name,
+                q_size,
+                param_name,
+                k_size,
+                param_name,
+                v_size,
+            )
+        return False
+    return True
+
+
+def _check_for_seq_len_0_nested(params: SDPAParams, debug=False) -> bool:
+    # When this function is called we are assured that the nt is dim==4
+    q_is_safe = (
+        _check_for_seq_len_0_and_consistent_head_dim_nested_helper(
+            params.query, "query", debug
+        )
+        if params.query.is_nested
+        else True
+    )
+    # short circuit if any is unsafe
+    if not q_is_safe:
+        return False
+
+    k_is_safe = (
+        _check_for_seq_len_0_and_consistent_head_dim_nested_helper(
+            params.key, "key", debug
+        )
+        if params.key.is_nested
+        else True
+    )
+    # short circuit if any is unsafe
+    if not k_is_safe:
+        return False
+
+    v_is_safe = (
+        _check_for_seq_len_0_and_consistent_head_dim_nested_helper(
+            params.value, "value", debug
+        )
+        if params.value.is_nested
+        else True
+    )
+    # short circuit if any is unsafe
+    if not v_is_safe:
+        return False
+
+    # We now know none of the inputs have ragged num_heads, so we can safely
+    # access .size(1)
+    q_num_heads = params.query.size(1)
+    k_num_heads = params.key.size(1)
+    v_num_heads = params.value.size(1)
+    same_num_heads = q_num_heads == k_num_heads and q_num_heads == v_num_heads
+
+    if not same_num_heads:
+        if (
+            params.query.requires_grad
+            or params.key.requires_grad
+            or params.value.requires_grad
+        ):
+            if debug:
+                log.warning(
+                    "Both fused kernels do not support training with broadcasted NT inputs."
+                )
+            return False
+        return _try_broadcast_param_size(
+            q_num_heads, k_num_heads, v_num_heads, "num heads", debug
+        )
+    return True
+
+
+def _can_use_flash_sdpa_jagged(params: SDPAParams, debug=False) -> bool:
+    constraints = (
+        _check_batch_size_nested,
+        _check_head_dim_size_flash_nested,
+        _check_for_seq_len_0_nested,
+    )
+    for constraint in constraints:
+        if not constraint(params, debug):
+            return False
+    return True
+
+
+def _can_use_efficient_sdpa_jagged(params: SDPAParams, debug=False) -> bool:
+    constraints = (
+        _check_batch_size_nested,
+        _check_for_seq_len_0_nested,
+    )
+    for constraint in constraints:
+        if not constraint(params, debug):
+            return False
+    return True
+
+
+def _can_use_math_sdpa_jagged(params: SDPAParams, debug=False) -> bool:
+    if (
+        not params.query.transpose(1, 2).is_contiguous()
+        or not params.key.transpose(1, 2).is_contiguous()
+        or not params.value.transpose(1, 2).is_contiguous()
+    ):
+        if debug:
+            log.warning(
+                "If inputs are nested tensors they must be contiguous after transposing."
+            )
+        return False
+    if params.is_causal:
+        if debug:
+            log.warning(
+                "Nested tensors for query / key are not supported when is_causal=True."
+            )
+        return False
+    return True
+
+
+def _select_sdp_backend(query, key, value, attn_mask, dropout, is_causal, enable_gqa):
+    if (
+        not flash_sdp_enabled()
+        and not mem_efficient_sdp_enabled()
+        and not math_sdp_enabled()
+        and not cudnn_sdp_enabled()
+    ):
+        return SDPBackend.ERROR
+
+    ordering = (
+        SDPBackend.FLASH_ATTENTION,
+        SDPBackend.EFFICIENT_ATTENTION,
+        SDPBackend.MATH,
+        SDPBackend.CUDNN_ATTENTION,
+    )
+
+    params = SDPAParams(query, key, value, attn_mask, dropout, is_causal, enable_gqa)
+
+    for backend in ordering:
+        if backend == SDPBackend.CUDNN_ATTENTION:
+            if can_use_cudnn_attention(params):
+                return SDPBackend.CUDNN_ATTENTION
+        if backend == SDPBackend.FLASH_ATTENTION:
+            if can_use_flash_attention(params) and _can_use_flash_sdpa_jagged(params):
+                return SDPBackend.FLASH_ATTENTION
+        if backend == SDPBackend.EFFICIENT_ATTENTION:
+            if can_use_efficient_attention(params) and _can_use_efficient_sdpa_jagged(
+                params
+            ):
+                return SDPBackend.EFFICIENT_ATTENTION
+        if backend == SDPBackend.MATH:
+            if math_sdp_enabled() and _can_use_math_sdpa_jagged(params):
+                return SDPBackend.MATH
+
+    log.warning("Memory efficient kernel not used because:")
+    can_use_efficient_attention(params, debug=True)
+    _can_use_efficient_sdpa_jagged(params, debug=True)
+    log.warning("Flash attention kernel not used because:")
+    can_use_flash_attention(params, debug=True)
+    _can_use_flash_sdpa_jagged(params, debug=True)
+    log.warning("Math attention kernel not used because:")
+    _can_use_math_sdpa_jagged(params, debug=True)
+    log.warning("cuDNN attention kernel not used because:")
+    can_use_cudnn_attention(params, debug=True)
+    return SDPBackend.ERROR
+
+
+def _cumulative_and_max_seq_len_nnz(qkv: torch.Tensor) -> tuple[torch.Tensor, int, int]:
+    # This function is used to calculate two pieces of metadata that are needed
+    # for use with flash-attention and efficient_attention kernels. They are the
+    # cumulative sequence_length over a batch of sequences and the maximum
+    # sequence length.
+
+    # It returns a tuple of cumulative sequence lengths and the maximum sequence
+    # length, and the last element in the cumulative_sequence_lengths
+    if not isinstance(qkv, NestedTensor):
+        raise ValueError("QKV must be nested for flash cumulative_seq_len calculation.")
+
+    if qkv.lengths() is None:
+        # TODO: Explore performance impact of copying
+        cumulative_seqlen = qkv.offsets().to(dtype=torch.int32, device=qkv.device)
+        max_seqlen = qkv._get_max_seqlen()
+        n_elem = qkv.values().shape[0]
+    else:
+        # TODO: Explore performance impact of copying
+        cumulative_seqlen = (
+            qkv.lengths().cumsum(0).to(dtype=torch.int32, device=qkv.device)
+        )
+        max_seqlen = qkv._get_max_seqlen()
+        # TODO: Explore performance impact when compiling
+        n_elem = int(cumulative_seqlen[-1].item())
+    return cumulative_seqlen, max_seqlen, n_elem
+
+
+def _is_safe_to_get_storage_as_tensor(tensor: torch.Tensor) -> bool:
+    # This function checks if a nested tensor is valid for
+    # use with the flash-attention and efficient_attention kernels without
+    # needing to call contiguous on the nested tensor input.
+    # It checks that the storage offsets' adjacent_differences are a constant
+    # multiple of the previous tensor in the nested tensor and that the strides
+    # are monitonically decreasing. This check is done after calling transpose on
+    # the nested tensor resulting in a Nt of shape [bsz, {seq_len}, num_heads, dim]
+
+    # Returns a boolean indicating if contiguous needs to be called for input
+    if not isinstance(tensor, NestedTensor):
+        raise AssertionError("tensor must be a NestedTensor")
+    offsets = tensor.offsets()
+    strides = tensor._strides
+
+    n_tensors = offsets.size(0) - 1
+    if n_tensors <= 1:
+        return True
+
+    # Check initially that the tensor strides are in strictly descending order
+    prev_stride = strides[1]
+    for stride in strides[2:]:
+        if prev_stride <= stride:
+            # This would mean that the last stride is greater than the seq_len
+            # stride
+            return False
+        prev_stride = stride
+
+    # Congrats you made it!
+    return True
+
+
+def _view_as_dense(
+    tensor: torch.Tensor, Nnz: int, num_heads: int, head_dim: int
+) -> torch.Tensor:
+    if tensor.is_nested:
+        return tensor.values()
+    return tensor.view(Nnz, num_heads, head_dim)
+
+
+# TODO: Next iteration should add test cases and check it works
+# def _sdpa_nested_preprocessing_with_broadcast(query, key, value):
+#     # Query (Batch x Num_heads x {Q_seq_len}  x Dim_per_head)
+#     # Key   (Batch x Num_heads x {KV_seq_len} x Dim_per_head)
+#     # Value (Batch x Num_heads x {KV_seq_len} x Dim_per_head)
+#     q_batch_size = query.size(0)
+#     k_batch_size = key.size(0)
+#     v_batch_size = value.size(0)
+
+#     output_batch_size = max(q_batch_size, k_batch_size, v_batch_size)
+
+#     q_num_heads = query.size(1)
+#     k_num_heads = key.size(1)
+#     v_num_heads = value.size(1)
+
+#     output_num_heads = max(q_num_heads, k_num_heads, v_num_heads)
+
+#     head_dim_qk = query.size(3)
+#     head_dim_v = value.size(3)
+
+#     q_t = query.transpose(1, 2)
+#     k_t = key.transpose(1, 2)
+#     v_t = value.transpose(1, 2)
+
+#     # Checks in sdp_utils ensure that if {*}_batch_size/{*}_num_heads !=
+#     # output_batch_size/num_heads then they are 1
+#     q_batch_size_needs_broadcast = q_batch_size != output_batch_size
+#     k_batch_size_needs_broadcast = k_batch_size != output_batch_size
+#     v_batch_size_needs_broadcast = v_batch_size != output_batch_size
+
+#     # If {*}_batch_size_needs_broadcast, then
+#     # (1) max_seqlen_batch_{*} is given by {*}_t.size(1)
+#     #     this is because needs_broadcast indicates that the batch_size is 1
+#     #     and hence there is only 1 value for seq_len
+#     # (2) The cum_seq_lens are given by [0, {*}_t.size(1), 2 * {*}_t.size(1),
+#     # ..., output_batch_size * {*}_t.size(1)]
+#     # (3) Nnz_{*} is given by output_batch_size * {*}_t.size(1)
+
+#     if q_batch_size_needs_broadcast or not q_t.is_nested:
+#         max_seqlen_batch_q = q_t.size(1)
+#         cumulative_sequence_length_q = torch.arange(
+#             0,
+#             (output_batch_size + 1) * max_seqlen_batch_q,
+#             max_seqlen_batch_q,
+#             device=q_t.device,
+#             dtype=torch.int32,
+#         )
+#         Nnz_q = output_batch_size * max_seqlen_batch_q
+#     else:
+#         (
+#             cumulative_sequence_length_q,
+#             max_seqlen_batch_q,
+#             Nnz_q,
+#         ) = _cumulative_and_max_seq_len_nnz(q_t)
+
+#     if k_batch_size_needs_broadcast and v_batch_size_needs_broadcast:
+#         assert k_t.size(1) == v_t.size(1)
+#         max_seqlen_batch_kv = k_t.size(1)
+#         cumulative_sequence_length_kv = torch.arange(
+#             0,
+#             (output_batch_size + 1) * max_seqlen_batch_kv,
+#             max_seqlen_batch_kv,
+#             device=k_t.device,
+#             dtype=torch.int32,
+#         )
+#         Nnz_kv = output_batch_size * max_seqlen_batch_kv
+#     else:
+#         cumulative_sequence_length_kv, max_seqlen_batch_kv, Nnz_kv = (
+#             _cumulative_and_max_seq_len_nnz(v_t)
+#             if k_batch_size_needs_broadcast
+#             else _cumulative_and_max_seq_len_nnz(k_t)
+#         )
+
+#     q_num_heads_needs_broadcast = q_num_heads != output_num_heads
+#     k_num_heads_needs_broadcast = k_num_heads != output_num_heads
+#     v_num_heads_needs_broadcast = v_num_heads != output_num_heads
+
+#     if not q_t.is_nested:
+#         query_buffer_reshaped = q_t.expand(
+#             output_batch_size, q_t.size(1), output_num_heads, head_dim_qk
+#         )
+#         query_buffer_reshaped = query_buffer_reshaped.reshape(
+#             Nnz_q, output_num_heads, head_dim_qk
+#         )
+#     else:
+#         if not q_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(q_t):
+#             q_t = q_t.contiguous()
+#         # If we are broadcasting then Nnz_q will be the output_batch_size since
+#         # seq_len is 1
+#         effective_batch_size_q = (
+#             output_batch_size if q_batch_size_needs_broadcast else Nnz_q
+#         )
+#         query_buffer_reshaped = _view_as_dense(
+#             q_t, effective_batch_size_q, output_num_heads, head_dim_qk
+#         )
+
+#     # If the physical layout of the NestedTensor's storage
+#     # is not: batch, {seq_len}, num_heads, head_dim then we need
+#     # to call contiguous
+#     if not k_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(k_t):
+#         k_t = k_t.contiguous()
+#     if not v_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(v_t):
+#         v_t = v_t.contiguous()
+
+#     effective_batch_size_k = (
+#         output_batch_size if k_batch_size_needs_broadcast else Nnz_kv
+#     )
+#     key_buffer_reshaped = _view_as_dense(
+#         k_t, effective_batch_size_k, output_num_heads, head_dim_qk
+#     )
+
+#     effective_batch_size_v = (
+#         output_batch_size if v_batch_size_needs_broadcast else Nnz_kv
+#     )
+#     value_buffer_reshaped = _view_as_dense(
+#         v_t, effective_batch_size_v, output_num_heads, head_dim_v
+#     )
+
+#     if not q_batch_size_needs_broadcast:
+#         output_shape = q_t._size
+#         if head_dim_v != head_dim_qk:
+#             output_shape[-1] = head_dim_v
+#         if q_num_heads_needs_broadcast:
+#             output_shape[1] = output_num_heads
+#     else:
+#         output_shape = torch.empty(3, dtype=torch.int64, device=torch.device("cpu"))
+#         output_shape[0] = q_t.size(1)
+#         output_shape[1] = output_num_heads
+#         output_shape[2] = head_dim_v
+
+#     return (
+#         query_buffer_reshaped,
+#         key_buffer_reshaped,
+#         value_buffer_reshaped,
+#         cumulative_sequence_length_q,
+#         cumulative_sequence_length_kv,
+#         max_seqlen_batch_q,
+#         max_seqlen_batch_kv,
+#         output_shape,
+#     )
+
+
+def _sdpa_nested_preprocessing(query, key, value):
+    # Query (Batch x Num_heads x {Q_seq_len}  x Dim_per_head)
+    # Key   (Batch x Num_heads x {KV_seq_len} x Dim_per_head)
+    # Value (Batch x Num_heads x {KV_seq_len} x Dim_per_head)
+    q_batch_size = query.size(0)
+    k_batch_size = key.size(0)
+    v_batch_size = value.size(0)
+
+    q_num_heads = query.size(1)
+    k_num_heads = key.size(1)
+    v_num_heads = value.size(1)
+
+    if not (q_batch_size == k_batch_size and q_batch_size == v_batch_size) or not (
+        q_num_heads == k_num_heads and k_num_heads == v_num_heads
+    ):
+        raise RuntimeError(
+            "This path is currently not implemented for jagged layout NT."
+        )
+        # return _sdpa_nested_preprocessing_with_broadcast(query, key, value)
+
+    num_heads = query.size(1)
+    head_dim_qk = query.size(3)
+    head_dim_v = value.size(3)
+    q_t = query.transpose(1, 2)
+    k_t = key.transpose(1, 2)
+    v_t = value.transpose(1, 2)
+
+    (
+        cumulative_sequence_length_q,
+        max_seqlen_batch_q,
+        Nnz_q,
+    ) = _cumulative_and_max_seq_len_nnz(q_t)
+    (
+        cumulative_sequence_length_kv,
+        max_seqlen_batch_kv,
+        Nnz_kv,
+    ) = _cumulative_and_max_seq_len_nnz(k_t)
+
+    # [TODO] K and V have to have the same Nnz, should probably torch_check
+    # assume in order to not iterate over v
+
+    # If the physical layout of the NestedTensor's storage
+    # is not: batch, {seq_len}, num_heads, head_dim then we need
+    # to call contiguous
+    if not q_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(q_t):
+        q_t = q_t.contiguous()
+    if not k_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(k_t):
+        k_t = k_t.contiguous()
+    if not v_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(v_t):
+        v_t = v_t.contiguous()
+
+    query_buffer_reshaped = _view_as_dense(q_t, Nnz_q, num_heads, head_dim_qk)
+    key_buffer_reshaped = _view_as_dense(k_t, Nnz_kv, num_heads, head_dim_qk)
+    value_buffer_reshaped = _view_as_dense(v_t, Nnz_kv, num_heads, head_dim_v)
+
+    output_nt_info = {
+        "offsets": q_t.offsets(),
+        "lengths": q_t.lengths(),
+        "max_seqlen": q_t._get_max_seqlen(),
+        "min_seqlen": q_t._get_min_seqlen(),
+    }
+
+    return (
+        query_buffer_reshaped,
+        key_buffer_reshaped,
+        value_buffer_reshaped,
+        cumulative_sequence_length_q,
+        cumulative_sequence_length_kv,
+        max_seqlen_batch_q,
+        max_seqlen_batch_kv,
+        output_nt_info,
+    )
+
+
+def _pad_last_dim(
+    tensor: torch.Tensor, alignment_size: int, slice: bool
+) -> torch.Tensor:
+    # FlashAttentionV2 requires that head dimension be a multiple of 8
+    # This was previously done within the kernel, however
+    # This causes the kernel to maybe alias query, key, value
+    # So instead we pad the head_dimensions to be a multiple of 8
+    # in the composite region
+    last_dim_size = tensor.size(-1)
+    if last_dim_size % alignment_size == 0:
+        return tensor
+    pad_count = alignment_size - (last_dim_size % alignment_size)
+    tensor = torch.nn.functional.pad(tensor, [0, pad_count])
+    if slice:
+        return tensor[..., 0:last_dim_size]
+    return tensor
+
+
+# TODO: coalesce with torch/nn/utils/attention.py
+def _calculate_scale(query, scale):
+    # TODO: Investigate why math.sqrt() isn't properly handled by Dynamo?
+    softmax_scale = scale if scale is not None else torch.sym_sqrt(1.0 / query.size(-1))
+    return softmax_scale
+
+
+def _post_process_flash_output(out: torch.Tensor, og_size):
+    if not out.is_nested and out.size(-1) != og_size:
+        out = out[..., 0:og_size]
+    return out
+
+
+def _is_computing_meta_flops(x):
+    # Note: there's a use case of using meta tensors & the dispatch-based flop counter.
+    # We can use this function to check for this scenario in order to handle it specially.
+    if not torch.jit.is_scripting() and x.device.type == "meta":
+        torch_dispatch_mode_stack = (
+            torch.utils._python_dispatch._get_current_dispatch_mode_stack()
+        )
+        return any(
+            type(x) is torch.utils.flop_counter._FlopCounterMode
+            for x in torch_dispatch_mode_stack
+        )
+    return False
+
+
+def _autocast(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """
+    [Autocasting SDPA for NJT]
+
+    Normal autocasting doesn't work for NJT+SDPA right now:
+    * NJT intercepts the __torch_function__ call for scaled_dot_product_attention, which happens
+      before we get to any aten ops or dispatcher logic; then the torch_function logic calls into
+      efficient attention or flash attention. So, autocasting on the scaled_dot_product_attention
+      op won't work because we never see that aten op.
+    * If we put autocasting on `_flash_attention_forward`, then we'll get autocasting to run, but
+      the kernel selection logic in torch_function handling (ie. jagged_scaled_dot_product_attention)
+      won't work correctly: the kernel selection logic will run before autocasting, and choose
+      a kernel based on the un-autocasted dtypes; but then autocasting will run and the actual
+      attention computation will happen in a different dtype.
+
+    An alternative is to just change the backend selection logic for SDPA+NJT to be autocast-aware
+    and rely on autocasting to do the actual conversions for flash attention / efficient attention.
+    However, by manually doing the actual autocast before the backend selection, we ensure that the
+    autocast handling for backend selection doesn't diverge from the autocast handling for the
+    actual dtype conversions.
+    """
+    device_type = query.device.type
+    # meta device is not supported by autocast, so break early for it
+    if _is_computing_meta_flops(query) or not torch.is_autocast_enabled(device_type):
+        return query, key, value, attn_mask
+
+    def cvt(x):
+        if x is None:
+            return x
+        target_dtype = torch.get_autocast_dtype(device_type)
+        if (
+            (not x.dtype.is_floating_point)
+            or x.dtype == target_dtype
+            or x.dtype == torch.float64
+        ):
+            return x
+        return x.to(target_dtype)
+
+    return cvt(query), cvt(key), cvt(value), cvt(attn_mask)
+
+
+def jagged_scaled_dot_product_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None = None,
+    dropout_p=0.0,
+    is_causal=False,
+    scale=None,
+    enable_gqa=False,
+):
+    query, key, value, attn_mask = _autocast(query, key, value, attn_mask)
+    _validate_sdpa_input(query, key, value, attn_mask, dropout_p, is_causal, scale)
+    # for mypy, ugh
+    if not (
+        isinstance(query, NestedTensor)
+        and isinstance(key, NestedTensor)
+        and isinstance(value, NestedTensor)
+    ):
+        raise AssertionError("query, key, and value must all be NestedTensor instances")
+    from torch.nested._internal.nested_tensor import (
+        nested_view_from_values_offsets_lengths,
+    )
+
+    # Special path for non-ragged sequence length (e.g. for SAM where we have a ragged
+    # second batch dim instead). For this case, we can just send the dense buffers through
+    # vanilla SDPA.
+    if query.dim() > 3 and key.dim() > 3 and value.dim() > 3 and query._ragged_idx == 1:
+        output = F.scaled_dot_product_attention(
+            query.values(),
+            key.values(),
+            value.values(),
+            attn_mask=(
+                attn_mask.values() if isinstance(attn_mask, NestedTensor) else attn_mask
+            ),
+            dropout_p=dropout_p,
+            is_causal=is_causal,
+            scale=scale,
+        )
+        return nested_view_from_values_offsets_lengths(
+            output,
+            query.offsets(),
+            query.lengths(),
+            min_seqlen=query._maybe_min_seqlen,  # type: ignore[attr-defined]
+            max_seqlen=query._maybe_max_seqlen,  # type: ignore[attr-defined]
+        )
+
+    compute_logsumexp = query.requires_grad or key.requires_grad or value.requires_grad
+
+    backend_choice = _select_sdp_backend(
+        query, key, value, attn_mask, dropout_p, is_causal, enable_gqa
+    )
+
+    if _is_computing_meta_flops(query):
+        # Backend choice will probably not be correct if we have a meta device,
+        # because backend choice is device-aware. In this case, we mostly just
+        # want to avoid using math backend (which does a .item() call).
+        # Arbitrarily choose flash attention.
+        backend_choice = SDPBackend.FLASH_ATTENTION
+
+    if backend_choice == SDPBackend.FLASH_ATTENTION:
+        og_size = query.size(-1)
+        query_padded = _pad_last_dim(query, 8, False)
+        key_padded = _pad_last_dim(key, 8, False)
+        value_padded = _pad_last_dim(value, 8, False)
+        # We need to calculate the scale based off the OG head dim size
+        og_scale = _calculate_scale(query, scale)
+        (
+            query_buffer_reshaped,
+            key_buffer_reshaped,
+            value_buffer_reshaped,
+            cumulative_sequence_length_q,
+            cumulative_sequence_length_kv,
+            max_seqlen_batch_q,
+            max_seqlen_batch_kv,
+            output_nt_info,
+        ) = _sdpa_nested_preprocessing(query_padded, key_padded, value_padded)
+        (
+            attention,
+            _logsumexp,
+            _philox_seed,
+            _philox_offset,
+            _debug_attn_mask,
+        ) = torch.ops.aten._flash_attention_forward(
+            query_buffer_reshaped,
+            key_buffer_reshaped,
+            value_buffer_reshaped,
+            cumulative_sequence_length_q,
+            cumulative_sequence_length_kv,
+            max_seqlen_batch_q,
+            max_seqlen_batch_kv,
+            dropout_p,
+            is_causal,
+            False,
+            scale=og_scale,
+        )
+        # Reshape output to convert nnz to batch_size and seq_len
+        attention = nested_view_from_values_offsets_lengths(
+            attention,  # output from flash_attn is [total_q, num_heads, head_size_og]
+            **output_nt_info,
+        ).transpose(1, 2)
+        return _post_process_flash_output(attention, og_size)
+    elif backend_choice == SDPBackend.EFFICIENT_ATTENTION:
+        (
+            query_reshaped,
+            key_reshaped,
+            value_reshaped,
+            cumulative_sequence_length_q,
+            cumulative_sequence_length_kv,
+            max_seqlen_batch_q,
+            max_seqlen_batch_kv,
+            output_nt_info,
+        ) = _sdpa_nested_preprocessing(query, key, value)
+        (
+            attention,
+            log_sumexp,
+            seed,
+            offset,
+            max_seqlen_q,
+            max_seqlen_batch_kv,
+        ) = torch.ops.aten._efficient_attention_forward(
+            query_reshaped.unsqueeze(0),
+            key_reshaped.unsqueeze(0),
+            value_reshaped.unsqueeze(0),
+            None,
+            cumulative_sequence_length_q,
+            cumulative_sequence_length_kv,
+            max_seqlen_batch_q,
+            max_seqlen_batch_kv,
+            dropout_p,
+            int(is_causal),
+            compute_logsumexp,
+            scale=scale,
+        )
+        # Reshape output to convert nnz to batch_size and seq_len
+        return nested_view_from_values_offsets_lengths(
+            attention.squeeze(0),
+            **output_nt_info,
+        ).transpose(1, 2)
+    elif backend_choice == SDPBackend.CUDNN_ATTENTION:
+        (
+            query_reshaped,
+            key_reshaped,
+            value_reshaped,
+            cumulative_sequence_length_q,
+            cumulative_sequence_length_kv,
+            max_seqlen_batch_q,
+            max_seqlen_batch_kv,
+            output_nt_info,
+        ) = _sdpa_nested_preprocessing(query, key, value)
+        (
+            attention,
+            logsumexp,
+            cum_seqlen_q,
+            cum_seqlen_kv,
+            max_seqlen_q,
+            max_seqlen_kv,
+            seed,
+            offset,
+            _,
+        ) = torch.ops.aten._cudnn_attention_forward(
+            query_reshaped,
+            key_reshaped,
+            value_reshaped,
+            attn_mask,
+            cumulative_sequence_length_q,
+            cumulative_sequence_length_kv,
+            max_seqlen_batch_q,
+            max_seqlen_batch_kv,
+            compute_logsumexp,
+            dropout_p,
+            is_causal,
+            False,
+            scale=scale,
+        )
+        return nested_view_from_values_offsets_lengths(
+            attention,
+            **output_nt_info,
+        ).transpose(1, 2)
+    elif backend_choice == SDPBackend.MATH:
+        # save the offsets and shape of the inputs, so we can reshape the final output
+        # query @ key = attn: [B, D1, j0, D'] @ [B, D1, D' j1] = [B, D1, j0, j1]
+        # attn @ value = out: [B, D1, j0, j1] @ [B, D1, j1, D2] = [B, D1, j0, D2]
+        offsets = query.offsets()
+        q_lengths = query.lengths()
+        min_seqlen = query._maybe_min_seqlen
+        max_seqlen = query._maybe_max_seqlen
+        d1 = query._size[1]
+        d2 = value._size[-1]
+
+        # convert jagged layout Nested Tensor to strided layout Nested Tensor
+        # which support the math implementation of SDPA
+        def get_strided_layout_nested_tensor(jagged_layout_nt):
+            lengths = jagged_layout_nt._offsets[1:] - jagged_layout_nt._offsets[:-1]
+            transpose = torch.transpose(jagged_layout_nt, 1, 2)
+            tensor_list = transpose.values().split(list(lengths), dim=0)
+            strided_nt = torch.nested.as_nested_tensor(list(tensor_list))
+            strided_nt = strided_nt.transpose(1, 2).contiguous()
+            return strided_nt
+
+        query = get_strided_layout_nested_tensor(query)
+        key = get_strided_layout_nested_tensor(key)
+        value = get_strided_layout_nested_tensor(value)
+
+        attn_out = torch._scaled_dot_product_attention_math(
+            query, key, value, attn_mask, dropout_p, is_causal, scale=scale
+        )[0]
+
+        # convert strided layout Nested Tensor back to jagged layout Nested Tensor
+        attn_out = attn_out.transpose(1, 2).contiguous().values()
+        attn_out = attn_out.view(-1, d1, d2)
+        attn_out = nested_view_from_values_offsets_lengths(
+            attn_out,
+            offsets,
+            lengths=q_lengths,
+            min_seqlen=min_seqlen,
+            max_seqlen=max_seqlen,
+        ).transpose(1, 2)
+
+        return attn_out
+    else:
+        raise RuntimeError(
+            "No viable backend for scaled_dot_product_attention was found."
+        )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..36c3ae609bb0871214138888d29501c803705598
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/__init__.py
@@ -0,0 +1,62 @@
+# mypy: allow-untyped-defs
+from torch.nn.parameter import (  # usort: skip
+    Buffer as Buffer,
+    Parameter as Parameter,
+    UninitializedBuffer as UninitializedBuffer,
+    UninitializedParameter as UninitializedParameter,
+)
+from torch.nn.modules import *  # usort: skip # noqa: F403
+from torch.nn import (
+    attention as attention,
+    functional as functional,
+    init as init,
+    modules as modules,
+    parallel as parallel,
+    parameter as parameter,
+    utils as utils,
+)
+from torch.nn.parallel import DataParallel as DataParallel
+
+
+def factory_kwargs(kwargs):
+    r"""Return a canonicalized dict of factory kwargs.
+
+    Given kwargs, returns a canonicalized dict of factory kwargs that can be directly passed
+    to factory functions like torch.empty, or errors if unrecognized kwargs are present.
+
+    This function makes it simple to write code like this::
+
+        class MyModule(nn.Module):
+            def __init__(self, **kwargs):
+                factory_kwargs = torch.nn.factory_kwargs(kwargs)
+                self.weight = Parameter(torch.empty(10, **factory_kwargs))
+
+    Why should you use this function instead of just passing `kwargs` along directly?
+
+    1. This function does error validation, so if there are unexpected kwargs we will
+    immediately report an error, instead of deferring it to the factory call
+    2. This function supports a special `factory_kwargs` argument, which can be used to
+    explicitly specify a kwarg to be used for factory functions, in the event one of the
+    factory kwargs conflicts with an already existing argument in the signature (e.g.
+    in the signature ``def f(dtype, **kwargs)``, you can specify ``dtype`` for factory
+    functions, as distinct from the dtype argument, by saying
+    ``f(dtype1, factory_kwargs={"dtype": dtype2})``)
+    """
+    if kwargs is None:
+        return {}
+    simple_keys = {"device", "dtype", "memory_format"}
+    expected_keys = simple_keys | {"factory_kwargs"}
+    if not kwargs.keys() <= expected_keys:
+        raise TypeError(f"unexpected kwargs {kwargs.keys() - expected_keys}")
+
+    # guarantee no input kwargs is untouched
+    r = dict(kwargs.get("factory_kwargs", {}))
+    for k in simple_keys:
+        if k in kwargs:
+            if k in r:
+                raise TypeError(
+                    f"{k} specified twice, in **kwargs and in factory_kwargs"
+                )
+            r[k] = kwargs[k]
+
+    return r
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/_reduction.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/_reduction.py
new file mode 100644
index 0000000000000000000000000000000000000000..a3ca62929a3b53ec7108df75034d51ccb7a7949c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/_reduction.py
@@ -0,0 +1,60 @@
+import warnings
+
+
+# NB: Keep this file in sync with enums in aten/src/ATen/core/Reduction.h
+
+
+def get_enum(reduction: str) -> int:
+    if reduction == "none":
+        ret = 0
+    elif reduction == "mean":
+        ret = 1
+    elif reduction == "elementwise_mean":
+        warnings.warn(
+            "reduction='elementwise_mean' is deprecated. "
+            "Please use reduction='mean' instead.",
+            stacklevel=2,
+        )
+        ret = 1
+    elif reduction == "sum":
+        ret = 2
+    else:
+        ret = -1  # TODO: remove once JIT exceptions support control flow
+        raise ValueError(f"{reduction} is not a valid value for reduction")
+    return ret
+
+
+# In order to support previous versions, accept boolean size_average and reduce
+# and convert them into the new constants for now
+
+
+# We use these functions in torch/legacy as well, in which case we'll silence the warning
+def legacy_get_string(
+    size_average: bool | None,
+    reduce: bool | None,
+    emit_warning: bool = True,
+) -> str:
+    warning = "size_average and reduce args will be deprecated, please use reduction='{}' instead."
+
+    if size_average is None:
+        size_average = True
+    if reduce is None:
+        reduce = True
+
+    if size_average and reduce:
+        ret = "mean"
+    elif reduce:
+        ret = "sum"
+    else:
+        ret = "none"
+    if emit_warning:
+        warnings.warn(warning.format(ret), stacklevel=2)
+    return ret
+
+
+def legacy_get_enum(
+    size_average: bool | None,
+    reduce: bool | None,
+    emit_warning: bool = True,
+) -> int:
+    return get_enum(legacy_get_string(size_average, reduce, emit_warning))
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f3c0e56e0b74bce7a26cbd5b6421f61a98e23d4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/__init__.py
@@ -0,0 +1,197 @@
+# mypy: allow-untyped-defs
+"""This module contains functions and classes that alter the behavior of torch.nn.functional.scaled_dot_product_attention"""
+
+import contextlib
+from collections.abc import Iterable
+from typing import Union
+from warnings import warn
+
+import torch.backends.cuda
+from torch._C import _SDPBackend as SDPBackend
+from torch.backends.cuda import (
+    can_use_efficient_attention,
+    can_use_flash_attention,
+    SDPAParams,
+)
+
+
+__all__: list[str] = [
+    "SDPBackend",
+    "sdpa_kernel",
+    "WARN_FOR_UNFUSED_KERNELS",
+    "register_flash_attention_impl",
+    "activate_flash_attention_impl",
+    "list_flash_attention_impls",
+    "current_flash_attention_impl",
+    "restore_flash_attention_impl",
+]
+
+
+# Note: [SDPA warnings]
+# TODO: Consider using this for sdpa regardless of subclasses
+# This only effects users of bias subclasses
+# If this is set to True, we will warn the user if they are not using the fused kernels
+# As well, it will raise warnings for all the reasons why the fused kernels can't be run.
+# To set this to True, run
+# torch.nn.attention.WARN_FOR_UNFUSED_KERNELS = True
+WARN_FOR_UNFUSED_KERNELS = False
+
+
+r"""An enum-like class that contains the different backends for scaled dot product attention.
+    This backend class is designed to be used with the sdpa_kernel context manager.
+
+    The following Enums are available:
+        - ERROR: An error occurred when trying to determine the backend.
+        - MATH: The math backend for scaled dot product attention.
+        - FLASH_ATTENTION: The flash attention backend for scaled dot product attention.
+        - EFFICIENT_ATTENTION: The efficient attention backend for scaled dot product attention.
+        - CUDNN_ATTENTION: The cuDNN backend for scaled dot product attention.
+        - OVERRIDEABLE: The overridable backend for extension.
+
+    See :func:`torch.nn.attention.sdpa_kernel` for more details.
+
+    .. warning:: This class is in beta and subject to change.
+"""
+SDPBackend.__module__ = __name__
+SDPBackend.__name__ = "SDPBackend"
+
+
+def _raise_kernel_warnings(params: SDPAParams) -> None:
+    """
+    If WARN_FOR_UNFUSED_KERNELS is set to True, this will raise warnings
+    for all the reasons why the fused kernels can't be run. If using subclasses
+    """
+    if WARN_FOR_UNFUSED_KERNELS:
+        if not can_use_efficient_attention(params):
+            warn("Efficient attention can't be used because:", stacklevel=2)
+            can_use_efficient_attention(params, True)
+        if not can_use_flash_attention(params):
+            warn("Flash attention can't be used because:", stacklevel=2)
+            can_use_flash_attention(params, True)
+
+
+_backend_names = {
+    "cudnn": "CUDNN_ATTENTION",
+    "flash": "FLASH_ATTENTION",
+    "mem_efficient": "EFFICIENT_ATTENTION",
+    "math": "MATH",
+    "overrideable": "OVERRIDEABLE",
+}
+
+
+def _backend_from_string(name: str):
+    return getattr(SDPBackend, name)
+
+
+def _cur_sdpa_kernel_backends(with_priority: bool = False):
+    backends = []
+    for name, val in _backend_names.items():
+        if getattr(torch._C, f"_get_{name}_sdp_enabled")():
+            backends.append(getattr(SDPBackend, val))
+    if with_priority:
+        curr_priority = torch._C._get_sdp_priority_order()
+        backends = sorted(
+            backends, key=lambda backend: curr_priority.index(int(backend))
+        )
+    return backends
+
+
+def _sdpa_kernel(backends: Iterable, set_priority: bool = False) -> None:
+    for name, val in _backend_names.items():
+        enabled = getattr(SDPBackend, val) in backends
+        getattr(torch._C, f"_set_sdp_use_{name}")(enabled)
+    if set_priority:
+        # backends should be a unique list
+        user_priority = [int(backend) for backend in backends]
+        previous_priority = torch._C._get_sdp_priority_order()
+        for backend in previous_priority:
+            if backend not in user_priority:
+                user_priority.append(int(backend))
+        torch._C._set_sdp_priority_order(user_priority)
+
+
+@contextlib.contextmanager
+def sdpa_kernel(backends: list[SDPBackend] | SDPBackend, set_priority: bool = False):
+    r"""
+    Context manager to select which backend to use for scaled dot product attention.
+
+    .. warning:: This function is beta and subject to change.
+
+    Args:
+        backends (Union[List[SDPBackend], SDPBackend]): A backend or list of backends for scaled dot product attention.
+        set_priority (bool=False): Whether the ordering of the backends is interpreted as their priority order.
+
+    Example:
+
+    .. code-block:: python
+
+        from torch.nn.functional import scaled_dot_product_attention
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+
+        # Only enable flash attention backend
+        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+            scaled_dot_product_attention(...)
+
+        # Enable the Math or Efficient attention backends
+        with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
+            scaled_dot_product_attention(...)
+
+        # Enable the cuDNN or flash attention backends, and in that order
+        with sdpa_kernel(
+            [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION], set_priority=True
+        ):
+            scaled_dot_product_attention(...)
+
+    This context manager can be used to select which backend to use for scaled dot product attention.
+    Upon exiting the context manager, the previous state of the flags will be restored, enabling all backends.
+    """
+    if not isinstance(backends, (list, SDPBackend)):
+        raise AssertionError(
+            f"Backend must be an instance of SDPBackend or a list of SDPBackend instances, got {type(backends).__name__}"
+        )
+
+    if isinstance(backends, SDPBackend):
+        backends = [backends]
+
+    backends = list(dict.fromkeys(backends))
+
+    previous_backends = _cur_sdpa_kernel_backends(with_priority=set_priority)
+    try:
+        _sdpa_kernel(backends, set_priority)
+        yield {}
+    finally:
+        _sdpa_kernel(previous_backends, set_priority)
+
+
+# variadic version of sdpa_kernel for dynamo to use while reconstructing
+@contextlib.contextmanager
+def _sdpa_kernel_variadic(*backends: SDPBackend):
+    with sdpa_kernel(list(backends)):
+        yield
+
+
+def _get_flash_version() -> str:
+    """This returns the closest matching tag for the flash attention backend"""
+    return "2.5.7"
+
+
+from . import _registry
+
+
+# Re-export registry types and functions for public API
+_FlashAttentionImpl = _registry._FlashAttentionImpl
+_RegisterFn = _registry._RegisterFn
+register_flash_attention_impl = _registry.register_flash_attention_impl
+activate_flash_attention_impl = _registry.activate_flash_attention_impl
+list_flash_attention_impls = _registry.list_flash_attention_impls
+current_flash_attention_impl = _registry.current_flash_attention_impl
+restore_flash_attention_impl = _registry.restore_flash_attention_impl
+
+register_flash_attention_impl.__module__ = __name__
+activate_flash_attention_impl.__module__ = __name__
+list_flash_attention_impls.__module__ = __name__
+current_flash_attention_impl.__module__ = __name__
+restore_flash_attention_impl.__module__ = __name__
+
+# Import built-in implementations to trigger self-registration
+from . import _fa3, _fa4  # noqa: F401  # noqa: F401
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_fa3.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_fa3.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ea065134c4346106700d8228900930ba37191df
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_fa3.py
@@ -0,0 +1,763 @@
+"""
+PROTOTYPE!
+Flash Attention 3 implementation.
+For fp8: only supports forward pass right now.
+For fp16/bf16: supports forward and backward pass.
+"""
+# mypy: allow-untyped-defs
+
+from __future__ import annotations
+
+import importlib
+import warnings
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+from dataclasses import dataclass
+from functools import cache
+from typing_extensions import TypeVarTuple, Unpack
+
+import torch
+from torch.library import Library
+
+from . import _registry
+
+
+__all__ = [
+    "register_flash_attention_fa3",
+]
+
+
+_FA3_CUDA_FWD: Callable | None = None  # Cache for torch.ops.flash_attn_3.fwd
+_FA3_CUDA_BWD: Callable | None = None  # Cache for torch.ops.flash_attn_3.bwd
+
+
+@dataclass
+class _FA3Handle:
+    library: Library | None
+
+    def remove(self) -> None:
+        self.library = None
+        # Clear the C++ flag
+        torch._C._set_sdp_use_fa3(False)
+
+
+@cache
+def _get_device_major(device: torch.device) -> int:
+    major, _ = torch.cuda.get_device_capability(device)
+    return major
+
+
+def register_flash_attention_fa3(
+    module_path: str = "flash_attn_interface",
+) -> _FA3Handle:
+    """
+    Register FA3 flash attention kernels with the PyTorch dispatcher.
+
+    Args:
+        module_path: Python module path to the FA3 implementation.
+    """
+    _fa3_import_module(module_path)
+
+    # Expose FA3 registration status to C++
+    torch._C._set_sdp_use_fa3(True)
+
+    return _FA3Handle(_fa3_register_kernels())
+
+
+def _fa3_import_module(module_path: str) -> None:
+    importlib.import_module(module_path)
+    if not hasattr(torch.ops, "flash_attn_3"):
+        raise RuntimeError(f"Module '{module_path}' does not expose FA3 kernels")
+    if not hasattr(torch.ops.flash_attn_3, "fwd"):
+        raise RuntimeError(
+            f"Module '{module_path}' does not expose FA3 forward kernels"
+        )
+    if not hasattr(torch.ops.flash_attn_3, "bwd"):
+        raise RuntimeError(
+            f"Module '{module_path}' does not expose FA3 backward kernels"
+        )
+    global _FA3_CUDA_FWD, _FA3_CUDA_BWD
+    _FA3_CUDA_FWD = torch.ops.flash_attn_3.fwd
+    _FA3_CUDA_BWD = torch.ops.flash_attn_3.bwd
+
+
+def _fa3_register_kernels() -> Library:
+    lib = Library("aten", "IMPL", "CUDA")  # noqa: TOR901
+    lib.impl(
+        "_flash_attention_forward.quantized", _fa3_flash_attention_forward_impl, "CUDA"
+    )
+    lib.impl(
+        "_scaled_dot_product_flash_attention.quantized",
+        _fa3_scaled_dot_product_flash_attention_forward_impl,
+        "CUDA",
+    )
+    lib.impl(
+        "_flash_attention_forward", _fa3_flash_attention_forward_impl_default, "CUDA"
+    )
+    lib.impl(
+        "_flash_attention_forward_no_dropout_inplace",
+        _fa3_flash_attention_forward_no_dropout_inplace_impl,
+        "CUDA",
+    )
+    lib.impl(
+        "_scaled_dot_product_flash_attention",
+        _fa3_scaled_dot_product_flash_attention_forward_impl_default,
+        "CUDA",
+    )
+
+    lib.impl("_flash_attention_backward", _fa3_flash_attention_backward_impl, "CUDA")
+    lib.impl(
+        "_scaled_dot_product_flash_attention_backward",
+        _fa3_scaled_dot_product_flash_attention_backward_impl,
+        "CUDA",
+    )
+    return lib
+
+
+def _fa3_common_support_error(
+    query: torch.Tensor,
+    tensors: tuple[torch.Tensor, ...],
+    dropout_p: float,
+    cum_seq_q: torch.Tensor | None,
+    q_descale: torch.Tensor | None,
+    k_descale: torch.Tensor | None,
+    v_descale: torch.Tensor | None,
+) -> str | None:
+    if dropout_p != 0.0:
+        return "dropout_p must be 0"
+
+    if not all(t.is_cuda for t in tensors):
+        return "inputs must be CUDA tensors"
+    if len({t.device for t in tensors}) != 1:
+        return "inputs must share device"
+    if query.dtype == torch.float8_e4m3fn and (
+        q_descale is None or k_descale is None or v_descale is None
+    ):
+        warnings.warn(
+            "When using SDPA with fp8, descale tensor should always be used"
+            " for accurate dequantization. Please use "
+            "_scaled_dot_product_attention_quantized and "
+            "provide the descale tensors.",
+            UserWarning,
+        )
+    if cum_seq_q is None and query.dim() != 4:
+        return "dense query must be 4D"
+    if cum_seq_q is not None and query.dim() != 3:
+        return "ragged query must be 3D"
+    if not torch.cuda.is_available():
+        return "CUDA not available"
+    if _get_device_major(query.device) != 9:
+        return "FA3 requires compute capability 9.0"
+    return None
+
+
+def _fa3_forward_support_error(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    dropout_p: float,
+    return_debug_mask: bool,
+    alibi_slopes: torch.Tensor | None,
+    seqused_k: torch.Tensor | None,
+    cum_seq_q: torch.Tensor | None,
+    q_descale: torch.Tensor | None,
+    k_descale: torch.Tensor | None,
+    v_descale: torch.Tensor | None,
+) -> str | None:
+    if return_debug_mask:
+        return "return_debug_mask must be False"
+    if alibi_slopes is not None:
+        return "alibi_slopes not supported"
+    if seqused_k is not None:
+        if seqused_k.dtype != torch.int32:
+            return "seqused_k must be int32"
+        if not seqused_k.is_cuda:
+            return "seqused_k must be CUDA"
+    supported_dtypes = (torch.float8_e4m3fn, torch.float16, torch.bfloat16)
+    if not all(t.dtype in supported_dtypes for t in {query, key, value}):
+        return f"inputs must be one of {supported_dtypes}"
+    if len({t.dtype for t in {query, key, value}}) != 1:
+        return "all inputs must have the same dtype"
+    error = _fa3_common_support_error(
+        query,
+        (query, key, value),
+        dropout_p,
+        cum_seq_q,
+        q_descale,
+        k_descale,
+        v_descale,
+    )
+    if error is not None:
+        if error == "inputs must share device":
+            return "query, key, value must be on same device"
+        return error
+    return None
+
+
+def _fa3_backward_support_error(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    dropout_p: float,
+    cum_seq_q: torch.Tensor | None,
+    window_size_left: int | None,
+    window_size_right: int | None,
+) -> str | None:
+    # FA3 backward ONLY supports fp16/bf16, NOT fp8
+    if query.dtype == torch.float8_e4m3fn:
+        return (
+            "FA3 backward does not support fp8 - use inference only (torch.no_grad())"
+        )
+    if logsumexp.dtype != torch.float32:
+        return "logsumexp dtype must be float32"
+    supported_dtypes = (torch.float16, torch.bfloat16)
+    if not all(t.dtype in supported_dtypes for t in {grad_out, query, key, value, out}):
+        return f"inputs must be one of {supported_dtypes}"
+    if len({t.dtype for t in {grad_out, query, key, value, out}}) != 1:
+        return "all inputs must have the same dtype"
+    error = _fa3_common_support_error(
+        query,
+        (grad_out, query, key, value, out, logsumexp),
+        dropout_p,
+        cum_seq_q,
+        None,
+        None,
+        None,
+    )
+    if error is not None:
+        return error
+    return None
+
+
+Ts = TypeVarTuple("Ts")
+
+
+def _transpose_dense(*tensors: Unpack[Ts]) -> tuple[Unpack[Ts]]:
+    return tuple(t.transpose(1, 2) for t in tensors)  # type: ignore[attr-defined]
+
+
+def _maybe_contiguous(x: torch.Tensor | None) -> torch.Tensor | None:
+    """Ensure tensor is contiguous in the last dimension."""
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+def _fa3_run_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor | None,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    scale: float | None,
+    is_causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
+    seqused_k: torch.Tensor | None,
+    out: torch.Tensor | None = None,
+    q_descale: torch.Tensor | None = None,
+    k_descale: torch.Tensor | None = None,
+    v_descale: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Run the FA3 forward pass by calling the C++ kernel directly.
+    """
+    if _FA3_CUDA_FWD is None:
+        raise RuntimeError("FA3 not registered")
+    # Ensure contiguous in the last dimension
+    q = _maybe_contiguous(query)
+    k = _maybe_contiguous(key)
+    v = (
+        value.contiguous()
+        if value.dtype == torch.float8_e4m3fn
+        and value.stride(-1) != 1
+        and value.stride(-3) != 1
+        else _maybe_contiguous(value)
+    )
+
+    cu_seqlens_q = _maybe_contiguous(cu_seq_q)
+    cu_seqlens_k = _maybe_contiguous(cu_seq_k)
+    seqused_k = _maybe_contiguous(seqused_k)
+    block_table = _maybe_contiguous(block_table)
+
+    out, softmax_lse, out_accum, softmax_lse_accum = _FA3_CUDA_FWD(
+        q,
+        k,
+        v,
+        None,  # k_new
+        None,  # v_new
+        None,  # qv
+        out,  # out_ (pre-allocated output)
+        cu_seqlens_q,  # cu_seqlens_q
+        cu_seqlens_k,  # cu_seqlens_k
+        None,  # cu_seqlens_k_new
+        None,  # seqused_q
+        seqused_k,  # seqused_k
+        max_q,  # max_seqlen_q
+        max_k,  # max_seqlen_k
+        block_table,  # block_table,
+        None,  # kv_batch_idx,
+        None,  # leftpad_k,
+        None,  # rotary_cos,
+        None,  # rotary_sin,
+        None,  # seqlens_rotary,
+        q_descale,  # q_descale,
+        k_descale,  # k_descale,
+        v_descale,  # v_descale,
+        scale,  # softmax_scale,
+        is_causal,  # causal,
+        window_size_left if window_size_left is not None else -1,  # window_size_left
+        window_size_right if window_size_right is not None else -1,  # window_size_right
+        0,  # attention_chunk,
+        0.0,  # softcap,
+        True,  # rotary_interleaved,
+        None,  # scheduler_metadata,
+        num_splits
+        or (1 if torch.are_deterministic_algorithms_enabled() else 0),  # num_splits,
+        None,  # pack_gqa,
+        torch._C._get_sm_carveout_experimental() or 0,  # sm_margin,
+    )
+    return out, softmax_lse.contiguous()
+
+
+def _fa3_run_backward(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    cu_seq_q: torch.Tensor | None,
+    cu_seq_k: torch.Tensor | None,
+    max_seqlen_q: int | None,
+    max_seqlen_k: int | None,
+    scale: float | None,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    deterministic: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if _FA3_CUDA_BWD is None:
+        raise RuntimeError("FA3 not registered")
+
+    # Ensure contiguous
+    dout = _maybe_contiguous(grad_out)
+    q = query.contiguous() if query.stride(-1) != 1 else query
+    k = key.contiguous() if key.stride(-1) != 1 else key
+    v = value.contiguous() if value.stride(-1) != 1 else value
+    o = _maybe_contiguous(out)
+    lse = _maybe_contiguous(logsumexp)
+
+    # Pre-allocate gradient tensors
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+    _FA3_CUDA_BWD(
+        dout,
+        q,
+        k,
+        v,
+        o,
+        lse,
+        dq,
+        dk,
+        dv,
+        cu_seq_q,
+        cu_seq_k,
+        None,
+        None,
+        max_seqlen_q,
+        max_seqlen_k,
+        scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        0.0,
+        deterministic,
+        torch._C._get_sm_carveout_experimental() or 0,
+    )
+    return dq, dk, dv
+
+
+def _fa3_flash_attention_forward_impl(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    q_descale: torch.Tensor | None = None,
+    k_descale: torch.Tensor | None = None,
+    v_descale: torch.Tensor | None = None,
+    *,
+    scale: float | None = None,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    seqused_k: torch.Tensor | None = None,
+    alibi_slopes: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    compute_auxiliary: bool = True,
+    num_splits: int | None = None,
+):
+    error = _fa3_forward_support_error(
+        query,
+        key,
+        value,
+        dropout_p,
+        return_debug_mask,
+        alibi_slopes,
+        seqused_k,
+        cum_seq_q,
+        q_descale,
+        k_descale,
+        v_descale,
+    )
+    if error is not None:
+        raise RuntimeError(f"FA3 flash_attention forward unsupported: {error}")
+    out, lse = _fa3_run_forward(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        seqused_k,
+        out,
+        q_descale,
+        k_descale,
+        v_descale,
+        block_table,
+        num_splits,
+    )
+    if compute_auxiliary:
+        rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
+        philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
+        debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
+    else:
+        rng_state = None
+        philox_offset = None
+        debug_mask = None
+    return out, lse, rng_state, philox_offset, debug_mask
+
+
+def _fa3_flash_attention_forward_no_dropout_inplace_impl(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    *,
+    scale: float | None = None,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    seqused_k: torch.Tensor | None = None,
+    alibi_slopes: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+):
+    _, lse, _, _, _ = _fa3_flash_attention_forward_impl(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        None,
+        None,
+        None,
+        scale=scale,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        seqused_k=seqused_k,
+        alibi_slopes=alibi_slopes,
+        out=out,
+        block_table=block_table,
+        compute_auxiliary=False,
+        num_splits=num_splits,
+    )
+    return lse
+
+
+def _fa3_flash_attention_forward_impl_default(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    *,
+    scale: float | None = None,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    seqused_k: torch.Tensor | None = None,
+    alibi_slopes: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    num_splits: int | None = None,
+):
+    return _fa3_flash_attention_forward_impl(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        None,
+        None,
+        None,
+        scale=scale,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        seqused_k=seqused_k,
+        alibi_slopes=alibi_slopes,
+        out=out,
+        block_table=block_table,
+        num_splits=num_splits,
+    )
+
+
+def _fa3_flash_attention_backward_impl(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    rng_state: torch.Tensor,
+    unused: torch.Tensor,
+    *,
+    scale: float | None = None,
+    window_size_left: int | None = None,
+    window_size_right: int | None = None,
+):
+    """FA3 implementation of _flash_attention_backward."""
+    error = _fa3_backward_support_error(
+        grad_out,
+        query,
+        key,
+        value,
+        out,
+        logsumexp,
+        dropout_p,
+        cum_seq_q,
+        window_size_left,
+        window_size_right,
+    )
+
+    if error is not None:
+        raise RuntimeError(f"FA3 flash_attention backward unsupported: {error}")
+
+    deterministic = torch.are_deterministic_algorithms_enabled()
+
+    dq, dk, dv = _fa3_run_backward(
+        grad_out,
+        query,
+        key,
+        value,
+        out,
+        logsumexp,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        scale,
+        is_causal,
+        window_size_left if window_size_left is not None else -1,
+        window_size_right if window_size_right is not None else -1,
+        deterministic,
+    )
+    return dq, dk, dv
+
+
+def _fa3_scaled_dot_product_flash_attention_forward_impl(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    q_descale: torch.Tensor | None = None,
+    k_descale: torch.Tensor | None = None,
+    v_descale: torch.Tensor | None = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    return_debug_mask: bool = False,
+    *,
+    scale: float | None = None,
+):
+    error = _fa3_forward_support_error(
+        query,
+        key,
+        value,
+        dropout_p,
+        return_debug_mask,
+        None,
+        None,
+        None,
+        q_descale,
+        k_descale,
+        v_descale,
+    )
+    if error is not None:
+        raise RuntimeError(f"FA3 SDPA forward unsupported: {error}")
+    q, k, v = _transpose_dense(query, key, value)
+
+    # Pre-allocate output with query's strides (BHSD layout), then create
+    # a BSHD view for the kernel. This ensures the returned output has
+    # the same memory layout as the input query.
+    out_dtype = torch.bfloat16 if query.dtype == torch.float8_e4m3fn else query.dtype
+    out_bhsd = torch.empty_like(query, dtype=out_dtype)
+    out_bshd = out_bhsd.transpose(1, 2)
+
+    max_q_flash = q.size(1)
+    max_k_flash = k.size(1)
+    _, lse, rng_state, philox_offset, debug_mask = _fa3_flash_attention_forward_impl(
+        q,
+        k,
+        v,
+        None,
+        None,
+        max_q_flash,
+        max_k_flash,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+        out=out_bshd,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+    )
+    max_q = query.size(2)
+    max_k = key.size(2)
+    return (
+        out_bhsd,
+        lse,
+        None,
+        None,
+        max_q,
+        max_k,
+        rng_state,
+        philox_offset,
+        debug_mask,
+    )
+
+
+def _fa3_scaled_dot_product_flash_attention_forward_impl_default(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    return_debug_mask: bool = False,
+    *,
+    scale: float | None = None,
+):
+    return _fa3_scaled_dot_product_flash_attention_forward_impl(
+        query,
+        key,
+        value,
+        None,
+        None,
+        None,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+    )
+
+
+def _fa3_scaled_dot_product_flash_attention_backward_impl(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    philox_seed: torch.Tensor,
+    philox_offset: torch.Tensor,
+    *,
+    scale: float | None = None,
+):
+    """FA3 implementation of _scaled_dot_product_flash_attention_backward."""
+    error = _fa3_backward_support_error(
+        grad_out, query, key, value, out, logsumexp, dropout_p, None, None, None
+    )
+    if error is not None:
+        raise RuntimeError(f"FA3 SDPA backward unsupported: {error}")
+
+    # SDPA uses BHSD layout, FA3 uses BSHD - transpose
+    grad_out_t, q_t, k_t, v_t, out_t = _transpose_dense(
+        grad_out, query, key, value, out
+    )
+
+    dq, dk, dv = _fa3_flash_attention_backward_impl(
+        grad_out_t,
+        q_t,
+        k_t,
+        v_t,
+        out_t,
+        logsumexp,
+        None,  # cum_seq_q (dense attention)
+        None,  # cum_seq_k
+        max_q,  # max_seqlen_q
+        max_k,  # max_seqlen_k
+        dropout_p,
+        is_causal,
+        philox_seed,
+        philox_offset,
+        scale=scale,
+    )
+
+    # Transpose gradients back to BHSD layout
+    dq_out, dk_out, dv_out = _transpose_dense(dq, dk, dv)
+    return dq_out, dk_out, dv_out
+
+
+_registry.register_flash_attention_impl("FA3", register_fn=register_flash_attention_fa3)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_fa4.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_fa4.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2786d9dceefeb868c8f724a50e56ed9da123ce3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_fa4.py
@@ -0,0 +1,541 @@
+"""UBER PROTOTYPE!!!"""
+# mypy: allow-untyped-defs
+
+from __future__ import annotations
+
+import importlib
+from dataclasses import dataclass
+from functools import cache
+from typing import Any, TYPE_CHECKING
+from typing_extensions import TypeVarTuple, Unpack
+
+from . import _registry
+
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+import torch
+from torch.library import Library
+
+
+__all__ = [
+    "register_flash_attention_fa4",
+]
+
+
+_FA4_MODULE_PATH: str | None = None
+
+
+@dataclass
+class _FA4Handle:
+    library: Library | None
+
+    def remove(self) -> None:
+        self.library = None
+
+
+@cache
+def _get_device_major(device: torch.device) -> int:
+    major, _ = torch.cuda.get_device_capability(device)
+    return major
+
+
+def register_flash_attention_fa4(
+    module_path: str = "flash_attn.cute.interface",
+) -> _FA4Handle:
+    """
+    Register FA4 flash attention kernels with the PyTorch dispatcher.
+
+    Args:
+        module_path: Python module path to the FA4 implementation.
+    """
+    global _FA4_MODULE_PATH
+    _ = _fa4_import_module(module_path)
+    _FA4_MODULE_PATH = module_path
+    return _FA4Handle(_fa4_register_kernels())
+
+
+@cache
+def _fa4_import_module(module_path: str) -> ModuleType:
+    module = importlib.import_module(module_path)
+    if not hasattr(module, "_flash_attn_fwd") or not hasattr(module, "_flash_attn_bwd"):
+        raise RuntimeError(f"Module '{module_path}' does not expose FA4 kernels")
+    return module
+
+
+def _fa4_register_kernels() -> Library:
+    lib = Library("aten", "IMPL", "CUDA")  # noqa: TOR901
+    lib.impl("_flash_attention_forward", _fa4_flash_attention_forward_impl, "CUDA")
+    lib.impl(
+        "_flash_attention_forward_no_dropout_inplace",
+        _fa4_flash_attention_forward_no_dropout_inplace_impl,
+        "CUDA",
+    )
+    lib.impl("_flash_attention_backward", _fa4_flash_attention_backward_impl, "CUDA")
+    lib.impl(
+        "_scaled_dot_product_flash_attention",
+        _fa4_scaled_dot_product_flash_attention_forward_impl,
+        "CUDA",
+    )
+    lib.impl(
+        "_scaled_dot_product_flash_attention_backward",
+        _fa4_scaled_dot_product_flash_attention_backward_impl,
+        "CUDA",
+    )
+    return lib
+
+
+def _fa4_common_support_error(
+    query: torch.Tensor,
+    tensors: tuple[torch.Tensor, ...],
+    cum_seq_q: torch.Tensor | None,
+    require_fp32: tuple[tuple[str, torch.Tensor], ...] = (),
+) -> str | None:
+    if not all(t.is_cuda for t in tensors):
+        return "inputs must be CUDA tensors"
+    if len({t.device for t in tensors}) != 1:
+        return "inputs must share device"
+    if query.dtype not in (torch.float16, torch.bfloat16):
+        return "query dtype must be float16 or bfloat16"
+    for name, tensor in require_fp32:
+        if tensor.dtype != torch.float32:
+            return f"{name} dtype must be float32"
+    if cum_seq_q is None and query.dim() != 4:
+        return "dense query must be 4D"
+    if cum_seq_q is not None and query.dim() != 3:
+        return "ragged query must be 3D"
+    if not torch.cuda.is_available():
+        return "CUDA not available"
+    if _get_device_major(query.device) not in (9, 10):
+        return "FA4 requires compute capability 9.0 or 10.0"
+    return None
+
+
+def _fa4_forward_support_error(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    dropout_p: float,
+    return_debug_mask: bool,
+    alibi_slopes: torch.Tensor | None,
+    seqused_k: torch.Tensor | None,
+    cum_seq_q: torch.Tensor | None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> str | None:
+    if dropout_p != 0.0:
+        return "dropout_p must be 0"
+    if return_debug_mask:
+        return "return_debug_mask must be False"
+    if alibi_slopes is not None:
+        return "alibi_slopes not supported"
+    if seqused_k is not None:
+        if seqused_k.dtype != torch.int32:
+            return "seqused_k must be int32"
+        if not seqused_k.is_cuda:
+            return "seqused_k must be CUDA"
+    major = _get_device_major(query.device)
+    if block_table is not None and major != 10:
+        return f"paged KV (block_table) not supported on SM {major}0"
+    if num_splits is not None and num_splits > 1 and major != 10:
+        return f"SplitKV (num_splits > 1) not supported on SM {major}0"
+    error = _fa4_common_support_error(
+        query,
+        (query, key, value),
+        cum_seq_q,
+    )
+    if error is not None:
+        if error == "inputs must share device":
+            return "query, key, value must be on same device"
+        return error
+    return None
+
+
+def _fa4_backward_support_error(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    dropout_p: float,
+    cum_seq_q: torch.Tensor | None,
+) -> str | None:
+    if dropout_p != 0.0:
+        return "dropout_p must be 0"
+    error = _fa4_common_support_error(
+        query,
+        (grad_out, query, key, value, out, logsumexp),
+        cum_seq_q,
+        require_fp32=(("logsumexp", logsumexp),),
+    )
+    if error is not None:
+        return error
+    return None
+
+
+def _aten_to_fa4_window_size(val: int | None) -> int | None:
+    """need to convert -1 to None for FA4"""
+    return None if val == -1 else val
+
+
+Ts = TypeVarTuple("Ts")
+
+
+def _transpose_dense(*tensors: Unpack[Ts]) -> tuple[Unpack[Ts]]:
+    return tuple(t.transpose(1, 2) for t in tensors)  # type: ignore[attr-defined]
+
+
+def _fa4_run_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor | None,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int | None,
+    max_k: int | None,
+    scale: float | None,
+    is_causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
+    seqused_k: torch.Tensor | None,
+    out: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if _FA4_MODULE_PATH is None:
+        raise RuntimeError("FA4 not registered")
+    module = _fa4_import_module(_FA4_MODULE_PATH)
+
+    kwargs: dict[str, Any] = {
+        "softmax_scale": scale,
+        "causal": is_causal,
+        "window_size_left": _aten_to_fa4_window_size(window_size_left),
+        "window_size_right": _aten_to_fa4_window_size(window_size_right),
+        "return_lse": True,
+        "cu_seqlens_q": cu_seq_q,
+        "cu_seqlens_k": cu_seq_k,
+        "max_seqlen_q": max_q,
+        "max_seqlen_k": max_k,
+        "seqused_k": seqused_k.contiguous() if seqused_k is not None else None,
+        "page_table": block_table,
+        "num_splits": num_splits or 1,
+        "out": out,
+    }
+    out, lse = module._flash_attn_fwd(query, key, value, **kwargs)
+    return out, lse.contiguous()
+
+
+def _fa4_run_backward(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    cu_seq_q: torch.Tensor | None,
+    cu_seq_k: torch.Tensor | None,
+    scale: float | None,
+    is_causal: bool,
+    window_size_left: int | None,
+    window_size_right: int | None,
+    deterministic: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if _FA4_MODULE_PATH is None:
+        raise RuntimeError("FA4 not registered")
+    module = _fa4_import_module(_FA4_MODULE_PATH)
+    dq, dk, dv = module._flash_attn_bwd(
+        query,
+        key,
+        value,
+        out,
+        grad_out,
+        logsumexp.contiguous(),
+        softmax_scale=scale,
+        causal=is_causal,
+        window_size_left=_aten_to_fa4_window_size(window_size_left),
+        window_size_right=_aten_to_fa4_window_size(window_size_right),
+        cu_seqlens_q=cu_seq_q,
+        cu_seqlens_k=cu_seq_k,
+        deterministic=deterministic,
+    )
+    return dq, dk, dv
+
+
+def _fa4_flash_attention_forward_impl(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    *,
+    scale: float | None = None,
+    window_size_left: int | None = None,
+    window_size_right: int | None = None,
+    seqused_k: torch.Tensor | None = None,
+    alibi_slopes: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    compute_auxiliary: bool = True,
+    num_splits: int | None = None,
+):
+    error = _fa4_forward_support_error(
+        query,
+        key,
+        value,
+        dropout_p,
+        return_debug_mask,
+        alibi_slopes,
+        seqused_k,
+        cum_seq_q,
+        block_table,
+        num_splits,
+    )
+    if error is not None:
+        raise RuntimeError(f"FA4 flash_attention forward unsupported: {error}")
+    out, lse = _fa4_run_forward(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        seqused_k,
+        out,
+        block_table,
+        num_splits,
+    )
+    if compute_auxiliary:
+        rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
+        philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
+        debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
+    else:
+        rng_state = None
+        philox_offset = None
+        debug_mask = None
+    return out, lse, rng_state, philox_offset, debug_mask
+
+
+def _fa4_flash_attention_forward_no_dropout_inplace_impl(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    return_debug_mask: bool,
+    *,
+    scale: float | None = None,
+    window_size_left: int | None = None,
+    window_size_right: int | None = None,
+    seqused_k: torch.Tensor | None = None,
+    alibi_slopes: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+):
+    _, lse, _, _, _ = _fa4_flash_attention_forward_impl(
+        query,
+        key,
+        value,
+        cum_seq_q,
+        cum_seq_k,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        seqused_k=seqused_k,
+        alibi_slopes=alibi_slopes,
+        out=out,
+        block_table=block_table,
+        compute_auxiliary=False,
+        num_splits=num_splits,
+    )
+    return lse
+
+
+def _fa4_flash_attention_backward_impl(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    rng_state: torch.Tensor,
+    unused: torch.Tensor,
+    *,
+    scale: float | None = None,
+    window_size_left: int | None = None,
+    window_size_right: int | None = None,
+):
+    error = _fa4_backward_support_error(
+        grad_out,
+        query,
+        key,
+        value,
+        out,
+        logsumexp,
+        dropout_p,
+        cum_seq_q,
+    )
+    if error is not None:
+        raise RuntimeError(f"FA4 flash_attention backward unsupported: {error}")
+    deterministic = torch.are_deterministic_algorithms_enabled()
+    dq, dk, dv = _fa4_run_backward(
+        grad_out,
+        query,
+        key,
+        value,
+        out,
+        logsumexp,
+        cum_seq_q,
+        cum_seq_k,
+        scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        deterministic,
+    )
+    return dq, dk, dv
+
+
+def _fa4_scaled_dot_product_flash_attention_forward_impl(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    return_debug_mask: bool = False,
+    *,
+    scale: float | None = None,
+):
+    error = _fa4_forward_support_error(
+        query,
+        key,
+        value,
+        dropout_p,
+        return_debug_mask,
+        None,
+        None,
+        None,
+    )
+    if error is not None:
+        raise RuntimeError(f"FA4 SDPA forward unsupported: {error}")
+    q, k, v = _transpose_dense(query, key, value)
+
+    # Pre-allocate output with query's strides (BHSD layout), then create
+    # a BSHD view for the kernel. This ensures the returned output has
+    # the same memory layout as the input query.
+    out_bhsd = torch.empty_like(query)
+    out_bshd = out_bhsd.transpose(1, 2)
+
+    max_q_flash = q.size(1)
+    max_k_flash = k.size(1)
+    _, lse, rng_state, philox_offset, debug_mask = _fa4_flash_attention_forward_impl(
+        q,
+        k,
+        v,
+        None,
+        None,
+        max_q_flash,
+        max_k_flash,
+        dropout_p,
+        is_causal,
+        return_debug_mask,
+        scale=scale,
+        out=out_bshd,
+    )
+    max_q = query.size(2)
+    max_k = key.size(2)
+    return (
+        out_bhsd,
+        lse,
+        None,
+        None,
+        max_q,
+        max_k,
+        rng_state,
+        philox_offset,
+        debug_mask,
+    )
+
+
+def _fa4_scaled_dot_product_flash_attention_backward_impl(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    logsumexp: torch.Tensor,
+    cum_seq_q: torch.Tensor | None,
+    cum_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    dropout_p: float,
+    is_causal: bool,
+    philox_seed: torch.Tensor,
+    philox_offset: torch.Tensor,
+    *,
+    scale: float | None = None,
+):
+    error = _fa4_backward_support_error(
+        grad_out,
+        query,
+        key,
+        value,
+        out,
+        logsumexp,
+        dropout_p,
+        None,
+    )
+    if error is not None:
+        raise RuntimeError(f"FA4 SDPA backward unsupported: {error}")
+    q, k, v, o, go = _transpose_dense(query, key, value, out, grad_out)
+    max_q = query.size(2)
+    max_k = key.size(2)
+    dq, dk, dv = _fa4_flash_attention_backward_impl(
+        go,
+        q,
+        k,
+        v,
+        o,
+        logsumexp,
+        None,
+        None,
+        max_q,
+        max_k,
+        dropout_p,
+        is_causal,
+        philox_seed,
+        philox_offset,
+        scale=scale,
+    )
+    dq, dk, dv = _transpose_dense(dq, dk, dv)
+    return dq, dk, dv
+
+
+_registry.register_flash_attention_impl("FA4", register_fn=register_flash_attention_fa4)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_registry.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..7034c8983916df191db0503b3cf0598a3a7dd480
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_registry.py
@@ -0,0 +1,137 @@
+# mypy: allow-untyped-defs
+"""Registry for flash attention implementations.
+
+This module contains the registration system for flash attention implementations.
+It has no torch dependencies to avoid circular imports during initialization.
+"""
+
+import logging
+from collections.abc import Callable
+from typing import Literal, Protocol
+
+
+logger = logging.getLogger(__name__)
+
+
+class FlashAttentionHandle(Protocol):
+    def remove(self) -> None: ...
+
+
+_RegisterFn = Callable[..., FlashAttentionHandle | None]
+_FlashAttentionImpl = Literal["FA3", "FA4"]
+
+_FLASH_ATTENTION_IMPLS: dict[str, _RegisterFn] = {}
+
+_FLASH_ATTENTION_ACTIVE: tuple[str, FlashAttentionHandle] | None = None
+
+
+def register_flash_attention_impl(
+    impl: str | _FlashAttentionImpl,
+    *,
+    register_fn: _RegisterFn,
+) -> None:
+    """
+    Register the callable that activates a flash attention impl.
+
+    .. note::
+        This function is intended for SDPA backend providers to register their
+        implementations. End users should use :func:`activate_flash_attention_impl`
+        to activate a registered implementation.
+
+    Args:
+        impl: Implementation identifier (e.g., ``"FA4"``).
+        register_fn: Callable that performs the actual dispatcher registration.
+            This function will be invoked by :func:`activate_flash_attention_impl`
+            and should register custom kernels with the PyTorch dispatcher.
+            It may optionally return a handle implementing
+            :class:`FlashAttentionHandle` to keep any necessary state alive.
+
+    Example:
+        >>> def my_impl_register(module_path: str = "my_flash_impl"):
+        ...     # Register custom kernels with torch dispatcher
+        ...     pass  # doctest: +SKIP
+        >>> register_flash_attention_impl(
+        ...     "MyImpl", register_fn=my_impl_register
+        ... )  # doctest: +SKIP
+    """
+    global _FLASH_ATTENTION_IMPLS
+    _FLASH_ATTENTION_IMPLS[impl] = register_fn
+
+
+def activate_flash_attention_impl(
+    impl: str | _FlashAttentionImpl,
+) -> None:
+    """
+    Activate into the dispatcher a previously registered flash attention impl.
+
+    .. note::
+        Backend providers should NOT automatically activate their implementation
+        on import. Users should explicitly opt-in by calling this function or via
+        environment variables to ensure multiple provider libraries can coexist.
+
+    Args:
+        impl: Implementation identifier to activate. See
+            :func:`~torch.nn.attention.list_flash_attention_impls` for available
+            implementations.
+            If the backend's :func:`register_flash_attention_impl` callable
+            returns a :class:`FlashAttentionHandle`, the registry keeps that
+            handle alive for the lifetime of the process (until explicit
+            uninstall support exists).
+
+    Example:
+        >>> activate_flash_attention_impl("FA4")  # doctest: +SKIP
+    """
+    global _FLASH_ATTENTION_ACTIVE, _FLASH_ATTENTION_IMPLS
+
+    restore_flash_attention_impl(
+        _raise_warn=False
+    )  # first restore any prev overrides (if any) to default
+
+    register_fn = _FLASH_ATTENTION_IMPLS.get(impl)
+    if register_fn is None:
+        raise ValueError(
+            f"Unknown flash attention impl '{impl}'. "
+            f"Available implementations: {list_flash_attention_impls()}"
+        )
+
+    handle = register_fn()
+    if handle is not None:
+        _FLASH_ATTENTION_ACTIVE = (impl, handle)
+
+
+def list_flash_attention_impls() -> list[str]:
+    """Return the names of all available flash attention implementations."""
+    return sorted(_FLASH_ATTENTION_IMPLS.keys())
+
+
+def current_flash_attention_impl() -> str | None:
+    """
+    Return the currently activated flash attention impl name, if any.
+
+    ``None`` indicates that no custom impl has been activated.
+    """
+    return (
+        _FLASH_ATTENTION_ACTIVE[0]
+        if _FLASH_ATTENTION_ACTIVE is not None
+        else _FLASH_ATTENTION_ACTIVE
+    )
+
+
+def restore_flash_attention_impl(_raise_warn: bool = True) -> None:
+    """
+    Restore the default FA2 implementation
+    """
+    global _FLASH_ATTENTION_ACTIVE
+
+    handle = None
+    if _FLASH_ATTENTION_ACTIVE is not None:
+        handle = _FLASH_ATTENTION_ACTIVE[1]
+
+    if handle is not None:
+        handle.remove()
+    elif _raise_warn:
+        logger.warning(
+            "Trying to restore default FA2 impl when no custom impl was activated"
+        )
+
+    _FLASH_ATTENTION_ACTIVE = None  # default
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d3e9538034eef6f4292bb7521fdba57391c4c156
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/_utils.py
@@ -0,0 +1,61 @@
+# mypy: allow-untyped-defs
+"""Defines utilities for interacting with scaled_dot_product_attention"""
+
+import math
+
+import torch
+
+
+__all__: list[str] = []
+
+
+def _input_requires_grad(*tensors: torch.Tensor) -> bool:
+    """Returns True if any of the tensors requires grad"""
+    return any(t.requires_grad for t in tensors)
+
+
+def _postprocess_flash_output(inpt_tensor: torch.Tensor, og_size: int) -> torch.Tensor:
+    """Handles the unpad of the last dimension"""
+    if inpt_tensor.size(-1) != og_size:
+        return inpt_tensor[..., :og_size]
+    return inpt_tensor
+
+
+def _calculate_scale(head_dim_size: int, scale: float | None) -> float:
+    """
+    For FlashAttention we pad the head dimension to be a multiple of 8 so we need to scale the output
+    by the original head size and not the padded.
+    """
+    if scale is not None:
+        return scale
+    return 1.0 / math.sqrt(head_dim_size)
+
+
+def _validate_sdpa_input(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None = None,
+    dropout_p=0.0,
+    is_causal=False,
+    scale=None,
+    allow_lowp_kv=False,
+) -> None:
+    if not allow_lowp_kv:
+        if query.dtype != key.dtype or query.dtype != value.dtype:
+            raise ValueError(
+                f"Expected query, key, and value to have the same dtype, "
+                f"but got query.dtype: {query.dtype}, key.dtype: {key.dtype}, "
+                f"and value.dtype: {value.dtype} instead."
+            )
+    if query.device != key.device or query.device != value.device:
+        raise ValueError(
+            f"Expected query, key, and value to have the same device type, "
+            f"but got query.device: {query.device}, key.device: {key.device}, "
+            f"and value.device: {value.device} instead."
+        )
+    if query.dim() < 2 or key.dim() < 2 or value.dim() < 2:
+        raise ValueError(
+            f"Expected query, key, and value to all be  at least 2 dimensional, but got query.dim: "
+            f"{query.dim()}, key.dim: {key.dim()} and value.dim: {value.dim()} instead."
+        )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/bias.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/bias.py
new file mode 100644
index 0000000000000000000000000000000000000000..be03c6ca3efc9a7c88e9a89577db579b0bf12ee8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/bias.py
@@ -0,0 +1,376 @@
+# mypy: allow-untyped-defs
+"""Defines bias subclasses that work with scaled_dot_product_attention"""
+
+from enum import auto, IntEnum
+from warnings import warn
+
+import torch
+import torch.nn.functional as F
+from torch.backends.cuda import (
+    can_use_efficient_attention,
+    can_use_flash_attention,
+    is_flash_attention_available,
+    SDPAParams,
+)
+from torch.nn.attention import _raise_kernel_warnings
+from torch.nn.attention._utils import (
+    _calculate_scale,
+    _input_requires_grad,
+    _postprocess_flash_output,
+    _validate_sdpa_input,
+)
+
+
+__all__ = ["causal_upper_left", "causal_lower_right", "CausalVariant", "CausalBias"]
+
+
+torch._dynamo.allow_in_graph(is_flash_attention_available)
+torch._dynamo.allow_in_graph(can_use_flash_attention)
+torch._dynamo.allow_in_graph(can_use_efficient_attention)
+torch._dynamo.allow_in_graph(SDPAParams)
+
+
+class CausalVariant(IntEnum):
+    r"""
+    Enum for causal variants used in attention mechanisms.
+
+    Defines two types of causal biases:
+
+    ``UPPER_LEFT``: Represents upper-left triangular bias for standard causal attention.
+    The equivalent pytorch code for constructing this bias is:
+
+    .. code-block:: python
+
+        torch.tril(torch.ones(size, dtype=torch.bool))
+
+    For instance, with ``shape=(3,4)``, the materialized bias tensor will be:
+
+    .. code-block:: text
+
+        [[1, 0, 0, 0],
+         [1, 1, 0, 0],
+         [1, 1, 1, 0]]
+
+
+    ``LOWER_RIGHT``: Represents lower-right triangular bias, the include values are aligned to the lower
+    right corner of the matrix.
+
+    The equivalent pytorch code for constructing this bias is:
+
+    .. code-block:: python
+
+        diagonal_offset = size[1] - size[0]
+        torch.tril(
+            torch.ones(size, dtype=torch.bool),
+            diagonal=diagonal_offset,
+        )
+
+    For instance, with ``shape=(3,4)``, the materialized bias tensor will be:
+
+    .. code-block:: text
+
+        [[1, 1, 0, 0],
+         [1, 1, 1, 0],
+         [1, 1, 1, 1]]
+
+    Note that these variants are equivalent to each other when the sequence lengths of the query and key/value
+    tensors are equal since the triangular matrix is square.
+
+    .. warning:: This enum is a prototype and subject to change.
+    """
+
+    UPPER_LEFT = auto()
+    LOWER_RIGHT = auto()
+
+
+class CausalBias(torch.Tensor):
+    """
+    A bias representing causal attention patterns. For an overview of the bias structure, see the :class:`CausalVariant` enum.
+
+    This class is used for defining causal (triangular) attention biases. For construing the bias, there exist
+    two factory functions: :func:`causal_upper_left` and :func:`causal_lower_right`.
+
+    Example:
+
+    .. code-block:: python
+
+        from torch.nn.attention.bias import causal_lower_right
+
+        bsz, num_heads, seqlen_q, seqlen_kv, head_dim = 32, 8, 4, 12, 8
+
+        # Create a lower-right causal bias
+        attn_bias = causal_lower_right(seqlen_q, seqlen_kv)
+
+        q = torch.randn(
+            bsz, num_heads, seqlen_q, head_dim, device="cuda", dtype=torch.float16
+        )
+        k = torch.randn(
+            bsz, num_heads, seqlen_kv, head_dim, device="cuda", dtype=torch.float16
+        )
+        v = torch.randn(
+            bsz, num_heads, seqlen_kv, head_dim, device="cuda", dtype=torch.float16
+        )
+
+        out = F.scaled_dot_product_attention(q, k, v, attn_bias)
+
+    .. warning:: This class is a prototype and subject to change.
+    """
+
+    def __init__(self, variant: CausalVariant, seq_len_q: int, seq_len_kv: int) -> None:
+        """
+        Initializes the CausalBias instance with a specified variant and sequence lengths.
+
+        Args:
+            variant (CausalVariant): The type of causal bias to use (either UPPER_LEFT or LOWER_RIGHT).
+            seq_len_q (int): The sequence length of the query tensor.
+            seq_len_kv (int): The sequence length of the key/value tensor.
+
+        Raises a warning if the LOWER_RIGHT variant is used with seq_len_q > seq_len_kv, as it may produce NaNs.
+        """
+        if not isinstance(variant, CausalVariant):
+            raise AssertionError(
+                f"variant must be a CausalVariant, got {type(variant).__name__}"
+            )
+        super().__init__()
+        self.variant = variant
+        self.seq_len_q = seq_len_q
+        self.seq_len_kv = seq_len_kv
+        if seq_len_q > seq_len_kv and variant == CausalVariant.LOWER_RIGHT:
+            warn(
+                "Lower right causal bias will produce NaNs in the output when seq_len_q > seq_len_kv!",
+                stacklevel=2,
+            )
+
+    def _upper_left(self, device: torch.device) -> torch.Tensor:
+        """Upper left causal bias"""
+        return torch.tril(
+            torch.ones(self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool)
+        )
+
+    def _lower_right(self, device: torch.device) -> torch.Tensor:
+        """Lower right causal bias"""
+        diagonal_offset = self.seq_len_kv - self.seq_len_q
+        return torch.tril(
+            torch.ones(
+                self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool
+            ),
+            diagonal=diagonal_offset,
+        )
+
+    # pyrefly: ignore [bad-return]
+    def _materialize(self, device: torch.device | None = None) -> torch.Tensor:
+        """
+        Materializes the causal bias into a tensor form.
+
+        Depending on the variant, this method generates either an upper-left or lower-right
+        triangular matrix to represent the causal bias.
+
+        Args:
+            device (Optional[torch.device]): The device on which to create the tensor. Defaults to CPU.
+
+        Returns:
+            torch.Tensor: The materialized bias tensor.
+        """
+        if device is None:
+            device = torch.device("cpu")
+        if self.variant == CausalVariant.UPPER_LEFT:
+            return self._upper_left(device)
+        elif self.variant == CausalVariant.LOWER_RIGHT:
+            return self._lower_right(device)
+
+    @staticmethod
+    def _dispatch(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: "CausalBias",
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        scale: float | None = None,
+        enable_gqa: bool = False,
+    ) -> torch.Tensor:
+        r"""
+        Handles the logic for computing attention with the specified causal bias.
+
+        Args:
+            query (Tensor): Query tensor; shape :math:`(N, ..., L, E)`.
+            key (Tensor): Key tensor; shape :math:`(N, ..., S, E)`.
+            value (Tensor): Value tensor; shape :math:`(N, ..., S, Ev)`.
+            attn_mask (CausalBias): The type of causal attention to apply.
+                A boolean mask where a value of True indicates that the element *should* take part in attention.
+                A float mask of the same type as query, key, value that is added to the attention score.
+            dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied
+            is_causal (bool): If true, assumes upper left causal attention masking and errors if both attn_mask and is_causal
+                are set.
+            scale (optional float): Scaling factor applied prior to softmax. If None, the default value is set
+                to :math:`\frac{1}{\sqrt{E}}`.
+            enable_gqa (optional bool): If set to True, Grouped Query Attention (GQA) is enabled, by default it is set to False.
+
+        Returns:
+            output (Tensor): Attention output; shape :math:`(N, ..., L, Ev)`.
+
+        Raises:
+            ValueError: If the causal bias variant is not a CausalVariant type.
+
+        """
+        if is_causal:
+            raise ValueError("CausalBias should not be used with causal=True")
+
+        if (
+            attn_mask.seq_len_q == attn_mask.seq_len_kv
+            or attn_mask.variant == CausalVariant.UPPER_LEFT
+        ):
+            return F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=None,
+                dropout_p=dropout_p,
+                is_causal=True,
+                scale=scale,
+                enable_gqa=enable_gqa,
+            )
+        elif attn_mask.variant == CausalVariant.LOWER_RIGHT:
+            _validate_sdpa_input(query, key, value, None, dropout_p, is_causal, scale)
+            sdpa_params = SDPAParams(
+                query, key, value, None, dropout_p, is_causal, enable_gqa
+            )
+            if can_use_flash_attention(sdpa_params):
+                alignment = 64 if query.device.type == "xpu" else 8
+                og_head_size = query.size(-1)
+                og_scale = _calculate_scale(og_head_size, scale)
+                needs_padding = og_head_size % alignment != 0
+                if needs_padding:
+                    pad_len = alignment - (og_head_size % alignment)
+                    query = torch.nn.functional.pad(query, (0, pad_len))
+                    key = torch.nn.functional.pad(key, (0, pad_len))
+                    value = torch.nn.functional.pad(value, (0, pad_len))
+                out = torch.ops.aten._scaled_dot_product_flash_attention(
+                    query,
+                    key,
+                    value,
+                    dropout_p,
+                    is_causal=True,  # TODO: Flash accepts causal = True and for this particular op it means lower right
+                    return_debug_mask=False,
+                    scale=og_scale,
+                )[0]
+                return _postprocess_flash_output(out, og_head_size)
+            if can_use_efficient_attention(sdpa_params):
+                compute_log_sumexp = False
+                if _input_requires_grad(query, key, value):
+                    compute_log_sumexp = True
+                return torch.ops.aten._efficient_attention_forward(
+                    query.transpose(1, 2),
+                    key.transpose(1, 2),
+                    value.transpose(1, 2),
+                    bias=None,
+                    cu_seqlens_q=None,
+                    cu_seqlens_k=None,
+                    max_seqlen_q=None,
+                    max_seqlen_k=None,
+                    dropout_p=dropout_p,
+                    custom_mask_type=int(attn_mask.variant),
+                    compute_log_sumexp=compute_log_sumexp,
+                    scale=scale,
+                    seqlen_k=None,
+                )[0].transpose(1, 2)
+            else:
+                _raise_kernel_warnings(sdpa_params)
+                # We can't use efficient attention the only support for lower right is via materialization
+                return F.scaled_dot_product_attention(
+                    query,
+                    key,
+                    value,
+                    attn_mask=attn_mask._materialize(query.device),
+                    dropout_p=dropout_p,
+                    is_causal=False,
+                    scale=scale,
+                    enable_gqa=enable_gqa,
+                )
+        else:
+            raise ValueError(
+                f"CausalBias.variant must be a CausalVariant type, but found: {attn_mask.variant}"
+            )
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        """Defines the behavior of torch.nn.functional.scaled_dot_product_attention when the attn_bias is an AttnBias"""
+        if kwargs is None:
+            kwargs = {}
+        if func is torch.nn.functional.scaled_dot_product_attention:
+            return cls._dispatch(*args, **kwargs)
+        return super().__torch_function__(func, types, args, kwargs)
+
+    def __repr__(self) -> str:  # type:ignore[override]
+        return self._materialize().__repr__()
+
+
+def causal_upper_left(*size) -> CausalBias:
+    """
+    Creates an upper-left triangular causal bias.
+
+    This function generates a upper-left triangular matrix to represent causal attention bias with a
+    diagonal offset set so that the inclusive values are aligned to the upper left corner of the matrix.
+    This equivalent to the `is_causal=True` argument in `scaled_dot_product_attention`.
+
+    The equivalent pytorch code for constructing this bias is:
+
+    .. code-block:: python
+
+        torch.tril(torch.ones(size, dtype=torch.bool))
+
+    For instance, with `shape=(3,4)`, the materialized bias tensor will be:
+
+    .. code-block:: text
+
+        [[1, 0, 0, 0],
+         [1, 1, 0, 0],
+         [1, 1, 1, 0]]
+
+    Args:
+        size: The size of the bias matrix.
+
+    Returns:
+        CausalBias: The UPPER_LEFT triangular causal bias variant.
+    """
+    if len(size) != 2:
+        raise AssertionError("causal_upper_left only supports 2D tensors")
+    seq_len_q, seq_len_kv = size
+    return CausalBias(CausalVariant.UPPER_LEFT, seq_len_q, seq_len_kv)
+
+
+def causal_lower_right(*size) -> CausalBias:
+    """
+    Creates a lower-right triangular causal bias.
+
+    This function generates a lower-right triangular matrix to represent causal attention bias with a
+    diagonal offset set so that the inclusive values are aligned to the lower right corner of the matrix.
+
+    The equivalent pytorch code for constructing this bias is:
+
+    .. code-block:: python
+
+        diagonal_offset = size[1] - size[0]
+        torch.tril(
+            torch.ones(size, dtype=torch.bool),
+            diagonal=diagonal_offset,
+        )
+
+    For instance, with `shape=(3,4)`, the materialized bias tensor will be:
+
+    .. code-block:: text
+
+        [[1, 1, 0, 0],
+         [1, 1, 1, 0],
+         [1, 1, 1, 1]]
+
+    Args:
+        size: The size of the bias matrix.
+
+    Returns:
+        CausalBias: The LOWER_RIGHT triangular causal bias variant.
+    """
+    if len(size) != 2:
+        raise AssertionError("causal_lower_right only supports 2D tensors")
+    seq_len_q, seq_len_kv = size
+    return CausalBias(CausalVariant.LOWER_RIGHT, seq_len_q, seq_len_kv)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a6694bbe3990bacac6025e8c8bd4ab86e80d2e9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/__init__.py
@@ -0,0 +1,2 @@
+# Experimental features are not mature yet and are subject to change.
+# We do not provide any BC/FC guarantees
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/_paged_attention.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/_paged_attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..98e990c81f9238de2ade88ee0077234a4a9acb23
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/_paged_attention.py
@@ -0,0 +1,358 @@
+# mypy: allow-untyped-defs
+"""
+This module implements Paged Attention on top of flex_attention.
+This module is experimental and subject to change.
+"""
+
+import torch
+from torch.nn.attention.flex_attention import (
+    _identity,
+    _mask_mod_signature,
+    _score_mod_signature,
+    BlockMask,
+    noop_mask,
+)
+
+
+__all__ = ["PagedAttention"]
+
+
+def _cdiv(x: int | float | torch.Tensor, multiple: int | float | torch.Tensor):
+    return (x + multiple - 1) // multiple
+
+
+class PagedAttention:
+    """
+    PagedAttention supports flex attention inference with a large batch size.
+    With PagedAttention, a batch of key/value tensors with varying kv length
+    is split into tensor blocks of fixed length and cached in a compact way.
+    Thus we can avoid redundant memory consumption due to varying kv length and
+    support a larger batch size.
+    """
+
+    def __init__(
+        self,
+        n_pages: int,
+        page_size: int,
+        max_batch_size: int,
+        device: str = "cuda",
+    ) -> None:
+        # number of pages
+        self.n_pages = n_pages
+
+        # number of tokens per page
+        self.page_size = page_size
+
+        # page table: [batch, logical_block_idx] -> physical_page_idx
+        self.page_table = -torch.ones(
+            (max_batch_size, self.n_pages), dtype=torch.int64, device=device
+        )
+
+        # capacity: batch_idx -> allocated sequence length
+        self.capacity = torch.zeros(max_batch_size, dtype=torch.int64, device=device)
+
+        # index of empty pages that is available for allocation
+        self.empty_pages = list(range(n_pages - 1, -1, -1))
+
+        # mapping from physical page index to logical page index
+        self.physical_to_logical = -torch.ones(
+            (max_batch_size, n_pages), dtype=torch.int64, device=device
+        )
+
+    def reserve(self, batch_idx: torch.Tensor, seq_len: torch.Tensor) -> None:
+        """
+        Requests the capacity of a given batch to be at least enough to
+        hold `seq_len` elements.
+
+        Args:
+            batch_idx (Tensor): batch index to be reserved; shape :math:`(1)`.
+            seq_len (Tensor): minimum capacity for the given batch; shape :math:`(1)`.
+        """
+
+        if seq_len <= self.capacity[batch_idx]:
+            return
+
+        num_pages_to_allocate = _cdiv(
+            seq_len - self.capacity[batch_idx], self.page_size
+        )
+
+        if len(self.empty_pages) < num_pages_to_allocate:
+            raise AssertionError(
+                f"requested {num_pages_to_allocate.item()} pages "
+                f"but there are only {len(self.empty_pages)} empty pages"
+            )
+
+        start_page_idx = self.capacity[batch_idx] // self.page_size
+        end_page_idx = start_page_idx + num_pages_to_allocate
+
+        # find empty physical pages
+        allocated_pages = torch.tensor(
+            self.empty_pages[-num_pages_to_allocate:],
+            device=num_pages_to_allocate.device,
+        )
+        self.empty_pages = self.empty_pages[:-num_pages_to_allocate]
+
+        # update page table
+        self.page_table[
+            batch_idx,
+            start_page_idx:end_page_idx,
+        ] = allocated_pages
+
+        # update metadata
+        self.physical_to_logical[batch_idx, allocated_pages] = torch.arange(
+            start_page_idx.item(),
+            end_page_idx.item(),
+            device=num_pages_to_allocate.device,
+        )
+        self.capacity[batch_idx] += num_pages_to_allocate * self.page_size
+
+    def erase(self, batch_idx: torch.Tensor) -> None:
+        """
+        Removes a single batch from paged attention.
+
+        Args:
+            batch_idx (Tensor): batch index to be removed; shape :math:`(1)`.
+        """
+
+        # find allocated pages
+        allocated_page_idx = self.page_table[batch_idx] != -1
+        allocated_pages = self.page_table[batch_idx][allocated_page_idx]
+
+        # clean metadata
+        self.capacity[batch_idx] = 0
+        self.empty_pages += allocated_pages.tolist()
+        self.physical_to_logical[batch_idx][:, allocated_pages] = -1
+        self.page_table[batch_idx] = -1
+
+    def assign(
+        self,
+        batch_idx: torch.Tensor,
+        input_pos: torch.Tensor,
+        k_val: torch.Tensor,
+        v_val: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+    ) -> None:
+        """
+        Assigns new contents `val` to the storage `cache` at the location
+        `batch_idx` and `input_pos`.
+
+        Args:
+            batch_idx (Tensor): batch index; shape :math:`(B)`.
+            input_pos (Tensor): input positions to be assigned for the given batch; shape :math:`(B, S)`.
+            val (Tensor): value to be assigned; shape :math:`(B, H, S, D)`
+            cache (Tensor): the cache to store the values; shape:`(1, H, MAX_S, D)`
+        """
+        if k_val.requires_grad:
+            raise RuntimeError("val must not require gradient")
+
+        B, H, S, K_D = k_val.shape
+        V_D = v_val.shape[3]
+        if B != batch_idx.shape[0]:
+            raise RuntimeError(
+                f"Expect val and batch_idx have the same batch size "
+                f"but got B={B} and B={batch_idx.shape[0]}."
+            )
+        if H != k_cache.shape[1]:
+            raise RuntimeError(
+                f"Expect val and cache has the same number of heads "
+                f"but got H={H} and H={k_cache.shape[1]}."
+            )
+        if S != input_pos.shape[1]:
+            raise RuntimeError(
+                f"Expect val and input_pos has the same length "
+                f"but got S={S} and S={input_pos.shape[0]}."
+            )
+        if K_D != k_cache.shape[3]:
+            raise RuntimeError(
+                f"Expect k_val and k_cache has the same hidden dim "
+                f"but got D={K_D} and D={k_cache.shape[3]}."
+            )
+        if V_D != v_cache.shape[3]:
+            raise RuntimeError(
+                f"Expect v_val and v_cache has the same hidden dim "
+                f"but got D={V_D} and D={v_cache.shape[3]}."
+            )
+
+        # find address
+        logical_block_idx = input_pos // self.page_size  # [B, S]
+        logical_block_offset = input_pos % self.page_size  # [B, S]
+        physical_block_idx = torch.gather(
+            self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64)
+        ).to(torch.int32)  # [B, S]
+
+        addr = (physical_block_idx * self.page_size + logical_block_offset).view(
+            -1
+        )  # [B*S]
+
+        k_val = k_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, K_D)
+        v_val = v_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, V_D)
+
+        k_cache[:, :, addr, :] = k_val
+        v_cache[:, :, addr, :] = v_val
+
+    def convert_logical_block_mask(
+        self,
+        block_mask: BlockMask,
+        batch_idx: torch.Tensor | None = None,
+        kv_len: torch.Tensor | None = None,
+    ) -> BlockMask:
+        """
+        Converts a logical block mask by mapping its logical kv indices to the corresponding
+        physical kv indices.
+
+        Args:
+            block_mask (BlockMask): logical block mask;
+                kv_indices shape :math:`(B, H, ROWS, MAX_BLOCKS_IN_COL)`.
+            batch_idx (Tensor): batch index corresponding to the block_mask
+                batch dimension. This provides flexibility to convert a
+                block mask with smaller batch size than the page table;
+                shape :math:`(B)`.
+            kv_len (Optional[Tensor]): actual KV sequence length for upper bound check;
+                shape :math:`(B,)` to handle multiple batches.
+        """
+        B, H, ROWS, MAX_BLOCKS_IN_COL = block_mask.kv_indices.shape
+
+        if block_mask.BLOCK_SIZE[1] != self.page_size:
+            raise RuntimeError(
+                f"Expect block_mask has the same column block size as page_size"
+                f"but got size={block_mask.BLOCK_SIZE[1]} and size={self.page_size}"
+            )
+
+        # Increase the num columns of converted block mask from logical block mask's
+        # num columns to n_pages, since a) the converted block mask
+        # may have larger indices values; and b) `_ordered_to_dense` realizes
+        # a dense tensor with these converted indices. There would be an IndexError
+        # if using the logical block mask's num columns.
+
+        device = block_mask.kv_num_blocks.device
+
+        if batch_idx is None:
+            batch_idx = torch.arange(B, device=device)
+        page_table = self.page_table[batch_idx]
+
+        new_kv_num_blocks = block_mask.kv_num_blocks.clone()
+
+        new_kv_indices = torch.zeros(
+            (B, H, ROWS, self.n_pages), dtype=torch.int32, device=device
+        )
+        new_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
+            torch.gather(
+                page_table, 1, block_mask.kv_indices.view(B, -1).to(torch.int64)
+            )
+            .view(block_mask.kv_indices.shape)
+            .to(torch.int32)
+        )
+
+        new_full_kv_indices, new_full_kv_num_blocks = None, None
+        if block_mask.full_kv_num_blocks is not None:
+            if block_mask.full_kv_indices is None:
+                raise AssertionError(
+                    "block_mask.full_kv_indices must not be None when full_kv_num_blocks is not None"
+                )
+            new_full_kv_num_blocks = block_mask.full_kv_num_blocks.clone()
+            new_full_kv_indices = torch.zeros(
+                (B, H, ROWS, self.n_pages), dtype=torch.int32, device=device
+            )
+            new_full_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
+                torch.gather(
+                    page_table,
+                    1,
+                    block_mask.full_kv_indices.view(B, -1).to(torch.int64),
+                )
+                .view(block_mask.full_kv_indices.shape)
+                .to(torch.int32)
+            )
+
+        new_mask_mod = self.get_mask_mod(block_mask.mask_mod, kv_len)
+
+        seq_lengths = (block_mask.seq_lengths[0], self.n_pages * self.page_size)
+        return BlockMask.from_kv_blocks(
+            new_kv_num_blocks,
+            new_kv_indices,
+            new_full_kv_num_blocks,
+            new_full_kv_indices,
+            block_mask.BLOCK_SIZE,
+            new_mask_mod,
+            seq_lengths=seq_lengths,
+        )
+
+    def get_mask_mod(
+        self,
+        mask_mod: _mask_mod_signature | None,
+        kv_len: torch.Tensor | None = None,
+    ) -> _mask_mod_signature:
+        """
+        Converts a mask_mod based on mapping from the physical block index to the logical
+        block index.
+
+        Args:
+            mask_mod (_mask_mod_signature): mask_mod based on the logical block index.
+            kv_len (Optional[torch.Tensor]): actual KV sequence length for upper bound check.
+        """
+        if mask_mod is None:
+            mask_mod = noop_mask
+
+        def new_mask_mod(
+            b: torch.Tensor,
+            h: torch.Tensor,
+            q_idx: torch.Tensor,
+            physical_kv_idx: torch.Tensor,
+        ):
+            physical_kv_block = physical_kv_idx // self.page_size
+            physical_kv_offset = physical_kv_idx % self.page_size
+            logical_block_idx = self.physical_to_logical[b, physical_kv_block]
+            logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
+            live_block = logical_block_idx >= 0
+            within_upper_bound = (
+                logical_kv_idx < kv_len[b] if kv_len is not None else True
+            )
+            within_lower_bound = logical_kv_idx >= 0
+            is_valid = live_block & within_upper_bound & within_lower_bound
+
+            return torch.where(is_valid, mask_mod(b, h, q_idx, logical_kv_idx), False)
+
+        return new_mask_mod
+
+    def get_score_mod(
+        self,
+        score_mod: _score_mod_signature | None,
+        kv_len: torch.Tensor | None = None,
+    ) -> _score_mod_signature:
+        """
+        Converts a score_mod based on mapping from the physical block index to the logical
+        block index.
+
+        Args:
+            score_mod (_score_mod_signature): score_mod based on the logical block index.
+            `kv_len (Optional[torch.Tensor]): actual KV sequence length for upper bound check.
+
+        """
+        if score_mod is None:
+            score_mod = _identity
+
+        def new_score_mod(
+            score: torch.Tensor,
+            b: torch.Tensor,
+            h: torch.Tensor,
+            q_idx: torch.Tensor,
+            physical_kv_idx: torch.Tensor,
+        ):
+            physical_kv_block = physical_kv_idx // self.page_size
+            physical_kv_offset = physical_kv_idx % self.page_size
+            logical_block_idx = self.physical_to_logical[b, physical_kv_block]
+            logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
+            live_block = logical_block_idx >= 0
+            within_upper_bound = (
+                logical_kv_idx < kv_len[b] if kv_len is not None else True
+            )
+            within_lower_bound = logical_kv_idx >= 0
+            is_valid = live_block & within_upper_bound & within_lower_bound
+
+            return torch.where(
+                is_valid,
+                score_mod(score, b, h, q_idx, logical_kv_idx),
+                float("-inf"),
+            )
+
+        return new_score_mod
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py
new file mode 100644
index 0000000000000000000000000000000000000000..47f9b92670a6d301986ed3df391fa509ee42fbe1
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py
@@ -0,0 +1,154 @@
+# mypy: allow-untyped-defs
+"""
+This operator implements FP8 scaled dot product attention using Flash Attention 3.
+This operator is experimental and subject to change.
+"""
+
+import warnings
+from enum import IntEnum
+
+import torch
+from torch import Tensor
+
+
+class DescaleType(IntEnum):
+    """Describes the scaling granularity for FP8 descale tensors.
+
+    Used with _scaled_dot_product_attention_quantized to explicitly specify
+    how the descale factors are applied to the quantized inputs.
+
+    .. warning::
+        This enum is experimental and subject to change.
+    """
+
+    PER_HEAD = 0
+    """Per-head descaling. Descale tensor shape: (batch_size, num_kv_heads)."""
+
+
+def _validate_descale(
+    descale: Tensor | None,
+    name: str,
+    query: Tensor,
+    key: Tensor,
+    descale_type: DescaleType,
+) -> None:
+    """Validate descale tensor for the specified scaling type.
+
+    Args:
+        descale: The descale tensor to validate (may be None)
+        name: Name of the descale tensor ("q", "k", or "v") for error messages
+        query: Query tensor to get batch size
+        key: Key tensor to get num_kv_heads
+        descale_type: The scaling granularity being used
+
+    Raises:
+        ValueError: If the descale tensor has invalid dtype, device, or shape
+
+    Note:
+        All descale tensors (q, k, v) use num_kv_heads for the head dimension.
+        For GQA/MQA where num_query_heads > num_kv_heads, q_descale is broadcast
+        from (B, H_kv) to match the query heads internally.
+    """
+    if descale is None:
+        return
+
+    # Check dtype
+    if descale.dtype != torch.float32:
+        raise ValueError(f"{name}_descale must have dtype float32, got {descale.dtype}")
+
+    # Check device
+    if not descale.is_cuda:
+        raise ValueError(f"{name}_descale must be a CUDA tensor")
+
+    # Check shape based on descale type
+    if descale_type == DescaleType.PER_HEAD:
+        batch_size = query.size(0)
+        # All descale tensors use num_kv_heads, even q_descale (broadcast internally)
+        # For BHSD layout, num_kv_heads is at dim 1 of key
+        num_kv_heads = key.size(1)
+
+        if descale.dim() != 2:
+            raise ValueError(
+                f"{name}_descale must be a 2D tensor with shape (batch_size, num_kv_heads) "
+                f"for PER_HEAD descaling, got {descale.dim()}D tensor"
+            )
+
+        if descale.size(0) != batch_size:
+            raise ValueError(
+                f"{name}_descale batch dimension must match query batch size, "
+                f"expected {batch_size}, got {descale.size(0)}"
+            )
+
+        if descale.size(1) != num_kv_heads:
+            raise ValueError(
+                f"{name}_descale head dimension must match num_kv_heads, "
+                f"expected {num_kv_heads}, got {descale.size(1)}"
+            )
+
+
+def _scaled_dot_product_attention_quantized(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    is_causal: bool = False,
+    scale: float | None = None,
+    q_descale: Tensor | None = None,
+    k_descale: Tensor | None = None,
+    v_descale: Tensor | None = None,
+    q_descale_type: DescaleType = DescaleType.PER_HEAD,
+    k_descale_type: DescaleType = DescaleType.PER_HEAD,
+    v_descale_type: DescaleType = DescaleType.PER_HEAD,
+) -> Tensor:
+    r"""Scaled dot product attention for FP8 inputs.
+
+    This is a specialized version of scaled_dot_product_attention that supports
+    FP8 quantized inputs (float8_e4m3fn) with per-head descaling. Requires the
+    Flash Attention 3 backend to be activated.
+
+    .. warning::
+        This function is experimental and only supports forward pass.
+
+    Args:
+        query (Tensor): Query tensor; shape :math:`(N, H_q, L, E)` dtype float8_e4m3fn
+        key (Tensor): Key tensor; shape :math:`(N, H, S, E)` dtype float8_e4m3fn
+        value (Tensor): Value tensor; shape :math:`(N, H, S, E_v)` dtype float8_e4m3fn
+        is_causal (bool): Apply causal attention mask
+        scale (float, optional): Scaling factor for attention weights
+        q_descale (Tensor, optional): Query descale tensor; shape :math:`(N, H)` for PER_HEAD
+        k_descale (Tensor, optional): Key descale tensor; shape :math:`(N, H)` for PER_HEAD
+        v_descale (Tensor, optional): Value descale tensor; shape :math:`(N, H)` for PER_HEAD
+        q_descale_type (DescaleType): Specifies the descaling granularity for query. Default: PER_HEAD
+        k_descale_type (DescaleType): Specifies the descaling granularity for key. Default: PER_HEAD
+        v_descale_type (DescaleType): Specifies the descaling granularity for value. Default: PER_HEAD
+
+    Returns:
+        Tensor: Attention output; shape :math:`(N, H_q, L, E_v)` dtype bfloat16
+    """
+    # Validate descale tensors
+    _validate_descale(q_descale, "q", query, key, q_descale_type)
+    _validate_descale(k_descale, "k", query, key, k_descale_type)
+    _validate_descale(v_descale, "v", query, key, v_descale_type)
+
+    if torch.is_grad_enabled() and (
+        query.requires_grad or key.requires_grad or value.requires_grad
+    ):
+        warnings.warn(
+            "_scaled_dot_product_attention_quantized does not support backward pass. "
+            "Gradients will not be computed for query, key, or value.",
+            UserWarning,
+        )
+    # Directly call the internal flash attention operator which has descale support
+    # NOTE: This should be torch._scaled_dot_product_flash_attention, but it does not work with torch.compile
+    result = torch.ops.aten._scaled_dot_product_flash_attention.quantized(
+        query,
+        key,
+        value,
+        q_descale,
+        k_descale,
+        v_descale,
+        0.0,
+        is_causal,
+        False,
+        scale=scale,
+    )
+    return result[0]  # Return the output tensor, mirroring scaled_dot_product_attention
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..5cf2c42e2b93c7c226f541a074b4e4d69140346c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py
@@ -0,0 +1,2093 @@
+# mypy: allow-untyped-defs
+# flake8: noqa: B950
+"""This module implements the user facing API for flex_attention in PyTorch."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import itertools
+import math
+import operator
+import types
+import typing
+import warnings
+from collections.abc import Callable
+from enum import Enum
+from typing import Any, cast, Literal, NamedTuple, overload, TypeAlias, TypeVar
+from typing_extensions import deprecated, Never, NotRequired, Self, TypedDict
+
+import torch
+from torch import Tensor
+from torch._higher_order_ops.flex_attention import flex_attention as flex_attention_hop
+from torch._higher_order_ops.utils import setup_compilation_env
+from torch.nn.attention._utils import _validate_sdpa_input
+from torch.utils._pytree import (
+    GetAttrKey,
+    tree_flatten,
+    tree_map_only,
+    tree_unflatten,
+    TreeSpec,
+)
+
+
+if typing.TYPE_CHECKING:
+    from torch._prims_common import DeviceLikeType
+    from torch.fx.node import BaseArgumentTypes
+
+
+# Private debug flag to disable internal compilation wrapping for debugging purposes.
+# WARNING: This is intended ONLY for debugging score_mod and mask_mod functions.
+# When enabled, this bypasses the required internal compilation that ensures correctness
+# and performance. Only use this temporarily when you need to set breakpoints
+# in your score_mod/mask_mod functions during development.
+#
+# This flag only affects the internal compilation when flex_attention is called directly.
+# If you have already wrapped flex_attention in torch.compile(), this flag has no effect
+# and the user's compilation will still occur.
+#
+# Usage:
+#   import torch.nn.attention.flex_attention as fa
+#   fa._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+#   # Now you can set breakpoints in your score_mod/mask_mod
+#   output = fa.flex_attention(q, k, v, score_mod=my_score_mod)
+#
+_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = False
+
+_WARNINGS_SHOWN: set[str] = set()
+
+
+def _warn_once(
+    warning_id: str, message: str, category: type[Warning] = UserWarning
+) -> None:
+    """Helper to ensure each warning is shown only once per process."""
+    if warning_id not in _WARNINGS_SHOWN:
+        if not torch.compiler.is_compiling():
+            warnings.warn(message, category, stacklevel=2)
+        _WARNINGS_SHOWN.add(warning_id)
+
+
+__all__ = [
+    "BlockMask",
+    "flex_attention",
+    "AuxOutput",
+    "AuxRequest",
+    "FlexKernelOptions",
+    "create_block_mask",
+    "create_mask",
+    "or_masks",
+    "and_masks",
+    "noop_mask",
+]
+
+_score_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor, Tensor], Tensor]
+_mask_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]
+_Backend: TypeAlias = Literal["AUTO", "TRITON", "FLASH", "TRITON_DECODE"]
+_R = TypeVar("_R")
+
+
+class FlexKernelOptions(TypedDict, total=False):
+    """Options for controlling the behavior of FlexAttention kernels.
+
+    These options are passed to the underlying Triton kernels to control performance
+    and numerical behavior. Most users will not need to specify these options as the
+    default autotuning provides good performance.
+
+    The options can be prefixed with ``fwd_`` or ``bwd_`` to apply only to forward or
+    backward pass respectively. For example: ``fwd_BLOCK_M`` and ``bwd_BLOCK_M1``.
+
+    Note:
+      We currently do not provide any backward compatibility guarantees for these options.
+      That being said most of these have remained pretty stable since their introduction. But
+      We do not consider this part of the public API just yet. We think that some documentation
+      Is better than secret hidden flags, but we may change these options in the future.
+
+    Example Usage:
+        .. code-block:: python
+
+            # Using dictionary (backward compatible)
+            kernel_opts = {"BLOCK_M": 64, "BLOCK_N": 64, "PRESCALE_QK": True}
+            output = flex_attention(q, k, v, kernel_options=kernel_opts)
+
+            # Using TypedDict (recommended for type safety)
+            from torch.nn.attention.flex_attention import FlexKernelOptions
+
+            kernel_opts: FlexKernelOptions = {
+                "BLOCK_M": 64,
+                "BLOCK_N": 64,
+                "PRESCALE_QK": True,
+            }
+            output = flex_attention(q, k, v, kernel_options=kernel_opts)
+
+            # Forward/backward specific options
+            kernel_opts: FlexKernelOptions = {
+                "fwd_BLOCK_M": 64,
+                "bwd_BLOCK_M1": 32,
+                "PRESCALE_QK": False,
+            }
+            output = flex_attention(q, k, v, kernel_options=kernel_opts)
+    """
+
+    # Performance tuning options
+
+    num_warps: NotRequired[int]
+    """Number of warps to use in the CUDA kernel. Higher values may improve performance
+    but increase register pressure. Default is determined by autotuning."""
+
+    num_stages: NotRequired[int]
+    """Number of pipeline stages in the CUDA kernel. Higher values may improve performance
+    but increase shared memory usage. Default is determined by autotuning."""
+
+    BLOCK_M: NotRequired[int]
+    """Thread block size for the sequence length dimension of Q in forward pass.
+    Must be a power of 2. Common values: 16, 32, 64, 128. Default is determined by autotuning."""
+
+    BLOCK_N: NotRequired[int]
+    """Thread block size for the sequence length dimension of K/V in forward pass.
+    Must be a power of 2. Common values: 16, 32, 64, 128. Default is determined by autotuning."""
+
+    # Backward-specific block sizes (when prefixed with 'bwd_')
+
+    BLOCK_M1: NotRequired[int]
+    """Thread block size for Q dimension in backward pass. Use as 'bwd_BLOCK_M1'.
+    Default is determined by autotuning."""
+
+    BLOCK_N1: NotRequired[int]
+    """Thread block size for K/V dimension in backward pass. Use as 'bwd_BLOCK_N1'.
+    Default is determined by autotuning."""
+
+    BLOCK_M2: NotRequired[int]
+    """Thread block size for second Q dimension in backward pass. Use as 'bwd_BLOCK_M2'.
+    Default is determined by autotuning."""
+
+    BLOCK_N2: NotRequired[int]
+    """Thread block size for second K/V dimension in backward pass. Use as 'bwd_BLOCK_N2'.
+    Default is determined by autotuning."""
+
+    PRESCALE_QK: NotRequired[bool]
+    """Whether to pre-scale QK by 1/sqrt(d) and change of base. This is slightly faster but
+    may have more numerical error. Default: False."""
+
+    ROWS_GUARANTEED_SAFE: NotRequired[bool]
+    """If True, guarantees that at least one value in each row is not masked out.
+    Allows skipping safety checks for better performance. Only set this if you are certain
+    your mask guarantees this property. For example, causal attention is guaranteed safe
+    because each query has at least 1 key-value to attend to. Default: False."""
+
+    BLOCKS_ARE_CONTIGUOUS: NotRequired[bool]
+    """If True, guarantees that all blocks in the mask are contiguous.
+    Allows optimizing block traversal. For example, causal masks would satisfy this,
+    but prefix_lm + sliding window would not. Default: False."""
+
+    WRITE_DQ: NotRequired[bool]
+    """Controls whether gradient scatters are done in the DQ iteration loop of the backward pass.
+    Setting this to False will force this to happen in the DK loop which depending on your
+    specific score_mod and mask_mod might be faster. Default: True."""
+
+    FORCE_USE_FLEX_ATTENTION: NotRequired[bool]
+    """If True, forces the use of the flex attention kernel instead of potentially using
+    the more optimized flex-decoding kernel for short sequences. This can be a helpful
+    option for debugging. Default: False."""
+
+    USE_TMA: NotRequired[bool]
+    """Whether to use Tensor Memory Accelerator (TMA) on supported hardware.
+    This is experimental and may not work on all hardware, currently specific
+    to NVIDIA GPUs Hopper+. Default: False."""
+
+    # ROCm-specific options
+
+    kpack: NotRequired[int]
+    """ROCm-specific kernel packing parameter."""
+
+    matrix_instr_nonkdim: NotRequired[int]
+    """ROCm-specific matrix instruction non-K dimension."""
+
+    waves_per_eu: NotRequired[int]
+    """ROCm-specific waves per execution unit."""
+
+    BACKEND: NotRequired[_Backend]
+    """Selects a specific kernel backend.
+
+    Options:
+        - "AUTO": Use current heuristics (typically Triton-based kernels with
+          automatic selection between flex_attention and flex_decoding)
+        - "TRITON": Standard Triton flex_attention kernel
+        - "TRITON_DECODE": Triton flex_decoding kernel, only available for short sequence lengths with specific configurations
+        - "FLASH": Experimental: Flash Attention kernel (cute-dsl), user needs to have flash installed
+
+    This option cannot be combined with legacy knobs such as ``FORCE_USE_FLEX_ATTENTION``.
+    Raises an error if the requested backend cannot be used. Default: "AUTO"
+    """
+
+
+class _KernelOptionsWithInternals(FlexKernelOptions, total=False):
+    OUTPUT_LOGSUMEXP: bool
+    OUTPUT_MAX: bool
+
+
+class AuxRequest(NamedTuple):
+    """Request which auxiliary outputs to compute from flex_attention.
+
+    Each field is a boolean indicating whether that auxiliary output should be computed.
+    """
+
+    lse: bool = False
+    max_scores: bool = False
+
+
+class AuxOutput(NamedTuple):
+    """Auxiliary outputs from flex_attention operation.
+
+    Fields will be None if not requested, or contain the tensor if requested.
+    """
+
+    lse: Tensor | None = None
+    max_scores: Tensor | None = None
+
+
+class _ModificationType(Enum):
+    """Enum for the type of modification function.
+    - SCORE_MOD: score_mod function which accepts a score as the first argument
+    - mask_mod: mask function which does not accept a score and is only used for generating
+    block mask
+    """
+
+    SCORE_MOD = 1
+    MASK_MOD = 2
+    UNKNOWN = 3
+
+
+def _get_mod_type(
+    fn: _score_mod_signature | _mask_mod_signature | Callable[..., Any],
+) -> _ModificationType:
+    """Get the type of modification function.
+    This function inspects the number of positional arguments of the function to determine
+    the type of modification function. If the function has 5 positional arguments, it is
+    considered as a score_mod function. If the function has 4 positional arguments, it is
+    considered as a mask function.
+    """
+    if hasattr(fn, "__code__"):
+        code = fn.__code__
+        num_positional_total = code.co_argcount
+        defaults = ()
+        if hasattr(fn, "__defaults__"):
+            defaults = fn.__defaults__ or ()
+        num_defaults = len(defaults)
+        num_positional_args = num_positional_total - num_defaults
+    else:
+        num_positional_args = sum(
+            1
+            for param in inspect.signature(fn).parameters.values()
+            if param.default is inspect.Parameter.empty
+        )
+    if num_positional_args != 5 and num_positional_args != 4:
+        raise AssertionError(
+            f"Expected 4 or 5 positional args, got {num_positional_args}"
+        )
+    if num_positional_args == 5:
+        return _ModificationType.SCORE_MOD
+    elif num_positional_args == 4:
+        return _ModificationType.MASK_MOD
+    else:
+        return _ModificationType.UNKNOWN
+
+
+# Need to define it here so that Dynamo doesn't skip it
+def _vmap_for_bhqkv(
+    fn: Callable[..., _R],
+    prefix: tuple[int | None, ...],
+    suffix: tuple[int | None, ...] = (),
+    out_dims: int | list[int | None] = 0,
+    group_dim: bool = False,
+) -> Callable[..., _R]:
+    """Used to vmap both score_mods and mask_mods over 4-dimensional/5-dimension inputs.
+    Mapping over the [b, hq, q_idx, kv_idx] or [b, hkv, g, q_idx, kv_idx] dimensions.
+
+    Args:
+        fn (callable): The function to vmap.
+        prefix (tuple): The prefix of the vmap. For score mod functions,
+                        this should be set to (0,). For mask_mods = ()
+        suffix (tuple): We need to add (0,) if gradOut is being mapped over,
+                        and (None,) * len(other_buffers).
+        out_dims (tuple): For forward cases, keep this as the default 0 since
+                          we are only returning 1 output. For backwards, the joint
+                          graph returns grads for B, H, Q_idx, KV_idx and other_buffers,
+                          so we set this to (0, None, None, None, None) + (None,) * len(other_buffers).
+
+    Returns:
+        callable: The vmapped function.
+    """
+    # We vamp a function 4 times, broadcasting the [b, h, q_idx, kv_idx] dimensions
+    dimensions: list[tuple[int | None, int | None, int | None, int | None]] = []
+    dimensions = [
+        (None, None, None, 0),
+        (None, None, 0, None),
+        (None, 0, None, None),
+    ]
+
+    if group_dim:
+        dimensions += [
+            (None, 0, None, None),
+        ]
+
+    dimensions += [
+        (0, None, None, None),
+    ]
+
+    for dims in dimensions:
+        fn = torch.vmap(fn, in_dims=prefix + dims + suffix, out_dims=out_dims)  # type: ignore[arg-type]
+    return fn
+
+
+def _identity(
+    score: Tensor,
+    batch: Tensor,
+    head: Tensor,
+    token_q: Tensor,
+    token_kv: Tensor,
+) -> Tensor:
+    return score
+
+
+def noop_mask(
+    batch: Tensor,
+    head: Tensor,
+    token_q: Tensor,
+    token_kv: Tensor,
+) -> Tensor:
+    """Returns a noop mask_mod"""
+    return batch.new_ones(size=(), dtype=torch.bool, device=batch.device)
+
+
+def _sliced_mask_mod_error(
+    batch: Tensor,
+    head: Tensor,
+    token_q: Tensor,
+    token_kv: Tensor,
+) -> Never:
+    """
+    Raises helpful error when using mask_mod from a sliced BlockMask.
+
+    After slicing a BlockMask, the mask_mod is reset and cannot be used directly.
+    Users must reassign mask_mod from the original (unsliced) BlockMask.
+    """
+    raise RuntimeError(
+        "Cannot use mask_mod from a sliced BlockMask. "
+        "When you slice a BlockMask using [], the mask_mod attribute is reset. "
+        "You must set it from the original BlockMask's mask_mod."
+        "\n\nIncorrect usage:"
+        "\n  base_mask = create_block_mask(my_mask_fn, ...)"
+        "\n  sliced_mask = base_mask[:, :, block_idx]"
+        "\n  sliced_mask.mask_mod = apply_offset(sliced_mask.mask_mod, offset)  # WRONG!"
+        "\n\nCorrect usage:"
+        "\n  base_mask = create_block_mask(my_mask_fn, ...)"
+        "\n  sliced_mask = base_mask[:, :, block_idx]"
+        "\n  sliced_mask.mask_mod = apply_offset(base_mask.mask_mod, offset)  # Use base_mask!"
+    )
+
+
+_DEFAULT_SPARSE_BLOCK_SIZE = 128
+_LARGE_SPARSE_BLOCK_SIZE = 1 << 30
+
+
+def _ordered_to_dense(num_blocks_in_row: Tensor, col_indices: Tensor) -> Tensor:
+    num_rows = col_indices.shape[-2]
+    num_cols = col_indices.shape[-1]
+    batch_dims = num_blocks_in_row.shape[:-1]
+    device = num_blocks_in_row.device
+
+    def create_dense_one(kv_num_blocks, kv_indices):
+        dense_mask = kv_indices.new_zeros(num_rows, num_cols + 1, dtype=torch.int32)
+
+        row_indices = torch.arange(num_rows, dtype=torch.int, device=device).unsqueeze(
+            -1
+        )
+        col_range = torch.arange(num_cols, dtype=torch.int, device=device)
+        index_mask = col_range < kv_num_blocks.unsqueeze(-1)
+
+        # We write to one spot "out of bounds"
+        valid_indices = torch.where(index_mask, kv_indices, num_cols)
+
+        # set the values in 'a' to 1 where the indices are valid
+        dense_mask[row_indices, valid_indices] = dense_mask.new_ones(())
+        return dense_mask[:, :num_cols].contiguous()
+
+    create_dense_batched = create_dense_one
+    for _ in range(len(batch_dims)):
+        create_dense_batched = torch.vmap(create_dense_batched, in_dims=(0, 0))
+
+    out = create_dense_batched(num_blocks_in_row, col_indices)
+    return out
+
+
+def _dense_to_ordered(dense_mask: Tensor) -> tuple[Tensor, Tensor]:
+    dense_mask = dense_mask.to(dtype=torch.int32)
+    num_blocks_in_row = dense_mask.sum(dim=-1)
+    col_indices = torch.argsort(dense_mask, dim=-1, descending=True, stable=True)
+    return (
+        num_blocks_in_row.to(torch.int32, memory_format=torch.contiguous_format),
+        col_indices.to(torch.int32, memory_format=torch.contiguous_format),
+    )
+
+
+def _transpose_ordered(
+    num_blocks_in_row: Tensor, col_indices: Tensor
+) -> tuple[Tensor, Tensor]:
+    dense = _ordered_to_dense(num_blocks_in_row, col_indices)
+    return _dense_to_ordered(dense.transpose(-2, -1))
+
+
+def _adjust_num_blocks_and_indices(
+    num_blocks: Tensor,
+    indices: Tensor,
+    new_num_rows: int,
+    new_num_cols: int,
+) -> tuple[Tensor, Tensor]:
+    indices = indices[:, :, :new_num_rows, :new_num_cols]
+    num_blocks = num_blocks[:, :, :new_num_rows]
+    num_blocks = torch.where(num_blocks < new_num_cols, num_blocks, new_num_cols)
+    num_blocks = torch.sum(indices < num_blocks[:, :, :, None], dim=-1).to(torch.int32)
+    return num_blocks, indices
+
+
+# TreeSpec for an empty tuple — used as the sentinel when there are no closure leaves.
+_EMPTY_CLOSURE_SPEC = tree_flatten(())[1]
+
+
+class _ExtractedLeaf:
+    """Sentinel in _StrippedClosure.leaf_entries marking a position that is
+    filled from the extracted pytree leaves list during reconstruction."""
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "_EXTRACTED_LEAF"
+
+
+_EXTRACTED_LEAF = _ExtractedLeaf()
+
+
+class _FunctionLeaf(typing.NamedTuple):
+    """Entry in _StrippedClosure.leaf_entries for a recursively processed
+    function.  Stores enough information to reconstruct the function from
+    the extracted leaves during unflattening."""
+
+    stripped: _StrippedClosure | Callable[..., Any]
+    closure_spec: TreeSpec
+    n_extracted: int  # number of extracted leaves this function contributes
+
+
+class _StrippedClosure(typing.NamedTuple):
+    """Data container holding the parts of a function needed for reconstruction.
+
+    Created by _extract_closure_pytree when closure tensors are lifted into
+    pytree leaves.  Unlike a FunctionType with None-filled cells, this is not
+    callable — it is pure data stored in the pytree context.
+    """
+
+    code: types.CodeType
+    globals_dict: dict[str, Any]
+    name: str
+    qualname: str
+    defaults: tuple[Any, ...] | None
+    kwdefaults: dict[str, Any] | None
+    extra_dict: dict[str, Any]
+    # Per-position info for the closure's flattened leaves.
+    # _EXTRACTED_LEAF → position filled from the extracted leaves list.
+    # _FunctionLeaf   → recursively processed function (not a valid pytree leaf).
+    leaf_entries: tuple[_ExtractedLeaf | _FunctionLeaf, ...]
+
+
+def _extract_closure_pytree(
+    fn, _seen: set[int] | None = None
+) -> tuple[
+    tuple[BaseArgumentTypes, ...], TreeSpec, _StrippedClosure | Callable[..., Any]
+]:
+    """Extract closure contents as a flattened sub-pytree.
+
+    Returns (extracted_leaves, closure_spec, fn_or_stripped) where:
+    - extracted_leaves: flattened non-function contents from the closure,
+      plus any tensors/scalars recursively extracted from nested function
+      closures
+    - closure_spec: TreeSpec describing how to reconstruct the closure contents
+    - fn_or_stripped: either the original fn (no extraction) or a
+      _StrippedClosure carrying the function parts needed for reconstruction
+
+    Functions found among the closure leaves are recursively processed: their
+    own closure tensors are extracted into the leaves list, and their skeleton
+    is stored in _StrippedClosure.leaf_entries as a _FunctionLeaf.  All other
+    values (tensors, scalars, None, etc.) remain as extracted leaves.
+
+    If fn is not a plain function, has no closure, or has empty cells, returns
+    the original function unchanged with no closure leaves.
+
+    Skipped under Dynamo tracing (torch.compiler.is_compiling) because Dynamo
+    can't trace through closure cell introspection and handles freevars via its
+    own lifting mechanism.
+    """
+    if not inspect.isfunction(fn) or torch.compiler.is_compiling():
+        return (), _EMPTY_CLOSURE_SPEC, fn
+
+    # Cycle detection for self-referencing closures.
+    if _seen is None:
+        _seen = set()
+    if id(fn) in _seen:
+        return (), _EMPTY_CLOSURE_SPEC, fn
+    _seen.add(id(fn))
+
+    closure = fn.__closure__
+    if not closure:
+        return (), _EMPTY_CLOSURE_SPEC, fn
+
+    try:
+        contents = tuple(cell.cell_contents for cell in closure)
+    except ValueError:
+        # Empty cell (created but not yet assigned) — can't extract
+        return (), _EMPTY_CLOSURE_SPEC, fn
+
+    closure_leaves, closure_spec = tree_flatten(contents)
+
+    extracted: list[BaseArgumentTypes] = []
+    leaf_entries: list[_ExtractedLeaf | _FunctionLeaf] = []
+    for leaf in closure_leaves:
+        if inspect.isfunction(leaf):
+            child_extracted, child_spec, child_stripped = _extract_closure_pytree(
+                leaf, _seen
+            )
+            extracted.extend(child_extracted)
+            leaf_entries.append(
+                _FunctionLeaf(child_stripped, child_spec, len(child_extracted))
+            )
+        else:
+            extracted.append(leaf)
+            leaf_entries.append(_EXTRACTED_LEAF)
+
+    stripped = _StrippedClosure(
+        code=fn.__code__,
+        globals_dict=fn.__globals__,
+        name=fn.__name__,
+        qualname=fn.__qualname__,
+        defaults=fn.__defaults__,
+        kwdefaults=fn.__kwdefaults__,
+        extra_dict=dict(fn.__dict__) if fn.__dict__ else {},
+        leaf_entries=tuple(leaf_entries),
+    )
+
+    return tuple(extracted), closure_spec, stripped
+
+
+def _reconstruct_closure_fn(stripped, extracted_leaves, closure_spec):
+    """Rebuild a function from a _StrippedClosure and flattened extracted leaves."""
+    if not isinstance(stripped, _StrippedClosure):
+        return stripped
+
+    all_leaves: list[BaseArgumentTypes | Callable[..., Any]] = []
+    idx = 0
+    for entry in stripped.leaf_entries:
+        if isinstance(entry, _FunctionLeaf):
+            child_fn = _reconstruct_closure_fn(
+                entry.stripped,
+                extracted_leaves[idx : idx + entry.n_extracted],
+                entry.closure_spec,
+            )
+            all_leaves.append(child_fn)
+            idx += entry.n_extracted
+        else:
+            # _EXTRACTED_LEAF — take from extracted leaves
+            all_leaves.append(extracted_leaves[idx])
+            idx += 1
+
+    contents = tree_unflatten(all_leaves, closure_spec)
+    new_cells = tuple(types.CellType(v) for v in contents)
+
+    restored = types.FunctionType(
+        stripped.code,
+        stripped.globals_dict,
+        stripped.name,
+        stripped.defaults,
+        new_cells,
+    )
+    restored.__qualname__ = stripped.qualname
+    if stripped.kwdefaults:
+        restored.__kwdefaults__ = stripped.kwdefaults
+    if stripped.extra_dict:
+        restored.__dict__.update(stripped.extra_dict)
+
+    return restored
+
+
+class _MaskModWrapper:
+    """Wraps a mask_mod or _StrippedClosure with value-based equality.
+
+    BlockMask stores an arbitrary callable (mask_mod) in its pytree context.
+    The default __eq__ for functions uses identity comparison, which is too
+    strict when the same closure is recreated (e.g., defined inside forward()).
+
+    When closure tensors have been extracted (by _extract_closure_pytree), fn
+    is a _StrippedClosure (pure data, not callable).  Equality compares the
+    code objects + closure_spec — no tensor dispatch is triggered.
+
+    When extraction is skipped (e.g., under Dynamo), fn is the original
+    callable and equality compares code objects + closure contents (for plain
+    functions) or delegates to __eq__ (for callable objects).
+    """
+
+    __slots__ = ("fn", "closure_spec")
+
+    def __init__(self, fn, closure_spec=None) -> None:
+        self.fn = fn
+        self.closure_spec = closure_spec
+
+    def __call__(self, b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
+        if isinstance(self.fn, _StrippedClosure):
+            raise RuntimeError(
+                "_MaskModWrapper with _StrippedClosure is not callable — "
+                "use _reconstruct_closure_fn to rebuild the function first"
+            )
+        return self.fn(b, h, q_idx, kv_idx)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, _MaskModWrapper):
+            return False
+        if self.fn is other.fn and self.closure_spec is other.closure_spec:
+            return True
+        # Extracted case: _StrippedClosure — compare code + closure_spec
+        if isinstance(self.fn, _StrippedClosure) and isinstance(
+            other.fn, _StrippedClosure
+        ):
+            return (
+                self.fn.code == other.fn.code
+                and self.closure_spec == other.closure_spec
+            )
+        # Non-extracted plain functions: compare code + closure contents
+        if inspect.isfunction(self.fn) and inspect.isfunction(other.fn):
+            return self.fn.__code__ == other.fn.__code__
+        # Callable objects: delegate to their __eq__
+        if not isinstance(self.fn, _StrippedClosure) and not isinstance(
+            other.fn, _StrippedClosure
+        ):
+            return self.fn == other.fn
+        return False
+
+    def __hash__(self) -> int:
+        if isinstance(self.fn, _StrippedClosure):
+            return hash(self.fn.code)
+        if inspect.isfunction(self.fn):
+            return hash(self.fn.__code__)
+        return hash(self.fn)
+
+    def __repr__(self) -> str:
+        return f"_MaskModWrapper({self.fn})"
+
+
+class BlockMask:
+    r"""
+    BlockMask is our format for representing a block-sparse attention mask.
+    It is somewhat of a cross in-between BCSR and a non-sparse format.
+
+    **Basics**
+
+    A block-sparse mask means that instead of representing the sparsity of
+    individual elements in the mask, a KV_BLOCK_SIZE x Q_BLOCK_SIZE block is
+    considered sparse only if every element within that block is sparse.
+    This aligns well with hardware, which generally expects to perform
+    contiguous loads and computation.
+
+    This format is primarily optimized for 1. simplicity, and 2. kernel
+    efficiency. Notably, it is *not* optimized for size, as this mask is always
+    reduced by a factor of KV_BLOCK_SIZE * Q_BLOCK_SIZE. If the size is a
+    concern, the tensors can be reduced in size by increasing the block size.
+
+    The essentials of our format are:
+
+    num_blocks_in_row: Tensor[ROWS]:
+    Describes the number of blocks present in each row.
+
+    col_indices: Tensor[ROWS, MAX_BLOCKS_IN_COL]:
+    `col_indices[i]` is the sequence of block positions for row i. The values of
+    this row after `col_indices[i][num_blocks_in_row[i]]` are undefined.
+
+    For example, to reconstruct the original tensor from this format:
+
+    .. code-block:: python
+
+        dense_mask = torch.zeros(ROWS, COLS)
+        for row in range(ROWS):
+            for block_idx in range(num_blocks_in_row[row]):
+                dense_mask[row, col_indices[row, block_idx]] = 1
+
+    Notably, this format makes it easier to implement a reduction along the
+    *rows* of the mask.
+
+    **Details**
+
+    The basics of our format require only kv_num_blocks and kv_indices. But, we
+    have up to 8 tensors on this object. This represents 4 pairs:
+
+    1. (kv_num_blocks, kv_indices): Used for the forwards pass of attention, as
+    we reduce along the KV dimension.
+
+    2. [OPTIONAL] (full_kv_num_blocks, full_kv_indices): This is optional and
+    purely an optimization. As it turns out, applying masking to every block
+    is quite expensive! If we specifically know which blocks are "full" and
+    don't require masking at all, then we can skip applying mask_mod to these
+    blocks. This requires the user to split out a separate mask_mod from the
+    score_mod. For causal masks, this is about a 15% speedup.
+
+    3. [GENERATED] (q_num_blocks, q_indices): Required for the backwards pass,
+    as computing dKV requires iterating along the mask along the Q dimension. These are autogenerated from 1.
+
+    4. [GENERATED] (full_q_num_blocks, full_q_indices): Same as above, but for
+    the backwards pass. These are autogenerated from 2.
+    """
+
+    seq_lengths: tuple[int, int]
+    kv_num_blocks: Tensor
+    kv_indices: Tensor
+    full_kv_num_blocks: Tensor | None
+    full_kv_indices: Tensor | None
+    q_num_blocks: Tensor | None
+    q_indices: Tensor | None
+    full_q_num_blocks: Tensor | None
+    full_q_indices: Tensor | None
+    BLOCK_SIZE: tuple[int, int]
+    mask_mod: _mask_mod_signature
+
+    # Attribute lists for pytree flatten/unflatten
+    _TENSOR_ATTRS = [
+        "kv_num_blocks",
+        "kv_indices",
+        "full_kv_num_blocks",
+        "full_kv_indices",
+        "q_num_blocks",
+        "q_indices",
+        "full_q_num_blocks",
+        "full_q_indices",
+    ]
+
+    _CONTEXT_ATTRS = [
+        "seq_lengths",
+        "BLOCK_SIZE",
+        "mask_mod",
+    ]
+
+    def __init__(
+        self,
+        seq_lengths: tuple[int, int],
+        kv_num_blocks: Tensor,
+        kv_indices: Tensor,
+        full_kv_num_blocks: Tensor | None,
+        full_kv_indices: Tensor | None,
+        q_num_blocks: Tensor | None,
+        q_indices: Tensor | None,
+        full_q_num_blocks: Tensor | None,
+        full_q_indices: Tensor | None,
+        BLOCK_SIZE: tuple[int, int],
+        mask_mod: _mask_mod_signature,
+    ) -> None:
+        if kv_indices.dim() < 2:
+            raise RuntimeError("BlockMask must have at least 2 dimensions")
+        if kv_num_blocks is None:
+            raise AssertionError("kv_num_blocks must be provided")
+        if kv_indices is None:
+            raise AssertionError("kv_indices must be provided")
+        if (full_kv_num_blocks is None) != (full_kv_indices is None):
+            raise AssertionError(
+                "full_kv_num_blocks and full_kv_indices must be both provided or omitted"
+            )
+        if (full_q_num_blocks is None) != (full_q_indices is None):
+            raise AssertionError(
+                "full_q_num_blocks and full_q_indices must be both provided or omitted"
+            )
+
+        self.seq_lengths = seq_lengths
+        self.kv_num_blocks = kv_num_blocks
+        self.kv_indices = kv_indices
+        self.full_kv_num_blocks = full_kv_num_blocks
+        self.full_kv_indices = full_kv_indices
+        self.q_num_blocks = q_num_blocks
+        self.q_indices = q_indices
+        self.full_q_num_blocks = full_q_num_blocks
+        self.full_q_indices = full_q_indices
+        self.BLOCK_SIZE = BLOCK_SIZE
+        self.mask_mod = mask_mod
+
+    @classmethod
+    def from_kv_blocks(
+        cls,
+        kv_num_blocks: Tensor,
+        kv_indices: Tensor,
+        full_kv_num_blocks: Tensor | None = None,
+        full_kv_indices: Tensor | None = None,
+        BLOCK_SIZE: int | tuple[int, int] = _DEFAULT_SPARSE_BLOCK_SIZE,
+        mask_mod: _mask_mod_signature | None = None,
+        seq_lengths: tuple[int, int] | None = None,
+        compute_q_blocks: bool = True,
+    ) -> Self:
+        """
+        Creates a BlockMask instance from key-value block information.
+
+        Args:
+            kv_num_blocks (Tensor): Number of kv_blocks in each Q_BLOCK_SIZE row tile.
+            kv_indices (Tensor): Indices of key-value blocks in each Q_BLOCK_SIZE row tile.
+            full_kv_num_blocks (Optional[Tensor]): Number of full kv_blocks in each Q_BLOCK_SIZE row tile.
+            full_kv_indices (Optional[Tensor]): Indices of full key-value blocks in each Q_BLOCK_SIZE row tile.
+            BLOCK_SIZE (Union[int, tuple[int, int]]): Size of KV_BLOCK_SIZE x Q_BLOCK_SIZE tiles.
+            mask_mod (Optional[Callable]): Function to modify the mask.
+
+        Returns:
+            BlockMask: Instance with full Q information generated via _transposed_ordered
+
+        Raises:
+            RuntimeError: If kv_indices has < 2 dimensions.
+            AssertionError: If only one of full_kv_* args is provided.
+        """
+        if kv_indices.dim() < 2:
+            raise RuntimeError("BlockMask must have at least 2 dimensions")
+
+        if (full_kv_num_blocks is None) != (full_kv_indices is None):
+            raise AssertionError(
+                "full_kv_num_blocks and full_kv_indices must be both provided or omitted"
+            )
+
+        # Generate q_num_blocks and q_indices
+        if compute_q_blocks:
+            q_num_blocks, q_indices = _transpose_ordered(kv_num_blocks, kv_indices)
+            if full_kv_num_blocks is not None:
+                if full_kv_indices is None:
+                    raise AssertionError("full_kv_indices must not be None")
+                full_q_num_blocks, full_q_indices = _transpose_ordered(
+                    full_kv_num_blocks, full_kv_indices
+                )
+            else:
+                full_q_num_blocks, full_q_indices = None, None
+        else:
+            q_num_blocks, q_indices = None, None
+            full_q_num_blocks, full_q_indices = None, None
+
+        if isinstance(BLOCK_SIZE, int):
+            BLOCK_SIZE = (BLOCK_SIZE, BLOCK_SIZE)
+
+        mask_mod = mask_mod if mask_mod is not None else noop_mask
+        if seq_lengths is None:
+            q_length = kv_indices.shape[-2] * BLOCK_SIZE[0]
+            kv_length = kv_indices.shape[-1] * BLOCK_SIZE[1]
+            seq_lengths = (q_length, kv_length)
+
+        return cls(
+            seq_lengths=seq_lengths,
+            kv_num_blocks=kv_num_blocks,
+            kv_indices=kv_indices,
+            full_kv_num_blocks=full_kv_num_blocks,
+            full_kv_indices=full_kv_indices,
+            q_num_blocks=q_num_blocks,
+            q_indices=q_indices,
+            full_q_num_blocks=full_q_num_blocks,
+            full_q_indices=full_q_indices,
+            BLOCK_SIZE=BLOCK_SIZE,
+            mask_mod=mask_mod,
+        )
+
+    @overload
+    def as_tuple(
+        self, flatten: Literal[True] = ...
+    ) -> tuple[
+        int,
+        int,
+        Tensor,
+        Tensor,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        int,
+        int,
+        _mask_mod_signature,
+    ]: ...
+
+    @overload
+    def as_tuple(
+        self, flatten: Literal[False]
+    ) -> tuple[
+        tuple[int, int],
+        Tensor,
+        Tensor,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        Tensor | None,
+        int | tuple[int, int],
+        _mask_mod_signature,
+    ]: ...
+
+    def as_tuple(self, flatten: bool = True) -> tuple[Any, ...]:
+        """
+        Returns a tuple of the attributes of the BlockMask.
+
+        Args:
+            flatten (bool): If True, it will flatten the tuple of (KV_BLOCK_SIZE, Q_BLOCK_SIZE)
+        """
+        if flatten:
+            block_size = (self.BLOCK_SIZE[0], self.BLOCK_SIZE[1])  # type: ignore[assignment]
+            seq_lengths = (self.seq_lengths[0], self.seq_lengths[1])  # type: ignore[assignment]
+        else:
+            block_size = (self.BLOCK_SIZE,)  # type: ignore[assignment]
+            seq_lengths = (self.seq_lengths,)  # type: ignore[assignment]
+
+        # pyrefly: ignore [not-iterable]
+        return (
+            *seq_lengths,
+            self.kv_num_blocks,
+            self.kv_indices,
+            self.full_kv_num_blocks,
+            self.full_kv_indices,
+            self.q_num_blocks,
+            self.q_indices,
+            self.full_q_num_blocks,
+            self.full_q_indices,
+            *block_size,
+            self.mask_mod,
+        )
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        *batch_dims, _, _ = self.kv_indices.shape
+        return tuple(batch_dims) + self.seq_lengths
+
+    def __str__(self) -> str:
+        s = f"BlockMask(shape={self.shape}, sparsity={self.sparsity():.2f}%, \n"
+        mask_str = self.to_string().strip()
+        s += mask_str
+        s += "\n)"
+        return s
+
+    def __getitem__(
+        self, index: int | slice | Tensor | tuple[int | slice | Tensor, ...]
+    ) -> Self:
+        """
+        Returns a new BlockMask instance by getting the mask for the given index position.
+
+        Args:
+            index: Index to apply to all attributes.
+
+        Example Usage:
+            .. code-block:: python
+
+                def causal_mask(b, h, q_idx, kv_idx):
+                    return q_idx >= kv_idx
+
+
+                block_mask = create_block_mask(
+                    causal_mask, 4, 2, 512, 512, device="cuda"
+                )
+                assert block_mask.kv_num_blocks.shape == (4, 2, 4)
+                assert block_mask.kv_indices.shape == (4, 2, 4, 4)
+
+                # Index on batch dimension
+                new_block_mask = block_mask[0]
+                assert new_block_mask.kv_num_blocks.shape == (2, 4)
+                assert new_block_mask.kv_indices.shape == (2, 4, 4)
+
+                # Index on batch and head dimension
+                new_block_mask = block_mask[0, 1]
+                assert new_block_mask.kv_num_blocks.shape == (4,)
+                assert new_block_mask.kv_indices.shape == (4, 4)
+
+                # slicing on batch and head dimension
+                new_block_mask = block_mask[0:2, 1:2]
+                assert new_block_mask.kv_num_blocks.shape == (2, 1, 4)
+                assert new_block_mask.kv_indices.shape == (2, 1, 4, 4)
+
+                # slicing on batch, head, and query dimension
+                new_block_mask = block_mask[
+                    0:2, 1:2, torch.tensor([1], dtype=torch.int32)
+                ]
+                assert new_block_mask.kv_num_blocks.shape == (2, 1, 1)
+                assert new_block_mask.kv_indices.shape == (2, 1, 1, 4)
+        """
+        index = (index,) if not isinstance(index, tuple) else index
+        padded = (*index, slice(None), slice(None), slice(None))[:3]
+        sizes = self.kv_num_blocks.shape[:3]
+        index = tuple(
+            (slice(i + n, i + n + 1) if -n <= i < 0 else slice(i, i + 1))
+            if isinstance(i, int)
+            else i
+            for i, n in zip(padded, sizes, strict=True)
+        )
+        new_kv_num_blocks = self.kv_num_blocks[index]
+        new_kv_indices = self.kv_indices[index]
+        if self.full_kv_num_blocks is not None:
+            if self.full_kv_indices is None:
+                raise AssertionError("full_kv_indices must not be None")
+            new_full_kv_num_blocks = self.full_kv_num_blocks[index]
+            new_full_kv_indices = self.full_kv_indices[index]
+        else:
+            new_full_kv_num_blocks = None
+            new_full_kv_indices = None
+        return BlockMask.from_kv_blocks(
+            new_kv_num_blocks,
+            new_kv_indices,
+            new_full_kv_num_blocks,
+            new_full_kv_indices,
+            BLOCK_SIZE=self.BLOCK_SIZE,
+            mask_mod=_sliced_mask_mod_error,
+            seq_lengths=self.seq_lengths,
+            compute_q_blocks=self.q_indices is not None,
+        )
+
+    def __repr__(self) -> str:
+        def shape_or_none(x: torch.Tensor | None):
+            return x.shape if x is not None else None
+
+        return (
+            f"BlockMask(\n"
+            f"    kv_num_blocks={self.kv_num_blocks.shape},\n"
+            f"    kv_indices={self.kv_indices.shape},\n"
+            f"    full_kv_num_blocks={shape_or_none(self.full_kv_num_blocks)},\n"
+            f"    full_kv_indices={shape_or_none(self.full_kv_indices)},\n"
+            f"    q_num_blocks={shape_or_none(self.q_num_blocks)},\n"
+            f"    q_indices={shape_or_none(self.q_indices)},\n"
+            f"    full_q_num_blocks={shape_or_none(self.full_q_num_blocks)},\n"
+            f"    full_q_indices={shape_or_none(self.full_q_indices)},\n"
+            f"    BLOCK_SIZE={self.BLOCK_SIZE},\n"
+            f"    shape={self.shape},\n"
+            f"    sparsity={self.sparsity():.2f}%,\n"
+            f"    mask_mod={self.mask_mod.__name__ if hasattr(self.mask_mod, '__name__') else self.mask_mod}\n"
+            f")"
+        )
+
+    def _adjust(self, new_q_len: int, new_kv_len: int) -> Self:
+        new_num_rows = (new_q_len + self.BLOCK_SIZE[0] - 1) // self.BLOCK_SIZE[0]
+        new_num_cols = (new_kv_len + self.BLOCK_SIZE[1] - 1) // self.BLOCK_SIZE[1]
+        new_kv_num_blocks, new_kv_indices = _adjust_num_blocks_and_indices(
+            self.kv_num_blocks, self.kv_indices, new_num_rows, new_num_cols
+        )
+        if self.full_kv_num_blocks is not None:
+            if self.full_kv_indices is None:
+                raise AssertionError("full_kv_indices must not be None")
+            (
+                new_full_kv_num_blocks,
+                new_full_kv_indices,
+            ) = _adjust_num_blocks_and_indices(
+                self.full_kv_num_blocks,
+                self.full_kv_indices,
+                new_num_rows,
+                new_num_cols,
+            )
+        else:
+            new_full_kv_num_blocks = None
+            new_full_kv_indices = None
+        return self.from_kv_blocks(
+            new_kv_num_blocks,
+            new_kv_indices,
+            new_full_kv_num_blocks,
+            new_full_kv_indices,
+            self.BLOCK_SIZE,
+            self.mask_mod,
+        )
+
+    def numel(self) -> int:
+        """Returns the number of elements (not accounting for sparsity) in the mask."""
+        shape = self.shape
+
+        def _prod(xs):
+            return functools.reduce(operator.mul, xs, 1)
+
+        return _prod(shape)
+
+    def sparsity(self) -> float:
+        """Computes the percentage of blocks that are sparse (i.e. not computed)"""
+        total_size = self.numel()
+        computed_blocks = self.kv_num_blocks.sum()
+        if self.full_kv_num_blocks is not None:
+            computed_blocks += self.full_kv_num_blocks.sum()
+
+        computed_size = computed_blocks.item() * self.BLOCK_SIZE[0] * self.BLOCK_SIZE[1]
+        dense_ratio = computed_size / total_size
+        return 100 * (1 - dense_ratio)
+
+    def to_dense(self) -> Tensor:
+        """Returns a dense block that is equivalent to the block mask."""
+        partial_dense = _ordered_to_dense(self.kv_num_blocks, self.kv_indices)
+        if self.full_kv_num_blocks is not None:
+            if self.full_kv_indices is None:
+                raise AssertionError("full_kv_indices must not be None")
+            # pyrefly: ignore [bad-return]
+            return partial_dense | _ordered_to_dense(
+                self.full_kv_num_blocks, self.full_kv_indices
+            )
+        return partial_dense
+
+    def to_string(
+        self, grid_size: int | tuple[int, int] = (20, 20), limit: int = 4
+    ) -> str:
+        """Returns a string representation of the block mask. Quite nifty.
+
+        If grid_size is -1, prints out an uncompressed version. Warning, it can be quite big!
+        """
+        dense_mask = self.to_dense()
+        *batch_dims, num_rows, num_cols = dense_mask.shape
+        if isinstance(grid_size, int):
+            max_rows = grid_size
+            max_cols = grid_size
+        elif grid_size == -1:
+            max_rows = num_rows
+            max_cols = num_cols
+        else:
+            max_rows, max_cols = grid_size
+
+        def create_block_vis(*batch_idx):
+            descriptors = []
+
+            descriptors.append(f"{batch_idx}")
+
+            vis = ", ".join(reversed(descriptors)) + "\n"
+
+            def summarize_section(section) -> str:
+                percentage = section.float().mean().item()
+                if percentage == 1:
+                    return "█"
+                elif percentage == 0:
+                    return " "
+                else:
+                    return "░"
+
+            def cdiv(a, b):
+                return (a + (b - 1)) // b
+
+            row_step = max(1, cdiv(num_rows, max_rows))
+            col_step = max(1, cdiv(num_cols, max_cols))
+
+            for r in range(0, num_rows, row_step):
+                for c in range(0, num_cols, col_step):
+                    cur_mask = dense_mask
+                    for idx in batch_idx:
+                        cur_mask = cur_mask[idx]
+                    char = summarize_section(
+                        cur_mask[r : r + row_step, c : c + col_step]
+                    )
+                    vis += char * 2
+                vis += "\n"
+            return vis
+
+        total_vis = []
+        for idx, batch_idx in enumerate(
+            itertools.product(*[range(i) for i in batch_dims])
+        ):
+            if idx == limit:
+                total_vis.append("...")
+                total_vis.append("To print out more, set BlockMask.to_string(limit=N)")
+                total_vis.append(
+                    "You can also index (BlockMask[batch, head]) to choose a specific batch or head"
+                )
+                break
+            block_vis = create_block_vis(*batch_idx)
+            total_vis.append(block_vis)
+
+        return "\n".join(total_vis)
+
+    def to(self, device: torch.device | str) -> BlockMask:
+        """Moves the BlockMask to the specified device.
+
+        Args:
+            device (torch.device or str): The target device to move the BlockMask to.
+                Can be a torch.device object or a string (e.g., 'cpu', 'cuda:0').
+
+        Returns:
+            BlockMask: A new BlockMask instance with all tensor components moved
+            to the specified device.
+
+        Note:
+            This method does not modify the original BlockMask in-place.
+            Instead, it returns a new BlockMask instance where individual tensor attributes
+            may or may not be moved to the specified device, depending on their
+            current device placement.
+        """
+        mapped_attributes = tree_map_only(
+            torch.Tensor,
+            lambda x: x.to(device),
+            self.as_tuple(flatten=False),
+        )
+        return BlockMask(*mapped_attributes)
+
+    @staticmethod
+    def _wrap_context_value(attr: str, value: Any) -> Any:
+        if attr == "mask_mod":
+            return _MaskModWrapper(value)
+        return value
+
+    @staticmethod
+    def _unwrap_context_value(attr: str, value: Any) -> Any:
+        if attr == "mask_mod":
+            if not isinstance(value, _MaskModWrapper):
+                raise AssertionError(f"Expected _MaskModWrapper, got {type(value)}")
+            return value.fn
+        return value
+
+    def _flatten(
+        self,
+    ) -> tuple[tuple[BaseArgumentTypes | None, ...], tuple[Any, ...]]:
+        """Flatten BlockMask into a list of tensors and context.
+
+        Closure tensors from mask_mod are extracted into the leaves via
+        _extract_closure_pytree so they are visible to the tracing
+        infrastructure (instead of being hidden in the pytree context).
+        """
+        tensors = tuple(getattr(self, attr) for attr in self._TENSOR_ATTRS)
+        closure_leaves, closure_spec, stripped = _extract_closure_pytree(self.mask_mod)
+        all_leaves = tensors + closure_leaves
+        context = tuple(
+            self._wrap_context_value(attr, getattr(self, attr))
+            if attr != "mask_mod"
+            else _MaskModWrapper(stripped, closure_spec)
+            for attr in self._CONTEXT_ATTRS
+        )
+        return all_leaves, context
+
+    @classmethod
+    def _unflatten(
+        cls,
+        leaves: tuple[Any, ...],
+        context: tuple[Any, ...],
+    ) -> Self:
+        """Unflatten leaves and context back into a BlockMask."""
+        n_regular = len(cls._TENSOR_ATTRS)
+        regular_leaves = leaves[:n_regular]
+        closure_leaves = leaves[n_regular:]
+        kwargs = {}
+        for attr, val in zip(cls._CONTEXT_ATTRS, context):
+            if attr == "mask_mod" and isinstance(val, _MaskModWrapper):
+                kwargs[attr] = _reconstruct_closure_fn(
+                    val.fn, closure_leaves, val.closure_spec
+                )
+            else:
+                kwargs[attr] = cls._unwrap_context_value(attr, val)
+        kwargs.update(zip(cls._TENSOR_ATTRS, regular_leaves))
+        return cls(**kwargs)
+
+    def _flatten_with_keys(
+        self,
+    ) -> tuple[tuple[tuple[GetAttrKey, Any], ...], tuple[tuple[GetAttrKey, Any], ...]]:
+        """Flatten BlockMask with keys for better tracing.
+
+        Closure tensors from mask_mod are extracted into the leaves via
+        _extract_closure_pytree so they are visible to the tracing
+        infrastructure (instead of being hidden in the pytree context).
+        """
+        tensors = tuple(
+            (GetAttrKey(attr), getattr(self, attr)) for attr in self._TENSOR_ATTRS
+        )
+        closure_leaves, closure_spec, stripped = _extract_closure_pytree(self.mask_mod)
+        closure_with_keys = tuple(
+            (GetAttrKey(f"_closure_{i}"), leaf) for i, leaf in enumerate(closure_leaves)
+        )
+        all_leaves = tensors + closure_with_keys
+        context = tuple(
+            (GetAttrKey(attr), self._wrap_context_value(attr, getattr(self, attr)))
+            if attr != "mask_mod"
+            else (GetAttrKey(attr), _MaskModWrapper(stripped, closure_spec))
+            for attr in self._CONTEXT_ATTRS
+        )
+        return all_leaves, context
+
+
+def _broadcast_to_dim(x: Tensor, dim: int) -> Tensor:
+    while x.dim() < dim:
+        x = x.unsqueeze(0)
+    return x
+
+
+def _round_up_to_multiple(x: int, multiple: int) -> int:
+    return (x + multiple - 1) // multiple * multiple
+
+
+def _convert_mask_to_block_mask(
+    mask: Tensor,
+    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    separate_full_blocks: bool = False,
+) -> tuple[Tensor, Tensor | None]:
+    if mask.dtype != torch.bool:
+        raise AssertionError(f"mask.dtype must be torch.bool, got {mask.dtype}")
+    mask = _broadcast_to_dim(mask, 4)
+
+    def padding_needed_for_multiple(x, multiple):
+        return _round_up_to_multiple(x, multiple) - x
+
+    mask = torch.nn.functional.pad(
+        mask,
+        (
+            0,
+            padding_needed_for_multiple(mask.shape[-1], KV_BLOCK_SIZE),
+            0,
+            padding_needed_for_multiple(mask.shape[-2], Q_BLOCK_SIZE),
+        ),
+    )
+    B, H, Q, KV = mask.shape
+    if Q % Q_BLOCK_SIZE != 0:
+        raise AssertionError(
+            f"Q ({Q}) must be divisible by Q_BLOCK_SIZE ({Q_BLOCK_SIZE})"
+        )
+    if KV % KV_BLOCK_SIZE != 0:
+        raise AssertionError(
+            f"KV ({KV}) must be divisible by KV_BLOCK_SIZE ({KV_BLOCK_SIZE})"
+        )
+    mask = mask.view(
+        B, H, Q // Q_BLOCK_SIZE, Q_BLOCK_SIZE, KV // KV_BLOCK_SIZE, KV_BLOCK_SIZE
+    )  # [B, H, Q//Q_BLOCK_SIZE, Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE, KV_BLOCK_SIZE]
+    mask = mask.permute(
+        0, 1, 2, 4, 3, 5
+    )  # [B, H, Q//Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE, Q_BLOCK_SIZE, KV_BLOCK_SIZE]
+    mask_block_sum = mask.sum(
+        dim=[-2, -1]
+    )  # [B, H, Q//Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE]
+    if separate_full_blocks:
+        full_block_sum = Q_BLOCK_SIZE * KV_BLOCK_SIZE
+        full_blocks = mask_block_sum == full_block_sum
+        partial_blocks = (mask_block_sum > 0) & (mask_block_sum < full_block_sum)
+        partial_blocks = partial_blocks.to(dtype=torch.int8)
+        full_blocks = full_blocks.to(dtype=torch.int8)
+        return partial_blocks, full_blocks
+    else:
+        partial_blocks = mask_block_sum > 0
+        partial_blocks = partial_blocks.to(dtype=torch.int8)
+        return partial_blocks, None
+
+
+def or_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:
+    """Returns a mask_mod that's the union of provided mask_mods"""
+    if not all(callable(arg) for arg in mask_mods):
+        raise RuntimeError(f"All inputs should be callable mask_mods: {mask_mods}")
+
+    def or_mask(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
+        result = b.new_zeros((), dtype=torch.bool)
+        for mask in mask_mods:
+            result = result | mask(b, h, q_idx, kv_idx)
+        return result
+
+    return or_mask
+
+
+def and_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:
+    """Returns a mask_mod that's the intersection of provided mask_mods"""
+    if not all(callable(arg) for arg in mask_mods):
+        raise RuntimeError(f"All inputs should be callable mask_mods: {mask_mods}")
+
+    def and_mask(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
+        result = b.new_ones((), dtype=torch.bool)
+        for mask in mask_mods:
+            result = result & mask(b, h, q_idx, kv_idx)
+        return result
+
+    return and_mask
+
+
+def _convert_block_mask_to_mask(
+    block_mask: Tensor,
+    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+) -> Tensor:
+    if block_mask.dim() != 4:
+        raise AssertionError(f"block_mask.dim() must be 4, got {block_mask.dim()}")
+    B, H, Q, KV = block_mask.shape
+    block_mask = block_mask.expand(Q_BLOCK_SIZE, KV_BLOCK_SIZE, *block_mask.shape)
+    block_mask = block_mask.permute(2, 3, 4, 0, 5, 1).reshape(
+        B, H, Q * Q_BLOCK_SIZE, KV * KV_BLOCK_SIZE
+    )
+    return block_mask
+
+
+def _create_sparse_block_from_block_mask(
+    block_mask: tuple[Tensor, Tensor | None],
+    mask_mod: _mask_mod_signature | None,
+    seq_lengths: tuple[int, int],
+    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
+) -> BlockMask:
+    partial_blocks, full_blocks = block_mask
+
+    partial_bm = _dense_to_ordered(partial_blocks)
+    if full_blocks is not None:
+        full_bm: tuple[Tensor | None, Tensor | None] = _dense_to_ordered(full_blocks)
+    else:
+        full_bm = (None, None)
+
+    return BlockMask.from_kv_blocks(
+        partial_bm[0],
+        partial_bm[1],
+        full_bm[0],
+        full_bm[1],
+        BLOCK_SIZE=(Q_BLOCK_SIZE, KV_BLOCK_SIZE),
+        mask_mod=mask_mod,
+        seq_lengths=seq_lengths,
+    )
+
+
+def create_mask(
+    mod_fn: _score_mod_signature | _mask_mod_signature,
+    B: int | None,
+    H: int | None,
+    Q_LEN: int,
+    KV_LEN: int,
+    device: DeviceLikeType | None = None,
+) -> Tensor:
+    r"""This function creates a mask tensor from a mod_fn function.
+
+    Args:
+        mod_fn (Union[_score_mod_signature, _mask_mod_signature]): Function to modify attention scores.
+        B (int): Batch size.
+        H (int): Number of query heads.
+        Q_LEN (int): Sequence length of query.
+        KV_LEN (int): Sequence length of key/value.
+        device (str): Device to run the mask creation on.
+
+    Returns:
+        mask (Tensor): A mask tensor with shape (B, H, M, N).
+    """
+    if device is None:
+        device = torch.accelerator.current_accelerator() or "cpu"
+    if B is None:
+        B = 1
+    if H is None:
+        H = 1
+    b = torch.arange(0, B, device=device)
+    h = torch.arange(0, H, device=device)
+    m = torch.arange(0, Q_LEN, device=device)
+    n = torch.arange(0, KV_LEN, device=device)
+    mod_type = _get_mod_type(mod_fn)
+
+    from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
+
+    with TransformGetItemToIndex():
+        if mod_type == _ModificationType.SCORE_MOD:
+            score_mod = mod_fn
+            score_mod = _vmap_for_bhqkv(score_mod, prefix=(0,))  # first input is score
+            out = score_mod(torch.zeros(B, H, Q_LEN, KV_LEN, device=device), b, h, m, n)
+            mask = torch.where(torch.isneginf(out), False, True)
+            return mask
+        elif mod_type == _ModificationType.MASK_MOD:
+            mask_mod = mod_fn
+            mask_mod = _vmap_for_bhqkv(mask_mod, prefix=())
+            mask = mask_mod(b, h, m, n)
+            return mask
+        else:
+            raise AssertionError
+
+
+def create_block_mask(
+    mask_mod: _mask_mod_signature,
+    B: int | None,
+    H: int | None,
+    Q_LEN: int,
+    KV_LEN: int,
+    device: DeviceLikeType | None = None,
+    BLOCK_SIZE: int | tuple[int, int] = _DEFAULT_SPARSE_BLOCK_SIZE,
+    _compile=False,
+) -> BlockMask:
+    r"""This function creates a block mask tuple from a mask_mod function.
+
+    Args:
+        mask_mod (Callable): mask_mod function. This is a callable that defines the
+            masking pattern for the attention mechanism. It takes four arguments:
+            b (batch size), h (number of heads), q_idx (query index), and kv_idx (key/value index).
+            It should return a boolean tensor indicating which attention connections are allowed (True)
+            or masked out (False).
+        B (int): Batch size.
+        H (int): Number of query heads.
+        Q_LEN (int): Sequence length of query.
+        KV_LEN (int): Sequence length of key/value.
+        device (str): Device to run the mask creation on.
+        BLOCK_SIZE (int or tuple[int, int]): Block size for the block mask. If a single int is provided it is used for both query and key/value.
+
+    Returns:
+        BlockMask:  A BlockMask object that contains the block mask information.
+
+    Example Usage:
+        .. code-block:: python
+
+            def causal_mask(b, h, q_idx, kv_idx):
+                return q_idx >= kv_idx
+
+
+            block_mask = create_block_mask(causal_mask, 1, 1, 8192, 8192, device="cuda")
+            query = torch.randn(1, 1, 8192, 64, device="cuda", dtype=torch.float16)
+            key = torch.randn(1, 1, 8192, 64, device="cuda", dtype=torch.float16)
+            value = torch.randn(1, 1, 8192, 64, device="cuda", dtype=torch.float16)
+            output = flex_attention(query, key, value, block_mask=block_mask)
+    """
+    if device is None:
+        device = torch.accelerator.current_accelerator() or "cpu"
+    mod_type = _get_mod_type(mask_mod)
+    if mod_type != _ModificationType.MASK_MOD:
+        raise AssertionError(
+            f"create-block_mask requires a mask_mod function! Got {mask_mod}"
+        )
+    if B is None:
+        B = 1
+    if H is None:
+        H = 1
+    if isinstance(BLOCK_SIZE, int):
+        Q_BLOCK_SIZE = BLOCK_SIZE
+        KV_BLOCK_SIZE = BLOCK_SIZE
+    else:
+        Q_BLOCK_SIZE, KV_BLOCK_SIZE = BLOCK_SIZE
+
+    if _compile:
+        warnings.warn(
+            "_compile flag on create_block_mask was originally added to work around a torch.compile limitation. That limitation has since been addressed. So, to compile create_block_mask, we suggest doing torch.compile(create_block_mask). This still works for now, but will be removed in the future.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return torch.compile(create_block_mask)(
+            mask_mod, B, H, Q_LEN, KV_LEN, device, BLOCK_SIZE
+        )
+
+    mask_tensor = create_mask(mask_mod, B, H, Q_LEN, KV_LEN, device)
+    partial_block_mask, full_block_mask = _convert_mask_to_block_mask(
+        mask_tensor,
+        Q_BLOCK_SIZE=Q_BLOCK_SIZE,
+        KV_BLOCK_SIZE=KV_BLOCK_SIZE,
+        separate_full_blocks=True,
+    )
+    block_mask = _create_sparse_block_from_block_mask(
+        (partial_block_mask, full_block_mask),
+        mask_mod,
+        (Q_LEN, KV_LEN),
+        Q_BLOCK_SIZE,
+        KV_BLOCK_SIZE,
+    )
+    return block_mask
+
+
+def _create_empty_block_mask(query: Tensor, key: Tensor) -> BlockMask:
+    r"""Default block mask for flex attention.
+    If users don't specify any block sparse mask info, we create this
+    empty block sparse mask. Which creates a BlockMask with 1 block that is the full length
+    of the query and key tensors.
+    """
+    device = query.device
+    return BlockMask.from_kv_blocks(
+        kv_num_blocks=torch.ones([1, 1, 1], dtype=torch.int32, device=device),
+        kv_indices=torch.zeros([1, 1, 1, 1], dtype=torch.int32, device=device),
+        BLOCK_SIZE=_LARGE_SPARSE_BLOCK_SIZE,
+        seq_lengths=(1, 1),
+    )
+
+
+def _apply_kernel_options(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    return_lse: bool,
+    kernel_options: FlexKernelOptions | None,
+    return_aux: AuxRequest | None = None,
+) -> _KernelOptionsWithInternals:
+    kernel_options = cast(
+        _KernelOptionsWithInternals,
+        {} if kernel_options is None else dict(kernel_options),
+    )
+
+    if "BACKEND" in kernel_options and kernel_options.get(
+        "FORCE_USE_FLEX_ATTENTION", False
+    ):
+        # TODO: remove FORCE_USE_FLEX_ATTENTION once BACKEND is fully adopted.
+        raise RuntimeError(
+            "BACKEND cannot be combined with legacy FORCE_USE_FLEX_ATTENTION. "
+            "BACKEND supersedes the legacy knob; please drop FORCE_USE_FLEX_ATTENTION "
+            "and only specify the desired BACKEND."
+        )
+
+    if "BACKEND" in kernel_options:
+        valid_backends = typing.get_args(_Backend)
+        if kernel_options["BACKEND"] not in valid_backends:
+            raise ValueError(
+                f"Invalid BACKEND value '{kernel_options['BACKEND']}'. "
+                f"Must be one of {valid_backends}"
+            )
+
+    kernel_options.setdefault("BACKEND", "AUTO")
+    kernel_options.setdefault("PRESCALE_QK", False)
+    kernel_options.setdefault("ROWS_GUARANTEED_SAFE", False)
+    kernel_options.setdefault("BLOCKS_ARE_CONTIGUOUS", False)
+    # This forces all biases grad scatters to be done in the DQ iteration loop of the backwards
+    kernel_options.setdefault("WRITE_DQ", True)
+
+    any_inputs_on_cpu_device = (
+        query.device.type == "cpu"
+        or key.device.type == "cpu"
+        or value.device.type == "cpu"
+    )
+
+    # Determine what auxiliary outputs are needed
+    output_lse = return_lse
+    output_max = False
+
+    if return_aux is not None:
+        # New API takes precedence over legacy parameters
+        output_lse = return_aux.lse
+        output_max = return_aux.max_scores
+
+    # If forward kernel needs to return logsumexp is decided by this rule internally.
+    if "OUTPUT_LOGSUMEXP" in kernel_options:
+        raise AssertionError("OUTPUT_LOGSUMEXP must not be in kernel_options")
+    kernel_options["OUTPUT_LOGSUMEXP"] = True
+    if not output_lse:
+        # We used to check if q,k,v required grads but since captured buffers can require grad
+        # we always write unless in no_grad
+        kernel_options["OUTPUT_LOGSUMEXP"] = torch.is_grad_enabled()
+        if any_inputs_on_cpu_device:
+            # CPU with torch.compile now supports inference, and will not return lse
+            # TODO: support CPU for training and return lse
+            kernel_options["OUTPUT_LOGSUMEXP"] = False
+
+    # If forward kernel needs to return max is decided by this rule internally.
+    if "OUTPUT_MAX" in kernel_options:
+        raise AssertionError("OUTPUT_MAX must not be in kernel_options")
+    if kernel_options["BACKEND"] == "FLASH" and output_max:
+        raise NotImplementedError(
+            "Returning max scores is not supported with BACKEND='FLASH'. "
+            "Use return_aux=AuxRequest(lse=True) or omit max_scores."
+        )
+    kernel_options["OUTPUT_MAX"] = output_max
+    if any_inputs_on_cpu_device and output_max:
+        # CPU doesn't support returning max yet
+        # TODO: support CPU for returning max
+        raise NotImplementedError("Returning max scores is not supported on CPU.")
+        kernel_options["OUTPUT_MAX"] = False
+
+    return kernel_options
+
+
+def _validate_embed_dim(query: Tensor, key: Tensor, value: Tensor) -> None:
+    if query.size(-1) != key.size(-1):
+        raise ValueError(
+            f"Expect query and key/value to have the same embedding dimension "
+            f"but got E={query.size(-1)} and E={key.size(-1)}."
+        )
+
+
+def _validate_device(query: Tensor, key: Tensor, value: Tensor) -> None:
+    """TODO: Remove once non cuda/cpu devices support is added
+    We only need to check query since we have already that q,k,v are on the same device
+    """
+    if query.device.type == "cpu" and (
+        query.requires_grad or key.requires_grad or value.requires_grad
+    ):
+        raise NotImplementedError(
+            "FlexAttention does not support backward on CPU. Please set the input requires_grad to False or use another device."
+        )
+    supported_devices = {"cuda", "cpu", "xpu", "hpu"}
+    if query.device.type not in supported_devices:
+        raise ValueError(
+            "FlexAttention is only supported on CUDA, CPU or HPU devices. "
+            f"Found input tensors on {query.device.type} device."
+        )
+
+
+def _enforce_mem_layouts(
+    query: Tensor, key: Tensor, value: Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Enforce memory layouts for query, key, and value tensors.
+
+    For non-FP8 dtypes, no action is taken.
+
+    For FP8 dtypes, we enforce the following memory layouts:
+    - Query tensor must be in row-major memory layout, as it will be the left-operand in the FP8 GEMM `q @ k.T`.
+    - Key tensor must be in row-major memory layout, as it will be transposed when used as the right-operand
+      in the FP8 GEMM `q @ k.T`, meaning it will correctly be in column-major memory layout for the GEMM.
+    - Value tensor must be in column-major memory layout, as it will be the right-operand in the FP8 GEMM `softmax_scores @ v`.
+
+    Returns the query, key, and value tensors with the enforced memory layouts.
+    """
+
+    def is_row_major(tensor: Tensor) -> bool:
+        return tensor.stride()[-1] == 1
+
+    def is_col_major(tensor: Tensor) -> bool:
+        return tensor.stride()[-2] == 1
+
+    # These memory layout constraint are only for FP8 GEMMs on NVIDIA GPU architectures >= SM89 and < SM100.
+    # This is because GPU arch < SM89 does not not support FP8 GEMMs, and
+    # SM100 has support for TN, NT, TT, NN layouts for FP8 GEMMs
+    # (i.e., left and right operands can be in row or column major layouts)
+    # so this check is only needed for older architectures.
+    # See: https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/blackwell_functionality.md
+    fp8_dtypes = (
+        torch.float8_e4m3fn,
+        torch.float8_e5m2,
+    )
+    gemm_precision = query.dtype
+
+    should_enforce_mem_layout = (
+        gemm_precision in fp8_dtypes
+        and torch.version.cuda is not None
+        and torch.cuda.get_device_capability("cuda") >= (8, 9)
+        and torch.cuda.get_device_capability("cuda") < (10, 0)
+    )
+    if not should_enforce_mem_layout:
+        return query, key, value
+
+    # Query must be in row-major memory layout as the left-operand in the FP8 GEMM `q @ k.T`
+    if not is_row_major(query):
+        query = query.contiguous()
+
+    # Key must be in row-major memory layout as it will be transposed when used as the right-operand
+    # in the FP8 GEMM `q @ k.T`, meaning it will correctly be in column-major memory layout for the GEMM.
+    if not is_row_major(key):
+        key = key.contiguous()
+
+    # Value must be in column-major memory layout as the right-operand in the FP8 GEMM `softmax_scores @ v`
+    if not is_col_major(value):
+        value = value.transpose(-2, -1).contiguous().transpose(-2, -1)
+    return query, key, value
+
+
+@overload
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: Literal[False] = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: None = ...,
+) -> Tensor: ...
+
+
+@overload
+@deprecated(
+    "return_lse is deprecated and will be removed in v2.10. "
+    "Use return_aux=AuxRequest(lse=True) instead.",
+    category=FutureWarning,
+)
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: Literal[True] = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: None = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+@overload
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: bool = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: AuxRequest,
+) -> tuple[Tensor, AuxOutput]: ...
+
+
+@overload
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = ...,
+    block_mask: BlockMask | None = ...,
+    scale: float | None = ...,
+    enable_gqa: bool = ...,
+    return_lse: Literal[True] = ...,
+    kernel_options: FlexKernelOptions | None = ...,
+    *,
+    return_aux: AuxRequest,
+) -> Never: ...
+
+
+def flex_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    score_mod: _score_mod_signature | None = None,
+    block_mask: BlockMask | None = None,
+    scale: float | None = None,
+    enable_gqa: bool = False,
+    return_lse: bool = False,
+    kernel_options: FlexKernelOptions | None = None,
+    *,
+    return_aux: AuxRequest | None = None,
+) -> Tensor | tuple[Tensor, Tensor] | tuple[Tensor, AuxOutput]:
+    r"""This function implements scaled dot product attention with an arbitrary attention score modification function
+    described in the `Flex Attention <https://arxiv.org/abs/2412.05496>`_ paper. See also the
+    `blog post <https://pytorch.org/blog/flexattention/>`_.
+
+    This function computes the scaled dot product attention between query, key, and value tensors with a user-defined
+    attention score modification function. The attention score modification function will be applied after the attention
+    scores have been calculated between the query and key tensors. The attention scores are calculated as follows:
+
+    The ``score_mod`` function should have the following signature:
+
+    .. code-block:: python
+
+        def score_mod(
+            score: Tensor,
+            batch: Tensor,
+            head: Tensor,
+            q_idx: Tensor,
+            k_idx: Tensor
+        ) -> Tensor:
+
+    Where:
+        - ``score``: A scalar tensor representing the attention score,
+          with the same data type and device as the query, key, and value tensors.
+        - ``batch``, ``head``, ``q_idx``, ``k_idx``: Scalar tensors indicating
+          the batch index, query head index, query index, and key/value index, respectively.
+          These should have the ``torch.int`` data type and be located on the same device as the score tensor.
+
+    Args:
+        query (Tensor): Query tensor; shape :math:`(B, Hq, L, E)`. For FP8 dtypes, should be in row-major memory layout for optimal performance.
+        key (Tensor): Key tensor; shape :math:`(B, Hkv, S, E)`. For FP8 dtypes, should be in row-major memory layout for optimal performance.
+        value (Tensor): Value tensor; shape :math:`(B, Hkv, S, Ev)`. For FP8 dtypes, should be in column-major memory layout for optimal performance.
+        score_mod (Optional[Callable]): Function to modify attention scores. By default no score_mod is applied.
+        block_mask (Optional[BlockMask]): BlockMask object that controls the blocksparsity pattern of the attention.
+        scale (Optional[float]): Scaling factor applied prior to softmax. If none, the default value is set to :math:`\frac{1}{\sqrt{E}}`.
+        enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA) and broadcasts key/value heads to query heads.
+        return_lse (bool): Whether to return the logsumexp of the attention scores. Default is False. **Deprecated**: Use ``return_aux=AuxRequest(lse=True)`` instead.
+        kernel_options (Optional[FlexKernelOptions]):
+            Options to control the behavior of the underlying Triton kernels.
+            See :class:`FlexKernelOptions` for available options and usage examples.
+        return_aux (Optional[AuxRequest]): Specifies which auxiliary outputs to compute and return.
+            If None, only the attention output is returned. Use ``AuxRequest(lse=True, max_scores=True)``
+            to request both auxiliary outputs.
+
+    Returns:
+        output (Tensor): Attention output; shape :math:`(B, Hq, L, Ev)`.
+
+        When ``return_aux`` is not None:
+            aux (AuxOutput): Auxiliary outputs with requested fields populated.
+
+        When ``return_aux`` is None (deprecated paths):
+            lse (Tensor): Log-sum-exp of attention scores; shape :math:`(B, Hq, L)`. Only returned if ``return_lse=True``.
+
+    Shape legend:
+        - :math:`N: \text{Batch size} ... : \text{Any number of other batch dimensions (optional)}`
+        - :math:`S: \text{Source sequence length}`
+        - :math:`L: \text{Target sequence length}`
+        - :math:`E: \text{Embedding dimension of the query and key}`
+        - :math:`Ev: \text{Embedding dimension of the value}`
+
+    .. warning::
+        `torch.nn.attention.flex_attention` is a prototype feature in PyTorch.
+        Please look forward to a more stable implementation in a future version of PyTorch.
+        Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype
+
+    """
+    # Some basic input validation
+    _validate_sdpa_input(query, key, value, allow_lowp_kv=True)
+    _validate_embed_dim(query, key, value)
+    _validate_device(query, key, value)
+    query, key, value = _enforce_mem_layouts(query, key, value)
+    if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:
+        raise NotImplementedError("NYI: query, key, and value must be 4D tensors")
+    if (not enable_gqa) and query.size(-3) != key.size(-3):
+        raise ValueError(
+            f"Expect query and key/value to have the same number of heads "
+            f"but got Hq={query.size(-3)} and Hkv={key.size(-3)}. "
+            f"Try setting enable_gqa=True for GQA."
+        )
+    if enable_gqa:
+        Hq = query.size(1)
+        Hkv = key.size(1)
+        if Hq % Hkv != 0:
+            raise ValueError(
+                f"Expect number of query heads to be a multiple of kv heads for GQA "
+                f"but got Hq={Hq} and Hkv={Hkv}."
+            )
+    if query.size(0) != key.size(0):
+        if block_mask is None:
+            raise ValueError(
+                f"Expect query and key/value to have the same batch size, "
+                f"or non-none block_mask, "
+                f"but got block_mask=None, Bq={query.size(0)}, and Bkv={key.size(0)}."
+            )
+
+        if block_mask.kv_num_blocks.size(0) != query.size(0):
+            raise ValueError(
+                f"Expect query and key/value to have the same batch size, "
+                f"or block_mask and query to have the same batch size, "
+                f"but got Bq={query.size(0)}, Bkv={key.size(0)}, B_block_mask={block_mask.kv_num_blocks.size(0)}."
+            )
+
+    if score_mod is None:
+        score_mod = _identity
+
+    if block_mask is None:
+        block_mask = _create_empty_block_mask(query, key)
+
+    # If BlockMask was sliced, its mask_mod is intentionally replaced with an error-raising stub.
+    # This guard ensures we surface the intended error message before any shape-based checks.
+    if getattr(block_mask, "mask_mod", None) is _sliced_mask_mod_error:
+        raise RuntimeError("Cannot use mask_mod from a sliced BlockMask")
+
+    if (
+        block_mask.BLOCK_SIZE[0] == _LARGE_SPARSE_BLOCK_SIZE
+        and block_mask.BLOCK_SIZE[1] == _LARGE_SPARSE_BLOCK_SIZE
+    ):
+        # This corresponds to the case where we essentially have a "no-op" block mask.
+        pass
+    else:
+        block_mask_q_len = block_mask.shape[-2]
+        block_mask_kv_len = block_mask.shape[-1]
+        if query.size(-2) > block_mask_q_len or key.size(-2) > block_mask_kv_len:
+            raise ValueError(
+                f"block_mask was created for block_mask.shape={block_mask.shape} but got q_len={query.size(-2)} and kv_len={key.size(-2)}. "
+                "As the block mask was created for a smaller length than you're using it for, you likely need to create a new block mask."
+            )
+        elif (
+            query.size(-2) < block_mask_q_len and key.size(-2) <= block_mask_kv_len
+        ) or (query.size(-2) <= block_mask_q_len and key.size(-2) < block_mask_kv_len):
+            raise ValueError(
+                f"block_mask was created for block_mask.shape={block_mask.shape} but got q_len={query.size(-2)} and kv_len={key.size(-2)}. "
+                "As the block mask was created for a larger length than you're using it for, you can either 1. create a new block mask with the correct length, or 2. 'adjust' the existing block mask to the correct length by calling block_mask._adjust(q_len, kv_len). This essentially 'crops' the block mask to the upper left corner, which does not work for all mask_mods!"
+            )
+        if query.size(-2) != block_mask_q_len:
+            raise AssertionError(
+                f"query.size(-2) ({query.size(-2)}) != block_mask_q_len ({block_mask_q_len})"
+            )
+        if key.size(-2) != block_mask_kv_len:
+            raise AssertionError(
+                f"key.size(-2) ({key.size(-2)}) != block_mask_kv_len ({block_mask_kv_len})"
+            )
+
+    if scale is None:
+        scale = 1.0 / math.sqrt(query.size(-1))
+
+    if query.device != block_mask.kv_num_blocks.device:  # type: ignore[union-attr]
+        raise RuntimeError(
+            f"Expect q/k/v and block_mask to be on the same device "
+            f"but got {query.device} and {block_mask.kv_num_blocks.device}."  # type: ignore[union-attr]
+        )
+
+    # Handle deprecation warnings for old parameters
+    if return_lse and return_aux is not None:
+        raise ValueError(
+            "Cannot specify both return_lse and return_aux. "
+            "return_lse is deprecated, please use return_aux=AuxRequest(lse=True) instead."
+        )
+    elif return_lse and return_aux is None:
+        _warn_once(
+            "deprecated_return_lse",
+            "return_lse is deprecated and will be removed in v2.10. "
+            "Please use return_aux=AuxRequest(lse=True) instead.",
+            category=FutureWarning,
+        )
+
+    kernel_options = _apply_kernel_options(
+        query,
+        key,
+        value,
+        return_lse,
+        kernel_options,
+        return_aux,
+    )
+
+    def _finalize_outputs(
+        out,
+        lse,
+        max_scores,
+        *,
+        return_aux: AuxRequest | None,
+        return_lse: bool,
+    ):
+        """Normalize stats and build return value (aux-aware, legacy-compatible)."""
+        ln2 = math.log(2.0)
+        return_lse = return_lse or return_aux is not None and return_aux.lse
+        return_max = return_aux is not None and return_aux.max_scores
+
+        lse_scaled = lse * ln2 if (return_lse and lse.numel() > 0) else None
+        max_scaled = (
+            max_scores * ln2 if (return_max and max_scores.numel() > 0) else None
+        )
+
+        if return_aux is not None:
+            return out, AuxOutput(
+                lse=lse_scaled,
+                max_scores=max_scaled,
+            )
+
+        if return_lse:
+            return out, lse_scaled
+
+        return out
+
+    if torch.compiler.is_dynamo_compiling():
+        # mark head_dim and number of heads to be static
+        for x in [query, key, value]:
+            torch._dynamo.mark_static(x, -3)
+            torch._dynamo.mark_static(x, -1)
+
+        out, lse, max_scores = flex_attention_hop(
+            query,
+            key,
+            value,
+            score_mod,
+            block_mask.as_tuple(),
+            scale,
+            kernel_options,  # type: ignore[union-attr]
+        )
+        return _finalize_outputs(
+            out, lse, max_scores, return_aux=return_aux, return_lse=return_lse
+        )
+
+    if not _FLEX_ATTENTION_DISABLE_COMPILE_DEBUG:
+        _warn_once(
+            warning_id="flex_attention_performance",
+            message=(
+                "flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.\n\n"
+                "SOLUTION: Use torch.compile(flex_attention)(...)\n\n"
+                "If you want to debug your score_mod/mask_mod, you can set:\n"
+                "torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True\n\n"
+                "This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results."
+            ),
+        )
+
+    if not torch._dynamo.is_dynamo_supported():
+        raise RuntimeError("flex_attention requires dynamo support")
+
+    # Dynamo is expecting a callable with "__code__" attribute.
+    # We cannot directly pass hop to it. So we wrap it in a dummy function.
+    def _flex_attention_hop_wrapper(*args, **kwargs):
+        return flex_attention_hop(*args, **kwargs)
+
+    with setup_compilation_env() as backend:
+        if _FLEX_ATTENTION_DISABLE_COMPILE_DEBUG:
+            flex_fn = _flex_attention_hop_wrapper
+        else:
+            flex_fn = torch.compile(
+                _flex_attention_hop_wrapper, backend=backend, fullgraph=True
+            )
+
+        out, lse, max_scores = flex_fn(
+            query,
+            key,
+            value,
+            score_mod,
+            block_mask.as_tuple(),
+            scale,
+            kernel_options,
+        )
+    return _finalize_outputs(
+        out, lse, max_scores, return_aux=return_aux, return_lse=return_lse
+    )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/varlen.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/varlen.py
new file mode 100644
index 0000000000000000000000000000000000000000..96cfc332b8699229e407935b723e17dfcf0c8725
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/attention/varlen.py
@@ -0,0 +1,671 @@
+"""
+Variable-length attention implementation using Flash Attention.
+
+This module provides a high-level Python interface for variable-length attention
+that calls into the optimized Flash Attention kernels.
+"""
+
+import logging
+from functools import lru_cache
+from typing import Any, NamedTuple
+
+import torch
+
+
+log = logging.getLogger(__name__)
+
+__all__ = ["varlen_attn", "varlen_attn_out", "AuxRequest"]
+
+
+def _normalize_window_size(window_size: list[int] | None) -> list[int]:
+    if window_size is None:
+        window_size = [-1, -1]
+
+    if len(window_size) != 2:
+        raise ValueError(f"window_size must have length 2, got {len(window_size)}")
+    return window_size
+
+
+@lru_cache(maxsize=8)
+def _should_use_cudnn(device_index: int) -> bool:
+    """Cache device capability check to avoid repeated CUDA calls."""
+    return False
+
+
+class AuxRequest(NamedTuple):
+    """
+    Request which auxiliary outputs to compute from varlen_attn.
+
+    Each field is a boolean indicating whether that auxiliary output should be computed.
+    """
+
+    lse: bool = False
+
+
+@torch.library.custom_op("torch_attn::_varlen_attn", mutates_args={})
+def _varlen_attn(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    is_causal: bool = False,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Private custom op for variable-length attention.
+
+    This is the internal implementation. Users should use the public varlen_attn function instead.
+    """
+    window_size = _normalize_window_size(window_size)
+
+    use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
+
+    if use_cudnn:
+        log.info("Using cuDNN backend for varlen_attn")
+
+        if enable_gqa:
+            # TODO: check this
+            raise RuntimeError("GQA is not supported with the cuDNN backend.")
+        if num_splits is not None:
+            # TODO: check this
+            raise RuntimeError("num_splits is not supported with the cuDNN backend.")
+        if window_size[0] != -1 or window_size[1] != -1:
+            raise RuntimeError(
+                "cuDNN backend does not support window attention. Please use Flash Attention backend."
+            )
+        if seqused_k is not None or block_table is not None:
+            # TODO: cuDNN supports per-sequence KV lengths via SEQ_LEN_KV + padding_mask,
+            # but _cudnn_attention_forward doesn't expose it yet.
+            raise RuntimeError(
+                "seqused_k/block_table is not yet supported with the cuDNN backend."
+            )
+
+        result = torch.ops.aten._cudnn_attention_forward(
+            query,
+            key,
+            value,
+            None,  # attn_bias
+            cu_seq_q,
+            cu_seq_k,
+            max_q,
+            max_k,
+            True,  # compute_log_sumexp
+            0.0,  # dropout_p hardcoded to 0.0
+            is_causal,
+            False,  # return_debug_mask
+            scale=scale,
+        )
+        # cuDNN returns: (output, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask)
+        output, softmax_lse, rng_state = result[0], result[1], result[6]
+    else:
+        log.info("Using Flash Attention backend for varlen_attn")
+        output, softmax_lse, rng_state, _, _ = torch.ops.aten._flash_attention_forward(
+            query,
+            key,
+            value,
+            cu_seq_q,
+            cu_seq_k,
+            max_q,
+            max_k,
+            0.0,  # dropout_p hardcoded to 0.0
+            is_causal,
+            return_debug_mask=False,
+            scale=scale,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            seqused_k=seqused_k,
+            block_table=block_table,
+            num_splits=num_splits,
+        )
+
+    rng_state_ = torch.zeros(
+        (2,), dtype=torch.uint64, device=query.device
+    )  # hardcoded since dropout is hardcoded to 0
+    return output, softmax_lse, rng_state_
+
+
+@_varlen_attn.register_fake
+def _varlen_attn_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    is_causal: bool = False,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Fake implementation for meta tensor computation and tracing.
+
+    Based on the 3D varlen path from meta__flash_attention_forward:
+    - query shape: (total, num_heads, head_dim)
+    - logsumexp shape: (num_heads, total_q)
+    """
+    window_size = _normalize_window_size(window_size)
+
+    # Output has same shape as query
+    output = torch.empty_like(query)
+
+    # For varlen path: logsumexp shape is (num_heads, total_q)
+    total_q = query.size(0)
+    num_heads = query.size(1)
+    logsumexp = torch.empty(
+        (num_heads, total_q), dtype=torch.float, device=query.device
+    )
+
+    if torch.version.hip:
+        preferred = torch._C._get_rocm_fa_preferred_backend()
+        if preferred == torch._C._ROCmFABackend.AOTriton:
+            # AOTriton ROCm path uses batched 3D
+            batch_size = cu_seq_q.size(0) - 1
+            logsumexp = torch.empty(
+                (batch_size, num_heads, max_q), dtype=torch.float, device=query.device
+            )
+
+    rng_state = torch.empty((2,), dtype=torch.uint64, device=query.device)
+
+    return output, logsumexp, rng_state
+
+
+def varlen_attn(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    *,
+    return_aux: AuxRequest | None = None,
+    scale: float | None = None,
+    window_size: tuple[int, int] = (-1, -1),
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    r"""Compute variable-length attention using Flash Attention.
+
+    This function is similar to scaled_dot_product_attention but optimized for
+    variable-length sequences using cumulative sequence position tensors.
+
+    Args:
+        query (Tensor): Query tensor; shape :math:`(T_q, H_q, D)`
+        key (Tensor): Key tensor; shape :math:`(T_k, H_{kv}, D)`, or
+            :math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
+        value (Tensor): Value tensor; shape :math:`(T_k, H_{kv}, D)`, or
+            :math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
+        cu_seq_q (Tensor): Cumulative sequence positions for queries; shape :math:`(N+1,)`
+        cu_seq_k (Tensor): Cumulative sequence positions for keys/values; shape :math:`(N+1,)`
+        max_q (int): Maximum query sequence length in the batch.
+        max_k (int): Maximum key/value sequence length in the batch.
+        return_aux (Optional[AuxRequest]): If not None and ``return_aux.lse`` is True, also returns the logsumexp tensor.
+        scale (float, optional): Scaling factor for attention scores
+        window_size (tuple[int, int], optional): Window size for sliding window attention as (left, right).
+            Use (-1, -1) for full attention (default), (-1, 0) for causal attention,
+            or (W, 0) for causal attention with sliding window of size W.
+        enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA)
+            and allows key/value to have fewer heads than query.
+            Each KV head is shared by a group of :math:`H_q / H_{kv}` query heads,
+            so :math:`H_q` must be divisible by :math:`H_{kv}`.
+            Default is False.
+        seqused_k (Tensor, optional): Number of valid KV tokens per batch element; shape :math:`(N,)`.
+            When set, only the first ``seqused_k[i]`` tokens in the key/value sequence for batch
+            element *i* participate in attention. Useful for KV-cache decoding where the cache slot
+            is larger than the actual sequence. Inference-only (not supported in backward).
+        block_table (Tensor, optional): Block table for paged KV cache; shape
+            :math:`(N, \text{max\_pages\_per\_seq})`, dtype ``int32``.
+            Requires ``seqused_k``. Inference-only (not supported in backward).
+
+            When ``block_table`` is provided, ``key`` and ``value`` are a "pool" of
+            pages of tokens of KV data and the pages belong to any sequence/order.
+            The ``block_table`` is what maps each sequence's logical chunks
+            back to physical pages in this pool.
+
+            ``seqused_k[i]`` tells the kernel how many tokens in sequence *i* are
+            actually valid, since the last page is typically only partially filled.
+        num_splits (int, optional): Number of splits for split-KV. Set to ``1``
+            to disable split-KV which enables batch invariance. Split-KV
+            parallelizes the key/value sequence dimension across multiple thread
+            blocks and combines partial results. The split decision depends
+            on ``max_k`` (the longest sequence in the batch), so different batch
+            compositions can change the reduction order and produce different
+            floating-point results for the same sequence. When this is disabled,
+            bitwise identical outputs are guaranteed for a given sequence
+            regardless of what other sequences are in the batch, at the
+            cost of lower GPU utilization when there are few queries. When
+            ``None`` (default), the kernel chooses automatically.
+
+    Returns:
+        output (Tensor): Output tensor from attention computation; shape :math:`(T_q, H_q, D)`.
+
+        If ``return_aux`` is not None and ``return_aux.lse`` is True:
+            lse (Tensor): Log-sum-exp of attention scores; shape :math:`(T_q, H_q)`.
+
+    Shape legend:
+        - :math:`N`: Batch size
+        - :math:`T_q`: Total number of query tokens in the batch (sum of all query sequence lengths)
+        - :math:`T_k`: Total number of key/value tokens in the batch (sum of all key/value sequence lengths)
+        - :math:`H_q`: Number of query attention heads
+        - :math:`H_{kv}`: Number of key/value attention heads (equal to :math:`H_q` unless GQA is enabled)
+        - :math:`D`: Head dimension
+
+    Example::
+
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+        >>> batch_size, max_seq_len, embed_dim, num_heads = 2, 512, 1024, 16
+        >>> head_dim = embed_dim // num_heads
+        >>> seq_lengths = []
+        >>> for _ in range(batch_size):
+        ...     length = torch.randint(1, max_seq_len // 64 + 1, (1,)).item() * 64
+        ...     seq_lengths.append(min(length, max_seq_len))
+        >>> seq_lengths = torch.tensor(seq_lengths, device="cuda")
+        >>> total_tokens = seq_lengths.sum().item()
+        >>>
+        >>> # Create packed query, key, value tensors
+        >>> query = torch.randn(
+        ...     total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
+        ... )
+        >>> key = torch.randn(
+        ...     total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
+        ... )
+        >>> value = torch.randn(
+        ...     total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
+        ... )
+        >>>
+        >>> # Build cumulative sequence tensor
+        >>> cu_seq = torch.zeros(batch_size + 1, device="cuda", dtype=torch.int32)
+        >>> cu_seq[1:] = seq_lengths.cumsum(0)
+        >>> max_len = seq_lengths.max().item()
+        >>>
+        >>> # Call varlen_attn
+        >>> output = varlen_attn(
+        ...     query, key, value, cu_seq, cu_seq, max_len, max_len
+        ... )
+    """
+
+    num_heads_q = query.size(1)
+    num_heads_k = key.size(2) if block_table is not None else key.size(1)
+    if not enable_gqa and num_heads_q != num_heads_k:
+        raise ValueError(
+            f"Expect query and key/value to have the same number of heads "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
+            f"Try setting enable_gqa=True for GQA."
+        )
+    if enable_gqa and num_heads_q % num_heads_k != 0:
+        raise ValueError(
+            f"Expect number of query heads to be a multiple of kv heads for GQA "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
+        )
+
+    is_causal = window_size == (-1, 0)
+    out, lse, _ = torch.ops.torch_attn._varlen_attn(
+        query,
+        key,
+        value,
+        cu_seq_q,
+        cu_seq_k,
+        max_q,
+        max_k,
+        is_causal,
+        scale,
+        list(window_size),
+        enable_gqa,
+        seqused_k,
+        block_table,
+        num_splits,
+    )
+    if return_aux is not None and return_aux.lse:
+        return out, lse
+    return out
+
+
+@torch.library.custom_op("torch_attn::_varlen_attn_out", mutates_args={"out"})
+def _varlen_attn_out(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    is_causal: bool = False,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> torch.Tensor:
+    """
+    Private custom op for variable-length attention with pre-allocated output.
+    Same as _varlen_attn but writes the attention output into the provided out tensor.
+    """
+    window_size = _normalize_window_size(window_size)
+
+    use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
+
+    if use_cudnn:
+        # TODO: look into this
+        raise RuntimeError("cuDNN backend does not support out variant.")
+
+    log.info("Using Flash Attention backend for varlen_attn_out")
+    softmax_lse = torch.ops.aten._flash_attention_forward_no_dropout_inplace(
+        out,
+        query,
+        key,
+        value,
+        cu_seq_q,
+        cu_seq_k,
+        max_q,
+        max_k,
+        0.0,  # dropout_p hardcoded to 0.0
+        is_causal,
+        False,  # return_debug_mask
+        scale=scale,
+        window_size_left=window_size[0],
+        window_size_right=window_size[1],
+        seqused_k=seqused_k,
+        block_table=block_table,
+        num_splits=num_splits,
+    )
+
+    return softmax_lse
+
+
+@_varlen_attn_out.register_fake
+def _varlen_attn_out_fake(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    is_causal: bool = False,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> torch.Tensor:
+    """
+    Fake implementation for meta tensor computation and tracing.
+    """
+    total_q = query.size(0)
+    num_heads = query.size(1)
+    logsumexp = torch.empty(
+        (num_heads, total_q), dtype=torch.float, device=query.device
+    )
+
+    if torch.version.hip:
+        preferred = torch._C._get_rocm_fa_preferred_backend()
+        if preferred == torch._C._ROCmFABackend.AOTriton:
+            batch_size = cu_seq_q.size(0) - 1
+            logsumexp = torch.empty(
+                (batch_size, num_heads, max_q), dtype=torch.float, device=query.device
+            )
+
+    return logsumexp
+
+
+def varlen_attn_out(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor | None,
+    max_q: int,
+    max_k: int,
+    *,
+    return_aux: AuxRequest | None = None,
+    scale: float | None = None,
+    window_size: tuple[int, int] = (-1, -1),
+    enable_gqa: bool = False,
+    seqused_k: torch.Tensor | None = None,
+    block_table: torch.Tensor | None = None,
+    num_splits: int | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    r"""Compute variable-length attention using Flash Attention with a pre-allocated output tensor.
+
+    Same as :func:`varlen_attn` but writes the attention output into the provided ``out`` tensor
+    instead of allocating a new one.
+
+    """
+    num_heads_q = query.size(1)
+    num_heads_k = key.size(2) if block_table is not None else key.size(1)
+    if not enable_gqa and num_heads_q != num_heads_k:
+        raise ValueError(
+            f"Expect query and key/value to have the same number of heads "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
+            f"Try setting enable_gqa=True for GQA."
+        )
+    if enable_gqa and num_heads_q % num_heads_k != 0:
+        raise ValueError(
+            f"Expect number of query heads to be a multiple of kv heads for GQA "
+            f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
+        )
+
+    is_causal = window_size == (-1, 0)
+    lse = torch.ops.torch_attn._varlen_attn_out(
+        out,
+        query,
+        key,
+        value,
+        cu_seq_q,
+        cu_seq_k,
+        max_q,
+        max_k,
+        is_causal,
+        scale,
+        list(window_size),
+        enable_gqa,
+        seqused_k,
+        block_table,
+        num_splits,
+    )
+    if return_aux is not None and return_aux.lse:
+        return out, lse
+    return out
+
+
+def _setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None:
+    (
+        query,
+        key,
+        value,
+        cu_seq_q,
+        cu_seq_k,
+        max_q,
+        max_k,
+        is_causal,
+        scale,
+        window_size,
+        enable_gqa,
+        seqused_k,
+        block_table,
+        num_splits,
+    ) = inputs
+    out, lse, rng_state = output
+
+    if seqused_k is not None:
+        raise RuntimeError("seqused_k is an inference-only parameter.")
+    if block_table is not None:
+        raise RuntimeError("block_table is an inference-only parameter.")
+
+    ctx.save_for_backward(query, key, value, cu_seq_q, cu_seq_k, out, lse, rng_state)
+
+    ctx.max_q = max_q
+    ctx.max_k = max_k
+    ctx.is_causal = is_causal
+    ctx.scale = scale
+    ctx.window_size = window_size
+
+
+@torch.library.custom_op("torch_attn::_varlen_attn_backward", mutates_args={})
+def _varlen_attn_backward(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor,
+    max_q: int,
+    max_k: int,
+    is_causal: bool,
+    rng_state: torch.Tensor,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    window_size = _normalize_window_size(window_size)
+
+    unused = torch.empty(0, device=query.device)
+
+    use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
+    if use_cudnn:
+        log.info("Using cuDNN backend for varlen_attn")
+        if window_size[0] != -1 or window_size[1] != -1:
+            raise RuntimeError(
+                "cuDNN backend does not support window attention. Please use Flash Attention backend."
+            )
+        dq, dk, dv = torch.ops.aten._cudnn_attention_backward(
+            grad_out,
+            query,
+            key,
+            value,
+            out,
+            lse,
+            cu_seq_q,
+            cu_seq_k,
+            max_q,
+            max_k,
+            0.0,
+            is_causal,
+            rng_state,
+            unused,
+            scale=scale,
+        )
+    else:
+        log.info("Using Flash Attention backend for varlen_attn")
+        dq, dk, dv = torch.ops.aten._flash_attention_backward(
+            grad_out,
+            query,
+            key,
+            value,
+            out,
+            lse,
+            cu_seq_q,
+            cu_seq_k,
+            max_q,
+            max_k,
+            0.0,
+            is_causal,
+            rng_state,
+            unused,
+            scale=scale,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+        )
+    return dq, dk, dv
+
+
+@_varlen_attn_backward.register_fake
+def _varlen_attn_backward_fake(
+    grad_out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seq_q: torch.Tensor,
+    cu_seq_k: torch.Tensor,
+    max_q: int,
+    max_k: int,
+    is_causal: bool,
+    rng_state: torch.Tensor,
+    scale: float | None = None,
+    window_size: list[int] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Fake implementation for meta tensor computation and tracing.
+    """
+    window_size = _normalize_window_size(window_size)
+
+    grad_query = torch.empty_like(query)
+    grad_key = torch.empty_like(key)
+    grad_value = torch.empty_like(value)
+
+    return grad_query, grad_key, grad_value
+
+
+def _backward(
+    ctx: Any, grad_out: torch.Tensor, grad_lse: torch.Tensor, grad_rng: torch.Tensor
+) -> tuple[torch.Tensor | None, ...]:
+    query, key, value, cu_seq_q, cu_seq_k, out, lse, rng_state = ctx.saved_tensors
+
+    max_q = ctx.max_q
+    max_k = ctx.max_k
+    is_causal = ctx.is_causal
+    scale = ctx.scale
+    window_size = ctx.window_size
+
+    dq, dk, dv = torch.ops.torch_attn._varlen_attn_backward(
+        grad_out,
+        query,
+        key,
+        value,
+        out,
+        lse,
+        cu_seq_q,
+        cu_seq_k,
+        max_q,
+        max_k,
+        is_causal,
+        rng_state,
+        scale,
+        window_size,
+    )
+    # cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, window_size, \
+    # enable_gqa, seqused_k, block_table, num_splits
+    num_params = 11
+    return (dq, dk, dv, *((None,) * num_params))
+
+
+_varlen_attn.register_autograd(_backward, setup_context=_setup_context)
+
+torch._dynamo.disallow_in_graph(
+    torch.ops.aten._flash_attention_forward_no_dropout_inplace
+)
+
+from torch.utils.flop_counter import (
+    _varlen_attn_backward_flop,
+    _varlen_attn_forward_flop,
+    _varlen_attn_out_flop,
+    flop_registry,
+)
+
+
+flop_registry[torch.ops.torch_attn._varlen_attn] = _varlen_attn_forward_flop
+flop_registry[torch.ops.torch_attn._varlen_attn_out] = _varlen_attn_out_flop
+flop_registry[torch.ops.torch_attn._varlen_attn_backward] = _varlen_attn_backward_flop
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/backends/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/backends/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/backends/thnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/backends/thnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..c56e923a84383a79c2a3f7ebddb3dfa1ce1f0953
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/backends/thnn.py
@@ -0,0 +1,6 @@
+# mypy: allow-untyped-defs
+# this is for historical pickle deserialization, it is not used otherwise
+
+
+def _get_thnn_function_backend() -> None:
+    pass
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/common_types.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/common_types.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1928414a396e9be5f4bc90e198b8c5eac77e23c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/common_types.py
@@ -0,0 +1,46 @@
+from typing import TypeAlias as _TypeAlias, TypeVar
+
+from torch import Tensor
+
+
+# ruff: noqa: PYI042,PYI047
+
+# Create some useful type aliases
+
+# Template for arguments which can be supplied as a tuple, or which can be a scalar which PyTorch will internally
+# broadcast to a tuple.
+# Comes in several variants: A tuple of unknown size, and a fixed-size tuple for 1d, 2d, or 3d operations.
+T = TypeVar("T")
+_scalar_or_tuple_any_t: _TypeAlias = T | tuple[T, ...]
+_scalar_or_tuple_1_t: _TypeAlias = T | tuple[T]
+_scalar_or_tuple_2_t: _TypeAlias = T | tuple[T, T]
+_scalar_or_tuple_3_t: _TypeAlias = T | tuple[T, T, T]
+_scalar_or_tuple_4_t: _TypeAlias = T | tuple[T, T, T, T]
+_scalar_or_tuple_5_t: _TypeAlias = T | tuple[T, T, T, T, T]
+_scalar_or_tuple_6_t: _TypeAlias = T | tuple[T, T, T, T, T, T]
+
+# For arguments which represent size parameters (eg, kernel size, padding)
+_size_any_t: _TypeAlias = _scalar_or_tuple_any_t[int]
+_size_1_t: _TypeAlias = _scalar_or_tuple_1_t[int]
+_size_2_t: _TypeAlias = _scalar_or_tuple_2_t[int]
+_size_3_t: _TypeAlias = _scalar_or_tuple_3_t[int]
+_size_4_t: _TypeAlias = _scalar_or_tuple_4_t[int]
+_size_5_t: _TypeAlias = _scalar_or_tuple_5_t[int]
+_size_6_t: _TypeAlias = _scalar_or_tuple_6_t[int]
+
+# For arguments which represent optional size parameters (eg, adaptive pool parameters)
+_size_any_opt_t: _TypeAlias = _scalar_or_tuple_any_t[int | None]
+_size_2_opt_t: _TypeAlias = _scalar_or_tuple_2_t[int | None]
+_size_3_opt_t: _TypeAlias = _scalar_or_tuple_3_t[int | None]
+
+# For arguments that represent a ratio to adjust each dimension of an input with (eg, upsampling parameters)
+_ratio_2_t: _TypeAlias = _scalar_or_tuple_2_t[float]
+_ratio_3_t: _TypeAlias = _scalar_or_tuple_3_t[float]
+_ratio_any_t: _TypeAlias = _scalar_or_tuple_any_t[float]
+
+_tensor_list_t: _TypeAlias = _scalar_or_tuple_any_t[Tensor]
+
+# For the return value of max pooling operations that may or may not return indices.
+# With the proposed 'Literal' feature to Python typing, it might be possible to
+# eventually eliminate this.
+_maybe_indices_t: _TypeAlias = _scalar_or_tuple_2_t[Tensor]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/cpp.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/cpp.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4ffd188cd39a43c2c1099eab81f620629fff1d6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/cpp.py
@@ -0,0 +1,90 @@
+# mypy: allow-untyped-defs
+"""Functionality for Python <-> C++ frontend inter-op."""
+
+from torch import nn
+
+
+class OrderedDictWrapper:
+    """A wrapper around a C++ OrderedDict.
+
+    It dynamically evaluates the OrderedDict getter on a bound C++ module, such
+    that new changes on the C++ side are picked up. Otherwise accessing e.g.
+    ``cpp_module._parameters`` just once would get a frozen copy of the parameters
+    at the time of access. ``torch.nn.Module`` accesses ``_parameters`` et al. via ``self.__dict__``
+    so using properties does not work.
+    """
+
+    def __init__(self, cpp_module, attr) -> None:
+        self.cpp_module = cpp_module
+        self.attr = attr
+
+    @property
+    def cpp_dict(self):
+        return getattr(self.cpp_module, self.attr)
+
+    # Magic methods cannot be assigned dynamically and bypass ``getattr``, so we
+    # must manually override them.
+
+    def items(self):
+        return self.cpp_dict.items()
+
+    def keys(self):
+        return self.cpp_dict.keys()
+
+    def values(self):
+        return self.cpp_dict.values()
+
+    def __iter__(self):
+        return self.cpp_dict.__iter__()
+
+    def __len__(self) -> int:
+        return self.cpp_dict.__len__()
+
+    def __contains__(self, key) -> bool:
+        return self.cpp_dict.__contains__(key)
+
+    def __getitem__(self, key):
+        return self.cpp_dict.__getitem__(key)
+
+
+class ModuleWrapper(nn.Module):
+    """A subclass of ``torch.nn.Module`` that wraps a C++ frontend module and delegates all access."""
+
+    def __init__(self, cpp_module) -> None:
+        # Assign before the super class constructor so ``self.training`` can be
+        # assigned to in the super class constructor.
+        self.cpp_module = cpp_module
+        super().__init__()
+        self._parameters = OrderedDictWrapper(cpp_module, "_parameters")  # type: ignore[assignment]
+        self._buffers: OrderedDictWrapper = OrderedDictWrapper(cpp_module, "_buffers")  # type: ignore[assignment]
+        self._modules: OrderedDictWrapper = OrderedDictWrapper(cpp_module, "_modules")  # type: ignore[assignment]
+        for attr in dir(cpp_module):
+            # Skip magic methods and the three attributes above.
+            if not attr.startswith("_"):
+                setattr(self, attr, getattr(self.cpp_module, attr))
+
+    def _apply(self, fn, recurse=True):
+        for param in self.parameters():
+            # Tensors stored in modules are graph leaves, and we don't
+            # want to create copy nodes, so we have to unpack the data.
+            param.data = fn(param.data)
+            if param._grad is not None:
+                param._grad.data = fn(param._grad.data)
+
+        for buf in self.buffers():
+            buf.data = fn(buf.data)
+
+        return self
+
+    # nn.Module defines training as a boolean
+    @property  # type: ignore[override]
+    # pyrefly: ignore [bad-override]
+    def training(self):
+        return self.cpp_module.training
+
+    @training.setter
+    def training(self, mode) -> None:
+        self.cpp_module.train(mode)
+
+    def __repr__(self) -> str:
+        return self.cpp_module.__repr__()
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/functional.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/functional.py
new file mode 100644
index 0000000000000000000000000000000000000000..0906e86ee623cc2ab7cb75f17b8d6446b4cd49c4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/functional.py
@@ -0,0 +1,6953 @@
+"""Functional interface."""
+
+import importlib
+import math
+import warnings
+from collections.abc import Callable
+from typing import Any as _Any, Optional, TYPE_CHECKING
+
+import torch
+from torch import _VF, sym_int as _sym_int, Tensor
+from torch._C import (
+    _add_docstr,
+    _infer_size,
+    _ScalingType as ScalingType,  # pyrefly: ignore [missing-module-attribute]
+    _SwizzleType as SwizzleType,  # pyrefly: ignore [missing-module-attribute]
+)
+from torch._jit_internal import (
+    _overload,
+    boolean_dispatch,
+    BroadcastingList1,
+    BroadcastingList2,  # pyrefly: ignore [missing-module-attribute]
+    BroadcastingList3,  # pyrefly: ignore [missing-module-attribute]
+)
+from torch._torch_docs import reproducibility_notes, sparse_support_notes, tf32_notes
+from torch.nn import _reduction as _Reduction, grad  # noqa: F401
+from torch.nn.modules.utils import _list_with_default, _pair, _single, _triple
+from torch.overrides import (
+    handle_torch_function,
+    has_torch_function,
+    has_torch_function_unary,
+    has_torch_function_variadic,
+)
+
+
+# Set visibility of the bound enums to this module
+ScalingType.__module__ = "torch.nn.functional"
+SwizzleType.__module__ = "torch.nn.functional"
+
+if TYPE_CHECKING:
+    from torch.types import _dtype as DType
+else:
+    # The JIT doesn't understand Union, nor torch.dtype here
+    DType = int
+
+try:
+    import numpy as np
+except ModuleNotFoundError:
+    np = None
+
+
+conv1d = _add_docstr(
+    torch.conv1d,
+    r"""
+conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor
+
+Applies a 1D convolution over an input signal composed of several input
+planes.
+
+{tf32_note}
+
+See :class:`~torch.nn.Conv1d` for details and output shape.
+
+Note:
+    {cudnn_reproducibility_note}
+
+Note:
+    This operator supports complex data types i.e. ``complex32, complex64, complex128``.
+""".format(**reproducibility_notes, **tf32_notes)
+    + r"""
+
+Args:
+    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`
+    weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kW)`
+    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: ``None``
+    stride: the stride of the convolving kernel. Can be a single number or
+      a one-element tuple `(sW,)`. Default: 1
+    padding: implicit paddings on both sides of the input. Can be a string {'valid', 'same'},
+      single number or a one-element tuple `(padW,)`. Default: 0
+      ``padding='valid'`` is the same as no padding. ``padding='same'`` pads
+      the input so the output has the same shape as the input. However, this mode
+      doesn't support any stride values other than 1.
+
+      .. warning::
+          For ``padding='same'``, if the ``weight`` is even-length and
+          ``dilation`` is odd in any dimension, a full :func:`pad` operation
+          may be needed internally. Lowering performance.
+    dilation: the spacing between kernel elements. Can be a single number or
+      a one-element tuple `(dW,)`. Default: 1
+    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by
+      the number of groups. Default: 1
+
+Examples::
+
+    >>> inputs = torch.randn(33, 16, 30)
+    >>> filters = torch.randn(20, 16, 5)
+    >>> F.conv1d(inputs, filters)
+""",
+)
+
+conv2d = _add_docstr(
+    torch.conv2d,
+    r"""
+conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor
+
+Applies a 2D convolution over an input image composed of several input
+planes.
+
+{tf32_note}
+
+See :class:`~torch.nn.Conv2d` for details and output shape.
+
+Note:
+    {cudnn_reproducibility_note}
+
+Note:
+    This operator supports complex data types i.e. ``complex32, complex64, complex128``.
+""".format(**reproducibility_notes, **tf32_notes)
+    + r"""
+
+Args:
+    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
+    weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kH , kW)`
+    bias: optional bias tensor of shape :math:`(\text{out\_channels})`. Default: ``None``
+    stride: the stride of the convolving kernel. Can be a single number or a
+      tuple `(sH, sW)`. Default: 1
+    padding: implicit paddings on both sides of the input. Can be a string {'valid', 'same'},
+      single number or a tuple `(padH, padW)`. Default: 0
+      ``padding='valid'`` is the same as no padding. ``padding='same'`` pads
+      the input so the output has the same shape as the input. However, this mode
+      doesn't support any stride values other than 1.
+
+      .. warning::
+          For ``padding='same'``, if the ``weight`` is even-length and
+          ``dilation`` is odd in any dimension, a full :func:`pad` operation
+          may be needed internally. Lowering performance.
+
+    dilation: the spacing between kernel elements. Can be a single number or
+      a tuple `(dH, dW)`. Default: 1
+    groups: split input into groups, both :math:`\text{in\_channels}` and :math:`\text{out\_channels}`
+      should be divisible by the number of groups. Default: 1
+
+Examples::
+
+    >>> # With square kernels and equal stride
+    >>> filters = torch.randn(8, 4, 3, 3)
+    >>> inputs = torch.randn(1, 4, 5, 5)
+    >>> F.conv2d(inputs, filters, padding=1)
+""",
+)  # noqa: E501
+
+conv3d = _add_docstr(
+    torch.conv3d,
+    r"""
+conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor
+
+Applies a 3D convolution over an input image composed of several input
+planes.
+
+{tf32_note}
+
+See :class:`~torch.nn.Conv3d` for details and output shape.
+
+Note:
+    {cudnn_reproducibility_note}
+
+Note:
+    This operator supports complex data types i.e. ``complex32, complex64, complex128``.
+""".format(**reproducibility_notes, **tf32_notes)
+    + r"""
+
+Args:
+    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iT , iH , iW)`
+    weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kT , kH , kW)`
+    bias: optional bias tensor of shape :math:`(\text{out\_channels})`. Default: None
+    stride: the stride of the convolving kernel. Can be a single number or a
+      tuple `(sT, sH, sW)`. Default: 1
+    padding: implicit paddings on both sides of the input. Can be a string {'valid', 'same'},
+      single number or a tuple `(padT, padH, padW)`. Default: 0
+      ``padding='valid'`` is the same as no padding. ``padding='same'`` pads
+      the input so the output has the same shape as the input. However, this mode
+      doesn't support any stride values other than 1.
+
+      .. warning::
+          For ``padding='same'``, if the ``weight`` is even-length and
+          ``dilation`` is odd in any dimension, a full :func:`pad` operation
+          may be needed internally. Lowering performance.
+
+    dilation: the spacing between kernel elements. Can be a single number or
+      a tuple `(dT, dH, dW)`. Default: 1
+    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by
+      the number of groups. Default: 1
+
+Examples::
+
+    >>> filters = torch.randn(33, 16, 3, 3, 3)
+    >>> inputs = torch.randn(20, 16, 50, 10, 20)
+    >>> F.conv3d(inputs, filters)
+""",
+)  # noqa: E501
+
+conv_transpose1d = _add_docstr(
+    torch.conv_transpose1d,
+    r"""
+conv_transpose1d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor
+
+Applies a 1D transposed convolution operator over an input signal
+composed of several input planes, sometimes also called "deconvolution".
+
+{tf32_note}
+
+See :class:`~torch.nn.ConvTranspose1d` for details and output shape.
+
+Note:
+    {cudnn_reproducibility_note}
+""".format(**reproducibility_notes, **tf32_notes)
+    + r"""
+
+Args:
+    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`
+    weight: filters of shape :math:`(\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kW)`
+    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: None
+    stride: the stride of the convolving kernel. Can be a single number or a
+      tuple ``(sW,)``. Default: 1
+    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both
+      sides of each dimension in the input. Can be a single number or a tuple
+      ``(padW,)``. Default: 0
+    output_padding: additional size added to one side of each dimension in the
+      output shape. Can be a single number or a tuple ``(out_padW)``. Default: 0
+    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the
+      number of groups. Default: 1
+    dilation: the spacing between kernel elements. Can be a single number or
+      a tuple ``(dW,)``. Default: 1
+
+Examples::
+
+    >>> inputs = torch.randn(20, 16, 50)
+    >>> weights = torch.randn(16, 33, 5)
+    >>> F.conv_transpose1d(inputs, weights)
+""",
+)
+
+conv_transpose2d = _add_docstr(
+    torch.conv_transpose2d,
+    r"""
+conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor
+
+Applies a 2D transposed convolution operator over an input image
+composed of several input planes, sometimes also called "deconvolution".
+
+{tf32_note}
+
+See :class:`~torch.nn.ConvTranspose2d` for details and output shape.
+
+Note:
+    {cudnn_reproducibility_note}
+""".format(**reproducibility_notes, **tf32_notes)
+    + r"""
+
+Args:
+    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
+    weight: filters of shape :math:`(\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kH , kW)`
+    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: None
+    stride: the stride of the convolving kernel. Can be a single number or a
+      tuple ``(sH, sW)``. Default: 1
+    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both
+      sides of each dimension in the input. Can be a single number or a tuple
+      ``(padH, padW)``. Default: 0
+    output_padding: additional size added to one side of each dimension in the
+      output shape. Can be a single number or a tuple ``(out_padH, out_padW)``.
+      Default: 0
+    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the
+      number of groups. Default: 1
+    dilation: the spacing between kernel elements. Can be a single number or
+      a tuple ``(dH, dW)``. Default: 1
+
+Examples::
+
+    >>> # With square kernels and equal stride
+    >>> inputs = torch.randn(1, 4, 5, 5)
+    >>> weights = torch.randn(4, 8, 3, 3)
+    >>> F.conv_transpose2d(inputs, weights, padding=1)
+""",
+)  # noqa: E501
+
+conv_transpose3d = _add_docstr(
+    torch.conv_transpose3d,
+    r"""
+conv_transpose3d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor
+
+Applies a 3D transposed convolution operator over an input image
+composed of several input planes, sometimes also called "deconvolution"
+
+{tf32_note}
+
+See :class:`~torch.nn.ConvTranspose3d` for details and output shape.
+
+Note:
+    {cudnn_reproducibility_note}
+""".format(**reproducibility_notes, **tf32_notes)
+    + r"""
+
+Args:
+    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iT , iH , iW)`
+    weight: filters of shape :math:`(\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kT , kH , kW)`
+    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: None
+    stride: the stride of the convolving kernel. Can be a single number or a
+      tuple ``(sT, sH, sW)``. Default: 1
+    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both
+      sides of each dimension in the input. Can be a single number or a tuple
+      ``(padT, padH, padW)``. Default: 0
+    output_padding: additional size added to one side of each dimension in the
+      output shape. Can be a single number or a tuple
+      ``(out_padT, out_padH, out_padW)``. Default: 0
+    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the
+      number of groups. Default: 1
+    dilation: the spacing between kernel elements. Can be a single number or
+      a tuple `(dT, dH, dW)`. Default: 1
+
+Examples::
+
+    >>> inputs = torch.randn(20, 16, 50, 10, 20)
+    >>> weights = torch.randn(16, 33, 3, 3, 3)
+    >>> F.conv_transpose3d(inputs, weights)
+""",
+)  # noqa: E501
+
+conv_tbc = _add_docstr(
+    torch.conv_tbc,
+    r"""
+Applies a 1-dimensional sequence convolution over an input sequence.
+Input and output dimensions are (Time, Batch, Channels) - hence TBC.
+
+Args:
+    input: input tensor of shape :math:`(\text{sequence length} \times batch \times \text{in\_channels})`
+    weight: filter of shape (:math:`\text{kernel width} \times \text{in\_channels} \times \text{out\_channels}`)
+    bias: bias of shape (:math:`\text{out\_channels}`)
+    pad: number of timesteps to pad. Default: 0
+""",
+)
+
+
+# Pooling
+avg_pool1d = _add_docstr(
+    torch.avg_pool1d,
+    r"""
+avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) -> Tensor
+
+Applies a 1D average pooling over an input signal composed of several
+input planes.
+
+See :class:`~torch.nn.AvgPool1d` for details and output shape.
+
+Args:
+    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`
+    kernel_size: the size of the window. Can be a single number or a
+      tuple `(kW,)`
+    stride: the stride of the window. Can be a single number or a tuple
+      `(sW,)`. Default: :attr:`kernel_size`
+    padding: implicit zero paddings on both sides of the input. Can be a single
+      number or a tuple `(padW,)`. Should be at most half of effective kernel
+      size, that is :math:`((kernelSize - 1) * dilation + 1) / 2`. Default: 0
+    ceil_mode: when True, will use `ceil` instead of `floor` to compute the
+        output shape. Default: ``False``
+    count_include_pad: when True, will include the zero-padding in the
+        averaging calculation. Default: ``True``
+
+Examples::
+
+    >>> # pool of square window of size=3, stride=2
+    >>> input = torch.tensor([[[1, 2, 3, 4, 5, 6, 7]]], dtype=torch.float32)
+    >>> F.avg_pool1d(input, kernel_size=3, stride=2)
+    tensor([[[ 2.,  4.,  6.]]])
+
+""",
+)
+
+
+avg_pool2d = _add_docstr(
+    torch._C._nn.avg_pool2d,
+    r"""
+avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -> Tensor
+
+Applies 2D average-pooling operation in :math:`kH \times kW` regions by step size
+:math:`sH \times sW` steps. The number of output features is equal to the number of
+input planes.
+
+See :class:`~torch.nn.AvgPool2d` for details and output shape.
+
+Args:
+    input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
+    kernel_size: size of the pooling region. Can be a single number, a single-element tuple or a
+      tuple `(kH, kW)`
+    stride: stride of the pooling operation. Can be a single number, a single-element tuple or a
+      tuple `(sH, sW)`. Default: :attr:`kernel_size`
+    padding: implicit zero paddings on both sides of the input. Can be a
+      single number, a single-element tuple or a tuple `(padH, padW)`.
+      Should be at most half of effective kernel size, that
+      is :math:`((kernelSize - 1) * dilation + 1) / 2`. Default: 0
+    ceil_mode: when True, will use `ceil` instead of `floor` in the formula
+        to compute the output shape. Default: ``False``
+    count_include_pad: when True, will include the zero-padding in the
+        averaging calculation. Default: ``True``
+    divisor_override: if specified, it will be used as divisor, otherwise
+         size of the pooling region will be used. Default: None
+""",
+)
+
+avg_pool3d = _add_docstr(
+    torch._C._nn.avg_pool3d,
+    r"""
+avg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -> Tensor
+
+Applies 3D average-pooling operation in :math:`kT \times kH \times kW` regions by step
+size :math:`sT \times sH \times sW` steps. The number of output features is equal to
+:math:`\lfloor\frac{\text{input planes}}{sT}\rfloor`.
+
+See :class:`~torch.nn.AvgPool3d` for details and output shape.
+
+Args:
+    input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iT \times iH , iW)`
+    kernel_size: size of the pooling region. Can be a single number or a
+      tuple `(kT, kH, kW)`
+    stride: stride of the pooling operation. Can be a single number or a
+      tuple `(sT, sH, sW)`. Default: :attr:`kernel_size`
+    padding: implicit zero paddings on both sides of the input. Can be a
+      single number or a tuple `(padT, padH, padW)`. Should be at most half
+      of effective kernel size, that is :math:`((kernelSize - 1) * dilation + 1) / 2`.
+      Default: 0
+    ceil_mode: when True, will use `ceil` instead of `floor` in the formula
+        to compute the output shape
+    count_include_pad: when True, will include the zero-padding in the
+        averaging calculation
+    divisor_override: if specified, it will be used as divisor, otherwise
+        size of the pooling region will be used. Default: None
+""",
+)
+
+
+def fractional_max_pool2d_with_indices(
+    input: Tensor,
+    kernel_size: BroadcastingList2[int],
+    output_size: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList2[float]] = None,  # noqa: UP045
+    return_indices: bool = False,
+    _random_samples: Tensor | None = None,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""
+    fractional_max_pool2d(input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None)
+
+    Applies 2D fractional max pooling over an input signal composed of several input planes.
+
+    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham
+
+    The max-pooling operation is applied in :math:`kH \times kW` regions by a stochastic
+    step size determined by the target output size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        kernel_size: the size of the window to take a max over.
+                     Can be a single number :math:`k` (for a square kernel of :math:`k \times k`)
+                     or a tuple `(kH, kW)`
+        output_size: the target output size of the image of the form :math:`oH \times oW`.
+                     Can be a tuple `(oH, oW)` or a single number :math:`oH` for a square image :math:`oH \times oH`
+        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.
+                      This has to be a number or tuple in the range (0, 1)
+        return_indices: if ``True``, will return the indices along with the outputs.
+                        Useful to pass to :func:`~torch.nn.functional.max_unpool2d`.
+
+    Examples::
+        >>> input = torch.randn(20, 16, 50, 32)
+        >>> # pool of square window of size=3, and target output size 13x12
+        >>> F.fractional_max_pool2d(input, 3, output_size=(13, 12))
+        >>> # pool of square window and target output size being half of input image size
+        >>> F.fractional_max_pool2d(input, 3, output_ratio=(0.5, 0.5))
+
+    .. _Fractional MaxPooling:
+        http://arxiv.org/abs/1412.6071
+    """
+    if has_torch_function_variadic(input, _random_samples):
+        return handle_torch_function(
+            fractional_max_pool2d_with_indices,
+            (input, _random_samples),
+            input,
+            kernel_size,
+            output_size=output_size,
+            output_ratio=output_ratio,
+            return_indices=return_indices,
+            _random_samples=_random_samples,
+        )
+    if output_size is None and output_ratio is None:
+        raise ValueError(
+            "fractional_max_pool2d requires specifying either an output_size or an output_ratio"
+        )
+    if output_size is None:
+        if output_ratio is None:
+            raise AssertionError("output_ratio is unexpectedly None")
+        if len(output_ratio) > 2:
+            raise ValueError(
+                "fractional_max_pool2d requires output_ratio to either be a single Int or tuple of Ints."
+            )
+        _output_ratio = _pair(output_ratio)
+        output_size = [
+            int(input.size(-2) * _output_ratio[0]),
+            int(input.size(-1) * _output_ratio[1]),
+        ]
+
+    if _random_samples is None:
+        n_batch = 1 if input.dim() == 3 else input.size(0)
+        _random_samples = torch.rand(
+            n_batch, input.size(-3), 2, dtype=input.dtype, device=input.device
+        )
+    return torch._C._nn.fractional_max_pool2d(
+        input, kernel_size, output_size, _random_samples
+    )
+
+
+def _fractional_max_pool2d(
+    input: Tensor,
+    kernel_size: BroadcastingList2[int],
+    output_size: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList2[float]] = None,  # noqa: UP045
+    return_indices: bool = False,
+    _random_samples: Tensor | None = None,
+) -> Tensor:
+    if has_torch_function_variadic(input, _random_samples):
+        return handle_torch_function(
+            fractional_max_pool2d,
+            (input, _random_samples),
+            input,
+            kernel_size,
+            output_size=output_size,
+            output_ratio=output_ratio,
+            return_indices=return_indices,
+            _random_samples=_random_samples,
+        )
+    return fractional_max_pool2d_with_indices(
+        input, kernel_size, output_size, output_ratio, return_indices, _random_samples
+    )[0]
+
+
+fractional_max_pool2d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=4,
+    default=False,
+    if_true=fractional_max_pool2d_with_indices,
+    if_false=_fractional_max_pool2d,
+    module_name=__name__,
+    func_name="fractional_max_pool2d",
+)
+
+
+def fractional_max_pool3d_with_indices(
+    input: Tensor,
+    kernel_size: BroadcastingList3[int],
+    output_size: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList3[float]] = None,  # noqa: UP045
+    return_indices: bool = False,
+    _random_samples: Tensor | None = None,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""
+    fractional_max_pool3d(input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None)
+
+    Applies 3D fractional max pooling over an input signal composed of several input planes.
+
+    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham
+
+    The max-pooling operation is applied in :math:`kT \times kH \times kW` regions by a stochastic
+    step size determined by the target output size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        kernel_size: the size of the window to take a max over.
+                     Can be a single number :math:`k` (for a square kernel of :math:`k \times k \times k`)
+                     or a tuple `(kT, kH, kW)`
+        output_size: the target output size of the form :math:`oT \times oH \times oW`.
+                     Can be a tuple `(oT, oH, oW)` or a single number :math:`oH` for a cubic output
+                     :math:`oH \times oH \times oH`
+        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.
+                      This has to be a number or tuple in the range (0, 1)
+        return_indices: if ``True``, will return the indices along with the outputs.
+                        Useful to pass to :func:`~torch.nn.functional.max_unpool3d`.
+
+    Shape:
+        - Input: :math:`(N, C, T_{in}, H_{in}, W_{in})` or :math:`(C, T_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, T_{out}, H_{out}, W_{out})` or :math:`(C, T_{out}, H_{out}, W_{out})`, where
+          :math:`(T_{out}, H_{out}, W_{out})=\text{output\_size}` or
+          :math:`(T_{out}, H_{out}, W_{out})=\text{output\_ratio} \times (T_{in}, H_{in}, W_{in})`
+
+    Examples::
+        >>> input = torch.randn(20, 16, 50, 32, 16)
+        >>> # pool of cubic window of size=3, and target output size 13x12x11
+        >>> F.fractional_max_pool3d(input, 3, output_size=(13, 12, 11))
+        >>> # pool of cubic window and target output size being half of input size
+        >>> F.fractional_max_pool3d(input, 3, output_ratio=(0.5, 0.5, 0.5))
+
+    .. _Fractional MaxPooling:
+        http://arxiv.org/abs/1412.6071
+    """
+    if has_torch_function_variadic(input, _random_samples):
+        return handle_torch_function(
+            fractional_max_pool3d_with_indices,
+            (input, _random_samples),
+            input,
+            kernel_size,
+            output_size=output_size,
+            output_ratio=output_ratio,
+            return_indices=return_indices,
+            _random_samples=_random_samples,
+        )
+    if output_size is None and output_ratio is None:
+        raise ValueError(
+            "fractional_max_pool3d requires specifying either an output_size or an output_ratio"
+        )
+    if output_size is None:
+        if output_ratio is None:
+            raise AssertionError("output_ratio is unexpectedly None")
+        _output_ratio = _triple(output_ratio)
+        output_size = [
+            int(input.size(-3) * _output_ratio[0]),
+            int(input.size(-2) * _output_ratio[1]),
+            int(input.size(-1) * _output_ratio[2]),
+        ]
+
+    if _random_samples is None:
+        n_batch = 1 if input.dim() == 4 else input.size(0)
+        _random_samples = torch.rand(
+            n_batch, input.size(-4), 3, dtype=input.dtype, device=input.device
+        )
+    return torch._C._nn.fractional_max_pool3d(
+        input, kernel_size, output_size, _random_samples
+    )
+
+
+def _fractional_max_pool3d(
+    input: Tensor,
+    kernel_size: BroadcastingList3[int],
+    output_size: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    output_ratio: Optional[BroadcastingList3[float]] = None,  # noqa: UP045
+    return_indices: bool = False,
+    _random_samples: Tensor | None = None,
+) -> Tensor:
+    if has_torch_function_variadic(input, _random_samples):
+        return handle_torch_function(
+            fractional_max_pool3d,
+            (input, _random_samples),
+            input,
+            kernel_size,
+            output_size=output_size,
+            output_ratio=output_ratio,
+            return_indices=return_indices,
+            _random_samples=_random_samples,
+        )
+    return fractional_max_pool3d_with_indices(
+        input, kernel_size, output_size, output_ratio, return_indices, _random_samples
+    )[0]
+
+
+fractional_max_pool3d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=4,
+    default=False,
+    if_true=fractional_max_pool3d_with_indices,
+    if_false=_fractional_max_pool3d,
+    module_name=__name__,
+    func_name="fractional_max_pool3d",
+)
+
+
+def max_pool1d_with_indices(
+    input: Tensor,
+    kernel_size: BroadcastingList1[int],
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
+    padding: BroadcastingList1[int] = 0,
+    dilation: BroadcastingList1[int] = 1,
+    ceil_mode: bool = False,
+    return_indices: bool = False,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""
+    max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)
+
+    Applies a 1D max pooling over an input signal composed of several input
+    planes.
+
+    .. note::
+        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from
+        what seen in :class:`~torch.nn.MaxPool1d`, and will change in a future release.
+
+    See :class:`~torch.nn.MaxPool1d` for details.
+
+    Args:
+        input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`, minibatch dim optional.
+        kernel_size: the size of the window. Can be a single number or a
+            tuple `(kW,)`
+        stride: the stride of the window. Can be a single number or a tuple
+            `(sW,)`. Default: :attr:`kernel_size`
+        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.
+        dilation: The stride between elements within a sliding window, must be > 0.
+        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This
+                   ensures that every element in the input tensor is covered by a sliding window.
+        return_indices: If ``True``, will return the argmax along with the max values.
+                        Useful for :class:`torch.nn.functional.max_unpool1d` later
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_pool1d_with_indices,
+            (input,),
+            input,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            ceil_mode=ceil_mode,
+            return_indices=return_indices,
+        )
+    if stride is None:
+        stride = torch.jit.annotate(list[int], [])
+    return torch.max_pool1d_with_indices(
+        input, kernel_size, stride, padding, dilation, ceil_mode
+    )
+
+
+def _max_pool1d(
+    input: Tensor,
+    kernel_size: BroadcastingList1[int],
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
+    padding: BroadcastingList1[int] = 0,
+    dilation: BroadcastingList1[int] = 1,
+    ceil_mode: bool = False,
+    return_indices: bool = False,
+) -> Tensor:
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_pool1d,
+            (input,),
+            input,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            ceil_mode=ceil_mode,
+            return_indices=return_indices,
+        )
+    if stride is None:
+        stride = torch.jit.annotate(list[int], [])
+    return torch.max_pool1d(input, kernel_size, stride, padding, dilation, ceil_mode)
+
+
+max_pool1d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=6,
+    default=False,
+    if_true=max_pool1d_with_indices,
+    if_false=_max_pool1d,
+    module_name=__name__,
+    func_name="max_pool1d",
+)
+
+
+def max_pool2d_with_indices(
+    input: Tensor,
+    kernel_size: BroadcastingList2[int],
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    padding: BroadcastingList2[int] = 0,
+    dilation: BroadcastingList2[int] = 1,
+    ceil_mode: bool = False,
+    return_indices: bool = False,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""
+    max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)
+
+    Applies a 2D max pooling over an input signal composed of several input
+    planes.
+
+    .. note::
+        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from
+        what seen in :class:`~torch.nn.MaxPool2d`, and will change in a future release.
+
+    See :class:`~torch.nn.MaxPool2d` for details.
+
+    Args:
+        input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`, minibatch dim optional.
+        kernel_size: size of the pooling region. Can be a single number or a
+            tuple `(kH, kW)`
+        stride: stride of the pooling operation. Can be a single number or a
+            tuple `(sH, sW)`. Default: :attr:`kernel_size`
+        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.
+        dilation: The stride between elements within a sliding window, must be > 0.
+        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This
+                   ensures that every element in the input tensor is covered by a sliding window.
+        return_indices: If ``True``, will return the argmax along with the max values.
+                        Useful for :class:`torch.nn.functional.max_unpool2d` later
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_pool2d_with_indices,
+            (input,),
+            input,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            ceil_mode=ceil_mode,
+            return_indices=return_indices,
+        )
+    if stride is None:
+        stride = torch.jit.annotate(list[int], [])
+    return torch._C._nn.max_pool2d_with_indices(
+        input, kernel_size, stride, padding, dilation, ceil_mode
+    )
+
+
+def _max_pool2d(
+    input: Tensor,
+    kernel_size: BroadcastingList2[int],
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    padding: BroadcastingList2[int] = 0,
+    dilation: BroadcastingList2[int] = 1,
+    ceil_mode: bool = False,
+    return_indices: bool = False,
+) -> Tensor:
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_pool2d,
+            (input,),
+            input,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            ceil_mode=ceil_mode,
+            return_indices=return_indices,
+        )
+    if stride is None:
+        stride = torch.jit.annotate(list[int], [])
+    return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
+
+
+max_pool2d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=6,
+    default=False,
+    if_true=max_pool2d_with_indices,
+    if_false=_max_pool2d,
+    module_name=__name__,
+    func_name="max_pool2d",
+)
+
+
+def max_pool3d_with_indices(
+    input: Tensor,
+    kernel_size: BroadcastingList3[int],
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    padding: BroadcastingList3[int] = 0,
+    dilation: BroadcastingList3[int] = 1,
+    ceil_mode: bool = False,
+    return_indices: bool = False,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""
+    max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)
+
+    Applies a 3D max pooling over an input signal composed of several input
+    planes.
+
+    .. note::
+        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from
+        what seen in :class:`~torch.nn.MaxPool3d`, and will change in a future release.
+
+    See :class:`~torch.nn.MaxPool3d` for details.
+
+    Args:
+        input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iD, iH , iW)`, minibatch dim optional.
+        kernel_size: size of the pooling region. Can be a single number or a
+                     tuple `(kT, kH, kW)`
+        stride: stride of the pooling operation. Can be a single number or a
+                tuple `(sT, sH, sW)`. Default: :attr:`kernel_size`
+        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.
+        dilation: The stride between elements within a sliding window, must be > 0.
+        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This
+                   ensures that every element in the input tensor is covered by a sliding window.
+        return_indices: If ``True``, will return the argmax along with the max values.
+                        Useful for :class:`torch.nn.functional.max_unpool3d` later
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_pool3d_with_indices,
+            (input,),
+            input,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            ceil_mode=ceil_mode,
+            return_indices=return_indices,
+        )
+    if stride is None:
+        stride = torch.jit.annotate(list[int], [])
+    return torch._C._nn.max_pool3d_with_indices(
+        input, kernel_size, stride, padding, dilation, ceil_mode
+    )
+
+
+def _max_pool3d(
+    input: Tensor,
+    kernel_size: BroadcastingList3[int],
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    padding: BroadcastingList3[int] = 0,
+    dilation: BroadcastingList3[int] = 1,
+    ceil_mode: bool = False,
+    return_indices: bool = False,
+) -> Tensor:
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_pool3d,
+            (input,),
+            input,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            ceil_mode=ceil_mode,
+            return_indices=return_indices,
+        )
+    if stride is None:
+        stride = torch.jit.annotate(list[int], [])
+    return torch.max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode)
+
+
+max_pool3d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=6,
+    default=False,
+    if_true=max_pool3d_with_indices,
+    if_false=_max_pool3d,
+    module_name=__name__,
+    func_name="max_pool3d",
+)
+
+
+def _unpool_output_size(
+    input: Tensor,
+    kernel_size: list[int],
+    stride: list[int],
+    padding: list[int],
+    output_size: list[int] | None,
+) -> list[int]:
+    input_size = input.size()
+    default_size = torch.jit.annotate(list[int], [])
+    for d in range(len(kernel_size)):
+        default_size.append(
+            (input_size[-len(kernel_size) + d] - 1) * stride[d]
+            + kernel_size[d]
+            - 2 * padding[d]
+        )
+    if output_size is None:
+        ret = default_size
+    else:
+        if len(output_size) == len(kernel_size) + 2:
+            output_size = output_size[2:]
+        if len(output_size) != len(kernel_size):
+            raise ValueError(
+                "output_size should be a sequence containing "
+                f"{len(kernel_size)} or {len(kernel_size) + 2} elements, but it has a length of '{len(output_size)}'"
+            )
+        for d in range(len(kernel_size)):
+            min_size = default_size[d] - stride[d]
+            max_size = default_size[d] + stride[d]
+            if not (min_size < output_size[d] < max_size):
+                raise ValueError(
+                    f'invalid output_size "{output_size}" (dim {d} must be between {min_size} and {max_size})'
+                )
+
+        ret = output_size
+    return ret
+
+
+def max_unpool1d(
+    input: Tensor,
+    indices: Tensor,
+    kernel_size: BroadcastingList1[int],
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
+    padding: BroadcastingList1[int] = 0,
+    output_size: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
+) -> Tensor:
+    r"""Compute a partial inverse of :class:`MaxPool1d`.
+
+    See :class:`~torch.nn.MaxUnpool1d` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_unpool1d,
+            (input,),
+            input,
+            indices,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            output_size=output_size,
+        )
+    kernel_size = _single(kernel_size)
+    if stride is not None:
+        _stride = _single(stride)
+    else:
+        _stride = kernel_size
+    padding = _single(padding)
+    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)
+    if isinstance(output_size, list):
+        output_size = output_size + [1]
+    else:
+        output_size = output_size + (1,)
+    return torch._C._nn.max_unpool2d(
+        input.unsqueeze(-1), indices.unsqueeze(-1), output_size
+    ).squeeze(-1)
+
+
+def max_unpool2d(
+    input: Tensor,
+    indices: Tensor,
+    kernel_size: BroadcastingList2[int],
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    padding: BroadcastingList2[int] = 0,
+    output_size: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+) -> Tensor:
+    r"""Compute a partial inverse of :class:`MaxPool2d`.
+
+    See :class:`~torch.nn.MaxUnpool2d` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_unpool2d,
+            (input,),
+            input,
+            indices,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            output_size=output_size,
+        )
+    kernel_size = _pair(kernel_size)
+    if stride is not None:
+        _stride = _pair(stride)
+    else:
+        _stride = kernel_size
+    padding = _pair(padding)
+    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)
+    return torch._C._nn.max_unpool2d(input, indices, output_size)
+
+
+def max_unpool3d(
+    input: Tensor,
+    indices: Tensor,
+    kernel_size: BroadcastingList3[int],
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    padding: BroadcastingList3[int] = 0,
+    output_size: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+) -> Tensor:
+    r"""Compute a partial inverse of :class:`MaxPool3d`.
+
+    See :class:`~torch.nn.MaxUnpool3d` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            max_unpool3d,
+            (input,),
+            input,
+            indices,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            output_size=output_size,
+        )
+    kernel_size = _triple(kernel_size)
+    if stride is not None:
+        _stride = _triple(stride)
+    else:
+        _stride = kernel_size
+    padding = _triple(padding)
+    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)
+    return torch._C._nn.max_unpool3d(input, indices, output_size, _stride, padding)
+
+
+def lp_pool3d(
+    input: Tensor,
+    norm_type: int | float,
+    kernel_size: BroadcastingList3[int],
+    stride: Optional[BroadcastingList3[int]] = None,  # noqa: UP045
+    ceil_mode: bool = False,
+) -> Tensor:
+    r"""
+    Apply a 3D power-average pooling over an input signal composed of several input planes.
+
+    If the sum of all inputs to the power of `p` is
+    zero, the gradient is set to zero as well.
+
+    When ``ceil_mode`` is ``True``, sliding windows may go off-bounds if they start within the left
+    padding or the input. Sliding windows that would start in the right padded region are ignored.
+
+    See :class:`~torch.nn.LPPool3d` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            lp_pool3d,
+            (input,),
+            input,
+            norm_type,
+            kernel_size,
+            stride=stride,
+            ceil_mode=ceil_mode,
+        )
+    kd, kw, kh = _triple(kernel_size)
+    if stride is not None:
+        out = avg_pool3d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)
+    else:
+        out = avg_pool3d(
+            input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode
+        )
+
+    return (
+        (torch.sign(out) * relu(torch.abs(out))).mul(kd * kw * kh).pow(1.0 / norm_type)
+    )
+
+
+def lp_pool2d(
+    input: Tensor,
+    norm_type: int | float,
+    kernel_size: BroadcastingList2[int],
+    stride: Optional[BroadcastingList2[int]] = None,  # noqa: UP045
+    ceil_mode: bool = False,
+) -> Tensor:
+    r"""
+    Apply a 2D power-average pooling over an input signal composed of several input planes.
+
+    If the sum of all inputs to the power of `p` is
+    zero, the gradient is set to zero as well.
+
+    When ``ceil_mode`` is ``True``, sliding windows may go off-bounds if they start within the left
+    padding or the input. Sliding windows that would start in the right padded region are ignored.
+
+    See :class:`~torch.nn.LPPool2d` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            lp_pool2d,
+            (input,),
+            input,
+            norm_type,
+            kernel_size,
+            stride=stride,
+            ceil_mode=ceil_mode,
+        )
+    kw, kh = _pair(kernel_size)
+    if stride is not None:
+        out = avg_pool2d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)
+    else:
+        out = avg_pool2d(
+            input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode
+        )
+
+    return (torch.sign(out) * relu(torch.abs(out))).mul(kw * kh).pow(1.0 / norm_type)
+
+
+def lp_pool1d(
+    input: Tensor,
+    norm_type: int | float,
+    kernel_size: int,
+    stride: Optional[BroadcastingList1[int]] = None,  # noqa: UP045
+    ceil_mode: bool = False,
+) -> Tensor:
+    r"""Apply a 1D power-average pooling over an input signal composed of several input planes.
+
+    If the sum of all inputs to the power of `p` is
+    zero, the gradient is set to zero as well.
+
+    When ``ceil_mode`` is ``True``, sliding windows may go off-bounds if they start within the left
+    padding or the input. Sliding windows that would start in the right padded region are ignored.
+
+    See :class:`~torch.nn.LPPool1d` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            lp_pool1d,
+            (input,),
+            input,
+            norm_type,
+            kernel_size,
+            stride=stride,
+            ceil_mode=ceil_mode,
+        )
+    if stride is not None:
+        out = avg_pool1d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)
+    else:
+        out = avg_pool1d(
+            input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode
+        )
+
+    return (
+        (torch.sign(out) * relu(torch.abs(out))).mul(kernel_size).pow(1.0 / norm_type)
+    )
+
+
+def adaptive_max_pool1d_with_indices(
+    input: Tensor,
+    output_size: BroadcastingList1[int],
+    return_indices: bool = False,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""
+    adaptive_max_pool1d(input, output_size, return_indices=False)
+
+    Applies a 1D adaptive max pooling over an input signal composed of
+    several input planes.
+
+    See :class:`~torch.nn.AdaptiveMaxPool1d` for details and output shape.
+
+    Args:
+        output_size: the target output size (single integer)
+        return_indices: whether to return pooling indices. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            adaptive_max_pool1d_with_indices,
+            (input,),
+            input,
+            output_size,
+            return_indices=return_indices,
+        )
+    return torch.adaptive_max_pool1d(input, output_size)
+
+
+def _adaptive_max_pool1d(
+    input: Tensor,
+    output_size: BroadcastingList1[int],
+    return_indices: bool = False,
+) -> Tensor:
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            adaptive_max_pool1d,
+            (input,),
+            input,
+            output_size,
+            return_indices=return_indices,
+        )
+    return adaptive_max_pool1d_with_indices(input, output_size)[0]
+
+
+adaptive_max_pool1d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=2,
+    default=False,
+    if_true=adaptive_max_pool1d_with_indices,
+    if_false=_adaptive_max_pool1d,
+    module_name=__name__,
+    func_name="adaptive_max_pool1d",
+)
+
+
+def adaptive_max_pool2d_with_indices(
+    input: Tensor,
+    output_size: BroadcastingList2[int],
+    return_indices: bool = False,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""adaptive_max_pool2d(input, output_size, return_indices=False)
+
+    Applies a 2D adaptive max pooling over an input signal composed of
+    several input planes.
+
+    See :class:`~torch.nn.AdaptiveMaxPool2d` for details and output shape.
+
+    Args:
+        output_size: the target output size (single integer or
+            double-integer tuple)
+        return_indices: whether to return pooling indices. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            adaptive_max_pool2d_with_indices,
+            (input,),
+            input,
+            output_size,
+            return_indices=return_indices,
+        )
+    # pyrefly: ignore [bad-argument-type]
+    output_size = _list_with_default(output_size, input.size())
+    return torch._C._nn.adaptive_max_pool2d(input, output_size)
+
+
+def _adaptive_max_pool2d(
+    input: Tensor,
+    output_size: BroadcastingList2[int],
+    return_indices: bool = False,
+) -> Tensor:
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            adaptive_max_pool2d,
+            (input,),
+            input,
+            output_size,
+            return_indices=return_indices,
+        )
+    return adaptive_max_pool2d_with_indices(input, output_size)[0]
+
+
+adaptive_max_pool2d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=2,
+    default=False,
+    if_true=adaptive_max_pool2d_with_indices,
+    if_false=_adaptive_max_pool2d,
+    module_name=__name__,
+    func_name="adaptive_max_pool2d",
+)
+
+
+def adaptive_max_pool3d_with_indices(
+    input: Tensor,
+    output_size: BroadcastingList3[int],
+    return_indices: bool = False,
+) -> tuple[Tensor, Tensor]:  # noqa: D400
+    r"""
+    adaptive_max_pool3d(input, output_size, return_indices=False)
+
+    Applies a 3D adaptive max pooling over an input signal composed of
+    several input planes.
+
+    See :class:`~torch.nn.AdaptiveMaxPool3d` for details and output shape.
+
+    Args:
+        output_size: the target output size (single integer or
+            triple-integer tuple)
+        return_indices: whether to return pooling indices. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            adaptive_max_pool3d_with_indices,
+            (input,),
+            input,
+            output_size,
+            return_indices=return_indices,
+        )
+    # pyrefly: ignore [bad-argument-type]
+    output_size = _list_with_default(output_size, input.size())
+    return torch._C._nn.adaptive_max_pool3d(input, output_size)
+
+
+def _adaptive_max_pool3d(
+    input: Tensor,
+    output_size: BroadcastingList3[int],
+    return_indices: bool = False,
+) -> Tensor:
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            adaptive_max_pool3d,
+            (input,),
+            input,
+            output_size,
+            return_indices=return_indices,
+        )
+    return adaptive_max_pool3d_with_indices(input, output_size)[0]
+
+
+adaptive_max_pool3d = boolean_dispatch(
+    arg_name="return_indices",
+    arg_index=2,
+    default=False,
+    if_true=adaptive_max_pool3d_with_indices,
+    if_false=_adaptive_max_pool3d,
+    module_name=__name__,
+    func_name="adaptive_max_pool3d",
+)
+
+
+adaptive_avg_pool1d = _add_docstr(
+    torch.adaptive_avg_pool1d,
+    r"""
+adaptive_avg_pool1d(input, output_size) -> Tensor
+
+Applies a 1D adaptive average pooling over an input signal composed of
+several input planes.
+
+See :class:`~torch.nn.AdaptiveAvgPool1d` for details and output shape.
+
+Args:
+    output_size: the target output size (single integer)
+""",
+)
+
+
+def adaptive_avg_pool2d(input: Tensor, output_size: BroadcastingList2[int]) -> Tensor:
+    r"""Apply a 2D adaptive average pooling over an input signal composed of several input planes.
+
+    See :class:`~torch.nn.AdaptiveAvgPool2d` for details and output shape.
+
+    Args:
+        output_size: the target output size (single integer or
+            double-integer tuple)
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(adaptive_avg_pool2d, (input,), input, output_size)
+    # pyrefly: ignore [bad-argument-type]
+    _output_size = _list_with_default(output_size, input.size())
+    return torch._C._nn.adaptive_avg_pool2d(input, _output_size)
+
+
+def adaptive_avg_pool3d(input: Tensor, output_size: BroadcastingList3[int]) -> Tensor:
+    r"""Apply a 3D adaptive average pooling over an input signal composed of several input planes.
+
+    See :class:`~torch.nn.AdaptiveAvgPool3d` for details and output shape.
+
+    Args:
+        output_size: the target output size (single integer or
+            triple-integer tuple)
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(adaptive_avg_pool3d, (input,), input, output_size)
+    # pyrefly: ignore [bad-argument-type]
+    _output_size = _list_with_default(output_size, input.size())
+    return torch._C._nn.adaptive_avg_pool3d(input, _output_size)
+
+
+# Activation functions
+def dropout(
+    input: Tensor,
+    p: float = 0.5,
+    training: bool = True,
+    inplace: bool = False,
+) -> Tensor:
+    r"""During training, randomly zeroes some elements of the input tensor with probability :attr:`p`.
+
+    Uses samples from a Bernoulli distribution.
+
+    See :class:`~torch.nn.Dropout` for details.
+
+    Args:
+        p: probability of an element to be zeroed. Default: 0.5
+        training: apply dropout if is ``True``. Default: ``True``
+        inplace: If set to ``True``, will do this operation in-place. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            dropout, (input,), input, p=p, training=training, inplace=inplace
+        )
+    if p < 0.0 or p > 1.0:
+        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+    return (
+        _VF.dropout_(input, p, training) if inplace else _VF.dropout(input, p, training)
+    )
+
+
+def alpha_dropout(
+    input: Tensor,
+    p: float = 0.5,
+    training: bool = False,
+    inplace: bool = False,
+) -> Tensor:
+    r"""Apply alpha dropout to the input.
+
+    See :class:`~torch.nn.AlphaDropout` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            alpha_dropout, (input,), input, p=p, training=training, inplace=inplace
+        )
+    if p < 0.0 or p > 1.0:
+        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+    return (
+        _VF.alpha_dropout_(input, p, training)
+        if inplace
+        else _VF.alpha_dropout(input, p, training)
+    )
+
+
+def dropout1d(
+    input: Tensor,
+    p: float = 0.5,
+    training: bool = True,
+    inplace: bool = False,
+) -> Tensor:
+    r"""Randomly zero out entire channels (a channel is a 1D feature map).
+
+    For example, the :math:`j`-th channel of the :math:`i`-th sample in the
+    batched input is a 1D tensor :math:`\text{input}[i, j]` of the input tensor.
+    Each channel will be zeroed out independently on every forward call with
+    probability :attr:`p` using samples from a Bernoulli distribution.
+
+    See :class:`~torch.nn.Dropout1d` for details.
+
+    Args:
+        p: probability of a channel to be zeroed. Default: 0.5
+        training: apply dropout if is ``True``. Default: ``True``
+        inplace: If set to ``True``, will do this operation in-place. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            dropout1d, (input,), input, p=p, training=training, inplace=inplace
+        )
+    if p < 0.0 or p > 1.0:
+        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+    inp_dim = input.dim()
+    if inp_dim not in (2, 3):
+        raise RuntimeError(
+            f"dropout1d: Expected 2D or 3D input, but received a {inp_dim}D input. "
+            "Note that dropout1d exists to provide channel-wise dropout on inputs with 1 "
+            "spatial dimension, a channel dimension, and an optional batch dimension "
+            "(i.e. 2D or 3D inputs)."
+        )
+
+    is_batched = inp_dim == 3
+    if not is_batched:
+        input = input.unsqueeze_(0) if inplace else input.unsqueeze(0)
+
+    result = (
+        _VF.feature_dropout_(input, p, training)
+        if inplace
+        else _VF.feature_dropout(input, p, training)
+    )
+
+    if not is_batched:
+        result = result.squeeze_(0) if inplace else result.squeeze(0)
+
+    return result
+
+
+def dropout2d(
+    input: Tensor,
+    p: float = 0.5,
+    training: bool = True,
+    inplace: bool = False,
+) -> Tensor:
+    r"""Randomly zero out entire channels (a channel is a 2D feature map).
+
+    For example, the :math:`j`-th channel of the :math:`i`-th sample in the
+    batched input is a 2D tensor :math:`\text{input}[i, j]` of the input tensor.
+    Each channel will be zeroed out independently on every forward call with
+    probability :attr:`p` using samples from a Bernoulli distribution.
+
+    See :class:`~torch.nn.Dropout2d` for details.
+
+    Args:
+        p: probability of a channel to be zeroed. Default: 0.5
+        training: apply dropout if is ``True``. Default: ``True``
+        inplace: If set to ``True``, will do this operation in-place. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            dropout2d, (input,), input, p=p, training=training, inplace=inplace
+        )
+    if p < 0.0 or p > 1.0:
+        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+    inp_dim = input.dim()
+    if inp_dim not in (3, 4):
+        warn_msg = (
+            f"dropout2d: Received a {inp_dim}-D input to dropout2d, which is deprecated "
+            "and will result in an error in a future release. To retain the behavior "
+            "and silence this warning, please use dropout instead. Note that dropout2d "
+            "exists to provide channel-wise dropout on inputs with 2 spatial dimensions, "
+            "a channel dimension, and an optional batch dimension (i.e. 3D or 4D inputs)."
+        )
+        warnings.warn(warn_msg, stacklevel=2)
+
+    # TODO: Properly support no-batch-dim inputs. For now, these are NOT supported; passing
+    # a 3D input will perform dropout1d behavior instead. This was done historically and the
+    # behavior is maintained here for now.
+    # See https://github.com/pytorch/pytorch/issues/77081
+    if inp_dim == 3:
+        warnings.warn(
+            "dropout2d: Received a 3D input to dropout2d and assuming that channel-wise "
+            "1D dropout behavior is desired - input is interpreted as shape (N, C, L), where C "
+            "is the channel dim. This behavior will change in a future release to interpret the "
+            "input as one without a batch dimension, i.e. shape (C, H, W). To maintain the 1D "
+            "channel-wise dropout behavior, please switch to using dropout1d instead.",
+            stacklevel=2,
+        )
+
+    result = (
+        _VF.feature_dropout_(input, p, training)
+        if inplace
+        else _VF.feature_dropout(input, p, training)
+    )
+
+    return result
+
+
+def dropout3d(
+    input: Tensor,
+    p: float = 0.5,
+    training: bool = True,
+    inplace: bool = False,
+) -> Tensor:
+    r"""Randomly zero out entire channels (a channel is a 3D feature map).
+
+    For example, the :math:`j`-th channel of the :math:`i`-th sample in the
+    batched input is a 3D tensor :math:`\text{input}[i, j]` of the input tensor.
+    Each channel will be zeroed out independently on every forward call with
+    probability :attr:`p` using samples from a Bernoulli distribution.
+
+    See :class:`~torch.nn.Dropout3d` for details.
+
+    Args:
+        p: probability of a channel to be zeroed. Default: 0.5
+        training: apply dropout if is ``True``. Default: ``True``
+        inplace: If set to ``True``, will do this operation in-place. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            dropout3d, (input,), input, p=p, training=training, inplace=inplace
+        )
+    if p < 0.0 or p > 1.0:
+        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+    inp_dim = input.dim()
+    if inp_dim not in (4, 5):
+        warn_msg = (
+            f"dropout3d: Received a {inp_dim}-D input to dropout3d, which is deprecated "
+            "and will result in an error in a future release. To retain the behavior "
+            "and silence this warning, please use dropout instead. Note that dropout3d "
+            "exists to provide channel-wise dropout on inputs with 3 spatial dimensions, "
+            "a channel dimension, and an optional batch dimension (i.e. 4D or 5D inputs)."
+        )
+        warnings.warn(warn_msg, stacklevel=2)
+
+    is_batched = inp_dim == 5
+    if not is_batched:
+        input = input.unsqueeze_(0) if inplace else input.unsqueeze(0)
+
+    result = (
+        _VF.feature_dropout_(input, p, training)
+        if inplace
+        else _VF.feature_dropout(input, p, training)
+    )
+
+    if not is_batched:
+        result = result.squeeze_(0) if inplace else result.squeeze(0)
+    return result
+
+
+def feature_alpha_dropout(
+    input: Tensor,
+    p: float = 0.5,
+    training: bool = False,
+    inplace: bool = False,
+) -> Tensor:
+    r"""Randomly masks out entire channels (a channel is a feature map).
+
+    For example, the :math:`j`-th channel of the :math:`i`-th sample in the batch input
+    is a tensor :math:`\text{input}[i, j]` of the input tensor. Instead of
+    setting activations to zero, as in regular Dropout, the activations are set
+    to the negative saturation value of the SELU activation function.
+
+    Each element will be masked independently on every forward call with
+    probability :attr:`p` using samples from a Bernoulli distribution.
+    The elements to be masked are randomized on every forward call, and scaled
+    and shifted to maintain zero mean and unit variance.
+
+    See :class:`~torch.nn.FeatureAlphaDropout` for details.
+
+    Args:
+        p: dropout probability of a channel to be zeroed. Default: 0.5
+        training: apply dropout if is ``True``. Default: ``True``
+        inplace: If set to ``True``, will do this operation in-place. Default: ``False``
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            feature_alpha_dropout,
+            (input,),
+            input,
+            p=p,
+            training=training,
+            inplace=inplace,
+        )
+    if p < 0.0 or p > 1.0:
+        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+    return (
+        _VF.feature_alpha_dropout_(input, p, training)
+        if inplace
+        else _VF.feature_alpha_dropout(input, p, training)
+    )
+
+
+def _threshold(
+    input: Tensor,
+    threshold: float,
+    value: float,
+    inplace: bool = False,
+) -> Tensor:
+    r"""Apply a threshold to each element of the input Tensor.
+
+    See :class:`~torch.nn.Threshold` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            _threshold, (input,), input, threshold, value, inplace=inplace
+        )
+    if inplace:
+        result = _VF.threshold_(input, threshold, value)
+    else:
+        result = _VF.threshold(input, threshold, value)
+    return result
+
+
+# We define this function as _threshold because it takes an argument
+# named threshold, which clobbers the recursive reference to the
+# function needed for __torch_function__ support
+threshold = _threshold
+
+threshold_ = _add_docstr(
+    _VF.threshold_,
+    r"""
+threshold_(input, threshold, value) -> Tensor
+
+In-place version of :func:`~threshold`.
+""",
+)
+
+
+def relu(input: Tensor, inplace: bool = False) -> Tensor:  # noqa: D400,D402
+    r"""relu(input, inplace=False) -> Tensor
+
+    Applies the rectified linear unit function element-wise. See
+    :class:`~torch.nn.ReLU` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(relu, (input,), input, inplace=inplace)
+    if inplace:
+        result = torch.relu_(input)
+    else:
+        result = torch.relu(input)
+    return result
+
+
+relu_ = _add_docstr(
+    torch.relu_,
+    r"""
+relu_(input) -> Tensor
+
+In-place version of :func:`~relu`.
+""",
+)
+
+
+def glu(input: Tensor, dim: int = -1) -> Tensor:  # noqa: D400,D402
+    r"""
+    glu(input, dim=-1) -> Tensor
+
+    The gated linear unit. Computes:
+
+    .. math ::
+        \text{GLU}(a, b) = a \otimes \sigma(b)
+
+    where `input` is split in half along `dim` to form `a` and `b`, :math:`\sigma`
+    is the sigmoid function and :math:`\otimes` is the element-wise product between matrices.
+
+    See `Language Modeling with Gated Convolutional Networks <https://arxiv.org/abs/1612.08083>`_.
+
+    Args:
+        input (Tensor): input tensor
+        dim (int): dimension on which to split the input. Default: -1
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(glu, (input,), input, dim=dim)
+    if input.dim() == 0:
+        raise RuntimeError(
+            "glu does not support scalars because halving size must be even"
+        )
+    return torch._C._nn.glu(input, dim)
+
+
+def hardtanh(
+    input: Tensor,
+    min_val: float = -1.0,
+    max_val: float = 1.0,
+    inplace: bool = False,
+) -> Tensor:  # noqa: D400,D402
+    r"""
+    hardtanh(input, min_val=-1., max_val=1., inplace=False) -> Tensor
+
+    Applies the HardTanh function element-wise. See :class:`~torch.nn.Hardtanh` for more
+    details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            hardtanh, (input,), input, min_val=min_val, max_val=max_val, inplace=inplace
+        )
+    if min_val > max_val:
+        raise ValueError("min_val cannot be greater than max_val")
+    if inplace:
+        result = torch._C._nn.hardtanh_(input, min_val, max_val)
+    else:
+        result = torch._C._nn.hardtanh(input, min_val, max_val)
+    return result
+
+
+hardtanh_ = _add_docstr(
+    torch._C._nn.hardtanh_,
+    r"""
+hardtanh_(input, min_val=-1., max_val=1.) -> Tensor
+
+In-place version of :func:`~hardtanh`.
+""",
+)
+
+
+def relu6(input: Tensor, inplace: bool = False) -> Tensor:  # noqa: D400,D402
+    r"""relu6(input, inplace=False) -> Tensor
+
+    Applies the element-wise function :math:`\text{ReLU6}(x) = \min(\max(0,x), 6)`.
+
+    See :class:`~torch.nn.ReLU6` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(relu6, (input,), input, inplace=inplace)
+    if inplace:
+        result = torch._C._nn.relu6_(input)
+    else:
+        result = torch._C._nn.relu6(input)
+    return result
+
+
+def elu(input: Tensor, alpha: float = 1.0, inplace: bool = False) -> Tensor:
+    r"""Apply the Exponential Linear Unit (ELU) function element-wise.
+
+    See :class:`~torch.nn.ELU` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(elu, (input,), input, alpha=alpha, inplace=inplace)
+    if inplace:
+        result = torch._C._nn.elu_(input, alpha)
+    else:
+        result = torch._C._nn.elu(input, alpha)
+    return result
+
+
+elu_ = _add_docstr(
+    torch._C._nn.elu_,
+    r"""
+elu_(input, alpha=1.) -> Tensor
+
+In-place version of :func:`~elu`.
+""",
+)
+
+
+def selu(input: Tensor, inplace: bool = False) -> Tensor:  # noqa: D400,D402
+    r"""selu(input, inplace=False) -> Tensor
+
+    Applies element-wise,
+    :math:`\text{SELU}(x) = scale * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))`,
+    with :math:`\alpha=1.6732632423543772848170429916717` and
+    :math:`scale=1.0507009873554804934193349852946`.
+
+    See :class:`~torch.nn.SELU` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(selu, (input,), input, inplace=inplace)
+    if inplace:
+        result = torch.selu_(input)
+    else:
+        result = torch.selu(input)
+    return result
+
+
+selu_ = _add_docstr(
+    torch.selu_,
+    r"""
+selu_(input) -> Tensor
+
+In-place version of :func:`~selu`.
+""",
+)
+
+
+def celu(
+    input: Tensor,
+    alpha: float = 1.0,
+    inplace: bool = False,
+) -> Tensor:  # noqa: D400,D402
+    r"""celu(input, alpha=1., inplace=False) -> Tensor
+
+    Applies element-wise,
+    :math:`\text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))`.
+
+    See :class:`~torch.nn.CELU` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            celu, (input,), input, alpha=alpha, inplace=inplace
+        )
+    if inplace:
+        result = torch.celu_(input, alpha)
+    else:
+        result = torch.celu(input, alpha)
+    return result
+
+
+celu_ = _add_docstr(
+    torch.celu_,
+    r"""
+celu_(input, alpha=1.) -> Tensor
+
+In-place version of :func:`~celu`.
+""",
+)
+
+
+def leaky_relu(
+    input: Tensor,
+    negative_slope: float = 0.01,
+    inplace: bool = False,
+) -> Tensor:  # noqa: D400,D402
+    r"""
+    leaky_relu(input, negative_slope=0.01, inplace=False) -> Tensor
+
+    Applies element-wise,
+    :math:`\text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)`
+
+    See :class:`~torch.nn.LeakyReLU` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            leaky_relu, (input,), input, negative_slope=negative_slope, inplace=inplace
+        )
+    if inplace:
+        result = torch._C._nn.leaky_relu_(input, negative_slope)
+    else:
+        result = torch._C._nn.leaky_relu(input, negative_slope)
+    return result
+
+
+leaky_relu_ = _add_docstr(
+    torch._C._nn.leaky_relu_,
+    r"""
+leaky_relu_(input, negative_slope=0.01) -> Tensor
+
+In-place version of :func:`~leaky_relu`.
+""",
+)
+
+
+prelu = _add_docstr(
+    torch.prelu,
+    r"""prelu(input, weight) -> Tensor
+
+Applies element-wise the function
+:math:`\text{PReLU}(x) = \max(0,x) + \text{weight} * \min(0,x)` where weight is a
+learnable parameter.
+
+.. note::
+    `weight` is expected to be a scalar or 1-D tensor. If `weight` is 1-D,
+    its size must match the number of input channels, determined by
+    `input.size(1)` when `input.dim() >= 2`, otherwise 1.
+    In the 1-D case, note that when `input` has dim > 2, `weight` can be expanded
+    to the shape of `input` in a way that is not possible using normal
+    :ref:`broadcasting semantics<broadcasting-semantics>`.
+
+See :class:`~torch.nn.PReLU` for more details.
+""",
+)
+
+
+def rrelu(
+    input: Tensor,
+    lower: float = 1.0 / 8,
+    upper: float = 1.0 / 3,
+    training: bool = False,
+    inplace: bool = False,
+) -> Tensor:  # noqa: D400,D402
+    r"""rrelu(input, lower=1./8, upper=1./3, training=False, inplace=False) -> Tensor
+
+    Randomized leaky ReLU.
+
+    See :class:`~torch.nn.RReLU` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            rrelu,
+            (input,),
+            input,
+            lower=lower,
+            upper=upper,
+            training=training,
+            inplace=inplace,
+        )
+    if inplace:
+        result = torch.rrelu_(input, lower, upper, training)
+    else:
+        result = torch.rrelu(input, lower, upper, training)
+    return result
+
+
+rrelu_ = _add_docstr(
+    torch.rrelu_,
+    r"""
+rrelu_(input, lower=1./8, upper=1./3, training=False) -> Tensor
+
+In-place version of :func:`~rrelu`.
+""",
+)
+
+logsigmoid = _add_docstr(
+    torch._C._nn.log_sigmoid,
+    r"""
+logsigmoid(input) -> Tensor
+
+Applies element-wise :math:`\text{LogSigmoid}(x_i) = \log \left(\frac{1}{1 + \exp(-x_i)}\right)`
+
+See :class:`~torch.nn.LogSigmoid` for more details.
+""",
+)
+
+gelu = _add_docstr(
+    torch._C._nn.gelu,
+    r"""
+gelu(input, approximate = 'none') -> Tensor
+
+When the approximate argument is 'none', it applies element-wise the function
+:math:`\text{GELU}(x) = x * \Phi(x)`
+
+where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.
+
+When the approximate argument is 'tanh', Gelu is estimated with
+
+.. math::
+    \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt{2 / \pi} * (x + 0.044715 * x^3)))
+
+See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_.
+""",
+)
+
+hardshrink = _add_docstr(
+    torch.hardshrink,
+    r"""
+hardshrink(input, lambd=0.5) -> Tensor
+
+Applies the hard shrinkage function element-wise
+
+See :class:`~torch.nn.Hardshrink` for more details.
+""",
+)
+
+
+def tanhshrink(input):  # noqa: D400,D402
+    r"""tanhshrink(input) -> Tensor
+
+    Applies element-wise, :math:`\text{Tanhshrink}(x) = x - \text{Tanh}(x)`
+
+    See :class:`~torch.nn.Tanhshrink` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(tanhshrink, (input,), input)
+    return input - input.tanh()
+
+
+def softsign(input):  # noqa: D400,D402
+    r"""softsign(input) -> Tensor
+
+    Applies element-wise, the function :math:`\text{SoftSign}(x) = \frac{x}{1 + |x|}`
+
+    See :class:`~torch.nn.Softsign` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(softsign, (input,), input)
+    return input / (input.abs() + 1)
+
+
+softplus = _add_docstr(
+    torch._C._nn.softplus,
+    r"""
+softplus(input, beta=1, threshold=20) -> Tensor
+
+Applies element-wise, the function :math:`\text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))`.
+
+For numerical stability the implementation reverts to the linear function
+when :math:`input \times \beta > threshold`.
+
+See :class:`~torch.nn.Softplus` for more details.
+""",
+)
+
+
+def _get_softmax_dim(name: str, ndim: int, stacklevel: int) -> int:
+    warnings.warn(
+        f"Implicit dimension choice for {name} has been deprecated. "
+        "Change the call to include dim=X as an argument.",
+        stacklevel=stacklevel,
+    )
+    if ndim == 0 or ndim == 1 or ndim == 3:
+        ret = 0
+    else:
+        ret = 1
+    return ret
+
+
+def softmin(
+    input: Tensor,
+    dim: int | None = None,
+    _stacklevel: int = 3,
+    dtype: DType | None = None,
+) -> Tensor:
+    r"""Apply a softmin function.
+
+    Note that :math:`\text{Softmin}(x) = \text{Softmax}(-x)`. See softmax definition for mathematical formula.
+
+    See :class:`~torch.nn.Softmin` for more details.
+
+    Args:
+        input (Tensor): input
+        dim (int): A dimension along which softmin will be computed (so every slice
+            along dim will sum to 1).
+        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.
+          If specified, the input tensor is casted to :attr:`dtype` before the operation
+          is performed. This is useful for preventing data type overflows. Default: None.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            softmin, (input,), input, dim=dim, _stacklevel=_stacklevel, dtype=dtype
+        )
+    if dim is None:
+        dim = _get_softmax_dim("softmin", input.dim(), _stacklevel)
+    if dtype is None:
+        ret = (-input).softmax(dim)
+    else:
+        ret = (-input).softmax(dim, dtype=dtype)
+    return ret
+
+
+def softmax(
+    input: Tensor,
+    dim: int | None = None,
+    _stacklevel: int = 3,
+    dtype: DType | None = None,
+) -> Tensor:
+    r"""Apply a softmax function.
+
+    Softmax is defined as:
+
+    :math:`\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}`
+
+    It is applied to all slices along dim, and will re-scale them so that the elements
+    lie in the range `[0, 1]` and sum to 1.
+
+    See :class:`~torch.nn.Softmax` for more details.
+
+    Args:
+        input (Tensor): input
+        dim (int): A dimension along which softmax will be computed.
+        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.
+          If specified, the input tensor is casted to :attr:`dtype` before the operation
+          is performed. This is useful for preventing data type overflows. Default: None.
+
+    .. note::
+        This function doesn't work directly with NLLLoss,
+        which expects the Log to be computed between the Softmax and itself.
+        Use log_softmax instead (it's faster and has better numerical properties).
+
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            softmax, (input,), input, dim=dim, _stacklevel=_stacklevel, dtype=dtype
+        )
+    if dim is None:
+        dim = _get_softmax_dim("softmax", input.dim(), _stacklevel)
+    if dtype is None:
+        ret = input.softmax(dim)
+    else:
+        ret = input.softmax(dim, dtype=dtype)
+    return ret
+
+
+def gumbel_softmax(
+    logits: Tensor,
+    tau: float = 1,
+    hard: bool = False,
+    eps: float = 1e-10,
+    dim: int = -1,
+) -> Tensor:
+    r"""
+    Sample from the Gumbel-Softmax distribution (`Link 1`_  `Link 2`_) and optionally discretize.
+
+    Args:
+      logits: `[..., num_features]` unnormalized log probabilities
+      tau: non-negative scalar temperature
+      hard: if ``True``, the returned samples will be discretized as one-hot vectors,
+            but will be differentiated as if it is the soft sample in autograd
+      dim (int): A dimension along which softmax will be computed. Default: -1.
+
+    Returns:
+      Sampled tensor of same shape as `logits` from the Gumbel-Softmax distribution.
+      If ``hard=True``, the returned samples will be one-hot, otherwise they will
+      be probability distributions that sum to 1 across `dim`.
+
+    .. note::
+      This function is here for legacy reasons, may be removed from nn.Functional in the future.
+
+    .. note::
+      The main trick for `hard` is to do  `y_hard - y_soft.detach() + y_soft`
+
+      It achieves two things:
+      - makes the output value exactly one-hot
+      (since we add then subtract y_soft value)
+      - makes the gradient equal to y_soft gradient
+      (since we strip all other gradients)
+
+    Examples::
+        >>> logits = torch.randn(20, 32)
+        >>> # Sample soft categorical using reparametrization trick:
+        >>> F.gumbel_softmax(logits, tau=1, hard=False)
+        >>> # Sample hard categorical using "Straight-through" trick:
+        >>> F.gumbel_softmax(logits, tau=1, hard=True)
+
+    .. _Link 1:
+        https://arxiv.org/abs/1611.00712
+    .. _Link 2:
+        https://arxiv.org/abs/1611.01144
+    """
+    if has_torch_function_unary(logits):
+        return handle_torch_function(
+            gumbel_softmax, (logits,), logits, tau=tau, hard=hard, eps=eps, dim=dim
+        )
+    if eps != 1e-10:
+        warnings.warn("`eps` parameter is deprecated and has no effect.", stacklevel=2)
+
+    gumbels = (
+        -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format)
+        .exponential_()
+        .log()
+    )  # ~Gumbel(0,1)
+    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
+    y_soft = gumbels.softmax(dim)
+
+    if hard:
+        # Straight through.
+        index = y_soft.max(dim, keepdim=True)[1]
+        y_hard = torch.zeros_like(
+            logits, memory_format=torch.legacy_contiguous_format
+        ).scatter_(dim, index, 1.0)
+        ret = y_hard - y_soft.detach() + y_soft
+    else:
+        # Reparameterization trick.
+        ret = y_soft
+    return ret
+
+
+def log_softmax(
+    input: Tensor,
+    dim: int | None = None,
+    _stacklevel: int = 3,
+    dtype: DType | None = None,
+) -> Tensor:
+    r"""Apply a softmax followed by a logarithm.
+
+    While mathematically equivalent to log(softmax(x)), doing these two
+    operations separately is slower and numerically unstable. This function
+    uses an alternative formulation to compute the output and gradient correctly.
+
+    See :class:`~torch.nn.LogSoftmax` for more details.
+
+    Args:
+        input (Tensor): input
+        dim (int): A dimension along which log_softmax will be computed.
+        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.
+          If specified, the input tensor is cast to :attr:`dtype` before the operation
+          is performed. This is useful for preventing data type overflows. Default: None.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            log_softmax, (input,), input, dim=dim, _stacklevel=_stacklevel, dtype=dtype
+        )
+    if dim is None:
+        dim = _get_softmax_dim("log_softmax", input.dim(), _stacklevel)
+    if dtype is None:
+        ret = input.log_softmax(dim)
+    else:
+        ret = input.log_softmax(dim, dtype=dtype)
+    return ret
+
+
+softshrink = _add_docstr(
+    torch._C._nn.softshrink,
+    r"""
+softshrink(input, lambd=0.5) -> Tensor
+
+Applies the soft shrinkage function elementwise
+
+See :class:`~torch.nn.Softshrink` for more details.
+""",
+)
+
+
+def tanh(input):  # noqa: D400,D402
+    r"""tanh(input) -> Tensor
+
+    Applies element-wise,
+    :math:`\text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)}{\exp(x) + \exp(-x)}`
+
+    See :class:`~torch.nn.Tanh` for more details.
+    """
+    return input.tanh()
+
+
+def sigmoid(input):  # noqa: D400,D402
+    r"""sigmoid(input) -> Tensor
+
+    Applies the element-wise function :math:`\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}`
+
+    See :class:`~torch.nn.Sigmoid` for more details.
+    """
+    return input.sigmoid()
+
+
+def hardsigmoid(input: Tensor, inplace: bool = False) -> Tensor:
+    r"""Apply the Hardsigmoid function element-wise.
+
+    .. math::
+        \text{Hardsigmoid}(x) = \begin{cases}
+            0 & \text{if~} x \le -3, \\
+            1 & \text{if~} x \ge +3, \\
+            x / 6 + 1 / 2 & \text{otherwise}
+        \end{cases}
+
+    Args:
+        inplace: If set to ``True``, will do this operation in-place. Default: ``False``
+
+    See :class:`~torch.nn.Hardsigmoid` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(hardsigmoid, (input,), input, inplace=inplace)
+    if inplace:
+        return torch._C._nn.hardsigmoid_(input)
+    return torch._C._nn.hardsigmoid(input)
+
+
+linear = _add_docstr(
+    torch._C._nn.linear,
+    r"""
+linear(input, weight, bias=None) -> Tensor
+
+Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.
+
+This operation supports 2-D :attr:`weight` with :ref:`sparse layout<sparse-docs>`
+
+{sparse_beta_warning}
+
+This operator supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+Shape:
+
+    - Input: :math:`(*, in\_features)` where `*` means any number of
+      additional dimensions, including none
+    - Weight: :math:`(out\_features, in\_features)` or :math:`(in\_features)`
+    - Bias: :math:`(out\_features)` or :math:`()`
+    - Output: :math:`(*, out\_features)` or :math:`(*)`, based on the shape of the weight
+""".format(**sparse_support_notes),
+)
+
+
+bilinear = _add_docstr(
+    torch.bilinear,
+    r"""
+bilinear(input1, input2, weight, bias=None) -> Tensor
+
+Applies a bilinear transformation to the incoming data:
+:math:`y = x_1^T A x_2 + b`
+
+Shape:
+
+    - input1: :math:`(N, *, H_{in1})` where :math:`H_{in1}=\text{in1\_features}`
+      and :math:`*` means any number of additional dimensions.
+      All but the last dimension of the inputs should be the same.
+    - input2: :math:`(N, *, H_{in2})` where :math:`H_{in2}=\text{in2\_features}`
+    - weight: :math:`(\text{out\_features}, \text{in1\_features},
+      \text{in2\_features})`
+    - bias: :math:`(\text{out\_features})`
+    - output: :math:`(N, *, H_{out})` where :math:`H_{out}=\text{out\_features}`
+      and all but the last dimension are the same shape as the input.
+""",
+)
+
+
+def silu(input: Tensor, inplace: bool = False) -> Tensor:
+    r"""Apply the Sigmoid Linear Unit (SiLU) function, element-wise.
+
+    The SiLU function is also known as the swish function.
+
+    .. math::
+        \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}
+
+    .. note::
+        See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_
+        where the SiLU (Sigmoid Linear Unit) was originally coined, and see
+        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation
+        in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:
+        a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_
+        where the SiLU was experimented with later.
+
+    See :class:`~torch.nn.SiLU` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(silu, (input,), input, inplace=inplace)
+    if inplace:
+        return torch._C._nn.silu_(input)
+    return torch._C._nn.silu(input)
+
+
+def mish(input: Tensor, inplace: bool = False) -> Tensor:
+    r"""Apply the Mish function, element-wise.
+
+    Mish: A Self Regularized Non-Monotonic Neural Activation Function.
+
+    .. math::
+        \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
+
+    .. note::
+        See `Mish: A Self Regularized Non-Monotonic Neural Activation Function <https://arxiv.org/abs/1908.08681>`_
+
+    See :class:`~torch.nn.Mish` for more details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(mish, (input,), input, inplace=inplace)
+    if inplace:
+        return torch._C._nn.mish_(input)
+    return torch._C._nn.mish(input)
+
+
+def hardswish(input: Tensor, inplace: bool = False) -> Tensor:
+    r"""Apply hardswish function, element-wise.
+
+    Follows implementation as described in the paper:
+    `Searching for MobileNetV3`_.
+
+    .. math::
+        \text{Hardswish}(x) = \begin{cases}
+            0 & \text{if~} x \le -3, \\
+            x & \text{if~} x \ge +3, \\
+            x \cdot (x + 3) /6 & \text{otherwise}
+        \end{cases}
+
+    See :class:`~torch.nn.Hardswish` for more details.
+
+    .. _`Searching for MobileNetV3`:
+        https://arxiv.org/abs/1905.02244
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(hardswish, (input,), input, inplace=inplace)
+    if inplace:
+        return torch._C._nn.hardswish_(input)
+    return torch._C._nn.hardswish(input)
+
+
+def _no_grad_embedding_renorm_(
+    weight: Tensor,
+    input: Tensor,
+    max_norm: float,
+    norm_type: float,
+    # pyrefly: ignore [bad-return]
+) -> tuple[Tensor, Tensor]:
+    torch.embedding_renorm_(weight.detach(), input, max_norm, norm_type)
+
+
+def embedding(
+    input: Tensor,
+    weight: Tensor,
+    padding_idx: int | None = None,
+    max_norm: float | None = None,
+    norm_type: float = 2.0,
+    scale_grad_by_freq: bool = False,
+    sparse: bool = False,
+) -> Tensor:
+    r"""Generate a simple lookup table that looks up embeddings in a fixed dictionary and size.
+
+    This module is often used to retrieve word embeddings using indices.
+    The input to the module is a list of indices, and the embedding matrix,
+    and the output is the corresponding word embeddings.
+
+    See :class:`torch.nn.Embedding` for more details.
+
+    .. note::
+        Note that the analytical gradients of this function with respect to
+        entries in :attr:`weight` at the row specified by :attr:`padding_idx`
+        are expected to differ from the numerical ones.
+
+    .. note::
+        Note that `:class:`torch.nn.Embedding` differs from this function in
+        that it initializes the row of :attr:`weight` specified by
+        :attr:`padding_idx` to all zeros on construction.
+
+    Args:
+        input (LongTensor): Tensor containing indices into the embedding matrix
+        weight (Tensor): The embedding matrix (must be 2-D) with number of rows equal to the maximum possible index + 1,
+            and number of columns equal to the embedding size
+        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;
+                                     therefore, the embedding vector at :attr:`padding_idx` is not updated during training,
+                                     i.e. it remains as a fixed "pad".
+        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`
+                                    is renormalized to have norm :attr:`max_norm`.
+                                    Note: this will modify :attr:`weight` in-place.
+        norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.
+        scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse of frequency of
+                                                the words in the mini-batch. Default ``False``.
+        sparse (bool, optional): If ``True``, gradient w.r.t. :attr:`weight` will be a sparse tensor. See Notes under
+                                 :class:`torch.nn.Embedding` for more details regarding sparse gradients.
+
+    Shape:
+        - Input: LongTensor of arbitrary shape containing the indices to extract
+        - Weight: Embedding matrix of floating point type with shape `(V, embedding_dim)`,
+          where V = maximum index + 1 and embedding_dim = the embedding size
+        - Output: `(*, embedding_dim)`, where `*` is the input shape
+
+    Examples::
+
+        >>> # a batch of 2 samples of 4 indices each
+        >>> input = torch.tensor([[1, 2, 4, 5], [4, 3, 2, 9]])
+        >>> # an embedding matrix containing 10 tensors of size 3
+        >>> embedding_matrix = torch.rand(10, 3)
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> F.embedding(input, embedding_matrix)
+        tensor([[[ 0.8490,  0.9625,  0.6753],
+                 [ 0.9666,  0.7761,  0.6108],
+                 [ 0.6246,  0.9751,  0.3618],
+                 [ 0.4161,  0.2419,  0.7383]],
+
+                [[ 0.6246,  0.9751,  0.3618],
+                 [ 0.0237,  0.7794,  0.0528],
+                 [ 0.9666,  0.7761,  0.6108],
+                 [ 0.3385,  0.8612,  0.1867]]])
+
+        >>> # example with padding_idx
+        >>> weights = torch.rand(10, 3)
+        >>> weights[0, :].zero_()
+        >>> embedding_matrix = weights
+        >>> input = torch.tensor([[0, 2, 0, 5]])
+        >>> F.embedding(input, embedding_matrix, padding_idx=0)
+        tensor([[[ 0.0000,  0.0000,  0.0000],
+                 [ 0.5609,  0.5384,  0.8720],
+                 [ 0.0000,  0.0000,  0.0000],
+                 [ 0.6262,  0.2438,  0.7471]]])
+    """
+    if has_torch_function_variadic(input, weight):
+        return handle_torch_function(
+            embedding,
+            (input, weight),
+            input,
+            weight,
+            padding_idx=padding_idx,
+            max_norm=max_norm,
+            norm_type=norm_type,
+            scale_grad_by_freq=scale_grad_by_freq,
+            sparse=sparse,
+        )
+    if padding_idx is not None:
+        if padding_idx > 0:
+            if padding_idx >= weight.size(0):
+                raise AssertionError("Padding_idx must be within num_embeddings")
+        elif padding_idx < 0:
+            if padding_idx < -weight.size(0):
+                raise AssertionError("Padding_idx must be within num_embeddings")
+            padding_idx = weight.size(0) + padding_idx
+    else:
+        padding_idx = -1
+    if max_norm is not None:
+        # Note [embedding_renorm contiguous]
+        # `embedding_renorm_` will call .contiguous() on input anyways, so we
+        # call it here and take advantage of the improved locality in the
+        # `embedding` call below too.
+        input = input.contiguous()
+        # Note [embedding_renorm set_grad_enabled]
+        # XXX: equivalent to
+        # with torch.no_grad():
+        #   torch.embedding_renorm_
+        # remove once script supports set_grad_enabled
+        _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
+    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
+
+
+def embedding_bag(
+    input: Tensor,
+    weight: Tensor,
+    offsets: Tensor | None = None,
+    max_norm: float | None = None,
+    norm_type: float = 2,
+    scale_grad_by_freq: bool = False,
+    mode: str = "mean",
+    sparse: bool = False,
+    per_sample_weights: Tensor | None = None,
+    include_last_offset: bool = False,
+    padding_idx: int | None = None,
+) -> Tensor:
+    r"""Compute sums, means or maxes of `bags` of embeddings.
+
+    Calculation is done without instantiating the intermediate embeddings.
+    See :class:`torch.nn.EmbeddingBag` for more details.
+
+    Note:
+        {backward_reproducibility_note}
+
+    Args:
+        input (LongTensor): Tensor containing bags of indices into the embedding matrix
+        weight (Tensor): The embedding matrix with number of rows equal to the maximum possible index + 1,
+            and number of columns equal to the embedding size
+        offsets (LongTensor, optional): Only used when :attr:`input` is 1D. :attr:`offsets` determines
+                             the starting index position of each bag (sequence) in :attr:`input`.
+        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`
+                                    is renormalized to have norm :attr:`max_norm`.
+                                    Note: this will modify :attr:`weight` in-place.
+        norm_type (float, optional): The ``p`` in the ``p``-norm to compute for the :attr:`max_norm` option.
+                                     Default ``2``.
+        scale_grad_by_freq (bool, optional): if given, this will scale gradients by the inverse of frequency of
+                                                the words in the mini-batch. Default ``False``.
+                                                Note: this option is not supported when ``mode="max"``.
+        mode (str, optional): ``"sum"``, ``"mean"`` or ``"max"``. Specifies the way to reduce the bag.
+                                 Default: ``"mean"``
+        sparse (bool, optional): if ``True``, gradient w.r.t. :attr:`weight` will be a sparse tensor. See Notes under
+                                 :class:`torch.nn.Embedding` for more details regarding sparse gradients.
+                                 Note: this option is not supported when ``mode="max"``.
+        per_sample_weights (Tensor, optional): a tensor of float / double weights, or None
+            to indicate all weights should be taken to be 1. If specified, :attr:`per_sample_weights`
+            must have exactly the same shape as input and is treated as having the same
+            :attr:`offsets`, if those are not None.
+
+        include_last_offset (bool, optional): if ``True``, the size of offsets is equal to the number of bags + 1.
+                                              The last element is the size of the input, or the ending index position
+                                              of the last bag (sequence). This matches the CSR format. Ignored when
+                                              input is 2D. Default ``False``.
+
+        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the
+                                     gradient; therefore, the embedding vector at :attr:`padding_idx` is not updated
+                                     during training, i.e. it remains as a fixed "pad". Note that the embedding
+                                     vector at :attr:`padding_idx` is excluded from the reduction.
+
+    Shape:
+        - :attr:`input` (LongTensor) and :attr:`offsets` (LongTensor, optional)
+
+          - If :attr:`input` is 2D of shape `(B, N)`, it will be treated as ``B`` bags (sequences)
+            each of fixed length ``N``, and this will return ``B`` values aggregated in a way
+            depending on the :attr:`mode`. :attr:`offsets` is ignored and required to be ``None`` in this case.
+
+          - If :attr:`input` is 1D of shape `(N)`, it will be treated as a concatenation of
+            multiple bags (sequences). :attr:`offsets` is required to be a 1D tensor containing
+            the starting index positions of each bag in :attr:`input`. Therefore, for :attr:`offsets`
+            of shape `(B)`, :attr:`input` will be viewed as having ``B`` bags.
+            Empty bags (i.e., having 0-length) will have returned vectors filled by zeros.
+
+        - :attr:`weight` (Tensor): the learnable weights of the module of shape `(num_embeddings, embedding_dim)`
+
+        - :attr:`per_sample_weights` (Tensor, optional). Has the same shape as :attr:`input`.
+
+        - :attr:`output`: aggregated embedding values of shape `(B, embedding_dim)`
+
+    Examples::
+
+        >>> # an Embedding module containing 10 tensors of size 3
+        >>> embedding_matrix = torch.rand(10, 3)
+        >>> # a batch of 2 samples of 4 indices each
+        >>> input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9])
+        >>> offsets = torch.tensor([0, 4])
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> F.embedding_bag(input, embedding_matrix, offsets)
+        tensor([[ 0.3397,  0.3552,  0.5545],
+                [ 0.5893,  0.4386,  0.5882]])
+
+        >>> # example with padding_idx
+        >>> embedding_matrix = torch.rand(10, 3)
+        >>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9])
+        >>> offsets = torch.tensor([0, 4])
+        >>> F.embedding_bag(input, embedding_matrix, offsets, padding_idx=2, mode='sum')
+        tensor([[ 0.0000,  0.0000,  0.0000],
+                [-0.7082,  3.2145, -2.6251]])
+    """
+    if has_torch_function_variadic(input, weight, offsets, per_sample_weights):
+        return handle_torch_function(
+            embedding_bag,
+            (input, weight, offsets, per_sample_weights),
+            input,
+            weight,
+            offsets=offsets,
+            max_norm=max_norm,
+            norm_type=norm_type,
+            scale_grad_by_freq=scale_grad_by_freq,
+            mode=mode,
+            sparse=sparse,
+            per_sample_weights=per_sample_weights,
+            include_last_offset=include_last_offset,
+            padding_idx=padding_idx,
+        )
+    # Check for backward compatibility.
+    # Used to be embedding_bag(weight, input, ...)
+    # Now is     embedding_bag(input, weight, ...)
+    if weight.dtype == torch.long and input.is_floating_point():
+        warnings.warn(
+            "Argument order of nn.functional.embedding_bag was changed. "
+            "Usage `embedding_bag(weight, input, ...)` is deprecated, "
+            "and should now be `embedding_bag(input, weight, ...)`.",
+            stacklevel=2,
+        )
+        weight, input = input, weight
+
+    if per_sample_weights is not None and input.size() != per_sample_weights.size():
+        raise ValueError(
+            f"embedding_bag: If per_sample_weights ({per_sample_weights.shape}) is not None, "
+            f"then it must have the same shape as the input ({input.shape})"
+        )
+
+    if not weight.dim() == 2:
+        raise ValueError(
+            f"weight has to be a 2D Tensor, but got Tensor of dimension {weight.dim()}"
+        )
+
+    if not torch.jit.is_scripting() and input.dim() == 2 and input.is_nested:
+        include_last_offset = True
+        # pyrefly: ignore [missing-attribute]
+        offsets = input.offsets()
+        input = input.values().reshape(-1)
+        if per_sample_weights is not None:
+            if not per_sample_weights.is_nested:
+                raise ValueError(
+                    "If input is nested, then per_sample_weights must be nested if specified"
+                )
+            per_sample_weights = per_sample_weights.values().reshape(-1)
+    elif input.dim() == 2:
+        if offsets is not None:
+            type_str = "<unknown>"
+            # TODO: Remove this once script supports type() calls
+            if not torch.jit.is_scripting():
+                type_str = str(type(offsets))
+            raise ValueError(
+                "if input is 2D, then offsets has to be None"
+                ", as input is treated is a mini-batch of"
+                " fixed length sequences. However, found "
+                f"offsets of type {type_str}"
+            )
+        offsets = torch.arange(
+            0, input.numel(), input.size(1), dtype=input.dtype, device=input.device
+        )
+        include_last_offset = False
+        input = input.reshape(-1)
+        if per_sample_weights is not None:
+            per_sample_weights = per_sample_weights.reshape(-1)
+    elif input.dim() == 1:
+        if offsets is None:
+            raise ValueError("offsets has to be a 1D Tensor but got None")
+        if offsets.dim() != 1:
+            raise ValueError("offsets has to be a 1D Tensor")
+    else:
+        raise ValueError(
+            f"input has to be 1D or 2D Tensor, but got Tensor of dimension {input.dim()}"
+        )
+    if mode == "sum":
+        mode_enum = 0
+    elif mode == "mean":
+        mode_enum = 1
+    elif mode == "max":
+        mode_enum = 2
+
+        if scale_grad_by_freq:
+            raise ValueError(
+                "max mode does not support scaling the gradient by the frequency"
+            )
+
+        if sparse:
+            raise ValueError("max mode does not support sparse weights")
+
+    else:
+        raise ValueError("mode has to be one of sum, mean or max")
+
+    if max_norm is not None:
+        # XXX: equivalent to
+        # with torch.no_grad():
+        #   torch.nembedding_renorm_
+        # remove once script supports set_grad_enabled
+        _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
+
+    if per_sample_weights is not None and mode != "sum":
+        raise NotImplementedError(
+            "embedding_bag: per_sample_weights was not None. "
+            "per_sample_weights is only supported for mode='sum' "
+            f"(got mode='{mode}'). Please open a feature request on GitHub."
+        )
+
+    ret, _, _, _ = torch.embedding_bag(
+        weight,
+        input,
+        offsets,
+        scale_grad_by_freq,
+        mode_enum,
+        sparse,
+        per_sample_weights,
+        include_last_offset,
+        padding_idx,
+    )
+    return ret
+
+
+if embedding_bag.__doc__:
+    embedding_bag.__doc__ = embedding_bag.__doc__.format(**reproducibility_notes)
+
+
+def _verify_batch_size(size: list[int]) -> None:
+    # XXX: JIT script does not support the reduce from functools, and mul op is a
+    # builtin, which cannot be used as a value to a func yet, so rewrite this size
+    # check to a simple equivalent for loop
+    #
+    # TODO: make use of reduce like below when JIT is ready with the missing features:
+    # from operator import mul
+    # from functools import reduce
+    #
+    #   if reduce(mul, size[2:], size[0]) == 1
+    size_prods = size[0]
+    for i in range(len(size) - 2):
+        size_prods *= size[i + 2]
+    if size_prods == 1:
+        raise ValueError(
+            f"Expected more than 1 value per channel when training, got input size {size}"
+        )
+
+
+def batch_norm(
+    input: Tensor,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    training: bool = False,
+    momentum: float = 0.1,
+    eps: float = 1e-5,
+) -> Tensor:
+    r"""Apply Batch Normalization for each channel across a batch of data.
+
+    See :class:`~torch.nn.BatchNorm1d`, :class:`~torch.nn.BatchNorm2d`,
+    :class:`~torch.nn.BatchNorm3d` for details.
+    """
+    if has_torch_function_variadic(input, running_mean, running_var, weight, bias):
+        return handle_torch_function(
+            batch_norm,
+            (input, running_mean, running_var, weight, bias),
+            input,
+            running_mean,
+            running_var,
+            weight=weight,
+            bias=bias,
+            training=training,
+            momentum=momentum,
+            eps=eps,
+        )
+    if training:
+        # pyrefly: ignore [bad-argument-type]
+        _verify_batch_size(input.size())
+
+    if training and eps <= 0.0:
+        raise ValueError(
+            f"batch_norm eps must be positive during training, but got {eps}"
+        )
+    elif eps < 0.0:
+        raise ValueError(f"batch_norm eps must be non-negative, but got {eps}")
+
+    return torch.batch_norm(
+        input,
+        weight,
+        bias,
+        running_mean,
+        running_var,
+        training,
+        momentum,
+        eps,
+        torch.backends.cudnn.enabled,
+    )
+
+
+def _verify_spatial_size(size: list[int]) -> None:
+    # Verify that there is > 1 spatial element for instance norm calculation.
+    size_prods = 1
+    for i in range(2, len(size)):
+        size_prods *= size[i]
+    if size_prods == 1:
+        raise ValueError(
+            f"Expected more than 1 spatial element when training, got input size {size}"
+        )
+
+
+def instance_norm(
+    input: Tensor,
+    running_mean: Tensor | None = None,
+    running_var: Tensor | None = None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    use_input_stats: bool = True,
+    momentum: float = 0.1,
+    eps: float = 1e-5,
+) -> Tensor:
+    r"""Apply Instance Normalization independently for each channel in every data sample within a batch.
+
+    See :class:`~torch.nn.InstanceNorm1d`, :class:`~torch.nn.InstanceNorm2d`,
+    :class:`~torch.nn.InstanceNorm3d` for details.
+    """
+    if has_torch_function_variadic(input, running_mean, running_var, weight, bias):
+        return handle_torch_function(
+            instance_norm,
+            (input, running_mean, running_var, weight, bias),
+            input,
+            running_mean=running_mean,
+            running_var=running_var,
+            weight=weight,
+            bias=bias,
+            use_input_stats=use_input_stats,
+            momentum=momentum,
+            eps=eps,
+        )
+    if use_input_stats:
+        # pyrefly: ignore [bad-argument-type]
+        _verify_spatial_size(input.size())
+    return torch.instance_norm(
+        input,
+        weight,
+        bias,
+        running_mean,
+        running_var,
+        use_input_stats,
+        momentum,
+        eps,
+        torch.backends.cudnn.enabled,
+    )
+
+
+def layer_norm(
+    input: Tensor,
+    normalized_shape: list[int],
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    eps: float = 1e-5,
+) -> Tensor:
+    r"""Apply Layer Normalization for last certain number of dimensions.
+
+    See :class:`~torch.nn.LayerNorm` for details.
+    """
+    if has_torch_function_variadic(input, weight, bias):
+        return handle_torch_function(
+            layer_norm,
+            (input, weight, bias),
+            input,
+            normalized_shape,
+            weight=weight,
+            bias=bias,
+            eps=eps,
+        )
+    return torch.layer_norm(
+        input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled
+    )
+
+
+def rms_norm(
+    input: Tensor,
+    normalized_shape: list[int],
+    weight: Tensor | None = None,
+    eps: float | None = None,
+) -> Tensor:
+    r"""Apply Root Mean Square Layer Normalization.
+
+    See :class:`~torch.nn.RMSNorm` for details.
+    """
+    if has_torch_function_variadic(input, weight):
+        return handle_torch_function(
+            rms_norm, (input, weight), input, normalized_shape, weight=weight, eps=eps
+        )
+    return torch.rms_norm(input, normalized_shape, weight, eps)
+
+
+def group_norm(
+    input: Tensor,
+    num_groups: int,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    eps: float = 1e-5,
+) -> Tensor:
+    r"""Apply Group Normalization for last certain number of dimensions.
+
+    See :class:`~torch.nn.GroupNorm` for details.
+    """
+    if has_torch_function_variadic(input, weight, bias):
+        return handle_torch_function(
+            group_norm,
+            (
+                input,
+                weight,
+                bias,
+            ),
+            input,
+            num_groups,
+            weight=weight,
+            bias=bias,
+            eps=eps,
+        )
+    if input.dim() < 2:
+        raise RuntimeError(
+            f"Expected at least 2 dimensions for input tensor but received {input.dim()}"
+        )
+    _verify_batch_size(
+        [input.size(0) * input.size(1) // num_groups, num_groups]
+        + list(input.size()[2:])
+    )
+    return torch.group_norm(
+        input, num_groups, weight, bias, eps, torch.backends.cudnn.enabled
+    )
+
+
+def local_response_norm(
+    input: Tensor,
+    size: int,
+    alpha: float = 1e-4,
+    beta: float = 0.75,
+    k: float = 1.0,
+) -> Tensor:
+    r"""Apply local response normalization over an input signal.
+
+    The input signal is composed of several input planes, where channels occupy the second dimension.
+    Normalization is applied across channels.
+
+    See :class:`~torch.nn.LocalResponseNorm` for details.
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            local_response_norm, (input,), input, size, alpha=alpha, beta=beta, k=k
+        )
+    dim = input.dim()
+    if dim < 3:
+        raise ValueError(
+            f"Expected 3D or higher dimensionality                          input (got {dim} dimensions)"
+        )
+
+    if input.numel() == 0:
+        return input
+
+    div = input.mul(input)
+    if dim == 3:
+        div = div.unsqueeze(1)
+        # pyrefly: ignore [bad-argument-type]
+        div = pad(div, (0, 0, size // 2, (size - 1) // 2))
+        div = avg_pool2d(div, (size, 1), stride=1).squeeze(1)
+    else:
+        sizes = input.size()
+        div = div.view(sizes[0], 1, sizes[1], sizes[2], -1)
+        # pyrefly: ignore [bad-argument-type]
+        div = pad(div, (0, 0, 0, 0, size // 2, (size - 1) // 2))
+        div = avg_pool3d(div, (size, 1, 1), stride=1).squeeze(1)
+        div = div.view(sizes)
+    div = div.mul(alpha).add(k).pow(beta)
+    return input / div
+
+
+# loss
+
+
+def ctc_loss(
+    log_probs: Tensor,
+    targets: Tensor,
+    input_lengths: Tensor,
+    target_lengths: Tensor,
+    blank: int = 0,
+    reduction: str = "mean",
+    zero_infinity: bool = False,
+) -> Tensor:
+    r"""Compute the Connectionist Temporal Classification loss.
+
+    See :class:`~torch.nn.CTCLoss` for details.
+
+    Note:
+        {cudnn_reproducibility_note}
+
+    Note:
+        {backward_reproducibility_note}
+
+    Args:
+        log_probs: :math:`(T, N, C)` or :math:`(T, C)` where `C = number of characters in alphabet including blank`,
+            `T = input length`, and `N = batch size`.
+            The logarithmized probabilities of the outputs
+            (e.g. obtained with :func:`torch.nn.functional.log_softmax`).
+        targets: :math:`(N, S)` or `(sum(target_lengths))`.
+                May be an empty tensor if all entries in `target_lengths` are zero.
+                In the second form, the targets are assumed to be concatenated.
+        input_lengths: :math:`(N)` or :math:`()`.
+            Lengths of the inputs (must each be :math:`\leq T`)
+        target_lengths: :math:`(N)` or :math:`()`.
+            Lengths of the targets
+        blank (int, optional):
+            Blank label. Default :math:`0`.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the output losses will be divided by the target lengths and
+            then the mean over the batch is taken, ``'sum'``: the output will be
+            summed. Default: ``'mean'``
+        zero_infinity (bool, optional):
+            Whether to zero infinite losses and the associated gradients.
+            Default: ``False``
+            Infinite losses mainly occur when the inputs are too short
+            to be aligned to the targets.
+
+    Example::
+
+        >>> log_probs = torch.randn(50, 16, 20).log_softmax(2).detach().requires_grad_()
+        >>> targets = torch.randint(1, 20, (16, 30), dtype=torch.long)
+        >>> input_lengths = torch.full((16,), 50, dtype=torch.long)
+        >>> target_lengths = torch.randint(10, 30, (16,), dtype=torch.long)
+        >>> loss = F.ctc_loss(log_probs, targets, input_lengths, target_lengths)
+        >>> loss.backward()
+    """
+    if has_torch_function_variadic(log_probs, targets, input_lengths, target_lengths):
+        return handle_torch_function(
+            ctc_loss,
+            (log_probs, targets, input_lengths, target_lengths),
+            log_probs,
+            targets,
+            input_lengths,
+            target_lengths,
+            blank=blank,
+            reduction=reduction,
+            zero_infinity=zero_infinity,
+        )
+    return torch.ctc_loss(
+        log_probs,
+        targets,
+        input_lengths,
+        target_lengths,
+        blank,
+        _Reduction.get_enum(reduction),
+        zero_infinity,
+    )
+
+
+if ctc_loss.__doc__:
+    ctc_loss.__doc__ = ctc_loss.__doc__.format(**reproducibility_notes)
+
+
+def nll_loss(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    ignore_index: int = -100,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:
+    r"""Compute the negative log likelihood loss.
+
+    See :class:`~torch.nn.NLLLoss` for details.
+
+    Args:
+        input: :math:`(N, C)` where `C = number of classes` or :math:`(N, C, H, W)`
+            in case of 2D Loss, or :math:`(N, C, d_1, d_2, ..., d_K)` where :math:`K \geq 1`
+            in the case of K-dimensional loss. `input` is expected to be log-probabilities.
+        target: :math:`(N)` where each value is :math:`0 \leq \text{targets}[i] \leq C-1`,
+            or :math:`(N, d_1, d_2, ..., d_K)` where :math:`K \geq 1` for
+            K-dimensional loss.
+        weight (Tensor, optional): A manual rescaling weight given to each
+            class. If given, has to be a Tensor of size `C`
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        ignore_index (int, optional): Specifies a target value that is ignored
+            and does not contribute to the input gradient. When :attr:`size_average` is
+            ``True``, the loss is averaged over non-ignored targets. Default: -100
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Example::
+
+        >>> # input is of size N x C = 3 x 5
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> # each element in target has to have 0 <= value < C
+        >>> target = torch.tensor([1, 0, 4])
+        >>> output = F.nll_loss(F.log_softmax(input, dim=1), target)
+        >>> output.backward()
+    """
+    if has_torch_function_variadic(input, target, weight):
+        return handle_torch_function(
+            nll_loss,
+            (input, target, weight),
+            input,
+            target,
+            weight=weight,
+            size_average=size_average,
+            ignore_index=ignore_index,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction = _Reduction.legacy_get_string(size_average, reduce)
+    return torch._C._nn.nll_loss_nd(
+        input,
+        target,
+        weight,
+        # pyrefly: ignore [bad-argument-type]
+        _Reduction.get_enum(reduction),
+        ignore_index,
+    )
+
+
+def poisson_nll_loss(
+    input: Tensor,
+    target: Tensor,
+    log_input: bool = True,
+    full: bool = False,
+    size_average: bool | None = None,
+    eps: float = 1e-8,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:
+    r"""Compute the Poisson negative log likelihood loss.
+
+    See :class:`~torch.nn.PoissonNLLLoss` for details.
+
+    Args:
+        input: Expectation of underlying Poisson distribution.
+        target: Random sample :math:`target \sim \text{Poisson}(input)`.
+        log_input: If ``True`` the loss is computed as
+            :math:`\exp(\text{input}) - \text{target} * \text{input}`, if ``False`` then loss is
+            :math:`\text{input} - \text{target} * \log(\text{input}+\text{eps})`. Default: ``True``
+        full: Whether to compute full loss, i. e. to add the Stirling
+            approximation term. Default: ``False``
+            :math:`\text{target} * \log(\text{target}) - \text{target} + 0.5 * \log(2 * \pi * \text{target})`.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        eps (float, optional): Small value to avoid evaluation of :math:`\log(0)` when
+            :attr:`log_input`\ =\ ``False``. Default: 1e-8
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    """
+    if has_torch_function_variadic(input, target):
+        return handle_torch_function(
+            poisson_nll_loss,
+            (input, target),
+            input,
+            target,
+            log_input=log_input,
+            full=full,
+            size_average=size_average,
+            eps=eps,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction = _Reduction.legacy_get_string(size_average, reduce)
+    if reduction != "none" and reduction != "mean" and reduction != "sum":
+        ret = input
+        raise ValueError(reduction + " is not a valid value for reduction")
+
+    ret = torch.poisson_nll_loss(
+        input, target, log_input, full, eps, _Reduction.get_enum(reduction)
+    )
+    return ret
+
+
+def gaussian_nll_loss(
+    input: Tensor,
+    target: Tensor,
+    var: Tensor | float,
+    full: bool = False,
+    eps: float = 1e-6,
+    reduction: str = "mean",
+) -> Tensor:
+    r"""Compute the Gaussian negative log likelihood loss.
+
+    See :class:`~torch.nn.GaussianNLLLoss` for details.
+
+    Args:
+        input: Expectation of the Gaussian distribution.
+        target: Sample from the Gaussian distribution.
+        var: Tensor of positive variance(s), one for each of the expectations
+            in the input (heteroscedastic), or a single one (homoscedastic),
+            or a positive scalar value to be used for all expectations.
+        full (bool, optional): Whether to include the constant term in the loss calculation. Default: ``False``.
+        eps (float, optional): Value added to var, for stability. Default: 1e-6.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the output is the average of all batch member losses,
+            ``'sum'``: the output is the sum of all batch member losses.
+            Default: ``'mean'``.
+    """
+    if has_torch_function_variadic(input, target, var):
+        return handle_torch_function(
+            gaussian_nll_loss,
+            (input, target, var),
+            input,
+            target,
+            var,
+            full=full,
+            eps=eps,
+            reduction=reduction,
+        )
+
+    # Entries of var must be non-negative
+    if isinstance(var, float):
+        if var < 0:
+            raise ValueError("var has negative entry/entries")
+        var = var * torch.ones_like(input)
+    elif torch.any(var < 0):
+        raise ValueError("var has negative entry/entries")
+
+    # Check var size
+    # If var.size == input.size, the case is heteroscedastic and no further checks are needed.
+    # Otherwise:
+    if var.size() != input.size():
+        # If var is one dimension short of input, but the sizes match otherwise, then this is a homoscedastic case.
+        # e.g. input.size = (10, 2, 3), var.size = (10, 2)
+        # -> unsqueeze var so that var.shape = (10, 2, 1)
+        # this is done so that broadcasting can happen in the loss calculation
+        if input.size()[:-1] == var.size():
+            var = torch.unsqueeze(var, -1)
+
+        # This checks if the var is broadcastable to the input and there is only one mismatched dimension.
+        # This is also a homoscedastic case.
+        # e.g. input.size = (10, 2, 3), var.size = (10, 2, 1)
+        # or  input.size = (4, 3, 32, 32), var.size = (4, 1, 32, 32)
+        elif (
+            input.ndim == var.ndim
+            and sum(y for x, y in zip(input.size(), var.size(), strict=True) if x != y)
+            == 1
+        ):  # Heteroscedastic case
+            pass
+
+        # If none of the above pass, then the size of var is incorrect.
+        else:
+            raise ValueError("var is of incorrect size")
+
+    # Check validity of reduction mode
+    if reduction != "none" and reduction != "mean" and reduction != "sum":
+        raise ValueError(reduction + " is not valid")
+
+    # Clamp for stability
+    var = var.clone()
+    with torch.no_grad():
+        var.clamp_(min=eps)
+
+    # Calculate the loss
+    loss = 0.5 * (torch.log(var) + (input - target) ** 2 / var)
+    if full:
+        loss += 0.5 * math.log(2 * math.pi)
+
+    if reduction == "mean":
+        return loss.mean()
+    elif reduction == "sum":
+        return loss.sum()
+    else:
+        return loss
+
+
+def kl_div(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+    log_target: bool = False,
+) -> Tensor:
+    r"""Compute the KL Divergence loss.
+
+    Refer - The `Kullback-Leibler divergence Loss
+    <https://en.wikipedia.org/wiki/Kullback-Leibler_divergence>`__
+
+    See :class:`~torch.nn.KLDivLoss` for details.
+
+    Args:
+        input: Tensor of arbitrary shape in log-probabilities.
+        target: Tensor of the same shape as input. See :attr:`log_target` for
+            the target's interpretation.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'batchmean'`` | ``'sum'`` | ``'mean'``.
+            ``'none'``: no reduction will be applied
+            ``'batchmean'``: the sum of the output will be divided by the batchsize
+            ``'sum'``: the output will be summed
+            ``'mean'``: the output will be divided by the number of elements in the output
+            Default: ``'mean'``
+        log_target (bool): A flag indicating whether ``target`` is passed in the log space.
+            It is recommended to pass certain distributions (like ``softmax``)
+            in the log space to avoid numerical issues caused by explicit ``log``.
+            Default: ``False``
+
+    .. note::
+        :attr:`size_average` and :attr:`reduce` are in the process of being deprecated,
+        and in the meantime, specifying either of those two args will override :attr:`reduction`.
+
+    .. warning::
+        :attr:`reduction` = ``'mean'`` doesn't return the true kl divergence value, please use
+        :attr:`reduction` = ``'batchmean'`` which aligns with KL math definition.
+    """
+    if has_torch_function_variadic(input, target):
+        return handle_torch_function(
+            kl_div,
+            (input, target),
+            input,
+            target,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+            log_target=log_target,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        if reduction == "mean":
+            warnings.warn(
+                "reduction: 'mean' divides the total loss by both the batch size and the support size."
+                "'batchmean' divides only by the batch size, and aligns with the KL div math definition."
+                "'mean' will be changed to behave the same as 'batchmean' in the next major release.",
+                stacklevel=2,
+            )
+
+        # special case for batchmean
+        if reduction == "batchmean":
+            reduction_enum = _Reduction.get_enum("sum")
+        else:
+            reduction_enum = _Reduction.get_enum(reduction)
+
+    reduced = torch.kl_div(input, target, reduction_enum, log_target=log_target)
+
+    if reduction == "batchmean" and input.dim() != 0:
+        reduced = reduced / input.size()[0]
+
+    return reduced
+
+
+def cross_entropy(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    ignore_index: int = -100,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+    label_smoothing: float = 0.0,
+) -> Tensor:
+    r"""Compute the cross entropy loss between input logits and target.
+
+    See :class:`~torch.nn.CrossEntropyLoss` for details.
+
+    Args:
+        input (Tensor) : Predicted unnormalized logits;
+            see Shape section below for supported shapes.
+        target (Tensor) : Ground truth class indices or class probabilities;
+            see Shape section below for supported shapes.
+        weight (Tensor, optional): a manual rescaling weight given to each
+            class. If given, has to be a Tensor of size `C`
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        ignore_index (int, optional): Specifies a target value that is ignored
+            and does not contribute to the input gradient. When :attr:`size_average` is
+            ``True``, the loss is averaged over non-ignored targets. Note that
+            :attr:`ignore_index` is only applicable when the target contains class indices.
+            Default: -100
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+        label_smoothing (float, optional): A float in [0.0, 1.0]. Specifies the amount
+            of smoothing when computing the loss, where 0.0 means no smoothing. The targets
+            become a mixture of the original ground truth and a uniform distribution as described in
+            `Rethinking the Inception Architecture for Computer Vision <https://arxiv.org/abs/1512.00567>`__. Default: :math:`0.0`.
+
+    Shape:
+        - Input: Shape :math:`(C)`, :math:`(N, C)` or :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \geq 1`
+          in the case of `K`-dimensional loss.
+        - Target: If containing class indices, shape :math:`()`, :math:`(N)` or :math:`(N, d_1, d_2, ..., d_K)` with
+          :math:`K \geq 1` in the case of K-dimensional loss where each value should be between :math:`[0, C)`.
+          If containing class probabilities, same shape as the input and each value should be between :math:`[0, 1]`.
+
+        where:
+
+        .. math::
+            \begin{aligned}
+                C ={} & \text{number of classes} \\
+                N ={} & \text{batch size} \\
+            \end{aligned}
+
+    Examples::
+
+        >>> # Example of target with class indices
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randint(5, (3,), dtype=torch.int64)
+        >>> loss = F.cross_entropy(input, target)
+        >>> loss.backward()
+        >>>
+        >>> # Example of target with class probabilities
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5).softmax(dim=1)
+        >>> loss = F.cross_entropy(input, target)
+        >>> loss.backward()
+    """
+    if has_torch_function_variadic(input, target, weight):
+        return handle_torch_function(
+            cross_entropy,
+            (input, target, weight),
+            input,
+            target,
+            weight=weight,
+            size_average=size_average,
+            ignore_index=ignore_index,
+            reduce=reduce,
+            reduction=reduction,
+            label_smoothing=label_smoothing,
+        )
+    if size_average is not None or reduce is not None:
+        reduction = _Reduction.legacy_get_string(size_average, reduce)
+    return torch._C._nn.cross_entropy_loss(
+        input,
+        target,
+        weight,
+        # pyrefly: ignore [bad-argument-type]
+        _Reduction.get_enum(reduction),
+        ignore_index,
+        label_smoothing,
+    )
+
+
+def binary_cross_entropy(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:
+    r"""Compute Binary Cross Entropy between the target and input probabilities.
+
+    See :class:`~torch.nn.BCELoss` for details.
+
+    Args:
+        input: Tensor of arbitrary shape as probabilities.
+        target: Tensor of the same shape as input with values between 0 and 1.
+        weight (Tensor, optional): a manual rescaling weight
+                if provided it's repeated to match input tensor shape
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Examples::
+
+        >>> input = torch.randn(3, 2, requires_grad=True)
+        >>> target = torch.rand(3, 2, requires_grad=False)
+        >>> loss = F.binary_cross_entropy(torch.sigmoid(input), target)
+        >>> loss.backward()
+    """
+    if has_torch_function_variadic(input, target, weight):
+        return handle_torch_function(
+            binary_cross_entropy,
+            (input, target, weight),
+            input,
+            target,
+            weight=weight,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    if target.size() != input.size():
+        raise ValueError(
+            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}) is deprecated. "
+            "Please ensure they have the same size."
+        )
+
+    if weight is not None:
+        new_size = _infer_size(target.size(), weight.size())
+        weight = weight.expand(new_size)
+
+    # pyrefly: ignore [bad-argument-type]
+    return torch._C._nn.binary_cross_entropy(input, target, weight, reduction_enum)
+
+
+def binary_cross_entropy_with_logits(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+    pos_weight: Tensor | None = None,
+) -> Tensor:
+    r"""Compute Binary Cross Entropy between target and input logits.
+
+    See :class:`~torch.nn.BCEWithLogitsLoss` for details.
+
+    Args:
+        input: Tensor of arbitrary shape as unnormalized scores (often referred to as logits).
+        target: Tensor of the same shape as input with values between 0 and 1
+        weight (Tensor, optional): a manual rescaling weight. If provided, the dimension
+           of weight supports :ref:`broadcasting to a common shape <broadcasting-semantics>`
+           with respect to the input (and target) shape.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+        pos_weight (Tensor, optional): a weight of positive examples to be broadcasted with target.
+            Must be a tensor with equal size along the class dimension to the number of classes.
+            Pay close attention to PyTorch's broadcasting semantics in order to achieve the desired
+            operations. For a target of size [B, C, H, W] (where B is batch size) pos_weight of
+            size [B, C, H, W] will apply different pos_weights to each element of the batch or
+            [C, H, W] the same pos_weights across the batch. To apply the same positive weight
+            along all spatial dimensions for a 2D multi-class target [C, H, W] use: [C, 1, 1].
+            Default: ``None``
+
+    Examples::
+
+         >>> input = torch.randn(3, requires_grad=True)
+         >>> target = torch.empty(3).random_(2)
+         >>> loss = F.binary_cross_entropy_with_logits(input, target)
+         >>> loss.backward()
+    """
+    if has_torch_function_variadic(input, target, weight, pos_weight):
+        return handle_torch_function(
+            binary_cross_entropy_with_logits,
+            (input, target, weight, pos_weight),
+            input,
+            target,
+            weight=weight,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+            pos_weight=pos_weight,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+
+    if not (target.size() == input.size()):
+        raise ValueError(
+            f"Target size ({target.size()}) must be the same as input size ({input.size()})"
+        )
+
+    return torch.binary_cross_entropy_with_logits(
+        input, target, weight, pos_weight, reduction_enum
+    )
+
+
+def smooth_l1_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+    beta: float = 1.0,
+) -> Tensor:
+    r"""Compute the Smooth L1 loss.
+
+    Function uses a squared term if the absolute
+    element-wise error falls below beta and an L1 term otherwise.
+
+    See :class:`~torch.nn.SmoothL1Loss` for details.
+
+    Args:
+        input (Tensor): Predicted values.
+        target (Tensor): Ground truth values.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+                                   'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                   'sum': the output will be summed. 'none': no reduction will be applied.
+                                   Default: 'mean'.
+        beta (float, optional): Specifies the threshold at which to change from the squared
+            term to the L1 term in the loss calculation. This value must be positive.
+            Default: 1.0.
+
+    Returns:
+        Tensor: L1 loss (optionally weighted).
+    """
+    if has_torch_function_variadic(input, target):
+        return handle_torch_function(
+            smooth_l1_loss,
+            (input, target),
+            input,
+            target,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+            beta=beta,
+        )
+    if target.size() != input.size():
+        warnings.warn(
+            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
+            "This will likely lead to incorrect results due to broadcasting. "
+            "Please ensure they have the same size.",
+            stacklevel=2,
+        )
+    if size_average is not None or reduce is not None:
+        reduction = _Reduction.legacy_get_string(size_average, reduce)
+
+    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
+
+    if beta == 0.0:
+        return torch._C._nn.l1_loss(
+            expanded_input,
+            expanded_target,
+            # pyrefly: ignore [bad-argument-type]
+            _Reduction.get_enum(reduction),
+        )
+    else:
+        return torch._C._nn.smooth_l1_loss(
+            expanded_input,
+            expanded_target,
+            # pyrefly: ignore [bad-argument-type]
+            _Reduction.get_enum(reduction),
+            beta,
+        )
+
+
+def huber_loss(
+    input: Tensor,
+    target: Tensor,
+    reduction: str = "mean",
+    delta: float = 1.0,
+    weight: Tensor | None = None,
+) -> Tensor:
+    r"""Compute the Huber loss, with optional weighting.
+
+    Function uses a squared term if the absolute
+    element-wise error falls below delta and a delta-scaled L1 term otherwise.
+
+    When delta equals 1, this loss is equivalent to SmoothL1Loss.
+    In general, Huber loss differs from SmoothL1Loss by a factor of delta (AKA beta in Smooth L1).
+
+    See :class:`~torch.nn.HuberLoss` for details.
+
+    Args:
+        input (Tensor): Predicted values.
+        target (Tensor): Ground truth values.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+                                   'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                   'sum': the output will be summed. 'none': no reduction will be applied.
+                                   Default: 'mean'.
+        delta (float, optional): The threshold at which to change between delta-scaled L1 and L2 loss. Default: 1.0.
+        weight (Tensor, optional): Weights for each sample. Default: None.
+
+    Returns:
+        Tensor: Huber loss (optionally weighted).
+    """
+    if has_torch_function_variadic(input, target, weight):
+        return handle_torch_function(
+            huber_loss,
+            (input, target, weight),
+            input,
+            target,
+            reduction=reduction,
+            delta=delta,
+            weight=weight,
+        )
+
+    if target.size() != input.size():
+        warnings.warn(
+            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
+            "This will likely lead to incorrect results due to broadcasting. "
+            "Please ensure they have the same size.",
+            stacklevel=2,
+        )
+
+    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
+
+    if weight is None:
+        # Use the optimized C++ backend for standard Huber loss
+        return torch._C._nn.huber_loss(
+            expanded_input,
+            expanded_target,
+            # pyrefly: ignore [bad-argument-type]
+            _Reduction.get_enum(reduction),
+            delta,
+        )
+    else:
+        if weight.size() != input.size():
+            raise ValueError("Weights and input must have the same size.")
+
+        # Calculate the unweighted loss first
+        unweighted_loss = torch._C._nn.huber_loss(
+            expanded_input,
+            expanded_target,
+            # pyrefly: ignore [bad-argument-type]
+            _Reduction.get_enum("none"),
+            delta,
+        )
+
+        # Apply weight to the unweighted loss
+        weighted_loss = unweighted_loss * weight
+
+        if reduction == "none":
+            return weighted_loss
+        elif reduction == "sum":
+            return torch.sum(weighted_loss)
+        elif reduction == "mean":
+            return weighted_loss.mean()
+        else:
+            raise ValueError(
+                f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', 'sum'."
+            )
+
+
+def l1_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+    weight: Tensor | None = None,
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the L1 loss, with optional weighting.
+
+    Function that takes the mean element-wise absolute value difference.
+
+    See :class:`~torch.nn.L1Loss` for details.
+
+    Args:
+        input (Tensor): Predicted values.
+        target (Tensor): Ground truth values.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+                                   'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                   'sum': the output will be summed. 'none': no reduction will be applied.
+                                   Default: 'mean'.
+        weight (Tensor, optional): Weights for each sample. Default: None.
+
+    Returns:
+        Tensor: L1 loss (optionally weighted).
+    """
+    if has_torch_function_variadic(input, target):
+        return handle_torch_function(
+            l1_loss,
+            (input, target, weight),
+            input,
+            target,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if target.size() != input.size():
+        warnings.warn(
+            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
+            "This will likely lead to incorrect results due to broadcasting. "
+            "Please ensure they have the same size.",
+            stacklevel=2,
+        )
+    if size_average is not None or reduce is not None:
+        reduction = _Reduction.legacy_get_string(size_average, reduce)
+
+    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
+
+    if weight is not None:
+        if weight.size() != input.size():
+            raise ValueError("Weights and input must have the same size.")
+
+        absolute_errors = torch.abs(expanded_input - expanded_target)
+        weighted_absolute_errors = absolute_errors * weight
+
+        if reduction == "none":
+            return weighted_absolute_errors
+        elif reduction == "sum":
+            return torch.sum(weighted_absolute_errors)
+        elif reduction == "mean":
+            return torch.sum(weighted_absolute_errors) / torch.sum(weight)
+        else:
+            raise ValueError(
+                f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', 'sum'."
+            )
+    else:
+        return torch._C._nn.l1_loss(
+            expanded_input,
+            expanded_target,
+            # pyrefly: ignore [bad-argument-type]
+            _Reduction.get_enum(reduction),
+        )
+
+
+def mse_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+    weight: Tensor | None = None,
+) -> Tensor:
+    r"""Compute the element-wise mean squared error, with optional weighting.
+
+    See :class:`~torch.nn.MSELoss` for details.
+
+    Args:
+        input (Tensor): Predicted values.
+        target (Tensor): Ground truth values.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+                                   'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                   'sum': the output will be summed. 'none': no reduction will be applied.
+                                   Default: 'mean'.
+        weight (Tensor, optional): Weights for each sample. Default: None.
+
+    Returns:
+        Tensor: Mean Squared Error loss (optionally weighted).
+    """
+    if has_torch_function_variadic(input, target, weight):
+        return handle_torch_function(
+            mse_loss,
+            (input, target, weight),
+            input,
+            target,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+            weight=weight,
+        )
+
+    if target.size() != input.size():
+        warnings.warn(
+            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
+            "This will likely lead to incorrect results due to broadcasting. "
+            "Please ensure they have the same size.",
+            stacklevel=2,
+        )
+
+    if size_average is not None or reduce is not None:
+        reduction = _Reduction.legacy_get_string(size_average, reduce)
+
+    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
+
+    if weight is not None:
+        if weight.size() != input.size():
+            raise ValueError("Weights and input must have the same size.")
+
+        # Perform weighted MSE loss manually
+        squared_errors = torch.pow(expanded_input - expanded_target, 2)
+        weighted_squared_errors = squared_errors * weight
+
+        if reduction == "none":
+            return weighted_squared_errors
+        elif reduction == "sum":
+            return torch.sum(weighted_squared_errors)
+        elif reduction == "mean":
+            return torch.sum(weighted_squared_errors) / torch.sum(weight)
+        else:
+            raise ValueError(
+                f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', 'sum'."
+            )
+    else:
+        return torch._C._nn.mse_loss(
+            expanded_input,
+            expanded_target,
+            # pyrefly: ignore [bad-argument-type]
+            _Reduction.get_enum(reduction),
+        )
+
+
+def margin_ranking_loss(
+    input1: Tensor,
+    input2: Tensor,
+    target: Tensor,
+    margin: float = 0,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the margin ranking loss.
+
+    See :class:`~torch.nn.MarginRankingLoss` for details.
+
+    Args:
+        input1 (Tensor): Predicted values.
+        input2 (Tensor): Predicted values.
+        target (Tensor): Ground truth values.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+                                   'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                   'sum': the output will be summed. 'none': no reduction will be applied.
+                                   Default: 'mean'.
+
+    Returns:
+        Tensor: Margin ranking loss.
+    """
+    if has_torch_function_variadic(input1, input2, target):
+        return handle_torch_function(
+            margin_ranking_loss,
+            (input1, input2, target),
+            input1,
+            input2,
+            target,
+            margin=margin,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    if input1.dim() != input2.dim() or input1.dim() != target.dim():
+        raise RuntimeError(
+            f"margin_ranking_loss : All input tensors should have same dimension but got sizes: "
+            f"input1: {input1.size()}, input2: {input2.size()}, target: {target.size()} "
+        )
+    return torch.margin_ranking_loss(input1, input2, target, margin, reduction_enum)
+
+
+def hinge_embedding_loss(
+    input: Tensor,
+    target: Tensor,
+    margin: float = 1.0,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the hinge embedding loss.
+
+    See :class:`~torch.nn.HingeEmbeddingLoss` for details.
+
+    Args:
+       input (Tensor): Predicted values.
+       target (Tensor): Ground truth values.
+       margin (float, optional): Margin for hinge loss. Has a default value of 1.
+       size_average (bool, optional): Deprecated (see :attr:`reduction`).
+       reduce (bool, optional): Deprecated (see :attr:`reduction`).
+       reduction (str, optional): Specifies the reduction to apply to the output:
+                                  'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                  'sum': the output will be summed. 'none': no reduction will be applied.
+                                  Default: 'mean'.
+
+    Returns:
+       Tensor: Hinge embedding loss.
+    """
+    if has_torch_function_variadic(input, target):
+        return handle_torch_function(
+            hinge_embedding_loss,
+            (input, target),
+            input,
+            target,
+            margin=margin,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    return torch.hinge_embedding_loss(input, target, margin, reduction_enum)
+
+
+def multilabel_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the multilabel margin loss.
+
+    See :class:`~torch.nn.MultiLabelMarginLoss` for details.
+
+    Args:
+       input (Tensor): Predicted values.
+       target (Tensor): Ground truth values.
+       size_average (bool, optional): Deprecated (see :attr:`reduction`).
+       reduce (bool, optional): Deprecated (see :attr:`reduction`).
+       reduction (str, optional): Specifies the reduction to apply to the output:
+                                  'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                  'sum': the output will be summed. 'none': no reduction will be applied.
+                                  Default: 'mean'.
+
+    Returns:
+       Tensor: Mutilabel margin loss.
+    """
+    if has_torch_function_variadic(input, target):
+        return handle_torch_function(
+            multilabel_margin_loss,
+            (input, target),
+            input,
+            target,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    # pyrefly: ignore [bad-argument-type]
+    return torch._C._nn.multilabel_margin_loss(input, target, reduction_enum)
+
+
+def soft_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the soft margin loss.
+
+    See :class:`~torch.nn.SoftMarginLoss` for details.
+
+    Args:
+       input (Tensor): Predicted values.
+       target (Tensor): Ground truth values.
+       size_average (bool, optional): Deprecated (see :attr:`reduction`).
+       reduce (bool, optional): Deprecated (see :attr:`reduction`).
+       reduction (str, optional): Specifies the reduction to apply to the output:
+                                  'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                  'sum': the output will be summed. 'none': no reduction will be applied.
+                                  Default: 'mean'.
+
+    Returns:
+       Tensor: Soft margin loss.
+    """
+    if has_torch_function_variadic(input, target):
+        return handle_torch_function(
+            soft_margin_loss,
+            (input, target),
+            input,
+            target,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    # pyrefly: ignore [bad-argument-type]
+    return torch._C._nn.soft_margin_loss(input, target, reduction_enum)
+
+
+def multilabel_soft_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the multilabel soft margin loss.
+
+    See :class:`~torch.nn.MultiLabelSoftMarginLoss` for details.
+
+    Args:
+       input (Tensor): Predicted values.
+       target (Tensor): Ground truth values.
+       size_average (bool, optional): Deprecated (see :attr:`reduction`).
+       reduce (bool, optional): Deprecated (see :attr:`reduction`).
+       reduction (str, optional): Specifies the reduction to apply to the output:
+                                  'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                  'sum': the output will be summed. 'none': no reduction will be applied.
+                                  Default: 'mean'.
+
+    Returns:
+       Tensor: Mutilabel soft margin loss.
+    """
+    if has_torch_function_variadic(input, target, weight):
+        return handle_torch_function(
+            multilabel_soft_margin_loss,
+            (input, target, weight),
+            input,
+            target,
+            weight=weight,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction = _Reduction.legacy_get_string(size_average, reduce)
+
+    loss = -(target * logsigmoid(input) + (1 - target) * logsigmoid(-input))
+
+    if weight is not None:
+        loss = loss * weight
+
+    class_dim = input.dim() - 1
+    C = input.size(class_dim)
+    loss = loss.sum(dim=class_dim) / C  # only return N loss values
+
+    if reduction == "none":
+        ret = loss
+    elif reduction == "mean":
+        ret = loss.mean()
+    elif reduction == "sum":
+        ret = loss.sum()
+    else:
+        ret = input
+        raise ValueError(reduction + " is not valid")
+    return ret
+
+
+def cosine_embedding_loss(
+    input1: Tensor,
+    input2: Tensor,
+    target: Tensor,
+    margin: float = 0,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the cosine embedding loss.
+
+    See :class:`~torch.nn.CosineEmbeddingLoss` for details.
+
+    Args:
+       input1 (Tensor): Predicted values.
+       input2 (Tensor): Predicted values.
+       target (Tensor): Ground truth values.
+       margin (float, optional): Margin for cosine embedding. Has a default value of 0.
+       size_average (bool, optional): Deprecated (see :attr:`reduction`).
+       reduce (bool, optional): Deprecated (see :attr:`reduction`).
+       reduction (str, optional): Specifies the reduction to apply to the output:
+                                  'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                  'sum': the output will be summed. 'none': no reduction will be applied.
+                                  Default: 'mean'.
+
+    Returns:
+       Tensor: Cosine embedding loss.
+    """
+    if has_torch_function_variadic(input1, input2, target):
+        return handle_torch_function(
+            cosine_embedding_loss,
+            (input1, input2, target),
+            input1,
+            input2,
+            target,
+            margin=margin,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    return torch.cosine_embedding_loss(input1, input2, target, margin, reduction_enum)
+
+
+def multi_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    p: int = 1,
+    margin: float = 1.0,
+    weight: Tensor | None = None,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:  # noqa: D400,D402
+    r"""Compute the multi margin loss, with optional weighting.
+
+    See :class:`~torch.nn.MultiMarginLoss` for details.
+
+    Args:
+        input (Tensor): Predicted values.
+        target (Tensor): Ground truth values.
+        p (int, optional): Has a default value of 1. 1 and 2 are the only supported values.
+        margin (float, optional): Margin for multi margin loss. Has a default value of 1.
+        weight (Tensor, optional): Weights for each sample. Default: None.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`).
+        reduce (bool, optional): Deprecated (see :attr:`reduction`).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+                                  'none' | 'mean' | 'sum'. 'mean': the mean of the output is taken.
+                                  'sum': the output will be summed. 'none': no reduction will be applied.
+                                  Default: 'mean'.
+
+    Returns:
+        Tensor: Multi margin loss (optionally weighted).
+    """
+    if has_torch_function_variadic(input, target, weight):
+        return handle_torch_function(
+            multi_margin_loss,
+            (input, target, weight),
+            input,
+            target,
+            p=p,
+            margin=margin,
+            weight=weight,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    if p != 1 and p != 2:
+        raise ValueError("only p == 1 and p == 2 supported")
+    if weight is not None:
+        if weight.dim() != 1:
+            raise ValueError("weight must be one-dimensional")
+
+    return torch._C._nn.multi_margin_loss(
+        input,
+        target,
+        p,
+        margin,
+        weight,
+        # pyrefly: ignore [bad-argument-type]
+        reduction_enum,
+    )
+
+
+pixel_shuffle = _add_docstr(
+    torch.pixel_shuffle,
+    r"""
+pixel_shuffle(input, upscale_factor) -> Tensor
+
+Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)` to a
+tensor of shape :math:`(*, C, H \times r, W \times r)`, where r is the :attr:`upscale_factor`.
+
+See :class:`~torch.nn.PixelShuffle` for details.
+
+Args:
+    input (Tensor): the input tensor
+    upscale_factor (int): factor to increase spatial resolution by
+
+Examples::
+
+    >>> input = torch.randn(1, 9, 4, 4)
+    >>> output = torch.nn.functional.pixel_shuffle(input, 3)
+    >>> print(output.size())
+    torch.Size([1, 1, 12, 12])
+""",
+)
+
+pixel_unshuffle = _add_docstr(
+    torch.pixel_unshuffle,
+    r"""
+pixel_unshuffle(input, downscale_factor) -> Tensor
+
+Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements in a
+tensor of shape :math:`(*, C, H \times r, W \times r)` to a tensor of shape
+:math:`(*, C \times r^2, H, W)`, where r is the :attr:`downscale_factor`.
+
+See :class:`~torch.nn.PixelUnshuffle` for details.
+
+Args:
+    input (Tensor): the input tensor
+    downscale_factor (int): factor to increase spatial resolution by
+
+Examples::
+
+    >>> input = torch.randn(1, 1, 12, 12)
+    >>> output = torch.nn.functional.pixel_unshuffle(input, 3)
+    >>> print(output.size())
+    torch.Size([1, 9, 4, 4])
+""",
+)
+
+channel_shuffle = _add_docstr(
+    torch.channel_shuffle,
+    r"""
+channel_shuffle(input, groups) -> Tensor
+
+Divide the channels in a tensor of shape :math:`(*, C , H, W)`
+into g groups and rearrange them as :math:`(*, C \frac g, g, H, W)`,
+while keeping the original tensor shape.
+
+See :class:`~torch.nn.ChannelShuffle` for details.
+
+Args:
+    input (Tensor): the input tensor
+    groups (int): number of groups to divide channels in and rearrange.
+
+Examples::
+
+    >>> input = torch.randn(1, 4, 2, 2)
+    >>> print(input)
+    [[[[1, 2],
+       [3, 4]],
+      [[5, 6],
+       [7, 8]],
+      [[9, 10],
+       [11, 12]],
+      [[13, 14],
+       [15, 16]],
+     ]]
+    >>> output = torch.nn.functional.channel_shuffle(input, 2)
+    >>> print(output)
+    [[[[1, 2],
+       [3, 4]],
+      [[9, 10],
+       [11, 12]],
+      [[5, 6],
+       [7, 8]],
+      [[13, 14],
+       [15, 16]],
+     ]]
+""",
+)
+
+native_channel_shuffle = _add_docstr(
+    torch.native_channel_shuffle,
+    r"""
+native_channel_shuffle(input, groups) -> Tensor
+
+Native kernel level implementation of the `channel_shuffle`.
+This function might become private in future releases, use with caution.
+
+Divide the channels in a tensor of shape :math:`(*, C , H, W)`
+into g groups and rearrange them as :math:`(*, C \frac g, g, H, W)`,
+while keeping the original tensor shape.
+
+See :class:`~torch.nn.ChannelShuffle` for details.
+
+Args:
+    input (Tensor): the input tensor
+    groups (int): number of groups to divide channels in and rearrange.
+
+Examples::
+
+    >>> input = torch.randn(1, 4, 2, 2)
+    >>> print(input)
+    [[[[1, 2],
+       [3, 4]],
+      [[5, 6],
+       [7, 8]],
+      [[9, 10],
+       [11, 12]],
+      [[13, 14],
+       [15, 16]],
+     ]]
+    >>> output = torch.nn.functional.native_channel_shuffle(input, 2)
+    >>> print(output)
+    [[[[1, 2],
+       [3, 4]],
+      [[9, 10],
+       [11, 12]],
+      [[5, 6],
+       [7, 8]],
+      [[13, 14],
+       [15, 16]],
+     ]]
+""",
+)
+
+
+@_overload
+def upsample(  # noqa: F811
+    input: Tensor,
+    size: int | None = None,
+    scale_factor: float | None = None,
+    mode: str = "nearest",
+    align_corners: bool | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:  # noqa: B950
+    pass
+
+
+@_overload
+def upsample(  # noqa: F811
+    input: Tensor,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
+    mode: str = "nearest",
+    align_corners: bool | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:  # noqa: B950
+    pass
+
+
+def upsample(  # noqa: F811
+    input,
+    size=None,
+    scale_factor=None,
+    mode="nearest",
+    align_corners=None,
+):
+    r"""Upsample input.
+
+    Provided tensor is upsampled to either the given :attr:`size` or the given
+    :attr:`scale_factor`
+
+    .. warning::
+        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.
+        This is equivalent with ``nn.functional.interpolate(...)``.
+
+    Note:
+        {backward_reproducibility_note}
+
+    The algorithm used for upsampling is determined by :attr:`mode`.
+
+    Currently temporal, spatial and volumetric upsampling are supported, i.e.
+    expected inputs are 3-D, 4-D or 5-D in shape.
+
+    The input dimensions are interpreted in the form:
+    `mini-batch x channels x [optional depth] x [optional height] x width`.
+
+    The modes available for upsampling are: `nearest`, `linear` (3D-only),
+    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `lanczos` (4D-only)
+
+    Args:
+        input (Tensor): the input tensor
+        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):
+            output spatial size.
+        scale_factor (float or Tuple[float]): multiplier for spatial size. Has to match input size if it is a tuple.
+        mode (str): algorithm used for upsampling:
+            ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
+            ``'trilinear'`` | ``'lanczos'``. Default: ``'nearest'``
+        align_corners (bool, optional): Geometrically, we consider the pixels of the
+            input and output as squares rather than points.
+            If set to ``True``, the input and output tensors are aligned by the
+            center points of their corner pixels, preserving the values at the corner pixels.
+            If set to ``False``, the input and output tensors are aligned by the corner
+            points of their corner pixels, and the interpolation uses edge value padding
+            for out-of-boundary values, making this operation *independent* of input size
+            when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`
+            is ``'linear'``, ``'bilinear'``, ``'bicubic'`` or ``'trilinear'``.
+            Default: ``False``
+
+    .. note::
+        With ``mode='bicubic'`` or ``mode='lanczos'``, it's possible to cause overshoot, in other words it can produce
+        negative values or values greater than 255 for images.
+        Explicitly call ``result.clamp(min=0, max=255)`` if you want to reduce the overshoot
+        when displaying the image.
+
+    .. warning::
+        With ``align_corners = True``, the linearly interpolating modes
+        (`linear`, `bilinear`, and `trilinear`) don't proportionally align the
+        output and input pixels, and thus the output values can depend on the
+        input size. This was the default behavior for these modes up to version
+        0.3.1. Since then, the default behavior is ``align_corners = False``.
+        See :class:`~torch.nn.Upsample` for concrete examples on how this
+        affects the outputs.
+
+    """
+    warnings.warn(
+        "`nn.functional.upsample` is deprecated. "
+        "Use `nn.functional.interpolate` instead.",
+        stacklevel=2,
+    )
+    return interpolate(input, size, scale_factor, mode, align_corners)
+
+
+if upsample.__doc__:
+    upsample.__doc__ = upsample.__doc__.format(**reproducibility_notes)
+
+
+def _is_integer(x) -> bool:
+    r"""Type check the input number is an integer.
+
+    Will return True for int, SymInt, Numpy integers and Tensors with integer elements.
+    """
+    if isinstance(x, (int, torch.SymInt)):
+        return True
+    if np is not None and isinstance(x, np.integer):
+        return True
+    return isinstance(x, Tensor) and not x.is_floating_point()
+
+
+@_overload
+def interpolate(  # noqa: F811
+    input: Tensor,
+    size: int | None = None,
+    scale_factor: list[float] | None = None,
+    mode: str = "nearest",
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
+    antialias: bool = False,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:  # noqa: B950
+    pass
+
+
+@_overload
+def interpolate(  # noqa: F811
+    input: Tensor,
+    size: list[int] | None = None,
+    scale_factor: list[float] | None = None,
+    mode: str = "nearest",
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
+    antialias: bool = False,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:  # noqa: B950
+    pass
+
+
+@_overload
+def interpolate(  # noqa: F811
+    input: Tensor,
+    size: int | None = None,
+    scale_factor: float | None = None,
+    mode: str = "nearest",
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
+    antialias: bool = False,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:  # noqa: B950
+    pass
+
+
+@_overload
+def interpolate(  # noqa: F811
+    input: Tensor,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
+    mode: str = "nearest",
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
+    antialias: bool = False,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:
+    pass
+
+
+def interpolate(  # noqa: F811
+    input: Tensor,
+    size: int | None = None,
+    scale_factor: list[float] | None = None,
+    mode: str = "nearest",
+    align_corners: bool | None = None,
+    recompute_scale_factor: bool | None = None,
+    antialias: bool = False,
+) -> Tensor:  # noqa: B950
+    r"""Down/up samples the input.
+
+    Tensor interpolated to either the given :attr:`size` or the given
+    :attr:`scale_factor`
+
+    The algorithm used for interpolation is determined by :attr:`mode`.
+
+    Currently temporal, spatial and volumetric sampling are supported, i.e.
+    expected inputs are 3-D, 4-D or 5-D in shape.
+
+    The input dimensions are interpreted in the form:
+    `mini-batch x channels x [optional depth] x [optional height] x width`.
+
+    The modes available for resizing are: `nearest`, `linear` (3D-only),
+    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `lanczos` (4D-only, CPU only), `area`, `nearest-exact`
+
+    Args:
+        input (Tensor): the input tensor
+        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):
+            output spatial size.
+        scale_factor (float or Tuple[float]): multiplier for spatial size. If `scale_factor` is a tuple,
+            its length has to match the number of spatial dimensions; `input.dim() - 2`.
+        mode (str): algorithm used for upsampling:
+            ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
+            ``'trilinear'`` | ``'lanczos'`` | ``'area'`` | ``'nearest-exact'``. Default: ``'nearest'``
+        align_corners (bool, optional): Geometrically, we consider the pixels of the
+            input and output as squares rather than points.
+            If set to ``True``, the input and output tensors are aligned by the
+            center points of their corner pixels, preserving the values at the corner pixels.
+            If set to ``False``, the input and output tensors are aligned by the corner
+            points of their corner pixels, and the interpolation uses edge value padding
+            for out-of-boundary values, making this operation *independent* of input size
+            when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`
+            is ``'linear'``, ``'bilinear'``, ``'bicubic'`` or ``'trilinear'``.
+            Default: ``None``. When ``None`` and :attr:`mode` is one of the linear modes,
+            it is treated as ``False``.
+        recompute_scale_factor (bool, optional): recompute the scale_factor for use in the
+            interpolation calculation. If `recompute_scale_factor` is ``True``, then
+            `scale_factor` must be passed in and `scale_factor` is used to compute the
+            output `size`. The computed output `size` will be used to infer new scales for
+            the interpolation. Note that when `scale_factor` is floating-point, it may differ
+            from the recomputed `scale_factor` due to rounding and precision issues.
+            If `recompute_scale_factor` is ``False``, then `size` or `scale_factor` will
+            be used directly for interpolation. Default: ``None``.
+        antialias (bool, optional): flag to apply anti-aliasing. Default: ``False``. Using anti-alias
+            option together with ``align_corners=False``, interpolation result would match Pillow
+            result for downsampling operation. Supported modes: ``'bilinear'``, ``'bicubic'``, ``'lanczos'``.
+
+    .. note::
+        With ``mode='bicubic'`` or ``mode='lanczos'``, it's possible to cause overshoot. For some dtypes, it can produce
+        negative values or values greater than 255 for images. Explicitly call ``result.clamp(min=0,max=255)``
+        if you want to reduce the overshoot when displaying the image.
+        For ``uint8`` inputs, it already performs saturating cast operation. So, no manual `clamp` operation is needed.
+
+    .. note::
+        Mode ``mode='lanczos'`` uses a Lanczos-3 windowed sinc filter (6 taps) and requires
+        ``antialias=True``. It only supports 4-D input (i.e. 2D spatial) and CPU. With ``antialias=True``
+        and ``align_corners=False``, the result matches PIL's ``Image.LANCZOS`` resampling filter.
+
+    .. note::
+        Mode ``mode='nearest-exact'`` matches Scikit-Image and PIL nearest neighbours interpolation
+        algorithms and fixes known issues with ``mode='nearest'``. This mode is introduced to keep
+        backward compatibility.
+        Mode ``mode='nearest'`` matches buggy OpenCV's ``INTER_NEAREST`` interpolation algorithm.
+
+    .. note::
+        The gradients for the dtype ``float16`` on CUDA may be inaccurate in the upsample operation
+        when using modes ``['linear', 'bilinear', 'bicubic', 'trilinear', 'area']``.
+        For more details, please refer to the discussion in
+        `issue#104157 <https://github.com/pytorch/pytorch/issues/104157>`_.
+
+    Note:
+        {backward_reproducibility_note}
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            interpolate,
+            (input,),
+            input,
+            size=size,
+            scale_factor=scale_factor,
+            mode=mode,
+            align_corners=align_corners,
+            recompute_scale_factor=recompute_scale_factor,
+            antialias=antialias,
+        )
+
+    if mode in ("nearest", "area", "nearest-exact"):
+        if align_corners is not None:
+            raise ValueError(
+                "align_corners option can only be set with the "
+                "interpolating modes: linear | bilinear | bicubic | trilinear"
+            )
+    else:
+        if align_corners is None:
+            align_corners = False
+
+    dim = input.dim() - 2  # Number of spatial dimensions.
+
+    # Process size and scale_factor.  Validate that exactly one is set.
+    # Validate its length if it is a list, or expand it if it is a scalar.
+    # After this block, exactly one of output_size and scale_factors will
+    # be non-None, and it will be a list (or tuple).
+    if size is not None and scale_factor is not None:
+        raise ValueError("only one of size or scale_factor should be defined")
+    elif size is not None:
+        if scale_factor is not None:
+            raise AssertionError("scale_factor must be None when size is specified")
+        scale_factors = None
+        if isinstance(size, (list, tuple)):
+            if len(size) != dim:
+                raise ValueError(
+                    "Input and output must have the same number of spatial dimensions, but got "
+                    f"input with spatial dimensions of {list(input.shape[2:])} and output size of {size}. "
+                    "Please provide input tensor in (N, C, d1, d2, ...,dK) format and "
+                    "output size in (o1, o2, ...,oK) format."
+                )
+            if not torch.jit.is_scripting():
+                if not all(_is_integer(x) for x in size):
+                    raise TypeError(
+                        "expected size to be one of int or Tuple[int] or Tuple[int, int] or "
+                        f"Tuple[int, int, int], but got size with types {[type(x) for x in size]}"
+                    )
+            output_size = size
+        else:
+            output_size = [size for _ in range(dim)]
+    elif scale_factor is not None:
+        if size is not None:
+            raise AssertionError("size must be None when scale_factor is specified")
+        output_size = None
+        if isinstance(scale_factor, (list, tuple)):
+            if len(scale_factor) != dim:
+                raise ValueError(
+                    "Input and scale_factor must have the same number of spatial dimensions, but "
+                    f"got input with spatial dimensions of {list(input.shape[2:])} and "
+                    f"scale_factor of shape {scale_factor}. "
+                    "Please provide input tensor in (N, C, d1, d2, ...,dK) format and "
+                    "scale_factor in (s1, s2, ...,sK) format."
+                )
+            scale_factors = scale_factor
+        else:
+            scale_factors = [scale_factor for _ in range(dim)]
+    else:
+        raise ValueError("either size or scale_factor should be defined")
+
+    if (
+        recompute_scale_factor is not None
+        and recompute_scale_factor
+        and size is not None
+    ):
+        raise ValueError(
+            "recompute_scale_factor is not meaningful with an explicit size."
+        )
+
+    # "area" mode always requires an explicit size rather than scale factor.
+    # Reuse the recompute_scale_factor code path.
+    if mode == "area" and output_size is None:
+        recompute_scale_factor = True
+
+    if recompute_scale_factor is not None and recompute_scale_factor:
+        # We compute output_size here, then un-set scale_factors.
+        # The C++ code will recompute it based on the (integer) output size.
+        if scale_factors is None:
+            raise AssertionError("scale_factors is unexpectedly None")
+        if not torch.jit.is_scripting() and torch._C._get_tracing_state():
+            # make scale_factor a tensor in tracing so constant doesn't get baked in
+            output_size = [
+                (
+                    torch.floor(
+                        (
+                            # pyrefly: ignore [missing-attribute]
+                            input.size(i + 2).float()
+                            * torch.tensor(scale_factors[i], dtype=torch.float32)
+                        ).float()
+                    )
+                )
+                for i in range(dim)
+            ]
+        elif torch.jit.is_scripting():
+            output_size = [
+                math.floor(float(input.size(i + 2)) * scale_factors[i])
+                for i in range(dim)
+            ]
+        else:
+            output_size = [
+                _sym_int(input.size(i + 2) * scale_factors[i]) for i in range(dim)
+            ]
+        scale_factors = None
+
+    if antialias and not (
+        mode in ("bilinear", "bicubic", "lanczos") and input.ndim == 4
+    ):
+        raise ValueError(
+            "Anti-alias option is restricted to bilinear, bicubic, and lanczos modes and requires a 4-D tensor as input"
+        )
+
+    if input.dim() == 3 and mode == "nearest":
+        # pyrefly: ignore [bad-argument-type]
+        return torch._C._nn.upsample_nearest1d(input, output_size, scale_factors)
+    if input.dim() == 4 and mode == "nearest":
+        # pyrefly: ignore [bad-argument-type]
+        return torch._C._nn.upsample_nearest2d(input, output_size, scale_factors)
+    if input.dim() == 5 and mode == "nearest":
+        # pyrefly: ignore [bad-argument-type]
+        return torch._C._nn.upsample_nearest3d(input, output_size, scale_factors)
+
+    if input.dim() == 3 and mode == "nearest-exact":
+        # pyrefly: ignore [bad-argument-type]
+        return torch._C._nn._upsample_nearest_exact1d(input, output_size, scale_factors)
+    if input.dim() == 4 and mode == "nearest-exact":
+        # pyrefly: ignore [bad-argument-type]
+        return torch._C._nn._upsample_nearest_exact2d(input, output_size, scale_factors)
+    if input.dim() == 5 and mode == "nearest-exact":
+        # pyrefly: ignore [bad-argument-type]
+        return torch._C._nn._upsample_nearest_exact3d(input, output_size, scale_factors)
+
+    if input.dim() == 3 and mode == "area":
+        if output_size is None:
+            raise AssertionError("output_size is unexpectedly None")
+        # pyrefly: ignore [bad-argument-type]
+        return adaptive_avg_pool1d(input, output_size)
+    if input.dim() == 4 and mode == "area":
+        if output_size is None:
+            raise AssertionError("output_size is unexpectedly None")
+        return adaptive_avg_pool2d(input, output_size)
+    if input.dim() == 5 and mode == "area":
+        if output_size is None:
+            raise AssertionError("output_size is unexpectedly None")
+        return adaptive_avg_pool3d(input, output_size)
+
+    if input.dim() == 3 and mode == "linear":
+        if align_corners is None:
+            raise AssertionError("align_corners is unexpectedly None")
+        return torch._C._nn.upsample_linear1d(
+            input,
+            # pyrefly: ignore [bad-argument-type]
+            output_size,
+            align_corners,
+            scale_factors,
+        )
+    if input.dim() == 4 and mode == "bilinear":
+        if align_corners is None:
+            raise AssertionError("align_corners is unexpectedly None")
+        if antialias:
+            return torch._C._nn._upsample_bilinear2d_aa(
+                input,
+                # pyrefly: ignore [bad-argument-type]
+                output_size,
+                align_corners,
+                scale_factors,
+            )
+        # Two levels are necessary to prevent TorchScript from touching
+        # are_deterministic_algorithms_enabled.
+        if not torch.jit.is_scripting():
+            if not input.is_cpu and torch.are_deterministic_algorithms_enabled():
+                # Use slow decomp whose backward will be in terms of index_put
+                # importlib is required because the import cannot be top level
+                # (cycle) and cannot be nested (TS doesn't support)
+                return importlib.import_module(
+                    "torch._decomp.decompositions"
+                )._upsample_linear_vec(input, output_size, align_corners, scale_factors)
+        return torch._C._nn.upsample_bilinear2d(
+            input,
+            # pyrefly: ignore [bad-argument-type]
+            output_size,
+            align_corners,
+            scale_factors,
+        )
+    if input.dim() == 5 and mode == "trilinear":
+        if align_corners is None:
+            raise AssertionError("align_corners is unexpectedly None")
+        # Two levels are necessary to prevent TorchScript from touching
+        # are_deterministic_algorithms_enabled.
+        if not torch.jit.is_scripting():
+            if not input.is_cpu and torch.are_deterministic_algorithms_enabled():
+                # Use slow decomp whose backward will be in terms of index_put
+                # importlib is required because the import cannot be top level
+                # (cycle) and cannot be nested (TS doesn't support)
+                return importlib.import_module(
+                    "torch._decomp.decompositions"
+                )._upsample_linear_vec(input, output_size, align_corners, scale_factors)
+        return torch._C._nn.upsample_trilinear3d(
+            input,
+            # pyrefly: ignore [bad-argument-type]
+            output_size,
+            align_corners,
+            scale_factors,
+        )
+    if input.dim() == 4 and mode == "bicubic":
+        if align_corners is None:
+            raise AssertionError("align_corners is unexpectedly None")
+        if antialias:
+            return torch._C._nn._upsample_bicubic2d_aa(
+                input,
+                # pyrefly: ignore [bad-argument-type]
+                output_size,
+                align_corners,
+                scale_factors,
+            )
+        return torch._C._nn.upsample_bicubic2d(
+            input,
+            # pyrefly: ignore [bad-argument-type]
+            output_size,
+            align_corners,
+            scale_factors,
+        )
+
+    if input.dim() == 4 and mode == "lanczos":
+        if align_corners is None:
+            raise AssertionError("align_corners is unexpectedly None")
+        if align_corners:
+            raise ValueError("Lanczos mode does not support align_corners=True")
+        if not antialias:
+            raise ValueError("Lanczos mode requires antialias=True")
+        return torch._C._nn._upsample_lanczos2d_aa(
+            input,
+            # pyrefly: ignore [bad-argument-type]
+            output_size,
+            align_corners,
+            scale_factors,
+        )
+
+    if input.dim() == 3 and mode == "bilinear":
+        raise NotImplementedError("Got 3D input, but bilinear mode needs 4D input")
+    if input.dim() == 3 and mode == "trilinear":
+        raise NotImplementedError("Got 3D input, but trilinear mode needs 5D input")
+    if input.dim() == 4 and mode == "linear":
+        raise NotImplementedError("Got 4D input, but linear mode needs 3D input")
+    if input.dim() == 4 and mode == "trilinear":
+        raise NotImplementedError("Got 4D input, but trilinear mode needs 5D input")
+    if input.dim() == 5 and mode == "linear":
+        raise NotImplementedError("Got 5D input, but linear mode needs 3D input")
+    if input.dim() == 5 and mode == "bilinear":
+        raise NotImplementedError("Got 5D input, but bilinear mode needs 4D input")
+
+    raise NotImplementedError(
+        "Input Error: Only 3D, 4D and 5D input Tensors supported"
+        f" (got {input.dim()}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | lanczos | area | nearest-exact"
+        f" (got {mode})"
+    )
+
+
+if interpolate.__doc__:
+    interpolate.__doc__ = interpolate.__doc__.format(**reproducibility_notes)
+
+
+@_overload
+def upsample_nearest(  # noqa: F811
+    input: Tensor,
+    size: int | None = None,
+    scale_factor: float | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:
+    pass
+
+
+@_overload
+def upsample_nearest(  # noqa: F811
+    input: Tensor,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:
+    pass
+
+
+def upsample_nearest(input, size=None, scale_factor=None):  # noqa: F811
+    r"""Upsamples the input, using nearest neighbours' pixel values.
+
+    .. warning::
+        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.
+        This is equivalent with ``nn.functional.interpolate(..., mode='nearest')``.
+
+    Currently spatial and volumetric upsampling are supported (i.e. expected
+    inputs are 4 or 5 dimensional).
+
+    Args:
+        input (Tensor): input
+        size (int or Tuple[int, int] or Tuple[int, int, int]): output spatia
+            size.
+        scale_factor (int): multiplier for spatial size. Has to be an integer.
+
+    Note:
+        {backward_reproducibility_note}
+    """
+    # DeprecationWarning is ignored by default
+    warnings.warn(
+        "`nn.functional.upsample_nearest` is deprecated. "
+        "Use `nn.functional.interpolate` instead.",
+        stacklevel=2,
+    )
+    return interpolate(input, size, scale_factor, mode="nearest")
+
+
+if upsample_nearest.__doc__:
+    upsample_nearest.__doc__ = upsample_nearest.__doc__.format(**reproducibility_notes)
+
+
+@_overload
+def upsample_bilinear(  # noqa: F811
+    input: Tensor,
+    size: int | None = None,
+    scale_factor: float | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:
+    pass
+
+
+@_overload
+def upsample_bilinear(  # noqa: F811
+    input: Tensor,
+    size: list[int] | None = None,
+    scale_factor: float | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:
+    pass
+
+
+@_overload
+def upsample_bilinear(  # noqa: F811
+    input: Tensor,
+    size: int | None = None,
+    scale_factor: list[float] | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:
+    pass
+
+
+@_overload
+def upsample_bilinear(  # noqa: F811
+    input: Tensor,
+    size: list[int] | None = None,
+    scale_factor: list[float] | None = None,
+    # pyrefly: ignore [bad-return]
+) -> Tensor:
+    pass
+
+
+def upsample_bilinear(input, size=None, scale_factor=None):  # noqa: F811
+    r"""Upsamples the input, using bilinear upsampling.
+
+    .. warning::
+        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.
+        This is equivalent with
+        ``nn.functional.interpolate(..., mode='bilinear', align_corners=True)``.
+
+    Expected inputs are spatial (4 dimensional). Use `upsample_trilinear` for
+    volumetric (5 dimensional) inputs.
+
+    Args:
+        input (Tensor): input
+        size (int or Tuple[int, int]): output spatial size.
+        scale_factor (int or Tuple[int, int]): multiplier for spatial size
+
+    Note:
+        {backward_reproducibility_note}
+    """
+    # DeprecationWarning is ignored by default
+    warnings.warn(
+        "`nn.functional.upsample_bilinear` is deprecated. "
+        "Use `nn.functional.interpolate` instead.",
+        stacklevel=2,
+    )
+    return interpolate(input, size, scale_factor, mode="bilinear", align_corners=True)
+
+
+if upsample_bilinear.__doc__:
+    upsample_bilinear.__doc__ = upsample_bilinear.__doc__.format(
+        **reproducibility_notes
+    )
+
+GRID_SAMPLE_INTERPOLATION_MODES = {
+    "bilinear": 0,
+    "nearest": 1,
+    "bicubic": 2,
+}
+
+GRID_SAMPLE_PADDING_MODES = {
+    "zeros": 0,
+    "border": 1,
+    "reflection": 2,
+}
+
+
+def grid_sample(
+    input: Tensor,
+    grid: Tensor,
+    mode: str = "bilinear",
+    padding_mode: str = "zeros",
+    align_corners: bool | None = None,
+) -> Tensor:
+    r"""Compute grid sample.
+
+    Given an :attr:`input` and a flow-field :attr:`grid`, computes the
+    ``output`` using :attr:`input` values and pixel locations from :attr:`grid`.
+
+    Currently, only spatial (4-D) and volumetric (5-D) :attr:`input` are
+    supported.
+
+    In the spatial (4-D) case, for :attr:`input` with shape
+    :math:`(N, C, H_\text{in}, W_\text{in})` and :attr:`grid` with shape
+    :math:`(N, H_\text{out}, W_\text{out}, 2)`, the output will have shape
+    :math:`(N, C, H_\text{out}, W_\text{out})`.
+
+    For each output location ``output[n, :, h, w]``, the size-2 vector
+    ``grid[n, h, w]`` specifies :attr:`input` pixel locations ``x`` and ``y``,
+    which are used to interpolate the output value ``output[n, :, h, w]``.
+    In the case of 5D inputs, ``grid[n, d, h, w]`` specifies the
+    ``x``, ``y``, ``z`` pixel locations for interpolating
+    ``output[n, :, d, h, w]``. :attr:`mode` argument specifies ``nearest`` or
+    ``bilinear`` interpolation method to sample the input pixels.
+
+    :attr:`grid` specifies the sampling pixel locations normalized by the
+    :attr:`input` spatial dimensions. Therefore, it should have most values in
+    the range of ``[-1, 1]``. For example, values ``x = -1, y = -1`` is the
+    left-top pixel of :attr:`input`, and values  ``x = 1, y = 1`` is the
+    right-bottom pixel of :attr:`input`.
+
+    If :attr:`grid` has values outside the range of ``[-1, 1]``, the corresponding
+    outputs are handled as defined by :attr:`padding_mode`. Options are
+
+        * ``padding_mode="zeros"``: use ``0`` for out-of-bound grid locations,
+        * ``padding_mode="border"``: use border values for out-of-bound grid locations,
+        * ``padding_mode="reflection"``: use values at locations reflected by
+          the border for out-of-bound grid locations. For location far away
+          from the border, it will keep being reflected until becoming in bound,
+          e.g., (normalized) pixel location ``x = -3.5`` reflects by border ``-1``
+          and becomes ``x' = 1.5``, then reflects by border ``1`` and becomes
+          ``x'' = -0.5``.
+
+    Note:
+        This function is often used in conjunction with :func:`affine_grid`
+        to build `Spatial Transformer Networks`_ .
+
+    Note:
+        When using the CUDA backend, this operation may induce nondeterministic
+        behaviour in its backward pass that is not easily switched off.
+        Please see the notes on :doc:`/notes/randomness` for background.
+
+    Note:
+        NaN values in :attr:`grid` would be interpreted as ``-1``.
+
+    Args:
+        input (Tensor): input of shape :math:`(N, C, H_\text{in}, W_\text{in})` (4-D case)
+                        or :math:`(N, C, D_\text{in}, H_\text{in}, W_\text{in})` (5-D case)
+        grid (Tensor): flow-field of shape :math:`(N, H_\text{out}, W_\text{out}, 2)` (4-D case)
+                       or :math:`(N, D_\text{out}, H_\text{out}, W_\text{out}, 3)` (5-D case)
+        mode (str): interpolation mode to calculate output values
+            ``'bilinear'`` | ``'nearest'`` | ``'bicubic'``. Default: ``'bilinear'``
+            Note: ``mode='bicubic'`` supports only 4-D input.
+            When ``mode='bilinear'`` and the input is 5-D, the interpolation mode
+            used internally will actually be trilinear. However, when the input is 4-D,
+            the interpolation mode will legitimately be bilinear.
+        padding_mode (str): padding mode for outside grid values
+            ``'zeros'`` | ``'border'`` | ``'reflection'``. Default: ``'zeros'``
+        align_corners (bool, optional): Geometrically, we consider the pixels of the
+            input  as squares rather than points.
+            If set to ``True``, the extrema (``-1`` and ``1``) are considered as referring
+            to the center points of the input's corner pixels. If set to ``False``, they
+            are instead considered as referring to the corner points of the input's corner
+            pixels, making the sampling more resolution agnostic.
+            This option parallels the ``align_corners`` option in
+            :func:`interpolate`, and so whichever option is used here
+            should also be used there to resize the input image before grid sampling.
+            Default: ``False``
+
+    Returns:
+        output (Tensor): output Tensor
+
+    .. _`Spatial Transformer Networks`:
+        https://arxiv.org/abs/1506.02025
+
+    .. warning::
+        When ``align_corners = True``, the grid positions depend on the pixel
+        size relative to the input image size, and so the locations sampled by
+        :func:`grid_sample` will differ for the same input given at different
+        resolutions (that is, after being upsampled or downsampled).
+        The default behavior up to version 1.2.0 was ``align_corners = True``.
+        Since then, the default behavior has been changed to ``align_corners = False``,
+        in order to bring it in line with the default for :func:`interpolate`.
+
+    .. note::
+        ``mode='bicubic'`` is implemented using the `cubic convolution algorithm`_ with :math:`\alpha=-0.75`.
+        The constant :math:`\alpha` might be different from packages to packages.
+        For example, `PIL`_ and `OpenCV`_ use -0.5 and -0.75 respectively.
+        This algorithm may "overshoot" the range of values it's interpolating.
+        For example, it may produce negative values or values greater than 255 when interpolating input in [0, 255].
+        Clamp the results with :func:`torch.clamp` to ensure they are within the valid range.
+    .. _`cubic convolution algorithm`: https://en.wikipedia.org/wiki/Bicubic_interpolation
+    .. _`PIL`: https://github.com/python-pillow/Pillow/blob/4634eafe3c695a014267eefdce830b4a825beed7/src/libImaging/Resample.c#L51
+    .. _`OpenCV`: https://github.com/opencv/opencv/blob/f345ed564a06178670750bad59526cfa4033be55/modules/imgproc/src/resize.cpp#L908
+    """
+    if has_torch_function_variadic(input, grid):
+        return handle_torch_function(
+            grid_sample,
+            (input, grid),
+            input,
+            grid,
+            mode=mode,
+            padding_mode=padding_mode,
+            align_corners=align_corners,
+        )
+    if mode != "bilinear" and mode != "nearest" and mode != "bicubic":
+        raise ValueError(
+            f"nn.functional.grid_sample(): expected mode to be 'bilinear', 'nearest' or 'bicubic', but got: '{mode}'"
+        )
+    if (
+        padding_mode != "zeros"
+        and padding_mode != "border"
+        and padding_mode != "reflection"
+    ):
+        raise ValueError(
+            "nn.functional.grid_sample(): expected padding_mode "
+            "to be 'zeros', 'border', or 'reflection', "
+            f"but got: '{padding_mode}'"
+        )
+
+    if mode == "bilinear":
+        mode_enum = 0
+    elif mode == "nearest":
+        mode_enum = 1
+    else:  # mode == 'bicubic'
+        mode_enum = 2
+
+    if padding_mode == "zeros":
+        padding_mode_enum = 0
+    elif padding_mode == "border":
+        padding_mode_enum = 1
+    else:  # padding_mode == 'reflection'
+        padding_mode_enum = 2
+
+    if align_corners is None:
+        warnings.warn(
+            "Default grid_sample and affine_grid behavior has changed "
+            "to align_corners=False since 1.3.0. Please specify "
+            "align_corners=True if the old behavior is desired. "
+            "See the documentation of grid_sample for details.",
+            stacklevel=2,
+        )
+        align_corners = False
+
+    return torch.grid_sampler(input, grid, mode_enum, padding_mode_enum, align_corners)
+
+
+def affine_grid(
+    theta: Tensor,
+    size: list[int],
+    align_corners: bool | None = None,
+) -> Tensor:
+    r"""Generate 2D or 3D flow field (sampling grid), given a batch of affine matrices :attr:`theta`.
+
+    .. note::
+        This function is often used in conjunction with :func:`grid_sample`
+        to build `Spatial Transformer Networks`_ .
+
+    Args:
+        theta (Tensor): input batch of affine matrices with shape
+            (:math:`N \times 2 \times 3`) for 2D or
+            (:math:`N \times 3 \times 4`) for 3D
+        size (torch.Size): the target output image size.
+            (:math:`N \times C \times H \times W` for 2D or
+            :math:`N \times C \times D \times H \times W` for 3D)
+            Example: torch.Size((32, 3, 24, 24))
+        align_corners (bool, optional): if ``True``, consider ``-1`` and ``1``
+            to refer to the centers of the corner pixels rather than the image corners.
+            Refer to :func:`grid_sample` for a more complete description.
+            A grid generated by :func:`affine_grid` should be passed to :func:`grid_sample`
+            with the same setting for this option.
+            Default: ``False``
+
+    Returns:
+        output (Tensor): output Tensor of size (:math:`N \times H \times W \times 2`)
+
+    .. _`Spatial Transformer Networks`:
+        https://arxiv.org/abs/1506.02025
+
+    .. warning::
+        When ``align_corners = True``, the grid positions depend on the pixel
+        size relative to the input image size, and so the locations sampled by
+        :func:`grid_sample` will differ for the same input given at different
+        resolutions (that is, after being upsampled or downsampled).
+        The default behavior up to version 1.2.0 was ``align_corners = True``.
+        Since then, the default behavior has been changed to ``align_corners = False``,
+        in order to bring it in line with the default for :func:`interpolate`.
+    .. warning::
+        When ``align_corners = True``, 2D affine transforms on 1D data and
+        3D affine transforms on 2D data (that is, when one of the spatial
+        dimensions has unit size) are ill-defined, and not an intended use case.
+        This is not a problem when ``align_corners = False``.
+        Up to version 1.2.0, all grid points along a unit dimension were
+        considered arbitrarily to be at ``-1``.
+        From version 1.3.0, under ``align_corners = True`` all grid points
+        along a unit dimension are considered to be at ``0``
+        (the center of the input image).
+    """
+    if has_torch_function_unary(theta):
+        return handle_torch_function(
+            affine_grid, (theta,), theta, size, align_corners=align_corners
+        )
+    if align_corners is None:
+        warnings.warn(
+            "Default grid_sample and affine_grid behavior has changed "
+            "to align_corners=False since 1.3.0. Please specify "
+            "align_corners=True if the old behavior is desired. "
+            "See the documentation of grid_sample for details.",
+            stacklevel=2,
+        )
+        align_corners = False
+
+    # enforce floating point dtype on theta
+    if not theta.is_floating_point():
+        raise ValueError(
+            f"Expected theta to have floating point type, but got {theta.dtype}"
+        )
+    # check that shapes and sizes match
+    if len(size) == 4:
+        if theta.dim() != 3 or theta.shape[-2] != 2 or theta.shape[-1] != 3:
+            raise ValueError(
+                f"Expected a batch of 2D affine matrices of shape Nx2x3 for size {size}. Got {theta.shape}."
+            )
+        spatial_size = size[-2:]  # spatial dimension sizes
+    elif len(size) == 5:
+        if theta.dim() != 3 or theta.shape[-2] != 3 or theta.shape[-1] != 4:
+            raise ValueError(
+                f"Expected a batch of 3D affine matrices of shape Nx3x4 for size {size}. Got {theta.shape}."
+            )
+        spatial_size = size[-3:]  # spatial dimension sizes
+    else:
+        raise NotImplementedError(
+            "affine_grid only supports 4D and 5D sizes, "
+            "for 2D and 3D affine transforms, respectively. "
+            f"Got size {size}."
+        )
+    # check for empty span
+    if align_corners and min(spatial_size) == 1:
+        warnings.warn(
+            "Since version 1.3.0, affine_grid behavior has changed "
+            "for unit-size grids when align_corners=True. "
+            "This is not an intended use case of affine_grid. "
+            "See the documentation of affine_grid for details.",
+            stacklevel=2,
+        )
+    elif min(size) <= 0:
+        raise ValueError(f"Expected non-zero, positive output size. Got {size}")
+
+    return torch.affine_grid_generator(theta, size, align_corners)
+
+
+def pad(
+    input: Tensor,
+    pad: list[int],
+    mode: str = "constant",
+    value: float | None = None,
+) -> Tensor:
+    r"""
+    pad(input, pad, mode="constant", value=None) -> Tensor
+
+    Pads tensor.
+
+    Padding size:
+        The padding size by which to pad some dimensions of :attr:`input`
+        are described starting from the last dimension and moving forward.
+        :math:`\left\lfloor\frac{\text{len(pad)}}{2}\right\rfloor` dimensions
+        of ``input`` will be padded.
+        For example, to pad only the last dimension of the input tensor, then
+        :attr:`pad` has the form
+        :math:`(\text{padding\_left}, \text{padding\_right})`;
+        to pad the last 2 dimensions of the input tensor, then use
+        :math:`(\text{padding\_left}, \text{padding\_right},`
+        :math:`\text{padding\_top}, \text{padding\_bottom})`;
+        to pad the last 3 dimensions, use
+        :math:`(\text{padding\_left}, \text{padding\_right},`
+        :math:`\text{padding\_top}, \text{padding\_bottom}`
+        :math:`\text{padding\_front}, \text{padding\_back})`.
+
+    Padding mode:
+        See :class:`torch.nn.CircularPad2d`, :class:`torch.nn.ConstantPad2d`,
+        :class:`torch.nn.ReflectionPad2d`, and :class:`torch.nn.ReplicationPad2d`
+        for concrete examples on how each of the padding modes works. Constant
+        padding is implemented for arbitrary dimensions. Circular, replicate and
+        reflection padding are implemented for padding the last 3 dimensions of a
+        4D or 5D input tensor, the last 2 dimensions of a 3D or 4D input tensor,
+        or the last dimension of a 2D or 3D input tensor.
+
+    Note:
+        When using the CUDA backend, this operation may induce nondeterministic
+        behaviour in its backward pass that is not easily switched off.
+        Please see the notes on :doc:`/notes/randomness` for background.
+
+    Args:
+        input (Tensor): N-dimensional tensor
+        pad (tuple): m-elements tuple, where
+            :math:`\frac{m}{2} \leq` input dimensions and :math:`m` is even.
+        mode: ``'constant'``, ``'reflect'``, ``'replicate'`` or ``'circular'``.
+            Default: ``'constant'``
+        value: fill value for ``'constant'`` padding. Default: ``0``
+
+    Examples::
+
+        >>> t4d = torch.empty(3, 3, 4, 2)
+        >>> p1d = (1, 1) # pad last dim by 1 on each side
+        >>> out = F.pad(t4d, p1d, "constant", 0)  # effectively zero padding
+        >>> print(out.size())
+        torch.Size([3, 3, 4, 4])
+        >>> p2d = (1, 1, 2, 2) # pad last dim by (1, 1) and 2nd to last by (2, 2)
+        >>> out = F.pad(t4d, p2d, "constant", 0)
+        >>> print(out.size())
+        torch.Size([3, 3, 8, 4])
+        >>> t4d = torch.empty(3, 3, 4, 2)
+        >>> p3d = (0, 1, 2, 1, 3, 3) # pad by (0, 1), (2, 1), and (3, 3)
+        >>> out = F.pad(t4d, p3d, "constant", 0)
+        >>> print(out.size())
+        torch.Size([3, 9, 7, 3])
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            torch.nn.functional.pad, (input,), input, pad, mode=mode, value=value
+        )
+    if not torch.jit.is_scripting():
+        if torch.are_deterministic_algorithms_enabled() and (
+            input.is_cuda or input.is_xpu
+        ):
+            if mode == "replicate":
+                # Use slow decomp whose backward will be in terms of index_put.
+                if torch.compiler.is_compiling():
+                    # nonstrict_trace makes Dynamo skip the function body
+                    # (which contains Dynamo-untraceable code) while
+                    # AOTAutograd still traces into it for the backward.
+                    return torch._dynamo.decorators.nonstrict_trace(
+                        torch._decomp.decompositions._replication_pad
+                    )(input, pad)
+                # importlib is required because the import cannot be top level
+                # (cycle) and cannot be nested (TS doesn't support)
+                return importlib.import_module(
+                    "torch._decomp.decompositions"
+                )._replication_pad(input, pad)
+    return torch._C._nn.pad(input, pad, mode, value)
+
+
+# TODO: Fix via https://github.com/pytorch/pytorch/issues/75798
+pad.__module__ = "torch.nn.functional"
+
+# distance
+
+
+pairwise_distance = _add_docstr(
+    torch.pairwise_distance,
+    r"""
+pairwise_distance(x1, x2, p=2.0, eps=1e-6, keepdim=False) -> Tensor
+
+See :class:`torch.nn.PairwiseDistance` for details
+""",
+)
+
+
+pdist = _add_docstr(
+    torch.pdist,
+    r"""
+pdist(input, p=2) -> Tensor
+
+Computes the p-norm distance between every pair of row vectors in the input.
+This is identical to the upper triangular portion, excluding the diagonal, of
+`torch.norm(input[:, None] - input, dim=2, p=p)`. This function will be faster
+if the rows are contiguous.
+
+If input has shape :math:`N \times M` then the output will have shape
+:math:`\frac{1}{2} N (N - 1)`.
+
+This function is equivalent to ``scipy.spatial.distance.pdist(input,
+'minkowski', p=p)`` if :math:`p \in (0, \infty)`. When :math:`p = 0` it is
+equivalent to ``scipy.spatial.distance.pdist(input, 'hamming') * M``.
+When :math:`p = \infty`, the closest scipy function is
+``scipy.spatial.distance.pdist(xn, lambda x, y: np.abs(x - y).max())``.
+
+Args:
+    input: input tensor of shape :math:`N \times M`.
+    p: p value for the p-norm distance to calculate between each vector pair
+        :math:`\in [0, \infty]`.
+""",
+)
+
+
+cosine_similarity = _add_docstr(
+    torch.cosine_similarity,
+    r"""
+cosine_similarity(x1, x2, dim=1, eps=1e-8) -> Tensor
+
+Returns cosine similarity between ``x1`` and ``x2``, computed along dim. ``x1`` and ``x2`` must be broadcastable
+to a common shape. ``dim`` refers to the dimension in this common shape. Dimension ``dim`` of the output is
+squeezed (see :func:`torch.squeeze`), resulting in the
+output tensor having 1 fewer dimension.
+
+.. math ::
+    \text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2, \epsilon) \cdot \max(\Vert x_2 \Vert _2, \epsilon)}
+
+Supports :ref:`type promotion <type-promotion-doc>`.
+
+Args:
+    x1 (Tensor): First input.
+    x2 (Tensor): Second input.
+    dim (int, optional): Dimension along which cosine similarity is computed. Default: 1
+    eps (float, optional): Small value to avoid division by zero.
+        Default: 1e-8
+
+Example::
+
+    >>> input1 = torch.randn(100, 128)
+    >>> input2 = torch.randn(100, 128)
+    >>> output = F.cosine_similarity(input1, input2)
+    >>> print(output)
+""",
+)
+
+
+one_hot = _add_docstr(
+    torch._C._nn.one_hot,
+    r"""
+one_hot(tensor, num_classes=-1) -> LongTensor
+
+Takes LongTensor with index values of shape ``(*)`` and returns a tensor
+of shape ``(*, num_classes)`` that have zeros everywhere except where the
+index of last dimension matches the corresponding value of the input tensor,
+in which case it will be 1.
+
+See also `One-hot on Wikipedia`_ .
+
+.. _One-hot on Wikipedia:
+    https://en.wikipedia.org/wiki/One-hot
+
+Arguments:
+    tensor (LongTensor): class values of any shape.
+    num_classes (int, optional):  Total number of classes. If set to -1, the number
+        of classes will be inferred as one greater than the largest class
+        value in the input tensor. Default: -1
+
+Returns:
+    LongTensor that has one more dimension with 1 values at the
+    index of last dimension indicated by the input, and 0 everywhere
+    else.
+
+Examples:
+    >>> F.one_hot(torch.arange(0, 5) % 3)
+    tensor([[1, 0, 0],
+            [0, 1, 0],
+            [0, 0, 1],
+            [1, 0, 0],
+            [0, 1, 0]])
+    >>> F.one_hot(torch.arange(0, 5) % 3, num_classes=5)
+    tensor([[1, 0, 0, 0, 0],
+            [0, 1, 0, 0, 0],
+            [0, 0, 1, 0, 0],
+            [1, 0, 0, 0, 0],
+            [0, 1, 0, 0, 0]])
+    >>> F.one_hot(torch.arange(0, 6).view(3,2) % 3)
+    tensor([[[1, 0, 0],
+             [0, 1, 0]],
+            [[0, 0, 1],
+             [1, 0, 0]],
+            [[0, 1, 0],
+             [0, 0, 1]]])
+""",
+)
+
+
+def triplet_margin_loss(
+    anchor: Tensor,
+    positive: Tensor,
+    negative: Tensor,
+    margin: float = 1.0,
+    p: float = 2,
+    eps: float = 1e-6,
+    swap: bool = False,
+    size_average: bool | None = None,
+    reduce: bool | None = None,
+    reduction: str = "mean",
+) -> Tensor:
+    r"""Compute the triplet loss between given input tensors and a margin greater than 0.
+
+    See :class:`~torch.nn.TripletMarginLoss` for details.
+    """
+    if has_torch_function_variadic(anchor, positive, negative):
+        return handle_torch_function(
+            triplet_margin_loss,
+            (anchor, positive, negative),
+            anchor,
+            positive,
+            negative,
+            margin=margin,
+            p=p,
+            eps=eps,
+            swap=swap,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+    if size_average is not None or reduce is not None:
+        reduction_enum = _Reduction.legacy_get_enum(size_average, reduce)
+    else:
+        reduction_enum = _Reduction.get_enum(reduction)
+    if margin <= 0:
+        raise ValueError(f"margin must be greater than 0, got {margin}")
+    return torch.triplet_margin_loss(
+        anchor, positive, negative, margin, p, eps, swap, reduction_enum
+    )
+
+
+def triplet_margin_with_distance_loss(
+    anchor: Tensor,
+    positive: Tensor,
+    negative: Tensor,
+    *,
+    distance_function: Callable[[Tensor, Tensor], Tensor] | None = None,
+    margin: float = 1.0,
+    swap: bool = False,
+    reduction: str = "mean",
+) -> Tensor:
+    r"""Compute the triplet margin loss for input tensors using a custom distance function.
+
+    See :class:`~torch.nn.TripletMarginWithDistanceLoss` for details.
+    """
+    if torch.jit.is_scripting():
+        raise NotImplementedError(
+            "F.triplet_margin_with_distance_loss does not support JIT scripting: "
+            "functions requiring Callables cannot be scripted."
+        )
+
+    if has_torch_function_variadic(anchor, positive, negative):
+        return handle_torch_function(
+            triplet_margin_with_distance_loss,
+            (anchor, positive, negative),
+            anchor,
+            positive,
+            negative,
+            distance_function=distance_function,
+            margin=margin,
+            swap=swap,
+            reduction=reduction,
+        )
+
+    # Check validity of reduction mode
+    if reduction not in ("mean", "sum", "none"):
+        raise ValueError(f"{reduction} is not a valid value for reduction")
+
+    # Check validity of margin
+    if margin <= 0:
+        raise ValueError(f"margin must be greater than 0, got {margin}")
+
+    # Check dimensions
+    a_dim = anchor.ndim
+    p_dim = positive.ndim
+    n_dim = negative.ndim
+    if not (a_dim == p_dim and p_dim == n_dim):
+        raise RuntimeError(
+            f"The anchor, positive, and negative tensors are expected to have "
+            f"the same number of dimensions, but got: anchor {a_dim}D, "
+            f"positive {p_dim}D, and negative {n_dim}D inputs"
+        )
+
+    # Calculate loss
+    if distance_function is None:
+        distance_function = torch.pairwise_distance
+
+    dist_pos = distance_function(anchor, positive)
+    dist_neg = distance_function(anchor, negative)
+    # The distance swap is described in the paper "Learning shallow
+    # convolutional feature descriptors with triplet losses" by V. Balntas, E.
+    # Riba et al.  If True, and if the positive example is closer to the
+    # negative example than the anchor is, swaps the positive example and the
+    # anchor in the loss computation.
+    if swap:
+        dist_swap = distance_function(positive, negative)
+        dist_neg = torch.minimum(dist_neg, dist_swap)
+    loss = torch.clamp_min(margin + dist_pos - dist_neg, 0)
+
+    # Apply reduction
+    if reduction == "sum":
+        return torch.sum(loss)
+    elif reduction == "mean":
+        return torch.mean(loss)
+    else:  # reduction == "none"
+        return loss
+
+
+def normalize(
+    input: Tensor,
+    p: float = 2.0,
+    dim: int = 1,
+    eps: float = 1e-12,
+    out: Tensor | None = None,
+) -> Tensor:
+    r"""Perform :math:`L_p` normalization of inputs over specified dimension.
+
+    For a tensor :attr:`input` of sizes :math:`(n_0, ..., n_{dim}, ..., n_k)`, each
+    :math:`n_{dim}` -element vector :math:`v` along dimension :attr:`dim` is transformed as
+
+    .. math::
+        v = \frac{v}{\max(\lVert v \rVert_p, \epsilon)}.
+
+    With the default arguments it uses the Euclidean norm over vectors along dimension :math:`1` for normalization.
+
+    Args:
+        input: input tensor of any shape
+        p (float): the exponent value in the norm formulation. Default: 2
+        dim (int or tuple of ints): the dimension to reduce. Default: 1
+        eps (float): small value to avoid division by zero. Default: 1e-12
+        out (Tensor, optional): the output tensor. If :attr:`out` is used, this
+                                operation won't be differentiable.
+    """
+    if has_torch_function_variadic(input, out):
+        return handle_torch_function(
+            normalize, (input, out), input, p=p, dim=dim, eps=eps, out=out
+        )
+    if out is None:
+        denom = input.norm(p, dim, keepdim=True).clamp_min(eps).expand_as(input)
+        return input / denom
+    else:
+        denom = input.norm(p, dim, keepdim=True).clamp_min_(eps).expand_as(input)
+        return torch.div(input, denom, out=out)
+
+
+def assert_int_or_pair(arg: list[int], arg_name: str, message: str) -> None:
+    if not (isinstance(arg, int) or len(arg) == 2):
+        raise AssertionError(message.format(arg_name))
+
+
+def unfold(
+    input: Tensor,
+    kernel_size: BroadcastingList2[int],
+    dilation: BroadcastingList2[int] = 1,
+    padding: BroadcastingList2[int] = 0,
+    stride: BroadcastingList2[int] = 1,
+) -> Tensor:
+    r"""Extract sliding local blocks from a batched input tensor.
+
+    .. warning::
+        Currently, only 4-D input tensors (batched image-like tensors) are
+        supported.
+
+    .. warning::
+
+        More than one element of the unfolded tensor may refer to a single
+        memory location. As a result, in-place operations (especially ones that
+        are vectorized) may result in incorrect behavior. If you need to write
+        to the tensor, please clone it first.
+
+
+    See :class:`torch.nn.Unfold` for details
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            unfold,
+            (input,),
+            input,
+            kernel_size,
+            dilation=dilation,
+            padding=padding,
+            stride=stride,
+        )
+    return torch._C._nn.im2col(
+        input, _pair(kernel_size), _pair(dilation), _pair(padding), _pair(stride)
+    )
+
+
+def fold(
+    input: Tensor,
+    output_size: BroadcastingList2[int],
+    kernel_size: BroadcastingList2[int],
+    dilation: BroadcastingList2[int] = 1,
+    padding: BroadcastingList2[int] = 0,
+    stride: BroadcastingList2[int] = 1,
+) -> Tensor:
+    r"""Combine an array of sliding local blocks into a large containing tensor.
+
+    .. warning::
+        Currently, only unbatched (3D) or batched (4D) image-like output tensors are supported.
+
+    See :class:`torch.nn.Fold` for details
+    """
+    if has_torch_function_unary(input):
+        return handle_torch_function(
+            fold,
+            (input,),
+            input,
+            output_size,
+            kernel_size,
+            dilation=dilation,
+            padding=padding,
+            stride=stride,
+        )
+    return torch._C._nn.col2im(
+        input,
+        _pair(output_size),
+        _pair(kernel_size),
+        _pair(dilation),
+        _pair(padding),
+        _pair(stride),
+    )
+
+
+#
+# multihead attention
+#
+
+
+def _in_projection_packed(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w: Tensor,
+    b: Tensor | None = None,
+) -> list[Tensor]:
+    r"""Perform the in-projection step of the attention operation, using packed weights.
+
+    Output is a triple containing projection tensors for query, key and value.
+
+    Args:
+        q, k, v: query, key and value tensors to be projected. For self-attention,
+            these are typically the same tensor; for encoder-decoder attention,
+            k and v are typically the same tensor. (We take advantage of these
+            identities for performance if they are present.) Regardless, q, k and v
+            must share a common embedding dimension; otherwise their shapes may vary.
+        w: projection weights for q, k and v, packed into a single tensor. Weights
+            are packed along dimension 0, in q, k, v order.
+        b: optional projection biases for q, k and v, packed into a single tensor
+            in q, k, v order.
+
+    Shape:
+        Inputs:
+        - q: :math:`(..., E)` where E is the embedding dimension
+        - k: :math:`(..., E)` where E is the embedding dimension
+        - v: :math:`(..., E)` where E is the embedding dimension
+        - w: :math:`(E * 3, E)` where E is the embedding dimension
+        - b: :math:`E * 3` where E is the embedding dimension
+
+        Output:
+        - in output list :math:`[q', k', v']`, each output tensor will have the
+            same shape as the corresponding input tensor.
+    """
+    E = q.size(-1)
+    if k is v:
+        if q is k:
+            # self-attention
+            proj = linear(q, w, b)
+            # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()
+            proj = (
+                proj.unflatten(-1, (3, E))
+                .unsqueeze(0)
+                .transpose(0, -2)
+                .squeeze(-2)
+                .contiguous()
+            )
+            # pyrefly: ignore [bad-return]
+            return proj[0], proj[1], proj[2]
+        else:
+            # encoder-decoder attention
+            w_q, w_kv = w.split([E, E * 2])
+            if b is None:
+                b_q = b_kv = None
+            else:
+                b_q, b_kv = b.split([E, E * 2])
+            q_proj = linear(q, w_q, b_q)
+            kv_proj = linear(k, w_kv, b_kv)
+            # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()
+            kv_proj = (
+                kv_proj.unflatten(-1, (2, E))
+                .unsqueeze(0)
+                .transpose(0, -2)
+                .squeeze(-2)
+                .contiguous()
+            )
+            # pyrefly: ignore [bad-return]
+            return (q_proj, kv_proj[0], kv_proj[1])
+    else:
+        w_q, w_k, w_v = w.chunk(3)
+        if b is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = b.chunk(3)
+        # pyrefly: ignore [bad-return]
+        return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
+
+
+def _in_projection(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w_q: Tensor,
+    w_k: Tensor,
+    w_v: Tensor,
+    b_q: Tensor | None = None,
+    b_k: Tensor | None = None,
+    b_v: Tensor | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    r"""Perform the in-projection step of the attention operation.
+
+    This is simply a triple of linear projections,
+    with shape constraints on the weights which
+    ensure embedding dimension uniformity in the projected outputs.
+    Output is a triple containing projection tensors for query, key and value.
+
+    Args:
+        q, k, v: query, key and value tensors to be projected.
+        w_q, w_k, w_v: weights for q, k and v, respectively.
+        b_q, b_k, b_v: optional biases for q, k and v, respectively.
+
+    Shape:
+        Inputs:
+        - q: :math:`(Qdims..., Eq)` where Eq is the query embedding dimension and Qdims are any
+            number of leading dimensions.
+        - k: :math:`(Kdims..., Ek)` where Ek is the key embedding dimension and Kdims are any
+            number of leading dimensions.
+        - v: :math:`(Vdims..., Ev)` where Ev is the value embedding dimension and Vdims are any
+            number of leading dimensions.
+        - w_q: :math:`(Eq, Eq)`
+        - w_k: :math:`(Eq, Ek)`
+        - w_v: :math:`(Eq, Ev)`
+        - b_q: :math:`(Eq)`
+        - b_k: :math:`(Eq)`
+        - b_v: :math:`(Eq)`
+
+        Output: in output triple :math:`(q', k', v')`,
+         - q': :math:`[Qdims..., Eq]`
+         - k': :math:`[Kdims..., Eq]`
+         - v': :math:`[Vdims..., Eq]`
+
+    """
+    Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
+    if w_q.shape != (Eq, Eq):
+        raise AssertionError(
+            f"expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}"
+        )
+    if w_k.shape != (Eq, Ek):
+        raise AssertionError(
+            f"expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}"
+        )
+    if w_v.shape != (Eq, Ev):
+        raise AssertionError(
+            f"expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}"
+        )
+    if b_q is not None and b_q.shape != (Eq,):
+        raise AssertionError(
+            f"expecting query bias shape of {(Eq,)}, but got {b_q.shape}"
+        )
+    if b_k is not None and b_k.shape != (Eq,):
+        raise AssertionError(
+            f"expecting key bias shape of {(Eq,)}, but got {b_k.shape}"
+        )
+    if b_v is not None and b_v.shape != (Eq,):
+        raise AssertionError(
+            f"expecting value bias shape of {(Eq,)}, but got {b_v.shape}"
+        )
+    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
+
+
+scaled_dot_product_attention = _add_docstr(
+    torch._C._nn.scaled_dot_product_attention,
+    r"""scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
+        is_causal=False, scale=None, enable_gqa=False) -> Tensor:
+
+    Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed,
+    and applying dropout if a probability greater than 0.0 is specified. The optional scale argument can only be
+    specified as a keyword argument.
+
+    .. code-block:: python
+
+        # Efficient implementation equivalent to the following:
+        def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
+                is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor:
+            L, S = query.size(-2), key.size(-2)
+            scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+            attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+            if is_causal:
+                assert attn_mask is None
+                temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0)
+                attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+
+            if attn_mask is not None:
+                if attn_mask.dtype == torch.bool:
+                    attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+                else:
+                    attn_bias = attn_mask + attn_bias
+
+            if enable_gqa:
+                key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
+                value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
+
+            attn_weight = query @ key.transpose(-2, -1) * scale_factor
+            attn_weight += attn_bias
+            attn_weight = torch.softmax(attn_weight, dim=-1)
+            attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+            return attn_weight @ value
+
+    .. warning::
+        This function is beta and subject to change.
+
+    .. warning::
+        This function always applies dropout according to the specified ``dropout_p`` argument.
+        To disable dropout during evaluation, be sure to pass a value of ``0.0`` when the module
+        that makes the function call is not in training mode.
+
+        For example:
+
+        .. code-block:: python
+
+            class MyModel(nn.Module):
+                def __init__(self, p=0.5):
+                    super().__init__()
+                    self.p = p
+
+                def forward(self, ...):
+                    return F.scaled_dot_product_attention(...,
+                        dropout_p=(self.p if self.training else 0.0))
+
+    Note:
+
+        The boolean mask semantics for ``attn_mask`` are the inverse of
+        :class:`~torch.nn.MultiheadAttention`'s ``key_padding_mask``.
+
+        In :func:`scaled_dot_product_attention`, ``True`` indicates values
+        to **participate** in attention.
+
+        In :class:`~torch.nn.MultiheadAttention`, ``True`` indicates values
+        to be **masked out** (padding).
+
+        If migrating from MHA, ensure you invert your boolean mask (e.g.,
+        using ``~mask`` or ``mask.logical_not()``).
+
+
+    Note:
+
+        There are currently three supported implementations of scaled dot product attention:
+
+            - `FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning`_
+            - `Memory-Efficient Attention`_
+            - A PyTorch implementation defined in C++ matching the above formulation
+
+        The function may call optimized kernels for improved performance when using the CUDA backend.
+        For all other backends, the PyTorch implementation will be used.
+
+        All implementations are enabled by default. Scaled dot product attention attempts to automatically select the
+        most optimal implementation based on the inputs. In order to provide more fine-grained control over what implementation
+        is used, the following functions are provided for enabling and disabling implementations.
+        The context manager is the preferred mechanism:
+
+            - :func:`torch.nn.attention.sdpa_kernel`: A context manager used to enable or disable any of the implementations.
+            - :func:`torch.backends.cuda.enable_flash_sdp`: Globally enables or disables FlashAttention.
+            - :func:`torch.backends.cuda.enable_mem_efficient_sdp`: Globally enables or disables  Memory-Efficient Attention.
+            - :func:`torch.backends.cuda.enable_math_sdp`: Globally enables or disables  the PyTorch C++ implementation.
+
+        Each of the fused kernels has specific input limitations. If the user requires the use of a specific fused implementation,
+        disable the PyTorch C++ implementation using :func:`torch.nn.attention.sdpa_kernel`.
+        In the event that a fused implementation is not available, a warning will be raised with the
+        reasons why the fused implementation cannot run.
+
+        Due to the nature of fusing floating point operations, the output of this function may be different
+        depending on what backend kernel is chosen.
+        The c++ implementation supports torch.float64 and can be used when higher precision is required.
+        For math backend, all intermediates are kept in torch.float if inputs are in torch.half or torch.bfloat16.
+    For more information please see :doc:`/notes/numerical_accuracy`
+
+        Grouped Query Attention (GQA) is an experimental feature. It currently works only for Flash_attention
+        and math kernel on CUDA tensor, and does not support Nested tensor.
+        Constraints for GQA:
+
+            - number_of_heads_query % number_of_heads_key_value == 0 and,
+            - number_of_heads_key == number_of_heads_value
+
+    Note:
+
+        {cudnn_reproducibility_note}
+    """.format(**reproducibility_notes)
+    + r"""
+    Args:
+        query (Tensor): Query tensor; shape :math:`(N, ..., Hq, L, E)`.
+        key (Tensor): Key tensor; shape :math:`(N, ..., H, S, E)`.
+        value (Tensor): Value tensor; shape :math:`(N, ..., H, S, Ev)`.
+        attn_mask (optional Tensor): Attention mask; shape must be broadcastable to the shape of attention weights,
+            which is :math:`(N,..., Hq, L, S)`. Two types of masks are supported.
+            A boolean mask where a value of True indicates that the element *should* take part in attention.
+            A float mask of the same type as query, key, value that is added to the attention score.
+        dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied
+        is_causal (bool): If set to true, the attention masking is a lower triangular matrix when the mask is a
+            square matrix. The attention masking has the form of the upper left causal bias due to the alignment
+            (see :class:`torch.nn.attention.bias.CausalBias`) when the mask is a non-square matrix.
+            An error is thrown if both attn_mask and is_causal are set.
+        scale (optional float, keyword-only): Scaling factor applied prior to softmax. If None, the default value is set
+            to :math:`\frac{1}{\sqrt{E}}`.
+        enable_gqa (bool): If set to True, Grouped Query Attention (GQA) is enabled, by default it is set to False.
+
+    Returns:
+        output (Tensor): Attention output; shape :math:`(N, ..., Hq, L, Ev)`.
+
+    Shape legend:
+        - :math:`N: \text{Batch size} ... : \text{Any number of other batch dimensions (optional)}`
+        - :math:`S: \text{Source sequence length}`
+        - :math:`L: \text{Target sequence length}`
+        - :math:`E: \text{Embedding dimension of the query and key}`
+        - :math:`Ev: \text{Embedding dimension of the value}`
+        - :math:`Hq: \text{Number of heads of query}`
+        - :math:`H: \text{Number of heads of key and value}`
+
+    Examples:
+
+        >>> # Optionally use the context manager to ensure one of the fused kernels is run
+        >>> query = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
+        >>> key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
+        >>> value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
+        >>> with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
+        >>>     F.scaled_dot_product_attention(query,key,value)
+
+
+        >>> # Sample for GQA for llama3
+        >>> query = torch.rand(32, 32, 128, 64, dtype=torch.float16, device="cuda")
+        >>> key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
+        >>> value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
+        >>> with sdpa_kernel(backends=[SDPBackend.MATH]):
+        >>>     F.scaled_dot_product_attention(query,key,value,enable_gqa=True)
+
+
+    .. _FlashAttention-2\: Faster Attention with Better Parallelism and Work Partitioning:
+        https://arxiv.org/abs/2307.08691
+    .. _Memory-Efficient Attention:
+        https://github.com/facebookresearch/xformers
+    .. _Grouped-Query Attention:
+        https://arxiv.org/pdf/2305.13245
+    """,
+)
+
+
+def _mha_shape_check(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    key_padding_mask: Tensor | None,
+    attn_mask: Tensor | None,
+    num_heads: int,
+):
+    # Verifies the expected shape for `query, `key`, `value`, `key_padding_mask` and `attn_mask`
+    # and returns if the input is batched or not.
+    # Raises an error if `query` is not 2-D (unbatched) or 3-D (batched) tensor.
+
+    # Shape check.
+    if query.dim() == 3:
+        # Batched Inputs
+        is_batched = True
+        if key.dim() != 3 or value.dim() != 3:
+            raise AssertionError(
+                "For batched (3-D) `query`, expected `key` and `value` to be 3-D"
+                f" but found {key.dim()}-D and {value.dim()}-D tensors respectively"
+            )
+        if key_padding_mask is not None:
+            if key_padding_mask.dim() != 2:
+                raise AssertionError(
+                    "For batched (3-D) `query`, expected `key_padding_mask` to be `None` or 2-D"
+                    f" but found {key_padding_mask.dim()}-D tensor instead"
+                )
+        if attn_mask is not None:
+            if attn_mask.dim() not in (2, 3):
+                raise AssertionError(
+                    "For batched (3-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D"
+                    f" but found {attn_mask.dim()}-D tensor instead"
+                )
+    elif query.dim() == 2:
+        # Unbatched Inputs
+        is_batched = False
+        if key.dim() != 2 or value.dim() != 2:
+            raise AssertionError(
+                "For unbatched (2-D) `query`, expected `key` and `value` to be 2-D"
+                f" but found {key.dim()}-D and {value.dim()}-D tensors respectively"
+            )
+
+        if key_padding_mask is not None:
+            if key_padding_mask.dim() != 1:
+                raise AssertionError(
+                    "For unbatched (2-D) `query`, expected `key_padding_mask` to be `None` or 1-D"
+                    f" but found {key_padding_mask.dim()}-D tensor instead"
+                )
+
+        if attn_mask is not None:
+            if attn_mask.dim() not in (2, 3):
+                raise AssertionError(
+                    "For unbatched (2-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D"
+                    f" but found {attn_mask.dim()}-D tensor instead"
+                )
+            if attn_mask.dim() == 3:
+                expected_shape = (num_heads, query.shape[0], key.shape[0])
+                if attn_mask.shape != expected_shape:
+                    raise AssertionError(
+                        f"Expected `attn_mask` shape to be {expected_shape} but got {attn_mask.shape}"
+                    )
+    else:
+        raise AssertionError(
+            f"query should be unbatched 2D or batched 3D tensor but received {query.dim()}-D query tensor"
+        )
+
+    return is_batched
+
+
+def _canonical_mask(
+    mask: Tensor | None,
+    mask_name: str,
+    other_type: DType | None,
+    other_name: str,
+    target_type: DType,
+    check_other: bool = True,
+) -> Tensor | None:
+    if mask is not None:
+        _mask_dtype = mask.dtype
+        _mask_is_float = torch.is_floating_point(mask)
+        if _mask_dtype != torch.bool and not _mask_is_float:
+            raise AssertionError(
+                f"only bool and floating types of {mask_name} are supported"
+            )
+        if check_other and other_type is not None:
+            if _mask_dtype != other_type:
+                warnings.warn(
+                    f"Support for mismatched {mask_name} and {other_name} "
+                    "is deprecated. Use same type for both instead.",
+                    stacklevel=2,
+                )
+        if not _mask_is_float:
+            mask = torch.zeros_like(mask, dtype=target_type).masked_fill_(
+                mask, float("-inf")
+            )
+    return mask
+
+
+def _none_or_dtype(input: Tensor | None) -> DType | None:
+    if input is None:
+        return None
+    elif isinstance(input, torch.Tensor):
+        return input.dtype
+    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")
+
+
+def _check_key_padding_mask(
+    key_padding_mask: torch.Tensor, src_len: int, bsz: int
+) -> None:
+    torch._check_with(
+        AssertionError,
+        key_padding_mask.shape[0] == bsz,
+        lambda: f"Expected key_padded_mask.shape[0] to be {bsz}, but got {key_padding_mask.shape[0]}",
+    )
+    torch._check_with(
+        AssertionError,
+        key_padding_mask.shape[1] == src_len,
+        lambda: f"Expected key_padded_mask.shape[1] to be {src_len}, but got {key_padding_mask.shape[1]}",
+    )
+
+
+def multi_head_attention_forward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    embed_dim_to_check: int,
+    num_heads: int,
+    in_proj_weight: Tensor | None,
+    in_proj_bias: Tensor | None,
+    bias_k: Tensor | None,
+    bias_v: Tensor | None,
+    add_zero_attn: bool,
+    dropout_p: float,
+    out_proj_weight: Tensor,
+    out_proj_bias: Tensor | None,
+    training: bool = True,
+    key_padding_mask: Tensor | None = None,
+    need_weights: bool = True,
+    attn_mask: Tensor | None = None,
+    use_separate_proj_weight: bool = False,
+    q_proj_weight: Tensor | None = None,
+    k_proj_weight: Tensor | None = None,
+    v_proj_weight: Tensor | None = None,
+    static_k: Tensor | None = None,
+    static_v: Tensor | None = None,
+    average_attn_weights: bool = True,
+    is_causal: bool = False,
+) -> tuple[Tensor, Tensor | None]:
+    r"""Forward method for MultiHeadAttention.
+
+    See :class:`torch.nn.MultiheadAttention` for details.
+
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        embed_dim_to_check: total dimension of the model.
+        num_heads: parallel attention heads.
+        in_proj_weight, in_proj_bias: input projection weight and bias.
+        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        dropout_p: probability of an element to be zeroed.
+        out_proj_weight, out_proj_bias: the output projection weight and bias.
+        training: apply dropout if is ``True``.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. This is an binary mask. When the value is True,
+            the corresponding value on the attention layer will be filled with -inf.
+        need_weights: output attn_output_weights.
+            Default: `True`
+            Note: `needs_weight` defaults to `True`, but should be set to `False`
+            For best performance when attention weights are not needed.
+            *Setting needs_weights to `True`
+            leads to a significant performance degradation.*
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+        is_causal: If specified, applies a causal mask as attention mask, and ignores
+            attn_mask for computing scaled dot product attention.
+            Default: ``False``.
+            .. warning::
+                is_causal is provides a hint that the attn_mask is the
+                causal mask.Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+        use_separate_proj_weight: the function accept the proj. weights for query, key,
+            and value in different forms. If false, in_proj_weight will be used, which is
+            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
+        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
+        static_k, static_v: static key and value used for attention operators.
+        average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across heads.
+            Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an effect
+            when ``need_weights=True.``. Default: True
+
+
+    Shape:
+        Inputs:
+        - query: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, E)` or :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, E)` or :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a FloatTensor is provided, it will be directly added to the value.
+          If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+          positions. If a BoolTensor is provided, positions with ``True``
+          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+
+        Outputs:
+        - attn_output: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: Only returned when ``need_weights=True``. If ``average_attn_weights=True``, returns
+          attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+          :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+          :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+          head of shape :math:`(num_heads, L, S)` when input is unbatched or :math:`(N, num_heads, L, S)`.
+    """
+    tens_ops = (
+        query,
+        key,
+        value,
+        in_proj_weight,
+        in_proj_bias,
+        bias_k,
+        bias_v,
+        out_proj_weight,
+        out_proj_bias,
+    )
+    if has_torch_function(tens_ops):
+        return handle_torch_function(
+            multi_head_attention_forward,
+            tens_ops,
+            query,
+            key,
+            value,
+            embed_dim_to_check,
+            num_heads,
+            in_proj_weight,
+            in_proj_bias,
+            bias_k,
+            bias_v,
+            add_zero_attn,
+            dropout_p,
+            out_proj_weight,
+            out_proj_bias,
+            training=training,
+            key_padding_mask=key_padding_mask,
+            need_weights=need_weights,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            use_separate_proj_weight=use_separate_proj_weight,
+            q_proj_weight=q_proj_weight,
+            k_proj_weight=k_proj_weight,
+            v_proj_weight=v_proj_weight,
+            static_k=static_k,
+            static_v=static_v,
+            average_attn_weights=average_attn_weights,
+        )
+
+    is_batched = _mha_shape_check(
+        query, key, value, key_padding_mask, attn_mask, num_heads
+    )
+
+    # For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input
+    # is batched, run the computation and before returning squeeze the
+    # batch dimension so that the output doesn't carry this temporary batch dimension.
+    if not is_batched:
+        # unsqueeze if the input is unbatched
+        query = query.unsqueeze(1)
+        key = key.unsqueeze(1)
+        value = value.unsqueeze(1)
+        if key_padding_mask is not None:
+            key_padding_mask = key_padding_mask.unsqueeze(0)
+
+    # set up shape vars
+    tgt_len, bsz, embed_dim = query.shape
+    src_len, _, _ = key.shape
+
+    key_padding_mask = _canonical_mask(
+        mask=key_padding_mask,
+        mask_name="key_padding_mask",
+        other_type=_none_or_dtype(attn_mask),
+        other_name="attn_mask",
+        target_type=query.dtype,
+    )
+
+    if is_causal and attn_mask is None:
+        raise RuntimeError(
+            "Need attn_mask if specifying the is_causal hint. "
+            "You may use the Transformer module method "
+            "`generate_square_subsequent_mask` to create this mask."
+        )
+
+    if is_causal and key_padding_mask is None and not need_weights:
+        # when we have a kpm or need weights, we need attn_mask
+        # Otherwise, we use the is_causal hint go as is_causal
+        # indicator to SDPA.
+        attn_mask = None
+    else:
+        attn_mask = _canonical_mask(
+            mask=attn_mask,
+            mask_name="attn_mask",
+            other_type=None,
+            other_name="",
+            target_type=query.dtype,
+            check_other=False,
+        )
+
+        if key_padding_mask is not None:
+            # We have the attn_mask, and use that to merge kpm into it.
+            # Turn off use of is_causal hint, as the merged mask is no
+            # longer causal.
+            is_causal = False
+
+    if embed_dim != embed_dim_to_check:
+        raise AssertionError(
+            f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
+        )
+    if isinstance(embed_dim, torch.Tensor):
+        # embed_dim can be a tensor when JIT tracing
+        head_dim = embed_dim.div(num_heads, rounding_mode="trunc")
+    else:
+        head_dim = embed_dim // num_heads
+    if head_dim * num_heads != embed_dim:
+        raise AssertionError(
+            f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
+        )
+    if use_separate_proj_weight:
+        # allow MHA to have different embedding dimensions when separate projection weights are used
+        if key.shape[:2] != value.shape[:2]:
+            raise AssertionError(
+                f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}"
+            )
+    else:
+        if key.shape != value.shape:
+            raise AssertionError(
+                f"key shape {key.shape} does not match value shape {value.shape}"
+            )
+
+    #
+    # compute in-projection
+    #
+    if not use_separate_proj_weight:
+        if in_proj_weight is None:
+            raise AssertionError(
+                "use_separate_proj_weight is False but in_proj_weight is None"
+            )
+        q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
+    else:
+        if q_proj_weight is None:
+            raise AssertionError(
+                "use_separate_proj_weight is True but q_proj_weight is None"
+            )
+        if k_proj_weight is None:
+            raise AssertionError(
+                "use_separate_proj_weight is True but k_proj_weight is None"
+            )
+        if v_proj_weight is None:
+            raise AssertionError(
+                "use_separate_proj_weight is True but v_proj_weight is None"
+            )
+        if in_proj_bias is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = in_proj_bias.chunk(3)
+        q, k, v = _in_projection(
+            query,
+            key,
+            value,
+            q_proj_weight,
+            k_proj_weight,
+            v_proj_weight,
+            b_q,
+            b_k,
+            b_v,
+        )
+
+    # prep attention mask
+
+    if attn_mask is not None:
+        # ensure attn_mask's dim is 3
+        if attn_mask.dim() == 2:
+            correct_2d_size = (tgt_len, src_len)
+            if attn_mask.shape != correct_2d_size:
+                raise RuntimeError(
+                    f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}."
+                )
+            attn_mask = attn_mask.unsqueeze(0)
+        elif attn_mask.dim() == 3:
+            correct_3d_size = (bsz * num_heads, tgt_len, src_len)
+            if attn_mask.shape != correct_3d_size:
+                raise RuntimeError(
+                    f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}."
+                )
+        else:
+            raise RuntimeError(
+                f"attn_mask's dimension {attn_mask.dim()} is not supported"
+            )
+
+    # add bias along batch dimension (currently second)
+    if bias_k is not None and bias_v is not None:
+        if static_k is not None:
+            raise AssertionError("bias cannot be added to static key.")
+        if static_v is not None:
+            raise AssertionError("bias cannot be added to static value.")
+        k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+        v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+        if attn_mask is not None:
+            # pyrefly: ignore [bad-argument-type]
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            # pyrefly: ignore [bad-argument-type]
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+    else:
+        if bias_k is not None:
+            raise AssertionError("bias_k is set but bias_v is None")
+        if bias_v is not None:
+            raise AssertionError("bias_v is set but bias_k is None")
+
+    #
+    # reshape q, k, v for multihead attention and make them batch first
+    #
+    # pyrefly: ignore [bad-argument-type, no-matching-overload]
+    q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    if static_k is None:
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
+        k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+    else:
+        # TODO finish disentangling control flow so we don't do in-projections when statics are passed
+        if static_k.size(0) != bsz * num_heads:
+            raise AssertionError(
+                f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
+            )
+        if static_k.size(2) != head_dim:
+            raise AssertionError(
+                f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
+            )
+        k = static_k
+    if static_v is None:
+        # pyrefly: ignore [bad-argument-type, no-matching-overload]
+        v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+    else:
+        # TODO finish disentangling control flow so we don't do in-projections when statics are passed
+        if static_v.size(0) != bsz * num_heads:
+            raise AssertionError(
+                f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
+            )
+        if static_v.size(2) != head_dim:
+            raise AssertionError(
+                f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
+            )
+        v = static_v
+
+    # add zero attention along batch dimension (now first)
+    if add_zero_attn:
+        zero_attn_shape = (bsz * num_heads, 1, head_dim)
+        k = torch.cat(
+            # pyrefly: ignore [no-matching-overload]
+            [k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)],
+            dim=1,
+        )
+        v = torch.cat(
+            # pyrefly: ignore [no-matching-overload]
+            [v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)],
+            dim=1,
+        )
+        if attn_mask is not None:
+            # pyrefly: ignore [bad-argument-type]
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            # pyrefly: ignore [bad-argument-type]
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+
+    # update source sequence length after adjustments
+    src_len = k.size(1)
+
+    # merge key padding and attention masks
+    if key_padding_mask is not None:
+        if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+            _check_key_padding_mask(key_padding_mask, src_len, bsz)
+
+        key_padding_mask = (
+            key_padding_mask.view(bsz, 1, 1, src_len)
+            .expand(-1, num_heads, -1, -1)
+            .reshape(bsz * num_heads, 1, src_len)
+        )
+        if attn_mask is None:
+            attn_mask = key_padding_mask
+        else:
+            attn_mask = attn_mask + key_padding_mask
+
+    # adjust dropout probability
+    if not training:
+        dropout_p = 0.0
+
+    #
+    # (deep breath) calculate attention and out projection
+    #
+
+    if need_weights:
+        _B, _Nt, E = q.shape
+        q_scaled = q * math.sqrt(1.0 / float(E))
+
+        if is_causal and attn_mask is None:
+            raise AssertionError("FIXME: is_causal not implemented for need_weights")
+
+        if attn_mask is not None:
+            attn_output_weights = torch.baddbmm(
+                attn_mask, q_scaled, k.transpose(-2, -1)
+            )
+        else:
+            attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
+        if not torch.jit.is_scripting():
+            del q_scaled, k
+        attn_output_weights = softmax(attn_output_weights, dim=-1)
+        if dropout_p > 0.0:
+            attn_output_weights = dropout(attn_output_weights, p=dropout_p)
+
+        attn_output = torch.bmm(attn_output_weights, v)
+        if not torch.jit.is_scripting():
+            del v
+
+        attn_output = (
+            attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
+        )
+        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+
+        # optionally average attention weights over heads
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        if average_attn_weights:
+            attn_output_weights = attn_output_weights.mean(dim=1)
+
+        if not is_batched:
+            # squeeze the output if input was unbatched
+            attn_output = attn_output.squeeze(1)
+            attn_output_weights = attn_output_weights.squeeze(0)
+        return attn_output, attn_output_weights
+    else:
+        # attn_mask can be either (L,S) or (N*num_heads, L, S)
+        # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
+        # in order to match the input for SDPA of (N, num_heads, L, S)
+        if attn_mask is not None:
+            if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                attn_mask = attn_mask.unsqueeze(0)
+            else:
+                attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)
+
+        # pyrefly: ignore [bad-argument-type]
+        q = q.view(bsz, num_heads, tgt_len, head_dim)
+        # pyrefly: ignore [bad-argument-type]
+        k = k.view(bsz, num_heads, src_len, head_dim)
+        # pyrefly: ignore [bad-argument-type]
+        v = v.view(bsz, num_heads, src_len, head_dim)
+
+        attn_output = scaled_dot_product_attention(
+            q, k, v, attn_mask, dropout_p, is_causal
+        )
+        # Free q, k, v and their backing projection storage before the
+        # .contiguous() call below allocates.  In self-attention the three
+        # tensors are views of a single packed projection, so releasing all
+        # references here lets the allocator reclaim that memory immediately.
+        if not torch.jit.is_scripting():
+            del q, k, v
+        attn_output = (
+            attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
+        )
+
+        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+        if not is_batched:
+            # squeeze the output if input was unbatched
+            attn_output = attn_output.squeeze(1)
+        return attn_output, None
+
+
+def grouped_mm(
+    mat_a: Tensor,
+    mat_b: Tensor,
+    *,
+    offs: Tensor | None = None,
+    bias: Tensor | None = None,
+    out_dtype: torch.dtype | None = None,
+) -> Tensor:
+    r"""
+    grouped_mm(mat_a, mat_b, *, offs=None, bias=None, out_dtype=None)
+
+    Computes a grouped matrix multiply that shares weight shapes across experts but
+    allows jagged token counts per expert, which is common in Mixture-of-Experts
+    (MoE) layers. Both ``mat_a`` and ``mat_b`` must be 2D or 3D tensors that already
+    satisfy the physical layout restrictions of grouped GEMM kernels (e.g., row-major
+    ``mat_a`` and column-major ``mat_b`` for FP8 inputs). Inputs are currently
+    expected to be ``torch.bfloat16`` values on CUDA devices with :math:`SM \ge 80`.
+
+    Args:
+        mat_a: Left operand. When 2D, its leading dimension is sliced into groups
+            according to ``offs``. When 3D, its first dimension enumerates the groups
+            directly and ``offs`` must be ``None``.
+        mat_b: Right operand. When both operands are 2D (e.g., MoE weight-gradient
+            updates), the trailing dimension of ``mat_a`` and the leading dimension of
+            ``mat_b`` are partitioned according to the same ``offs`` tensor. For the
+            common forward pass (``out = input @ weight.T``) ``mat_b`` is 3D with
+            shape ``(num_groups, N, K)``.
+        offs: Optional 1D tensor of monotonically increasing ``int32`` offsets that
+            delimit the jagged dimension of any 2D operand. ``offs[i]`` marks the end
+            of group ``i`` and ``offs[-1]`` must be strictly less than the total
+            length of that operand's sliced dimension; elements beyond ``offs[-1]``
+            are ignored.
+        bias: Optional tensor that is added to the grouped outputs. Bias is not
+            jagged and must be broadcastable to the result shape of each group.
+        out_dtype: Optional dtype that controls the accumulation/output dtype.
+            Passing ``torch.float32`` accumulates BF16 inputs in FP32 while keeping
+            the grouped GEMM API non-differentiable.
+
+    Returns:
+        A tensor containing the concatenated results of each per-group GEMM with
+        shape inferred from the operands and ``offs``.
+    """
+
+    return torch._grouped_mm(mat_a, mat_b, offs=offs, bias=bias, out_dtype=out_dtype)
+
+
+def scaled_mm(
+    mat_a: Tensor,
+    mat_b: Tensor,
+    scale_a: Tensor | list[Tensor],
+    scale_recipe_a: ScalingType | list[ScalingType],
+    scale_b: Tensor | list[Tensor],
+    scale_recipe_b: ScalingType | list[ScalingType],
+    swizzle_a: SwizzleType | list[SwizzleType] | None = None,
+    swizzle_b: SwizzleType | list[SwizzleType] | None = None,
+    bias: Tensor | None = None,
+    output_dtype: torch.dtype | None = torch.bfloat16,
+    contraction_dim: list[int] | tuple[int, ...] = (),
+    use_fast_accum: bool = False,
+) -> Tensor:
+    r"""
+    scaled_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a, swizzle_b, bias, output_dtype,
+              contraction_dim, use_fast_accum)
+
+    Applies a scaled matrix-multiply, mm(mat_a, mat_b) where the scaling of mat_a and mat_b are described by
+    scale_recipe_a and scale_recipe_b respectively.
+
+    Args:
+        scale_a: Tensor containing decoding scaling factors for mat_a
+        scale_recipe_a: Enum describing how mat_a has been scaled
+        scale_b: Tensor containing decoding scaling factors for mat_b
+        scale_recipe_b: Enum describing how mat_b has been scaled
+        swizzle_a: Enum describing the swizzling pattern (if any) of scale_a
+        swizzle_b: Enum describing the swizzling pattern (if any) of scale_b
+        bias: optional bias term to be added to the output
+        output_dtype: dtype used for the output tensor
+        contraction_dim: describe which dimensions are :math:`K` in the matmul.
+        use_fast_accum: enable/disable tensor-core fast accumulation (Hopper-GPUs only)
+    """
+
+    def expand_single_value(v: _Any | list[_Any] | None) -> list[_Any]:
+        if v is None:
+            return []
+        elif not isinstance(v, (list)):
+            return [
+                v,
+            ]
+        else:
+            return v
+
+    scale_a = expand_single_value(scale_a)
+    scale_recipe_a = expand_single_value(scale_recipe_a)
+    scale_b = expand_single_value(scale_b)
+    scale_recipe_b = expand_single_value(scale_recipe_b)
+    swizzle_a = expand_single_value(swizzle_a)
+    swizzle_b = expand_single_value(swizzle_b)
+
+    # native_functions has restrictions on what can be defined
+    # & passed through - std::optional<ArrayRef<Tensor>> for instance
+    # *cannot* be passed, but an empty vector (list) can.
+    # So, we need to convert None arguments for lists in python
+    # explicitly into empty lists.
+    def list_or_empty(l: list[_Any] | None) -> list[_Any]:
+        return l if l else []
+
+    def enum_list_as_int_list(l: _Any | list[_Any]) -> list[_Any]:
+        if not isinstance(l, list):
+            l = [
+                l,
+            ]
+        return [li.value for li in l]
+
+    out = torch._scaled_mm_v2(
+        mat_a,
+        mat_b,
+        scale_a,
+        enum_list_as_int_list(scale_recipe_a),
+        enum_list_as_int_list(list_or_empty(swizzle_a)),
+        scale_b,
+        enum_list_as_int_list(scale_recipe_b),
+        enum_list_as_int_list(list_or_empty(swizzle_b)),
+        bias,
+        output_dtype,
+        contraction_dim,
+        use_fast_accum,
+    )
+
+    return out
+
+
+def scaled_grouped_mm(
+    mat_a: Tensor,
+    mat_b: Tensor,
+    scale_a: Tensor | list[Tensor],
+    scale_recipe_a: ScalingType | list[ScalingType],
+    scale_b: Tensor | list[Tensor],
+    scale_recipe_b: ScalingType | list[ScalingType],
+    swizzle_a: SwizzleType | list[SwizzleType] | None = None,
+    swizzle_b: SwizzleType | list[SwizzleType] | None = None,
+    bias: Tensor | None = None,
+    offs: Tensor | None = None,
+    output_dtype: torch.dtype | None = torch.bfloat16,
+    contraction_dim: list[int] | tuple[int, ...] = (),
+    use_fast_accum: bool = False,
+) -> Tensor:
+    r"""
+    scaled_grouped_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a, swizzle_b, bias, offs,
+              output_dtype, use_fast_accum)
+
+    Applies a grouped scaled matrix-multiply, grouped_mm(mat_a, mat_b) where the scaling of mat_a and mat_b are described by
+    scale_recipe_a and scale_recipe_b respectively.
+
+    Args:
+        scale_a: Tensor containing decoding scaling factors for mat_a
+        scale_recipe_a: Enum describing how mat_a has been scaled
+        scale_b: Tensor containing decoding scaling factors for mat_b
+        scale_recipe_b: Enum describing how mat_b has been scaled
+        swizzle_a: Enum describing the swizzling pattern (if any) of scale_a
+        swizzle_b: Enum describing the swizzling pattern (if any) of scale_b
+        bias: optional bias term to be added to the output
+        offs: optional offsets into the source tensors denoting group start indices
+        output_dtype: dtype used for the output tensor
+        contraction_dim: describe which dimensions are :math:`K` in the matmul.
+        use_fast_accum: enable/disable tensor-core fast accumulation (Hopper-GPUs only)
+    """
+
+    def expand_single_value(v: _Any | list[_Any] | None) -> list[_Any]:
+        if v is None:
+            return []
+        elif not isinstance(v, (list)):
+            return [
+                v,
+            ]
+        else:
+            return v
+
+    scale_a = expand_single_value(scale_a)
+    scale_recipe_a = expand_single_value(scale_recipe_a)
+    scale_b = expand_single_value(scale_b)
+    scale_recipe_b = expand_single_value(scale_recipe_b)
+    swizzle_a = expand_single_value(swizzle_a)
+    swizzle_b = expand_single_value(swizzle_b)
+
+    # native_functions has restrictions on what can be defined
+    # & passed through - std::optional<ArrayRef<Tensor>> for instance
+    # *cannot* be passed, but an empty vector (list) can.
+    # So, we need to convert None arguments for lists in python
+    # explicitly into empty lists.
+    def list_or_empty(l: list[_Any] | None) -> list[_Any]:
+        return l if l else []
+
+    def enum_list_as_int_list(l: _Any | list[_Any]) -> list[_Any]:
+        if not isinstance(l, list):
+            l = [
+                l,
+            ]
+        return [li.value for li in l]
+
+    out = torch._scaled_grouped_mm_v2(
+        mat_a,
+        mat_b,
+        scale_a,
+        enum_list_as_int_list(scale_recipe_a),
+        enum_list_as_int_list(list_or_empty(swizzle_a)),
+        scale_b,
+        enum_list_as_int_list(scale_recipe_b),
+        enum_list_as_int_list(list_or_empty(swizzle_b)),
+        offs,
+        bias,
+        output_dtype,
+        contraction_dim,
+        use_fast_accum,
+    )
+
+    return out
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/functional.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/functional.pyi
new file mode 100644
index 0000000000000000000000000000000000000000..a724dc0c68b438af4f1f6c5143eff809acfa4c47
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/functional.pyi
@@ -0,0 +1,1200 @@
+# @generated by tools/pyi/gen_pyi.py from torch/nn/functional.pyi.in
+# mypy: allow-untyped-defs
+
+from collections.abc import Callable, Sequence
+from enum import Enum
+from typing import Any, Literal, overload, TypeAlias
+
+from torch import Tensor
+from torch.types import _dtype, _int, _size
+
+from .common_types import (
+    _ratio_any_t,
+    _size_1_t,
+    _size_2_opt_t,
+    _size_2_t,
+    _size_3_opt_t,
+    _size_3_t,
+    _size_any_t,
+)
+
+__all__ = [
+    "GRID_SAMPLE_INTERPOLATION_MODES",
+    "GRID_SAMPLE_PADDING_MODES",
+]
+
+# 'TypedDict' is a new accepted type that represents a dictionary with a fixed set of allowed keys.
+# It is standards-track but not in `typing` yet. We leave this hear to be uncommented once the feature
+# is wide-spread.
+
+# from mypy_extensions import TypedDict
+
+# GRID_SAMPLE_INTERPOLATION_MODES = TypedDict('GRID_SAMPLE_INTERPOLATION_MODES', {'bilinear': int, 'nearest': int})
+# GRID_SAMPLE_PADDING_MODES = TypedDict('GRID_SAMPLE_PADDING_MODES', {'zeros': int, 'border': int, 'reflection': int})
+
+GRID_SAMPLE_INTERPOLATION_MODES: TypeAlias = dict[str, int]
+GRID_SAMPLE_PADDING_MODES: TypeAlias = dict[str, int]
+
+# These stubs were generated by running stubgen (`stubgen --parse-only functional.py`), followed by manual cleaning.
+#
+# The 'BroadcastingList{1,2,3}' types were replaced by `_size` or _output_ratio, as appropriate.
+# This was necessary since the JIT uses BroadcastingList* types but static checking with mypy etc requires a `Sequence`
+# type. There is no way to express the expected lengths of these lists in the current Python typing system.
+#
+# Functions created via `_add_docstr` in `functional.py` where merely typed as `Any` by `stubgen`, so those were
+# deleted from the stub and replaced by generated declarations. See `gen_pyi` for the implementation of the code
+# generation logic for those functions. In the future, it might be worth looking into using the mypy plugin system
+# to encode the type semantics of `_add_docstr`, should that system ever become widespread.
+def _canonical_mask(
+    mask: Tensor | None,
+    mask_name: str,
+    other_type: _dtype | None,
+    other_name: str,
+    target_type: _dtype,
+    check_other: bool = True,
+) -> Tensor | None: ...
+
+__all__ += ["_canonical_mask"]
+
+def _none_or_dtype(input: Tensor | None) -> _dtype | None: ...
+
+__all__ += ["_none_or_dtype"]
+
+def adaptive_avg_pool2d(input: Tensor, output_size: _size_2_opt_t) -> Tensor: ...
+
+__all__ += ["adaptive_avg_pool2d"]
+
+def adaptive_avg_pool3d(input: Tensor, output_size: _size_3_opt_t) -> Tensor: ...
+
+__all__ += ["adaptive_avg_pool3d"]
+
+def adaptive_max_pool1d_with_indices(
+    input: Tensor,
+    output_size: _size,
+    return_indices: bool = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["adaptive_max_pool1d_with_indices"]
+
+def adaptive_max_pool2d_with_indices(
+    input: Tensor,
+    output_size: _size_2_opt_t,
+    return_indices: bool = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["adaptive_max_pool2d_with_indices"]
+
+def adaptive_max_pool3d_with_indices(
+    input: Tensor,
+    output_size: _size_3_opt_t,
+    return_indices: bool = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["adaptive_max_pool3d_with_indices"]
+
+def affine_grid(
+    theta: Tensor,
+    size: list[int],
+    align_corners: Any | None = ...,
+) -> Tensor: ...
+
+__all__ += ["affine_grid"]
+
+def alpha_dropout(
+    input: Tensor,
+    p: float = ...,
+    training: bool = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["alpha_dropout"]
+
+def assert_int_or_pair(arg: Any, arg_name: Any, message: Any) -> None: ...
+
+__all__ += ["assert_int_or_pair"]
+
+def batch_norm(
+    input: Tensor,
+    running_mean: Tensor | None,
+    running_var: Tensor | None,
+    weight: Tensor | None = ...,
+    bias: Tensor | None = ...,
+    training: bool = ...,
+    momentum: float = ...,
+    eps: float = ...,
+) -> Tensor: ...
+
+__all__ += ["batch_norm"]
+
+def binary_cross_entropy_with_logits(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+    pos_weight: Tensor | None = ...,
+) -> Tensor: ...
+
+__all__ += ["binary_cross_entropy_with_logits"]
+
+def binary_cross_entropy(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["binary_cross_entropy"]
+
+def celu(input: Tensor, alpha: float = ..., inplace: bool = ...) -> Tensor: ...
+
+__all__ += ["celu"]
+
+def cosine_embedding_loss(
+    input1: Tensor,
+    input2: Tensor,
+    target: Tensor,
+    margin: float = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["cosine_embedding_loss"]
+
+def cross_entropy(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = ...,
+    size_average: bool | None = ...,
+    ignore_index: int = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+    label_smoothing: float = ...,
+) -> Tensor: ...
+
+__all__ += ["cross_entropy"]
+
+def ctc_loss(
+    log_probs: Tensor,
+    targets: Tensor,
+    input_lengths: Tensor,
+    target_lengths: Tensor,
+    blank: int = ...,
+    reduction: str = ...,
+    zero_infinity: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["ctc_loss"]
+
+def dropout(
+    input: Tensor,
+    p: float = ...,
+    training: bool = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["dropout"]
+
+def dropout1d(
+    input: Tensor,
+    p: float = ...,
+    training: bool = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["dropout1d"]
+
+def dropout2d(
+    input: Tensor,
+    p: float = ...,
+    training: bool = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["dropout2d"]
+
+def dropout3d(
+    input: Tensor,
+    p: float = ...,
+    training: bool = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["dropout3d"]
+
+def elu(input: Tensor, alpha: float = ..., inplace: bool = ...) -> Tensor: ...
+
+__all__ += ["elu"]
+
+def embedding_bag(
+    input: Tensor,
+    weight: Tensor,
+    offsets: Tensor | None = ...,
+    max_norm: float | None = ...,
+    norm_type: float = ...,
+    scale_grad_by_freq: bool = ...,
+    mode: str = ...,
+    sparse: bool = ...,
+    per_sample_weights: Tensor | None = ...,
+    include_last_offset: bool = ...,
+    padding_idx: int | None = ...,
+) -> Tensor: ...
+
+__all__ += ["embedding_bag"]
+
+def embedding(
+    input: Tensor,
+    weight: Tensor,
+    padding_idx: int | None = ...,
+    max_norm: float | None = ...,
+    norm_type: float = ...,
+    scale_grad_by_freq: bool = ...,
+    sparse: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["embedding"]
+
+def feature_alpha_dropout(
+    input: Tensor,
+    p: float = ...,
+    training: bool = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["feature_alpha_dropout"]
+
+def fold(
+    input: Tensor,
+    output_size: _size_any_t,
+    kernel_size: _size_any_t,
+    dilation: _size_any_t = ...,
+    padding: _size_any_t = ...,
+    stride: _size_any_t = ...,
+) -> Tensor: ...
+
+__all__ += ["fold"]
+
+def fractional_max_pool2d_with_indices(
+    input: Tensor,
+    kernel_size: _size,
+    output_size: _size | None = ...,
+    output_ratio: _ratio_any_t | None = ...,
+    return_indices: bool = ...,
+    _random_samples: Tensor | None = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["fractional_max_pool2d_with_indices"]
+
+def fractional_max_pool3d_with_indices(
+    input: Tensor,
+    kernel_size: _size,
+    output_size: _size | None = ...,
+    output_ratio: _ratio_any_t | None = ...,
+    return_indices: bool = ...,
+    _random_samples: Tensor | None = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["fractional_max_pool3d_with_indices"]
+
+def gaussian_nll_loss(
+    input: Tensor,
+    target: Tensor,
+    var: Tensor | float,
+    full: bool | None = ...,
+    eps: float | None = ...,
+    reduction: str | None = ...,
+) -> Tensor: ...
+
+__all__ += ["gaussian_nll_loss"]
+
+def glu(input: Tensor, dim: int = ...) -> Tensor: ...
+
+__all__ += ["glu"]
+
+def grid_sample(
+    input: Tensor,
+    grid: Tensor,
+    mode: str = ...,
+    padding_mode: str = ...,
+    align_corners: Any | None = ...,
+) -> Tensor: ...
+
+__all__ += ["grid_sample"]
+
+def group_norm(
+    input: Tensor,
+    num_groups: int,
+    weight: Tensor | None = ...,
+    bias: Tensor | None = ...,
+    eps: float = ...,
+) -> Tensor: ...
+
+__all__ += ["group_norm"]
+
+def gumbel_softmax(
+    logits: Tensor,
+    tau: float = ...,
+    hard: bool = ...,
+    eps: float = ...,
+    dim: int = ...,
+) -> Tensor: ...
+
+__all__ += ["gumbel_softmax"]
+
+def hardsigmoid(input: Tensor, inplace: bool = False) -> Tensor: ...
+
+__all__ += ["hardsigmoid"]
+
+def hardswish(input: Tensor, inplace: bool = False) -> Tensor: ...
+
+__all__ += ["hardswish"]
+
+def hardtanh(
+    input: Tensor,
+    min_val: float = ...,
+    max_val: float = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["hardtanh"]
+
+def hinge_embedding_loss(
+    input: Tensor,
+    target: Tensor,
+    margin: float = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["hinge_embedding_loss"]
+
+def huber_loss(
+    input: Tensor,
+    target: Tensor,
+    reduction: str = ...,
+    delta: float = ...,
+) -> Tensor: ...
+
+__all__ += ["huber_loss"]
+
+def instance_norm(
+    input: Tensor,
+    running_mean: Tensor | None = ...,
+    running_var: Tensor | None = ...,
+    weight: Tensor | None = ...,
+    bias: Tensor | None = ...,
+    use_input_stats: bool = ...,
+    momentum: float = ...,
+    eps: float = ...,
+) -> Tensor: ...
+
+__all__ += ["instance_norm"]
+
+def interpolate(
+    input: Tensor,
+    size: int | Sequence[int] | None = ...,
+    scale_factor: float | Sequence[float] | None = ...,
+    mode: str = ...,
+    align_corners: bool | None = ...,
+    recompute_scale_factor: bool | None = ...,
+    antialias: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["interpolate"]
+
+def kl_div(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+    log_target: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["kl_div"]
+
+def l1_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["l1_loss"]
+
+def layer_norm(
+    input: Tensor,
+    normalized_shape: Sequence[int],
+    weight: Tensor | None = ...,
+    bias: Tensor | None = ...,
+    eps: float = ...,
+) -> Tensor: ...
+
+__all__ += ["layer_norm"]
+
+def leaky_relu(
+    input: Tensor,
+    negative_slope: float = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["leaky_relu"]
+
+def local_response_norm(
+    input: Tensor,
+    size: int,
+    alpha: float = ...,
+    beta: float = ...,
+    k: float = ...,
+) -> Tensor: ...
+
+__all__ += ["local_response_norm"]
+
+def log_softmax(
+    input: Tensor,
+    dim: int | None = ...,
+    _stacklevel: int = ...,
+    dtype: _dtype | None = ...,
+) -> Tensor: ...
+
+__all__ += ["log_softmax"]
+
+def lp_pool1d(
+    input: Tensor,
+    norm_type: float,
+    kernel_size: _size_1_t,
+    stride: _size | None | int = ...,
+    ceil_mode: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["lp_pool1d"]
+
+def lp_pool2d(
+    input: Tensor,
+    norm_type: float,
+    kernel_size: _size_2_t,
+    stride: _size | None | int = ...,
+    ceil_mode: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["lp_pool2d"]
+
+def lp_pool3d(
+    input: Tensor,
+    norm_type: float,
+    kernel_size: _size_3_t,
+    stride: _size | None | int = ...,
+    ceil_mode: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["lp_pool3d"]
+
+def margin_ranking_loss(
+    input1: Tensor,
+    input2: Tensor,
+    target: Tensor,
+    margin: float = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["margin_ranking_loss"]
+
+def max_pool1d_with_indices(
+    input: Tensor,
+    kernel_size: _size,
+    stride: _size | None = ...,
+    padding: _size = ...,
+    dilation: _size = ...,
+    ceil_mode: bool = ...,
+    return_indices: bool = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["max_pool1d_with_indices"]
+
+def max_pool2d_with_indices(
+    input: Tensor,
+    kernel_size: _size,
+    stride: _size | None = ...,
+    padding: _size = ...,
+    dilation: _size = ...,
+    ceil_mode: bool = ...,
+    return_indices: bool = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["max_pool2d_with_indices"]
+
+def max_pool3d_with_indices(
+    input: Tensor,
+    kernel_size: _size,
+    stride: _size | None = ...,
+    padding: _size = ...,
+    dilation: _size = ...,
+    ceil_mode: bool = ...,
+    return_indices: bool = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += ["max_pool3d_with_indices"]
+
+def max_unpool1d(
+    input: Tensor,
+    indices: Tensor,
+    kernel_size: _size,
+    stride: _size | None = ...,
+    padding: _size = ...,
+    output_size: _size | None = ...,
+) -> Tensor: ...
+
+__all__ += ["max_unpool1d"]
+
+def max_unpool2d(
+    input: Tensor,
+    indices: Tensor,
+    kernel_size: _size,
+    stride: _size | None = ...,
+    padding: _size = ...,
+    output_size: _size | None = ...,
+) -> Tensor: ...
+
+__all__ += ["max_unpool2d"]
+
+def max_unpool3d(
+    input: Tensor,
+    indices: Tensor,
+    kernel_size: _size,
+    stride: _size | None = ...,
+    padding: _size = ...,
+    output_size: _size | None = ...,
+) -> Tensor: ...
+
+__all__ += ["max_unpool3d"]
+
+def mish(input: Tensor, inplace: bool = False) -> Tensor: ...
+
+__all__ += ["mish"]
+
+def mse_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["mse_loss"]
+
+def multi_head_attention_forward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    embed_dim_to_check: int,
+    num_heads: int,
+    in_proj_weight: Tensor | None,
+    in_proj_bias: Tensor | None,
+    bias_k: Tensor | None,
+    bias_v: Tensor | None,
+    add_zero_attn: bool,
+    dropout_p: float,
+    out_proj_weight: Tensor,
+    out_proj_bias: Tensor | None,
+    training: bool = True,
+    key_padding_mask: Tensor | None = None,
+    need_weights: bool = True,
+    attn_mask: Tensor | None = None,
+    use_separate_proj_weight: bool = False,
+    q_proj_weight: Tensor | None = None,
+    k_proj_weight: Tensor | None = None,
+    v_proj_weight: Tensor | None = None,
+    static_k: Tensor | None = None,
+    static_v: Tensor | None = None,
+    average_attn_weights: bool = True,
+    is_causal: bool = False,
+) -> tuple[Tensor, Tensor | None]: ...
+
+__all__ += ["multi_head_attention_forward"]
+
+def multi_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    p: int = ...,
+    margin: float = ...,
+    weight: Tensor | None = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["multi_margin_loss"]
+
+def multilabel_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["multilabel_margin_loss"]
+
+def multilabel_soft_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["multilabel_soft_margin_loss"]
+
+def nll_loss(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = ...,
+    size_average: bool | None = ...,
+    ignore_index: int = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["nll_loss"]
+
+def normalize(
+    input: Tensor,
+    p: float = ...,
+    dim: int = ...,
+    eps: float = ...,
+    out: Tensor | None = ...,
+) -> Tensor: ...
+
+__all__ += ["normalize"]
+
+def poisson_nll_loss(
+    input: Tensor,
+    target: Tensor,
+    log_input: bool = ...,
+    full: bool = ...,
+    size_average: bool | None = ...,
+    eps: float = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["poisson_nll_loss"]
+
+def relu(input: Tensor, inplace: bool = ...) -> Tensor: ...
+
+__all__ += ["relu"]
+
+def relu6(input: Tensor, inplace: bool = ...) -> Tensor: ...
+
+__all__ += ["relu6"]
+
+def rms_norm(
+    input: Tensor,
+    normalized_shape: Sequence[int],
+    weight: Tensor | None = ...,
+    eps: float | None = ...,
+) -> Tensor: ...
+
+__all__ += ["rms_norm"]
+
+def rrelu(
+    input: Tensor,
+    lower: float = ...,
+    upper: float = ...,
+    training: bool = ...,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["rrelu"]
+
+def scaled_mm(
+    mat_a: Tensor,
+    mat_b: Tensor,
+    scale_a: Tensor | list[Tensor],
+    scale_recipe_a: ScalingType | list[ScalingType],
+    scale_b: Tensor | list[Tensor],
+    scale_recipe_b: ScalingType | list[ScalingType],
+    swizzle_a: SwizzleType | list[SwizzleType] | None = None,
+    swizzle_b: SwizzleType | list[SwizzleType] | None = None,
+    bias: Tensor | None = None,
+    output_dtype: _dtype = ...,
+    contraction_dim: list[int] | tuple[int, ...] = (),
+    use_fast_accum: bool = False,
+) -> Tensor: ...
+
+__all__ += ["scaled_mm"]
+
+def grouped_mm(
+    mat_a: Tensor,
+    mat_b: Tensor,
+    *,
+    offs: Tensor | None = None,
+    bias: Tensor | None = None,
+    out_dtype: _dtype | None = None,
+) -> Tensor: ...
+
+__all__ += ["grouped_mm"]
+
+class SwizzleType(Enum):
+    NO_SWIZZLE = 0
+    SWIZZLE_32_4_4 = 1
+
+__all__ += ["SwizzleType"]
+
+class ScalingType(Enum):
+    TensorWise = 0
+    RowWise = 1
+    BlockWise1x16 = 2
+    BlockWise1x32 = 3
+    BlockWise1x128 = 4
+    BlockWise128x128 = 5
+
+__all__ += ["ScalingType"]
+
+def selu(input: Tensor, inplace: bool = ...) -> Tensor: ...
+
+__all__ += ["selu"]
+
+def sigmoid(input: Tensor) -> Tensor: ...
+
+__all__ += ["sigmoid"]
+
+def silu(input: Tensor, inplace: bool = False) -> Tensor: ...
+
+__all__ += ["silu"]
+
+def smooth_l1_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+    beta: float = ...,
+) -> Tensor: ...
+
+__all__ += ["smooth_l1_loss"]
+
+def soft_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["soft_margin_loss"]
+
+def softmax(
+    input: Tensor,
+    dim: int | None = ...,
+    _stacklevel: int = ...,
+    dtype: _dtype | None = ...,
+) -> Tensor: ...
+
+__all__ += ["softmax"]
+
+def softmin(
+    input: Tensor,
+    dim: int | None = ...,
+    _stacklevel: int = ...,
+    dtype: _dtype | None = ...,
+) -> Tensor: ...
+
+__all__ += ["softmin"]
+
+def softsign(input: Any) -> Tensor: ...
+
+__all__ += ["softsign"]
+
+def tanh(input: Any) -> Tensor: ...
+
+__all__ += ["tanh"]
+
+def tanhshrink(input: Any) -> Tensor: ...
+
+__all__ += ["tanhshrink"]
+
+def threshold(
+    input: Tensor,
+    threshold: float,
+    value: float,
+    inplace: bool = ...,
+) -> Tensor: ...
+
+__all__ += ["threshold"]
+
+def triplet_margin_loss(
+    anchor: Tensor,
+    positive: Tensor,
+    negative: Tensor,
+    margin: float = ...,
+    p: float = ...,
+    eps: float = ...,
+    swap: bool = ...,
+    size_average: bool | None = ...,
+    reduce: bool | None = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["triplet_margin_loss"]
+
+def triplet_margin_with_distance_loss(
+    anchor: Tensor,
+    positive: Tensor,
+    negative: Tensor,
+    *,
+    distance_function: Callable[[Tensor, Tensor], Tensor] | None = ...,
+    margin: float = ...,
+    swap: bool = ...,
+    reduction: str = ...,
+) -> Tensor: ...
+
+__all__ += ["triplet_margin_with_distance_loss"]
+
+def unfold(
+    input: Tensor,
+    kernel_size: _size_any_t,
+    dilation: _size_any_t = ...,
+    padding: _size_any_t = ...,
+    stride: _size_any_t = ...,
+) -> Tensor: ...
+
+__all__ += ["unfold"]
+
+def upsample_bilinear(
+    input: Any,
+    size: Any | None = ...,
+    scale_factor: Any | None = ...,
+) -> Tensor: ...
+
+__all__ += ["upsample_bilinear"]
+
+def upsample_nearest(
+    input: Any,
+    size: Any | None = ...,
+    scale_factor: Any | None = ...,
+) -> Tensor: ...
+
+__all__ += ["upsample_nearest"]
+
+def upsample(
+    input: Any,
+    size: Any | None = ...,
+    scale_factor: Any | None = ...,
+    mode: str = ...,
+    align_corners: Any | None = ...,
+) -> Tensor: ...
+
+__all__ += ["upsample"]
+
+from torch import (
+    adaptive_avg_pool1d as adaptive_avg_pool1d,
+    avg_pool1d as avg_pool1d,
+    bilinear as bilinear,
+    celu_ as celu_,
+    channel_shuffle as channel_shuffle,
+    conv1d as conv1d,
+    conv2d as conv2d,
+    conv3d as conv3d,
+    conv_tbc as conv_tbc,
+    conv_transpose1d as conv_transpose1d,
+    conv_transpose2d as conv_transpose2d,
+    conv_transpose3d as conv_transpose3d,
+    cosine_similarity as cosine_similarity,
+    hardshrink as hardshrink,
+    native_channel_shuffle as native_channel_shuffle,
+    pairwise_distance as pairwise_distance,
+    pdist as pdist,
+    pixel_shuffle as pixel_shuffle,
+    pixel_unshuffle as pixel_unshuffle,
+    prelu as prelu,
+    relu_ as relu_,
+    rrelu_ as rrelu_,
+    selu_ as selu_,
+)
+from torch._C._nn import (
+    avg_pool2d as avg_pool2d,
+    avg_pool3d as avg_pool3d,
+    elu_ as elu_,
+    gelu as gelu,
+    hardtanh_ as hardtanh_,
+    leaky_relu_ as leaky_relu_,
+    linear as linear,
+    log_sigmoid as logsigmoid,
+    one_hot as one_hot,
+    pad as pad,
+    scaled_dot_product_attention as scaled_dot_product_attention,
+    softplus as softplus,
+    softshrink as softshrink,
+)
+
+@overload
+def adaptive_max_pool1d(
+    input: Tensor,
+    output_size: _int | _size,
+    return_indices: Literal[False] = False,
+) -> Tensor: ...
+@overload
+def adaptive_max_pool1d(
+    input: Tensor,
+    output_size: _int | _size,
+    return_indices: Literal[True],
+    /,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def adaptive_max_pool1d(
+    input: Tensor,
+    output_size: _int | _size,
+    *,
+    return_indices: Literal[True],
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def adaptive_max_pool2d(
+    input: Tensor,
+    output_size: _int | _size,
+    return_indices: Literal[False] = False,
+) -> Tensor: ...
+@overload
+def adaptive_max_pool2d(
+    input: Tensor,
+    output_size: _int | _size,
+    return_indices: Literal[True],
+    /,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def adaptive_max_pool2d(
+    input: Tensor,
+    output_size: _int | _size,
+    *,
+    return_indices: Literal[True],
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def adaptive_max_pool3d(
+    input: Tensor,
+    output_size: _int | _size,
+    return_indices: Literal[False] = False,
+) -> Tensor: ...
+@overload
+def adaptive_max_pool3d(
+    input: Tensor,
+    output_size: _int | _size,
+    return_indices: Literal[True],
+    /,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def adaptive_max_pool3d(
+    input: Tensor,
+    output_size: _int | _size,
+    *,
+    return_indices: Literal[True],
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def fractional_max_pool2d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    output_size: _int | _size | None = None,
+    output_ratio: _ratio_any_t | None = None,
+    return_indices: Literal[False] = False,
+    _random_samples: Tensor | None = None,
+) -> Tensor: ...
+@overload
+def fractional_max_pool2d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    output_size: _int | _size | None,
+    output_ratio: _ratio_any_t | None,
+    return_indices: Literal[True],
+    /,
+    _random_samples: Tensor | None = None,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def fractional_max_pool2d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    output_size: _int | _size | None = None,
+    output_ratio: _ratio_any_t | None = None,
+    *,
+    return_indices: Literal[True],
+    _random_samples: Tensor | None = None,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def fractional_max_pool3d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    output_size: _int | _size | None = None,
+    output_ratio: _ratio_any_t | None = None,
+    return_indices: Literal[False] = False,
+    _random_samples: Tensor | None = None,
+) -> Tensor: ...
+@overload
+def fractional_max_pool3d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    output_size: _int | _size | None,
+    output_ratio: _ratio_any_t | None,
+    return_indices: Literal[True],
+    /,
+    _random_samples: Tensor | None = None,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def fractional_max_pool3d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    output_size: _int | _size | None = None,
+    output_ratio: _ratio_any_t | None = None,
+    *,
+    return_indices: Literal[True],
+    _random_samples: Tensor | None = None,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def max_pool1d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None = None,
+    padding: _int | _size = 0,
+    dilation: _int | _size = 1,
+    ceil_mode: bool = False,
+    return_indices: Literal[False] = False,
+) -> Tensor: ...
+@overload
+def max_pool1d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None,
+    padding: _int | _size,
+    dilation: _int | _size,
+    ceil_mode: bool,
+    return_indices: Literal[True],
+    /,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def max_pool1d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None = None,
+    padding: _int | _size = 0,
+    dilation: _int | _size = 1,
+    ceil_mode: bool = False,
+    *,
+    return_indices: Literal[True],
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def max_pool2d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None = None,
+    padding: _int | _size = 0,
+    dilation: _int | _size = 1,
+    ceil_mode: bool = False,
+    return_indices: Literal[False] = False,
+) -> Tensor: ...
+@overload
+def max_pool2d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None,
+    padding: _int | _size,
+    dilation: _int | _size,
+    ceil_mode: bool,
+    return_indices: Literal[True],
+    /,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def max_pool2d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None = None,
+    padding: _int | _size = 0,
+    dilation: _int | _size = 1,
+    ceil_mode: bool = False,
+    *,
+    return_indices: Literal[True],
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def max_pool3d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None = None,
+    padding: _int | _size = 0,
+    dilation: _int | _size = 1,
+    ceil_mode: bool = False,
+    return_indices: Literal[False] = False,
+) -> Tensor: ...
+@overload
+def max_pool3d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None,
+    padding: _int | _size,
+    dilation: _int | _size,
+    ceil_mode: bool,
+    return_indices: Literal[True],
+    /,
+) -> tuple[Tensor, Tensor]: ...
+@overload
+def max_pool3d(
+    input: Tensor,
+    kernel_size: _int | _size,
+    stride: _int | _size | None = None,
+    padding: _int | _size = 0,
+    dilation: _int | _size = 1,
+    ceil_mode: bool = False,
+    *,
+    return_indices: Literal[True],
+) -> tuple[Tensor, Tensor]: ...
+
+__all__ += [
+    "adaptive_avg_pool1d",
+    "avg_pool1d",
+    "bilinear",
+    "celu_",
+    "channel_shuffle",
+    "conv_tbc",
+    "conv_transpose1d",
+    "conv_transpose2d",
+    "conv_transpose3d",
+    "conv1d",
+    "conv2d",
+    "conv3d",
+    "cosine_similarity",
+    "hardshrink",
+    "native_channel_shuffle",
+    "pairwise_distance",
+    "pdist",
+    "pixel_shuffle",
+    "pixel_unshuffle",
+    "prelu",
+    "relu_",
+    "rrelu_",
+    "selu_",
+    "avg_pool2d",
+    "avg_pool3d",
+    "elu_",
+    "gelu",
+    "hardtanh_",
+    "leaky_relu_",
+    "linear",
+    "logsigmoid",
+    "one_hot",
+    "pad",
+    "scaled_dot_product_attention",
+    "softplus",
+    "softshrink",
+    "max_pool1d",
+    "adaptive_max_pool1d",
+    "max_pool2d",
+    "fractional_max_pool2d",
+    "adaptive_max_pool2d",
+    "max_pool3d",
+    "fractional_max_pool3d",
+    "adaptive_max_pool3d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/grad.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/grad.py
new file mode 100644
index 0000000000000000000000000000000000000000..61e817dbed612e51f1ff426373fbef1d192e331e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/grad.py
@@ -0,0 +1,298 @@
+# mypy: allow-untyped-defs
+"""Gradient interface."""
+
+import torch
+from torch.nn.modules.utils import _pair, _single, _triple
+
+
+def conv1d_input(
+    input_size,
+    weight,
+    grad_output,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+):
+    r"""Compute the gradient of conv1d with respect to the input of the convolution.
+
+    This is same as the 1D transposed convolution operator under the hood but requires
+    the shape of the gradient w.r.t. input to be specified explicitly.
+
+    Args:
+        input_size : Shape of the input gradient tensor
+        weight: weight tensor (out_channels x in_channels/groups x kW)
+        grad_output : output gradient tensor (minibatch x out_channels x oW)
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+
+    Examples::
+
+        >>> input = torch.randn(1, 1, 3, requires_grad=True)
+        >>> weight = torch.randn(1, 1, 1, requires_grad=True)
+        >>> output = F.conv1d(input, weight)
+        >>> grad_output = torch.randn(output.shape)
+        >>> grad_input = torch.autograd.grad(output, input, grad_output)
+        >>> F.grad.conv1d_input(input.shape, weight, grad_output)
+
+    """
+    input = grad_output.new_empty(1).expand(input_size)
+
+    return torch.ops.aten.convolution_backward(
+        grad_output,
+        input,
+        weight,
+        None,
+        _single(stride),
+        _single(padding),
+        _single(dilation),
+        False,
+        [0],
+        groups,
+        (True, False, False),
+    )[0]
+
+
+def conv1d_weight(
+    input,
+    weight_size,
+    grad_output,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+):
+    r"""Compute the gradient of conv1d with respect to the weight of the convolution.
+
+    Args:
+        input: input tensor of shape (minibatch x in_channels x iW)
+        weight_size : Shape of the weight gradient tensor
+        grad_output : output gradient tensor (minibatch x out_channels x oW)
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+
+    Examples::
+
+        >>> input = torch.randn(1, 1, 3, requires_grad=True)
+        >>> weight = torch.randn(1, 1, 1, requires_grad=True)
+        >>> output = F.conv1d(input, weight)
+        >>> grad_output = torch.randn(output.shape)
+        >>> # xdoctest: +SKIP
+        >>> grad_weight = torch.autograd.grad(output, filter, grad_output)
+        >>> F.grad.conv1d_weight(input, weight.shape, grad_output)
+
+    """
+    weight = grad_output.new_empty(1).expand(weight_size)
+
+    return torch.ops.aten.convolution_backward(
+        grad_output,
+        input,
+        weight,
+        None,
+        _single(stride),
+        _single(padding),
+        _single(dilation),
+        False,
+        [0],
+        groups,
+        (False, True, False),
+    )[1]
+
+
+def conv2d_input(
+    input_size,
+    weight,
+    grad_output,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+):
+    r"""Compute the gradient of conv2d with respect to the input of the convolution.
+
+    This is same as the 2D transposed convolution operator under the hood but requires
+    the shape of the gradient w.r.t. input to be specified explicitly.
+
+    Args:
+        input_size : Shape of the input gradient tensor
+        weight: weight tensor (out_channels x in_channels/groups x kH x kW)
+        grad_output : output gradient tensor (minibatch x out_channels x oH x oW)
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+
+    Examples::
+
+        >>> input = torch.randn(1, 1, 3, 3, requires_grad=True)
+        >>> weight = torch.randn(1, 1, 1, 2, requires_grad=True)
+        >>> output = F.conv2d(input, weight)
+        >>> grad_output = torch.randn(output.shape)
+        >>> grad_input = torch.autograd.grad(output, input, grad_output)
+        >>> F.grad.conv2d_input(input.shape, weight, grad_output)
+
+    """
+    input = grad_output.new_empty(1).expand(input_size)
+
+    return torch.ops.aten.convolution_backward(
+        grad_output,
+        input,
+        weight,
+        None,
+        _pair(stride),
+        _pair(padding),
+        _pair(dilation),
+        False,
+        [0],
+        groups,
+        (True, False, False),
+    )[0]
+
+
+def conv2d_weight(
+    input,
+    weight_size,
+    grad_output,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+):
+    r"""Compute the gradient of conv2d with respect to the weight of the convolution.
+
+    Args:
+        input: input tensor of shape (minibatch x in_channels x iH x iW)
+        weight_size : Shape of the weight gradient tensor
+        grad_output : output gradient tensor (minibatch x out_channels x oH x oW)
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+
+    Examples::
+
+        >>> input = torch.randn(1, 1, 3, 3, requires_grad=True)
+        >>> weight = torch.randn(1, 1, 1, 2, requires_grad=True)
+        >>> output = F.conv2d(input, weight)
+        >>> grad_output = torch.randn(output.shape)
+        >>> # xdoctest: +SKIP
+        >>> grad_weight = torch.autograd.grad(output, filter, grad_output)
+        >>> F.grad.conv2d_weight(input, weight.shape, grad_output)
+
+    """
+    weight = grad_output.new_empty(1).expand(weight_size)
+
+    return torch.ops.aten.convolution_backward(
+        grad_output,
+        input,
+        weight,
+        None,
+        _pair(stride),
+        _pair(padding),
+        _pair(dilation),
+        False,
+        [0],
+        groups,
+        (False, True, False),
+    )[1]
+
+
+def conv3d_input(
+    input_size,
+    weight,
+    grad_output,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+):
+    r"""Compute the gradient of conv3d with respect to the input of the convolution.
+
+    This is same as the 3D transposed convolution operator under the hood but requires
+    the shape of the gradient w.r.t. input to be specified explicitly.
+
+    Args:
+        input_size : Shape of the input gradient tensor
+        weight: weights tensor (out_channels x in_channels/groups x kT x kH x kW)
+        grad_output : output gradient tensor (minibatch x out_channels x oT x oH x oW)
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+
+    Examples::
+
+        >>> input = torch.randn(2, 8, 10, 10, 20, requires_grad=True)
+        >>> weight = torch.randn(4, 8, 2, 3, 3, requires_grad=True)
+        >>> output = F.conv3d(input, weight)
+        >>> grad_output = torch.randn(output.shape)
+        >>> grad_input = torch.autograd.grad(output, input, grad_output)
+        >>> F.grad.conv3d_input(input.shape, weight, grad_output)
+
+    """
+    input = grad_output.new_empty(1).expand(input_size)
+
+    return torch.ops.aten.convolution_backward(
+        grad_output,
+        input,
+        weight,
+        None,
+        _triple(stride),
+        _triple(padding),
+        _triple(dilation),
+        False,
+        [0],
+        groups,
+        (True, False, False),
+    )[0]
+
+
+def conv3d_weight(
+    input,
+    weight_size,
+    grad_output,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+):
+    r"""Compute the gradient of conv3d with respect to the weight of the convolution.
+
+    Args:
+        input: input tensor of shape (minibatch x in_channels x iT x iH x iW)
+        weight_size : Shape of the weight gradient tensor
+        grad_output : output gradient tensor (minibatch x out_channels x oT x oH x oW)
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+
+    Examples::
+
+        >>> input = torch.randn(2, 8, 10, 10, 20, requires_grad=True)
+        >>> weight = torch.randn(4, 8, 2, 3, 3, requires_grad=True)
+        >>> output = F.conv3d(input, weight)
+        >>> grad_output = torch.randn(output.shape)
+        >>> grad_weight = torch.autograd.grad(output, weight, grad_output)
+        >>> F.grad.conv3d_weight(input, weight.shape, grad_output)
+
+    """
+    weight = grad_output.new_empty(1).expand(weight_size)
+
+    return torch.ops.aten.convolution_backward(
+        grad_output,
+        input,
+        weight,
+        None,
+        _triple(stride),
+        _triple(padding),
+        _triple(dilation),
+        False,
+        [0],
+        groups,
+        (False, True, False),
+    )[1]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/init.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/init.py
new file mode 100644
index 0000000000000000000000000000000000000000..982845a1cfec09fe51b83bcebdcd409bea4f8da2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/init.py
@@ -0,0 +1,798 @@
+"""This file contains utilities for initializing neural network parameters."""
+
+import math
+import warnings
+from collections.abc import Callable
+from typing import Literal, TypeVar
+from typing_extensions import ParamSpec
+
+import torch
+from torch import Tensor
+
+
+__all__ = [
+    "calculate_gain",
+    "uniform_",
+    "normal_",
+    "trunc_normal_",
+    "constant_",
+    "ones_",
+    "zeros_",
+    "eye_",
+    "dirac_",
+    "xavier_uniform_",
+    "xavier_normal_",
+    "kaiming_uniform_",
+    "kaiming_normal_",
+    "orthogonal_",
+    "sparse_",
+    # Deprecated aliases (for backward compatibility)
+    "uniform",
+    "normal",
+    "constant",
+    "eye",
+    "dirac",
+    "xavier_uniform",
+    "xavier_normal",
+    "kaiming_uniform",
+    "kaiming_normal",
+    "orthogonal",
+    "sparse",
+]
+
+
+_R = TypeVar("_R")
+_P = ParamSpec("_P")
+
+_NonlinearityType = Literal[
+    "linear",
+    "conv1d",
+    "conv2d",
+    "conv3d",
+    "conv_transpose1d",
+    "conv_transpose2d",
+    "conv_transpose3d",
+    "sigmoid",
+    "tanh",
+    "relu",
+    "leaky_relu",
+    "selu",
+]
+
+_FanMode = Literal["fan_in", "fan_out"]
+
+
+# These no_grad_* functions are necessary as wrappers around the parts of these
+# functions that use `with torch.no_grad()`. The JIT doesn't support context
+# managers, so these need to be implemented as builtins. Using these wrappers
+# lets us keep those builtins small and reusable.
+def _no_grad_uniform_(
+    tensor: Tensor, a: float, b: float, generator: torch.Generator | None = None
+) -> Tensor:
+    with torch.no_grad():
+        return tensor.uniform_(a, b, generator=generator)
+
+
+def _no_grad_normal_(
+    tensor: Tensor,
+    mean: float,
+    std: float,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    with torch.no_grad():
+        return tensor.normal_(mean, std, generator=generator)
+
+
+def _no_grad_trunc_normal_(
+    tensor: Tensor,
+    mean: float,
+    std: float,
+    a: float,
+    b: float,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    # Meta tensors have no storage, so sampling is a no-op.
+    if tensor.is_meta:
+        return tensor
+
+    def norm_cdf(x: float) -> float:
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    with torch.no_grad():
+        p = norm_cdf((b - mean) / std) - norm_cdf((a - mean) / std)
+
+        if p > 0.3:
+            # Cast bounds to tensor dtype so the rejection mask is consistent
+            # with the tensor's representable values.
+            lo = tensor.new_tensor(a, device="cpu").item()
+            hi = tensor.new_tensor(b, device="cpu").item()
+            result = tensor.normal_(mean, std, generator=generator)
+            while True:
+                mask = (result < lo) | (result > hi)
+                if not mask.any():
+                    break
+                result = torch.where(
+                    mask,
+                    torch.empty_like(result).normal_(mean, std, generator=generator),
+                    result,
+                )
+            if tensor is not result:
+                tensor.copy_(result)
+        else:
+            mode = max(a, min(mean, b))
+            log_peak = -0.5 * ((mode - mean) / std) ** 2
+
+            candidates = torch.empty_like(tensor)
+            accept_buf = torch.empty_like(tensor)
+
+            # First iteration: sample directly into tensor to avoid
+            # a where() + copy_() if all samples are accepted.
+            tensor.uniform_(a, b, generator=generator)
+            candidates.copy_(tensor)
+            # log_pdf = -0.5 * ((candidates - mean) / std) ** 2
+            candidates.sub_(mean).div_(std).pow_(2).mul_(-0.5).sub_(log_peak)
+            pending = accept_buf.uniform_(generator=generator).log_().gt(candidates)
+            if not pending.any():
+                pass
+            else:
+                result = tensor
+                while True:
+                    candidates.uniform_(a, b, generator=generator)
+                    result = torch.where(pending, candidates, result)
+                    # log_pdf = -0.5 * ((candidates - mean) / std) ** 2
+                    candidates.sub_(mean).div_(std).pow_(2).mul_(-0.5).sub_(log_peak)
+                    pending = torch.where(
+                        pending,
+                        accept_buf.uniform_(generator=generator).log_().gt(candidates),
+                        pending,
+                    )
+                    if not pending.any():
+                        break
+                tensor.copy_(result)
+
+        return tensor
+
+
+def _no_grad_fill_(tensor: Tensor, val: float) -> Tensor:
+    with torch.no_grad():
+        return tensor.fill_(val)
+
+
+def _no_grad_zero_(tensor: Tensor) -> Tensor:
+    with torch.no_grad():
+        return tensor.zero_()
+
+
+def calculate_gain(
+    nonlinearity: _NonlinearityType, param: int | float | None = None
+) -> float:
+    r"""Return the recommended gain value for the given nonlinearity function.
+
+    The values are as follows:
+
+    ================= ====================================================
+    nonlinearity      gain
+    ================= ====================================================
+    Linear / Identity :math:`1`
+    Conv{1,2,3}D      :math:`1`
+    Sigmoid           :math:`1`
+    Tanh              :math:`\frac{5}{3}`
+    ReLU              :math:`\sqrt{2}`
+    Leaky Relu        :math:`\sqrt{\frac{2}{1 + \text{negative\_slope}^2}}`
+    SELU              :math:`\frac{3}{4}`
+    ================= ====================================================
+
+    .. warning::
+        In order to implement `Self-Normalizing Neural Networks`_ ,
+        you should use ``nonlinearity='linear'`` instead of ``nonlinearity='selu'``.
+        This gives the initial weights a variance of ``1 / N``,
+        which is necessary to induce a stable fixed point in the forward pass.
+        In contrast, the default gain for ``SELU`` sacrifices the normalization
+        effect for more stable gradient flow in rectangular layers.
+
+    Args:
+        nonlinearity: the non-linear function (`nn.functional` name)
+        param: optional parameter for the non-linear function
+
+    Examples:
+        >>> gain = nn.init.calculate_gain(
+        ...     "leaky_relu", 0.2
+        ... )  # leaky_relu with negative_slope=0.2
+
+    .. _Self-Normalizing Neural Networks: https://papers.nips.cc/paper/2017/hash/5d44ee6f2c3f71b73125876103c8f6c4-Abstract.html
+    """
+    linear_fns = [
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+        "conv_transpose1d",
+        "conv_transpose2d",
+        "conv_transpose3d",
+    ]
+    if nonlinearity in linear_fns or nonlinearity == "sigmoid":
+        return 1
+    elif nonlinearity == "tanh":
+        return 5.0 / 3
+    elif nonlinearity == "relu":
+        return math.sqrt(2.0)
+    elif nonlinearity == "leaky_relu":
+        if param is None:
+            negative_slope = 0.01
+        elif (
+            not isinstance(param, bool)
+            and isinstance(param, int)
+            or isinstance(param, float)
+        ):
+            # True/False are instances of int, hence check above
+            negative_slope = param
+        else:
+            raise ValueError(f"negative_slope {param} not a valid number")
+        return math.sqrt(2.0 / (1 + negative_slope**2))
+    elif nonlinearity == "selu":
+        return (
+            3.0 / 4
+        )  # Value found empirically (https://github.com/pytorch/pytorch/pull/50664)
+    else:
+        raise ValueError(f"Unsupported nonlinearity {nonlinearity}")
+
+
+def uniform_(
+    tensor: Tensor,
+    a: float = 0.0,
+    b: float = 1.0,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input Tensor with values drawn from the uniform distribution.
+
+    :math:`\mathcal{U}(a, b)`.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        a: the lower bound of the uniform distribution
+        b: the upper bound of the uniform distribution
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.uniform_(w)
+    """
+    if torch.overrides.has_torch_function_variadic(tensor):
+        return torch.overrides.handle_torch_function(
+            uniform_, (tensor,), tensor=tensor, a=a, b=b, generator=generator
+        )
+    return _no_grad_uniform_(tensor, a, b, generator)
+
+
+def normal_(
+    tensor: Tensor,
+    mean: float = 0.0,
+    std: float = 1.0,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input Tensor with values drawn from the normal distribution.
+
+    :math:`\mathcal{N}(\text{mean}, \text{std}^2)`.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.normal_(w)
+    """
+    if torch.overrides.has_torch_function_variadic(tensor):
+        return torch.overrides.handle_torch_function(
+            normal_, (tensor,), tensor=tensor, mean=mean, std=std, generator=generator
+        )
+    return _no_grad_normal_(tensor, mean, std, generator)
+
+
+def trunc_normal_(
+    tensor: Tensor,
+    mean: float = 0.0,
+    std: float = 1.0,
+    a: float = -2.0,
+    b: float = 2.0,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input Tensor with values drawn from a truncated normal distribution.
+
+    The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+
+    For reduced-precision types (``torch.float16`` and ``torch.bfloat16``),
+    sampling quality depends on the underlying ``normal_()`` and ``uniform_()``
+    implementations which operate at higher internal precision to avoid
+    quantization artifacts.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b, generator=generator)
+
+
+def constant_(tensor: Tensor, val: float) -> Tensor:
+    r"""Fill the input Tensor with the value :math:`\text{val}`.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        val: the value to fill the tensor with
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.constant_(w, 0.3)
+    """
+    if torch.overrides.has_torch_function_variadic(tensor):
+        return torch.overrides.handle_torch_function(
+            constant_, (tensor,), tensor=tensor, val=val
+        )
+    return _no_grad_fill_(tensor, val)
+
+
+def ones_(tensor: Tensor) -> Tensor:
+    r"""Fill the input Tensor with the scalar value `1`.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.ones_(w)
+    """
+    return _no_grad_fill_(tensor, 1.0)
+
+
+def zeros_(tensor: Tensor) -> Tensor:
+    r"""Fill the input Tensor with the scalar value `0`.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.zeros_(w)
+    """
+    return _no_grad_zero_(tensor)
+
+
+def eye_(tensor: Tensor) -> Tensor:
+    r"""Fill the 2-dimensional input `Tensor` with the identity matrix.
+
+    Preserves the identity of the inputs in `Linear` layers, where as
+    many inputs are preserved as possible.
+
+    Args:
+        tensor: a 2-dimensional `torch.Tensor`
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.eye_(w)
+    """
+    if tensor.ndimension() != 2:
+        raise ValueError("Only tensors with 2 dimensions are supported")
+
+    with torch.no_grad():
+        torch.eye(*tensor.shape, out=tensor, requires_grad=tensor.requires_grad)
+    return tensor
+
+
+def dirac_(tensor: Tensor, groups: int = 1) -> Tensor:
+    r"""Fill the {3, 4, 5}-dimensional input `Tensor` with the Dirac delta function.
+
+    Preserves the identity of the inputs in `Convolutional`
+    layers, where as many input channels are preserved as possible. In case
+    of groups>1, each group of channels preserves identity
+
+    Args:
+        tensor: a {3, 4, 5}-dimensional `torch.Tensor`
+        groups (int, optional): number of groups in the conv layer (default: 1)
+    Examples:
+        >>> w = torch.empty(3, 16, 5, 5)
+        >>> nn.init.dirac_(w)
+        >>> w = torch.empty(3, 24, 5, 5)
+        >>> nn.init.dirac_(w, 3)
+    """
+    dimensions = tensor.ndimension()
+    if dimensions not in [3, 4, 5]:
+        raise ValueError("Only tensors with 3, 4, or 5 dimensions are supported")
+
+    sizes = tensor.size()
+
+    if sizes[0] % groups != 0:
+        raise ValueError("dim 0 must be divisible by groups")
+
+    if tensor.is_meta:
+        return tensor
+
+    out_chans_per_grp = sizes[0] // groups
+    min_dim = min(out_chans_per_grp, sizes[1])
+
+    with torch.no_grad():
+        tensor.zero_()
+
+        for g in range(groups):
+            for d in range(min_dim):
+                if dimensions == 3:  # Temporal convolution
+                    tensor[g * out_chans_per_grp + d, d, tensor.size(2) // 2] = 1
+                elif dimensions == 4:  # Spatial convolution
+                    tensor[
+                        g * out_chans_per_grp + d,
+                        d,
+                        tensor.size(2) // 2,
+                        tensor.size(3) // 2,
+                    ] = 1
+                else:  # Volumetric convolution
+                    tensor[
+                        g * out_chans_per_grp + d,
+                        d,
+                        tensor.size(2) // 2,
+                        tensor.size(3) // 2,
+                        tensor.size(4) // 2,
+                    ] = 1
+    return tensor
+
+
+def _calculate_fan_in_and_fan_out(tensor: Tensor) -> tuple[int, int]:
+    dimensions = tensor.dim()
+    if dimensions < 2:
+        raise ValueError(
+            "Fan in and fan out can not be computed for tensor with fewer than 2 dimensions"
+        )
+
+    num_input_fmaps = tensor.size(1)
+    num_output_fmaps = tensor.size(0)
+    receptive_field_size = 1
+    if tensor.dim() > 2:
+        # math.prod is not always available, accumulate the product manually
+        # we could use functools.reduce but that is not supported by TorchScript
+        for s in tensor.shape[2:]:
+            receptive_field_size *= s
+    fan_in = num_input_fmaps * receptive_field_size
+    fan_out = num_output_fmaps * receptive_field_size
+
+    return fan_in, fan_out
+
+
+def xavier_uniform_(
+    tensor: Tensor,
+    gain: float = 1.0,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input `Tensor` with values using a Xavier uniform distribution.
+
+    The method is described in `Understanding the difficulty of training
+    deep feedforward neural networks` - Glorot, X. & Bengio, Y. (2010).
+    The resulting tensor will have values sampled from
+    :math:`\mathcal{U}(-a, a)` where
+
+    .. math::
+        a = \text{gain} \times \sqrt{\frac{6}{\text{fan\_in} + \text{fan\_out}}}
+
+    Also known as Glorot initialization.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        gain: an optional scaling factor
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.xavier_uniform_(w, gain=nn.init.calculate_gain("relu"))
+    """
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
+    a = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
+
+    return _no_grad_uniform_(tensor, -a, a, generator)
+
+
+def xavier_normal_(
+    tensor: Tensor,
+    gain: float = 1.0,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input `Tensor` with values using a Xavier normal distribution.
+
+    The method is described in `Understanding the difficulty of training deep feedforward
+    neural networks` - Glorot, X. & Bengio, Y. (2010). The resulting tensor
+    will have values sampled from :math:`\mathcal{N}(0, \text{std}^2)` where
+
+    .. math::
+        \text{std} = \text{gain} \times \sqrt{\frac{2}{\text{fan\_in} + \text{fan\_out}}}
+
+    Also known as Glorot initialization.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        gain: an optional scaling factor
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.xavier_normal_(w)
+    """
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
+
+    return _no_grad_normal_(tensor, 0.0, std, generator)
+
+
+def _calculate_correct_fan(tensor: Tensor, mode: _FanMode) -> int:
+    # pyrefly: ignore [bad-assignment]
+    mode = mode.lower()
+    valid_modes = ["fan_in", "fan_out"]
+    if mode not in valid_modes:
+        raise ValueError(f"Mode {mode} not supported, please use one of {valid_modes}")
+
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    return fan_in if mode == "fan_in" else fan_out
+
+
+def kaiming_uniform_(
+    tensor: Tensor,
+    a: float = 0,
+    mode: _FanMode = "fan_in",
+    nonlinearity: _NonlinearityType = "leaky_relu",
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input `Tensor` with values using a Kaiming uniform distribution.
+
+    The method is described in `Delving deep into rectifiers: Surpassing
+    human-level performance on ImageNet classification` - He, K. et al. (2015).
+    The resulting tensor will have values sampled from
+    :math:`\mathcal{U}(-\text{bound}, \text{bound})` where
+
+    .. math::
+        \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}}
+
+    Also known as He initialization.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        a: the negative slope of the rectifier used after this layer (only
+            used with ``'leaky_relu'``)
+        mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``
+            preserves the magnitude of the variance of the weights in the
+            forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the
+            backwards pass.
+        nonlinearity: the non-linear function (`nn.functional` name),
+            recommended to use only with ``'relu'`` or ``'leaky_relu'`` (default).
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.kaiming_uniform_(w, mode="fan_in", nonlinearity="relu")
+
+    Note:
+        Be aware that ``fan_in`` and ``fan_out`` are calculated assuming
+        that the weight matrix is used in a transposed manner,
+        (i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).
+        This is important for correct initialization.
+        If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,
+        pass in a transposed weight matrix, i.e. ``nn.init.kaiming_uniform_(w.T, ...)``.
+    """
+    if torch.overrides.has_torch_function_variadic(tensor):
+        return torch.overrides.handle_torch_function(
+            kaiming_uniform_,
+            (tensor,),
+            tensor=tensor,
+            a=a,
+            mode=mode,
+            nonlinearity=nonlinearity,
+            generator=generator,
+        )
+
+    if 0 in tensor.shape:
+        warnings.warn("Initializing zero-element tensors is a no-op", stacklevel=2)
+        return tensor
+    fan = _calculate_correct_fan(tensor, mode)
+    gain = calculate_gain(nonlinearity, a)
+    std = gain / math.sqrt(fan)
+    bound = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
+    with torch.no_grad():
+        return tensor.uniform_(-bound, bound, generator=generator)
+
+
+def kaiming_normal_(
+    tensor: Tensor,
+    a: float = 0,
+    mode: _FanMode = "fan_in",
+    nonlinearity: _NonlinearityType = "leaky_relu",
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input `Tensor` with values using a Kaiming normal distribution.
+
+    The method is described in `Delving deep into rectifiers: Surpassing
+    human-level performance on ImageNet classification` - He, K. et al. (2015).
+    The resulting tensor will have values sampled from
+    :math:`\mathcal{N}(0, \text{std}^2)` where
+
+    .. math::
+        \text{std} = \frac{\text{gain}}{\sqrt{\text{fan\_mode}}}
+
+    Also known as He initialization.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        a: the negative slope of the rectifier used after this layer (only
+            used with ``'leaky_relu'``)
+        mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``
+            preserves the magnitude of the variance of the weights in the
+            forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the
+            backwards pass.
+        nonlinearity: the non-linear function (`nn.functional` name),
+            recommended to use only with ``'relu'`` or ``'leaky_relu'`` (default).
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.kaiming_normal_(w, mode="fan_out", nonlinearity="relu")
+
+    Note:
+        Be aware that ``fan_in`` and ``fan_out`` are calculated assuming
+        that the weight matrix is used in a transposed manner,
+        (i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).
+        This is important for correct initialization.
+        If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,
+        pass in a transposed weight matrix, i.e. ``nn.init.kaiming_normal_(w.T, ...)``.
+    """
+    if 0 in tensor.shape:
+        warnings.warn("Initializing zero-element tensors is a no-op", stacklevel=2)
+        return tensor
+    fan = _calculate_correct_fan(tensor, mode)
+    gain = calculate_gain(nonlinearity, a)
+    std = gain / math.sqrt(fan)
+    with torch.no_grad():
+        return tensor.normal_(0, std, generator=generator)
+
+
+def orthogonal_(
+    tensor: Tensor,
+    gain: float = 1,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the input `Tensor` with a (semi) orthogonal matrix.
+
+    Described in `Exact solutions to the nonlinear dynamics of learning in deep
+    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
+    at least 2 dimensions, and for tensors with more than 2 dimensions the
+    trailing dimensions are flattened.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
+        gain: optional scaling factor
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.orthogonal_(w)
+    """
+    if tensor.ndimension() < 2:
+        raise ValueError("Only tensors with 2 or more dimensions are supported")
+
+    if tensor.numel() == 0 or tensor.is_meta:
+        # no-op
+        return tensor
+    rows = tensor.size(0)
+    cols = tensor.numel() // rows
+    flattened = tensor.new_empty((rows, cols)).normal_(0, 1, generator=generator)
+
+    if rows < cols:
+        flattened.t_()
+
+    # Compute the qr factorization
+    q, r = torch.linalg.qr(flattened)
+    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
+    d = torch.diag(r, 0)
+    ph = d.sign()
+    q *= ph
+
+    if rows < cols:
+        q.t_()
+
+    with torch.no_grad():
+        tensor.view_as(q).copy_(q)
+        tensor.mul_(gain)
+    return tensor
+
+
+def sparse_(
+    tensor: Tensor,
+    sparsity: float,
+    std: float = 0.01,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    r"""Fill the 2D input `Tensor` as a sparse matrix.
+
+    The non-zero elements will be drawn from the normal distribution
+    :math:`\mathcal{N}(0, 0.01)`, as described in `Deep learning via
+    Hessian-free optimization` - Martens, J. (2010).
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        sparsity: The fraction of elements in each column to be set to zero
+        std: the standard deviation of the normal distribution used to generate
+            the non-zero values
+        generator: the torch Generator to sample from (default: None)
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.sparse_(w, sparsity=0.1)
+    """
+    if tensor.ndimension() != 2:
+        raise ValueError("Only tensors with 2 dimensions are supported")
+
+    if tensor.is_meta:
+        return tensor
+
+    rows, cols = tensor.shape
+    num_zeros = math.ceil(sparsity * rows)
+
+    with torch.no_grad():
+        tensor.normal_(0, std, generator=generator)
+        for col_idx in range(cols):
+            row_indices = torch.randperm(rows)
+            zero_indices = row_indices[:num_zeros]
+            tensor[zero_indices, col_idx] = 0
+    return tensor
+
+
+# for backward compatibility
+def _make_deprecate(meth: Callable[_P, _R]) -> Callable[_P, _R]:
+    new_name = meth.__name__
+    old_name = new_name[:-1]
+
+    def deprecated_init(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        warnings.warn(
+            f"`nn.init.{old_name}` is now deprecated in favor of `nn.init.{new_name}`.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        return meth(*args, **kwargs)
+
+    deprecated_init.__doc__ = rf"""
+    {old_name}(...)
+
+    .. warning::
+        This method is now deprecated in favor of :func:`torch.nn.init.{new_name}`.
+
+    See :func:`~torch.nn.init.{new_name}` for details."""
+    deprecated_init.__name__ = old_name
+    return deprecated_init
+
+
+uniform = _make_deprecate(uniform_)
+normal = _make_deprecate(normal_)
+constant = _make_deprecate(constant_)
+eye = _make_deprecate(eye_)
+dirac = _make_deprecate(dirac_)
+xavier_uniform = _make_deprecate(xavier_uniform_)
+xavier_normal = _make_deprecate(xavier_normal_)
+kaiming_uniform = _make_deprecate(kaiming_uniform_)
+kaiming_normal = _make_deprecate(kaiming_normal_)
+orthogonal = _make_deprecate(orthogonal_)
+sparse = _make_deprecate(sparse_)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fe9a09aa31464fd4e88b2e46b4210561a70e42e7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/__init__.py
@@ -0,0 +1,36 @@
+from torch.ao.nn.intrinsic import (
+    BNReLU2d,
+    BNReLU3d,
+    ConvBn1d,
+    ConvBn2d,
+    ConvBn3d,
+    ConvBnReLU1d,
+    ConvBnReLU2d,
+    ConvBnReLU3d,
+    ConvReLU1d,
+    ConvReLU2d,
+    ConvReLU3d,
+    LinearBn1d,
+    LinearReLU,
+)
+from torch.ao.nn.intrinsic.modules.fused import _FusedModule  # noqa: F401
+
+# Include the subpackages in case user imports from it directly
+from torch.nn.intrinsic import modules, qat, quantized  # noqa: F401
+
+
+__all__ = [
+    "ConvBn1d",
+    "ConvBn2d",
+    "ConvBn3d",
+    "ConvBnReLU1d",
+    "ConvBnReLU2d",
+    "ConvBnReLU3d",
+    "ConvReLU1d",
+    "ConvReLU2d",
+    "ConvReLU3d",
+    "LinearReLU",
+    "BNReLU2d",
+    "BNReLU3d",
+    "LinearBn1d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b1f0d85ab1aba10947e3fe0570a0ee49a4bb66e9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/modules/__init__.py
@@ -0,0 +1,33 @@
+from torch.nn.intrinsic.modules.fused import (
+    _FusedModule,
+    BNReLU2d,
+    BNReLU3d,
+    ConvBn1d,
+    ConvBn2d,
+    ConvBn3d,
+    ConvBnReLU1d,
+    ConvBnReLU2d,
+    ConvBnReLU3d,
+    ConvReLU1d,
+    ConvReLU2d,
+    ConvReLU3d,
+    LinearBn1d,
+    LinearReLU,
+)
+
+
+__all__ = [
+    "BNReLU2d",
+    "BNReLU3d",
+    "ConvBn1d",
+    "ConvBn2d",
+    "ConvBn3d",
+    "ConvBnReLU1d",
+    "ConvBnReLU2d",
+    "ConvBnReLU3d",
+    "ConvReLU1d",
+    "ConvReLU2d",
+    "ConvReLU3d",
+    "LinearBn1d",
+    "LinearReLU",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/modules/fused.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/modules/fused.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ba11c92f05d21d97473349eb13174bbca8ca4d2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/modules/fused.py
@@ -0,0 +1,33 @@
+from torch.ao.nn.intrinsic import (
+    BNReLU2d,
+    BNReLU3d,
+    ConvBn1d,
+    ConvBn2d,
+    ConvBn3d,
+    ConvBnReLU1d,
+    ConvBnReLU2d,
+    ConvBnReLU3d,
+    ConvReLU1d,
+    ConvReLU2d,
+    ConvReLU3d,
+    LinearBn1d,
+    LinearReLU,
+)
+from torch.ao.nn.intrinsic.modules.fused import _FusedModule  # noqa: F401
+
+
+__all__ = [
+    "BNReLU2d",
+    "BNReLU3d",
+    "ConvBn1d",
+    "ConvBn2d",
+    "ConvBn3d",
+    "ConvBnReLU1d",
+    "ConvBnReLU2d",
+    "ConvBnReLU3d",
+    "ConvReLU1d",
+    "ConvReLU2d",
+    "ConvReLU3d",
+    "LinearBn1d",
+    "LinearReLU",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..302ed488399965e7db72656cd4bc8e522085df3d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/__init__.py
@@ -0,0 +1 @@
+from torch.nn.intrinsic.qat.modules import *  # noqa: F403
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..cc439cd2fa6b439cec7b6ba0c9eec372c60a0caf
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/__init__.py
@@ -0,0 +1,32 @@
+from torch.nn.intrinsic.qat.modules.conv_fused import (
+    ConvBn1d,
+    ConvBn2d,
+    ConvBn3d,
+    ConvBnReLU1d,
+    ConvBnReLU2d,
+    ConvBnReLU3d,
+    ConvReLU1d,
+    ConvReLU2d,
+    ConvReLU3d,
+    freeze_bn_stats,
+    update_bn_stats,
+)
+from torch.nn.intrinsic.qat.modules.linear_fused import LinearBn1d
+from torch.nn.intrinsic.qat.modules.linear_relu import LinearReLU
+
+
+__all__ = [
+    "LinearReLU",
+    "LinearBn1d",
+    "ConvReLU1d",
+    "ConvReLU2d",
+    "ConvReLU3d",
+    "ConvBn1d",
+    "ConvBn2d",
+    "ConvBn3d",
+    "ConvBnReLU1d",
+    "ConvBnReLU2d",
+    "ConvBnReLU3d",
+    "update_bn_stats",
+    "freeze_bn_stats",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/conv_fused.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/conv_fused.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8dc1d49aad310819ef892b2ca6e984d1e392f06
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/conv_fused.py
@@ -0,0 +1,39 @@
+r"""Intrinsic QAT Modules.
+
+This file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/intrinsic/qat/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.intrinsic.qat import (
+    ConvBn1d,
+    ConvBn2d,
+    ConvBn3d,
+    ConvBnReLU1d,
+    ConvBnReLU2d,
+    ConvBnReLU3d,
+    ConvReLU1d,
+    ConvReLU2d,
+    ConvReLU3d,
+    freeze_bn_stats,
+    update_bn_stats,
+)
+
+
+__all__ = [
+    # Modules
+    "ConvBn1d",
+    "ConvBnReLU1d",
+    "ConvReLU1d",
+    "ConvBn2d",
+    "ConvBnReLU2d",
+    "ConvReLU2d",
+    "ConvBn3d",
+    "ConvBnReLU3d",
+    "ConvReLU3d",
+    # Utilities
+    "freeze_bn_stats",
+    "update_bn_stats",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/linear_fused.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/linear_fused.py
new file mode 100644
index 0000000000000000000000000000000000000000..79567d67bd1f930b89f6270232efbefec222ac61
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/linear_fused.py
@@ -0,0 +1,15 @@
+r"""Intrinsic QAT Modules.
+
+This file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/intrinsic/qat/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.intrinsic.qat import LinearBn1d
+
+
+__all__ = [
+    "LinearBn1d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/linear_relu.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/linear_relu.py
new file mode 100644
index 0000000000000000000000000000000000000000..71705320075efd08040418fb69c492882a6bd03d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/qat/modules/linear_relu.py
@@ -0,0 +1,15 @@
+r"""Intrinsic QAT Modules.
+
+This file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/intrinsic/qat/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.intrinsic.qat import LinearReLU
+
+
+__all__ = [
+    "LinearReLU",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c09fddf6e75f989349d206ee1dcab6c9eb25678
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/__init__.py
@@ -0,0 +1,14 @@
+# to ensure customers can use the module below
+# without importing it directly
+from torch.nn.intrinsic.quantized import dynamic, modules  # noqa: F401
+from torch.nn.intrinsic.quantized.modules import *  # noqa: F403
+
+
+__all__ = [
+    "BNReLU2d",
+    "BNReLU3d",
+    "ConvReLU1d",
+    "ConvReLU2d",
+    "ConvReLU3d",
+    "LinearReLU",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ce645703ddb584c34d80c15b562f0dacd14350b4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/__init__.py
@@ -0,0 +1 @@
+from torch.nn.intrinsic.quantized.dynamic.modules import *  # noqa: F403
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..49d148aa61a549518de0e98dad2a441bd837eba9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/modules/__init__.py
@@ -0,0 +1,6 @@
+from torch.nn.intrinsic.quantized.dynamic.modules.linear_relu import LinearReLU
+
+
+__all__ = [
+    "LinearReLU",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/modules/linear_relu.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/modules/linear_relu.py
new file mode 100644
index 0000000000000000000000000000000000000000..85df3dd4b441a342c8126dd4e8cea89f5bfa5c08
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/dynamic/modules/linear_relu.py
@@ -0,0 +1,6 @@
+from torch.ao.nn.intrinsic.quantized.dynamic import LinearReLU
+
+
+__all__ = [
+    "LinearReLU",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7cc79b04418f8a7a8128c52a64e938396a94ee94
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/__init__.py
@@ -0,0 +1,17 @@
+from torch.nn.intrinsic.quantized.modules.bn_relu import BNReLU2d, BNReLU3d
+from torch.nn.intrinsic.quantized.modules.conv_relu import (
+    ConvReLU1d,
+    ConvReLU2d,
+    ConvReLU3d,
+)
+from torch.nn.intrinsic.quantized.modules.linear_relu import LinearReLU
+
+
+__all__ = [
+    "LinearReLU",
+    "ConvReLU1d",
+    "ConvReLU2d",
+    "ConvReLU3d",
+    "BNReLU2d",
+    "BNReLU3d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/bn_relu.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/bn_relu.py
new file mode 100644
index 0000000000000000000000000000000000000000..f17253c93246f52c49b5310788028e91c61b92ad
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/bn_relu.py
@@ -0,0 +1,7 @@
+from torch.ao.nn.intrinsic.quantized import BNReLU2d, BNReLU3d
+
+
+__all__ = [
+    "BNReLU2d",
+    "BNReLU3d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/conv_relu.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/conv_relu.py
new file mode 100644
index 0000000000000000000000000000000000000000..5a72b4d5ebe8afd13d67eb6c1e9d227fce7aa6b2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/conv_relu.py
@@ -0,0 +1,8 @@
+from torch.ao.nn.intrinsic.quantized import ConvReLU1d, ConvReLU2d, ConvReLU3d
+
+
+__all__ = [
+    "ConvReLU1d",
+    "ConvReLU2d",
+    "ConvReLU3d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/linear_relu.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/linear_relu.py
new file mode 100644
index 0000000000000000000000000000000000000000..2d848c84d927669f2381edcbadd7e557ac429822
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/intrinsic/quantized/modules/linear_relu.py
@@ -0,0 +1,6 @@
+from torch.ao.nn.intrinsic.quantized import LinearReLU
+
+
+__all__ = [
+    "LinearReLU",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9e076cbe1b69e262cd91ed9f7f415c07c9acc9d8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/__init__.py
@@ -0,0 +1,335 @@
+from .module import Module  # usort: skip
+from .linear import Bilinear, Identity, LazyLinear, Linear  # usort: skip
+from .activation import (
+    CELU,
+    ELU,
+    GELU,
+    GLU,
+    Hardshrink,
+    Hardsigmoid,
+    Hardswish,
+    Hardtanh,
+    LeakyReLU,
+    LogSigmoid,
+    LogSoftmax,
+    Mish,
+    MultiheadAttention,
+    PReLU,
+    ReLU,
+    ReLU6,
+    RReLU,
+    SELU,
+    Sigmoid,
+    SiLU,
+    Softmax,
+    Softmax2d,
+    Softmin,
+    Softplus,
+    Softshrink,
+    Softsign,
+    Tanh,
+    Tanhshrink,
+    Threshold,
+)
+from .adaptive import AdaptiveLogSoftmaxWithLoss
+from .batchnorm import (
+    BatchNorm1d,
+    BatchNorm2d,
+    BatchNorm3d,
+    LazyBatchNorm1d,
+    LazyBatchNorm2d,
+    LazyBatchNorm3d,
+    SyncBatchNorm,
+)
+from .channelshuffle import ChannelShuffle
+from .container import (
+    Container,
+    ModuleDict,
+    ModuleList,
+    ParameterDict,
+    ParameterList,
+    Sequential,
+)
+from .conv import (
+    Conv1d,
+    Conv2d,
+    Conv3d,
+    ConvTranspose1d,
+    ConvTranspose2d,
+    ConvTranspose3d,
+    LazyConv1d,
+    LazyConv2d,
+    LazyConv3d,
+    LazyConvTranspose1d,
+    LazyConvTranspose2d,
+    LazyConvTranspose3d,
+)
+from .distance import CosineSimilarity, PairwiseDistance
+from .dropout import (
+    AlphaDropout,
+    Dropout,
+    Dropout1d,
+    Dropout2d,
+    Dropout3d,
+    FeatureAlphaDropout,
+)
+from .flatten import Flatten, Unflatten
+from .fold import Fold, Unfold
+from .instancenorm import (
+    InstanceNorm1d,
+    InstanceNorm2d,
+    InstanceNorm3d,
+    LazyInstanceNorm1d,
+    LazyInstanceNorm2d,
+    LazyInstanceNorm3d,
+)
+from .loss import (
+    BCELoss,
+    BCEWithLogitsLoss,
+    CosineEmbeddingLoss,
+    CrossEntropyLoss,
+    CTCLoss,
+    GaussianNLLLoss,
+    HingeEmbeddingLoss,
+    HuberLoss,
+    KLDivLoss,
+    L1Loss,
+    MarginRankingLoss,
+    MSELoss,
+    MultiLabelMarginLoss,
+    MultiLabelSoftMarginLoss,
+    MultiMarginLoss,
+    NLLLoss,
+    NLLLoss2d,
+    PoissonNLLLoss,
+    SmoothL1Loss,
+    SoftMarginLoss,
+    TripletMarginLoss,
+    TripletMarginWithDistanceLoss,
+)
+from .normalization import (
+    CrossMapLRN2d,
+    GroupNorm,
+    LayerNorm,
+    LocalResponseNorm,
+    RMSNorm,
+)
+from .padding import (
+    CircularPad1d,
+    CircularPad2d,
+    CircularPad3d,
+    ConstantPad1d,
+    ConstantPad2d,
+    ConstantPad3d,
+    ReflectionPad1d,
+    ReflectionPad2d,
+    ReflectionPad3d,
+    ReplicationPad1d,
+    ReplicationPad2d,
+    ReplicationPad3d,
+    ZeroPad1d,
+    ZeroPad2d,
+    ZeroPad3d,
+)
+from .pixelshuffle import PixelShuffle, PixelUnshuffle
+from .pooling import (
+    AdaptiveAvgPool1d,
+    AdaptiveAvgPool2d,
+    AdaptiveAvgPool3d,
+    AdaptiveMaxPool1d,
+    AdaptiveMaxPool2d,
+    AdaptiveMaxPool3d,
+    AvgPool1d,
+    AvgPool2d,
+    AvgPool3d,
+    FractionalMaxPool2d,
+    FractionalMaxPool3d,
+    LPPool1d,
+    LPPool2d,
+    LPPool3d,
+    MaxPool1d,
+    MaxPool2d,
+    MaxPool3d,
+    MaxUnpool1d,
+    MaxUnpool2d,
+    MaxUnpool3d,
+)
+from .rnn import GRU, GRUCell, LSTM, LSTMCell, RNN, RNNBase, RNNCell, RNNCellBase
+from .sparse import Embedding, EmbeddingBag
+from .transformer import (
+    Transformer,
+    TransformerDecoder,
+    TransformerDecoderLayer,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+from .upsampling import Upsample, UpsamplingBilinear2d, UpsamplingNearest2d
+
+
+__all__ = [
+    "AdaptiveAvgPool1d",
+    "AdaptiveAvgPool2d",
+    "AdaptiveAvgPool3d",
+    "AdaptiveLogSoftmaxWithLoss",
+    "AdaptiveMaxPool1d",
+    "AdaptiveMaxPool2d",
+    "AdaptiveMaxPool3d",
+    "AlphaDropout",
+    "AvgPool1d",
+    "AvgPool2d",
+    "AvgPool3d",
+    "BCELoss",
+    "BCEWithLogitsLoss",
+    "BatchNorm1d",
+    "BatchNorm2d",
+    "BatchNorm3d",
+    "Bilinear",
+    "CELU",
+    "CTCLoss",
+    "ChannelShuffle",
+    "CircularPad1d",
+    "CircularPad2d",
+    "CircularPad3d",
+    "ConstantPad1d",
+    "ConstantPad2d",
+    "ConstantPad3d",
+    "Container",
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "ConvTranspose1d",
+    "ConvTranspose2d",
+    "ConvTranspose3d",
+    "CosineEmbeddingLoss",
+    "CosineSimilarity",
+    "CrossEntropyLoss",
+    "CrossMapLRN2d",
+    "Dropout",
+    "Dropout1d",
+    "Dropout2d",
+    "Dropout3d",
+    "ELU",
+    "Embedding",
+    "EmbeddingBag",
+    "FeatureAlphaDropout",
+    "Flatten",
+    "Fold",
+    "FractionalMaxPool2d",
+    "FractionalMaxPool3d",
+    "GELU",
+    "GLU",
+    "GRU",
+    "GRUCell",
+    "GaussianNLLLoss",
+    "GroupNorm",
+    "Hardshrink",
+    "Hardsigmoid",
+    "Hardswish",
+    "Hardtanh",
+    "HingeEmbeddingLoss",
+    "HuberLoss",
+    "Identity",
+    "InstanceNorm1d",
+    "InstanceNorm2d",
+    "InstanceNorm3d",
+    "KLDivLoss",
+    "L1Loss",
+    "LPPool1d",
+    "LPPool2d",
+    "LPPool3d",
+    "LSTM",
+    "LSTMCell",
+    "LayerNorm",
+    "LazyBatchNorm1d",
+    "LazyBatchNorm2d",
+    "LazyBatchNorm3d",
+    "LazyConv1d",
+    "LazyConv2d",
+    "LazyConv3d",
+    "LazyConvTranspose1d",
+    "LazyConvTranspose2d",
+    "LazyConvTranspose3d",
+    "LazyInstanceNorm1d",
+    "LazyInstanceNorm2d",
+    "LazyInstanceNorm3d",
+    "LazyLinear",
+    "LeakyReLU",
+    "Linear",
+    "LocalResponseNorm",
+    "LogSigmoid",
+    "LogSoftmax",
+    "MSELoss",
+    "MarginRankingLoss",
+    "MaxPool1d",
+    "MaxPool2d",
+    "MaxPool3d",
+    "MaxUnpool1d",
+    "MaxUnpool2d",
+    "MaxUnpool3d",
+    "Mish",
+    "Module",
+    "ModuleDict",
+    "ModuleList",
+    "MultiLabelMarginLoss",
+    "MultiLabelSoftMarginLoss",
+    "MultiMarginLoss",
+    "MultiheadAttention",
+    "NLLLoss",
+    "NLLLoss2d",
+    "PReLU",
+    "PairwiseDistance",
+    "ParameterDict",
+    "ParameterList",
+    "PixelShuffle",
+    "PixelUnshuffle",
+    "PoissonNLLLoss",
+    "RMSNorm",
+    "RNN",
+    "RNNBase",
+    "RNNCell",
+    "RNNCellBase",
+    "RReLU",
+    "ReLU",
+    "ReLU6",
+    "ReflectionPad1d",
+    "ReflectionPad2d",
+    "ReflectionPad3d",
+    "ReplicationPad1d",
+    "ReplicationPad2d",
+    "ReplicationPad3d",
+    "SELU",
+    "Sequential",
+    "SiLU",
+    "Sigmoid",
+    "SmoothL1Loss",
+    "SoftMarginLoss",
+    "Softmax",
+    "Softmax2d",
+    "Softmin",
+    "Softplus",
+    "Softshrink",
+    "Softsign",
+    "SyncBatchNorm",
+    "Tanh",
+    "Tanhshrink",
+    "Threshold",
+    "Transformer",
+    "TransformerDecoder",
+    "TransformerDecoderLayer",
+    "TransformerEncoder",
+    "TransformerEncoderLayer",
+    "TripletMarginLoss",
+    "TripletMarginWithDistanceLoss",
+    "Unflatten",
+    "Unfold",
+    "Upsample",
+    "UpsamplingBilinear2d",
+    "UpsamplingNearest2d",
+    "ZeroPad1d",
+    "ZeroPad2d",
+    "ZeroPad3d",
+]
+
+# Please keep this list sorted
+if __all__ != sorted(__all__):
+    raise AssertionError("__all__ must be sorted")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/_functions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..59352264b3216b85f8dc238b4da714ae2167ad3d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/_functions.py
@@ -0,0 +1,320 @@
+# mypy: allow-untyped-defs
+import torch
+import torch.distributed as dist
+from torch.autograd.function import Function
+
+
+class SyncBatchNorm(Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(
+        self,
+        input,
+        weight,
+        bias,
+        running_mean,
+        running_var,
+        eps,
+        momentum,
+        process_group,
+        world_size,
+    ):
+        if not (
+            input.is_contiguous(memory_format=torch.channels_last)
+            or input.is_contiguous(memory_format=torch.channels_last_3d)
+        ):
+            input = input.contiguous()
+        if weight is not None:
+            weight = weight.contiguous()
+
+        size = int(input.numel() // input.size(1))
+        if size == 1 and world_size < 2:
+            raise ValueError(
+                f"Expected more than 1 value per channel when training, got input size {size}"
+            )
+
+        num_channels = input.shape[1]
+        if input.numel() > 0:
+            # calculate mean/invstd for input.
+            mean, invstd = torch.batch_norm_stats(input, eps)
+
+            count = torch.full(
+                (1,),
+                input.numel() // input.size(1),
+                dtype=mean.dtype,
+                device=mean.device,
+            )
+
+            # C, C, 1 -> (2C + 1)
+            combined = torch.cat([mean, invstd, count], dim=0)
+        else:
+            # for empty input, set stats and the count to zero. The stats with
+            # zero count will be filtered out later when computing global mean
+            # & invstd, but they still needs to participate the all_gather
+            # collective communication to unblock other peer processes.
+            combined = torch.zeros(
+                2 * num_channels + 1, dtype=input.dtype, device=input.device
+            )
+
+        # Use allgather instead of allreduce because count could be different across
+        # ranks, simple all reduce op can not give correct results.
+        # batch_norm_gather_stats_with_counts calculates global mean & invstd based on
+        # all gathered mean, invstd and count.
+        # for nccl backend, use the optimized version of all gather.
+        # The Gloo backend does not support `all_gather_into_tensor`.
+        if process_group._get_backend_name() != "gloo":
+            # world_size * (2C + 1)
+            combined_size = combined.numel()
+            combined_flat = torch.empty(
+                1,
+                combined_size * world_size,
+                dtype=combined.dtype,
+                device=combined.device,
+            )
+            dist.all_gather_into_tensor(
+                combined_flat, combined, process_group, async_op=False
+            )
+            combined = torch.reshape(combined_flat, (world_size, combined_size))
+            # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
+            mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)
+        else:
+            # world_size * (2C + 1)
+            combined_list = [torch.empty_like(combined) for _ in range(world_size)]
+            dist.all_gather(combined_list, combined, process_group, async_op=False)
+            combined = torch.stack(combined_list, dim=0)
+            # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
+            mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)
+
+        if not (torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()):
+            # The lines below force a synchronization between CUDA and CPU, because
+            # the shape of the result count_all depends on the values in mask tensor.
+            # Such synchronizations break CUDA Graph capturing.
+            # See https://github.com/pytorch/pytorch/issues/78549
+            # FIXME: https://github.com/pytorch/pytorch/issues/78656 describes
+            # a better longer-term solution.
+
+            # remove stats from empty inputs
+            mask = count_all.squeeze(-1) >= 1
+            count_all = count_all[mask]
+            mean_all = mean_all[mask]
+            invstd_all = invstd_all[mask]
+
+        # calculate global mean & invstd
+        counts = count_all.view(-1)
+        if running_mean is not None and counts.dtype != running_mean.dtype:
+            counts = counts.to(running_mean.dtype)
+        mean, invstd = torch.batch_norm_gather_stats_with_counts(
+            input,
+            mean_all,
+            invstd_all,
+            running_mean,
+            running_var,
+            momentum,
+            eps,
+            counts,
+        )
+
+        self.save_for_backward(input, weight, mean, invstd, count_all.to(torch.int32))
+        self.process_group = process_group
+
+        # apply element-wise normalization
+        if input.numel() > 0:
+            return torch.batch_norm_elemt(input, weight, bias, mean, invstd, eps)
+        else:
+            return torch.empty_like(input)
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(self, grad_output):
+        if not (
+            grad_output.is_contiguous(memory_format=torch.channels_last)
+            or grad_output.is_contiguous(memory_format=torch.channels_last_3d)
+        ):
+            grad_output = grad_output.contiguous()
+        saved_input, weight, mean, invstd, count_tensor = self.saved_tensors
+        grad_input = grad_weight = grad_bias = None
+        process_group = self.process_group
+
+        if saved_input.numel() > 0:
+            # calculate local stats as well as grad_weight / grad_bias
+            (
+                sum_dy,
+                sum_dy_xmu,
+                grad_weight,
+                grad_bias,
+            ) = torch.batch_norm_backward_reduce(
+                grad_output,
+                saved_input,
+                mean,
+                invstd,
+                weight,
+                self.needs_input_grad[0],
+                self.needs_input_grad[1],
+                self.needs_input_grad[2],
+            )
+
+            if self.needs_input_grad[0]:
+                # synchronizing stats used to calculate input gradient.
+                num_channels = sum_dy.shape[0]
+                combined = torch.cat([sum_dy, sum_dy_xmu], dim=0)
+                torch.distributed.all_reduce(
+                    combined,
+                    torch.distributed.ReduceOp.SUM,
+                    process_group,
+                    async_op=False,
+                )
+                sum_dy, sum_dy_xmu = torch.split(combined, num_channels)
+
+                # backward pass for gradient calculation
+                if weight is not None and weight.dtype != mean.dtype:
+                    weight = weight.to(mean.dtype)
+                grad_input = torch.batch_norm_backward_elemt(
+                    grad_output,
+                    saved_input,
+                    mean,
+                    invstd,
+                    weight,
+                    sum_dy,
+                    sum_dy_xmu,
+                    count_tensor,
+                )
+            # synchronizing of grad_weight / grad_bias is not needed as distributed
+            # training would handle all reduce.
+            if weight is None or not self.needs_input_grad[1]:
+                grad_weight = None
+
+            if weight is None or not self.needs_input_grad[2]:
+                grad_bias = None
+        else:
+            # This process got an empty input tensor in the forward pass.
+            # Although this process can directly set grad_input as an empty
+            # tensor of zeros, it still needs to participate in the collective
+            # communication to unblock its peers, as other peer processes might
+            # have received non-empty inputs.
+            num_channels = saved_input.shape[1]
+            if self.needs_input_grad[0]:
+                # launch all_reduce to unblock other peer processes
+                combined = torch.zeros(
+                    2 * num_channels, dtype=saved_input.dtype, device=saved_input.device
+                )
+                torch.distributed.all_reduce(
+                    combined,
+                    torch.distributed.ReduceOp.SUM,
+                    process_group,
+                    async_op=False,
+                )
+
+            # Leave grad_input, grad_weight and grad_bias as None, which will be
+            # interpreted by the autograd engine as Tensors full of zeros.
+
+        return grad_input, grad_weight, grad_bias, None, None, None, None, None, None
+
+
+class CrossMapLRN2d(Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, input, size, alpha=1e-4, beta=0.75, k=1):
+        ctx.size = size
+        ctx.alpha = alpha
+        ctx.beta = beta
+        ctx.k = k
+        ctx.scale = None
+
+        if input.dim() != 4:
+            raise ValueError(
+                f"CrossMapLRN2d: Expected input to be 4D, got {input.dim()}D instead."
+            )
+
+        ctx.scale = ctx.scale or input.new()
+        output = input.new()
+        channels = input.size(1)
+
+        output.resize_as_(input)
+        ctx.scale.resize_as_(input)
+
+        # use output storage as temporary buffer
+        input_square = output
+        torch.pow(input, 2, out=input_square)
+
+        pre_pad = int((ctx.size - 1) / 2 + 1)
+        pre_pad_crop = min(pre_pad, channels)
+
+        scale_first = ctx.scale.select(1, 0)
+        scale_first.zero_()
+        # compute first feature map normalization
+        for c in range(pre_pad_crop):
+            scale_first.add_(input_square.select(1, c))
+
+        # reuse computations for next feature maps normalization
+        # by adding the next feature map and removing the previous
+        for c in range(1, channels):
+            scale_previous = ctx.scale.select(1, c - 1)
+            scale_current = ctx.scale.select(1, c)
+            scale_current.copy_(scale_previous)
+            if c < channels - pre_pad + 1:
+                square_next = input_square.select(1, c + pre_pad - 1)
+                scale_current.add_(square_next, alpha=1)
+
+            if c > pre_pad:
+                square_previous = input_square.select(1, c - pre_pad)
+                scale_current.add_(square_previous, alpha=-1)
+
+        ctx.scale.mul_(ctx.alpha / ctx.size).add_(ctx.k)
+
+        torch.pow(ctx.scale, -ctx.beta, out=output)
+        output.mul_(input)
+
+        ctx.save_for_backward(input, output)
+        return output
+
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def backward(ctx, grad_output):
+        input, output = ctx.saved_tensors
+        grad_input = grad_output.new()
+
+        batch_size = input.size(0)
+        channels = input.size(1)
+        input_height = input.size(2)
+        input_width = input.size(3)
+
+        paddded_ratio = input.new(channels + ctx.size - 1, input_height, input_width)
+        accum_ratio = input.new(input_height, input_width)
+
+        cache_ratio_value = 2 * ctx.alpha * ctx.beta / ctx.size
+        inversePrePad = int(ctx.size - (ctx.size - 1) / 2)
+
+        grad_input.resize_as_(input)
+        torch.pow(ctx.scale, -ctx.beta, out=grad_input).mul_(grad_output)
+
+        paddded_ratio.zero_()
+        padded_ratio_center = paddded_ratio.narrow(0, inversePrePad, channels)
+        for n in range(batch_size):
+            torch.mul(grad_output[n], output[n], out=padded_ratio_center)
+            padded_ratio_center.div_(ctx.scale[n])
+            torch.sum(
+                paddded_ratio.narrow(0, 0, ctx.size - 1),
+                0,
+                keepdim=False,
+                out=accum_ratio,
+            )
+            for c in range(channels):
+                accum_ratio.add_(paddded_ratio[c + ctx.size - 1])
+                grad_input[n][c].addcmul_(
+                    input[n][c], accum_ratio, value=-cache_ratio_value
+                )
+                accum_ratio.add_(paddded_ratio[c], alpha=-1)
+
+        return grad_input, None, None, None, None
+
+
+class BackwardHookFunction(torch.autograd.Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, *args):
+        ctx.mark_non_differentiable(*[arg for arg in args if not arg.requires_grad])
+        return args
+
+    @staticmethod
+    def backward(ctx, *args):
+        return args
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/activation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/activation.py
new file mode 100644
index 0000000000000000000000000000000000000000..739b2a95fc3c1ffbdaeb4c11820f0dc57574975d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/activation.py
@@ -0,0 +1,1912 @@
+# mypy: allow-untyped-defs
+import warnings
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
+from torch.nn.parameter import Parameter
+
+from .linear import NonDynamicallyQuantizableLinear
+from .module import Module
+
+
+__all__ = [
+    "Threshold",
+    "ReLU",
+    "RReLU",
+    "Hardtanh",
+    "ReLU6",
+    "Sigmoid",
+    "Hardsigmoid",
+    "Tanh",
+    "SiLU",
+    "Mish",
+    "Hardswish",
+    "ELU",
+    "CELU",
+    "SELU",
+    "GLU",
+    "GELU",
+    "Hardshrink",
+    "LeakyReLU",
+    "LogSigmoid",
+    "Softplus",
+    "Softshrink",
+    "MultiheadAttention",
+    "PReLU",
+    "Softsign",
+    "Tanhshrink",
+    "Softmin",
+    "Softmax",
+    "Softmax2d",
+    "LogSoftmax",
+]
+
+
+class Threshold(Module):
+    r"""Thresholds each element of the input Tensor.
+
+    Threshold is defined as:
+
+    .. math::
+        y =
+        \begin{cases}
+        x, &\text{ if } x > \text{threshold} \\
+        \text{value}, &\text{ otherwise }
+        \end{cases}
+
+    Args:
+        threshold: The value to threshold at
+        value: The value to replace with
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Threshold.png
+
+    Examples::
+
+        >>> m = nn.Threshold(0, 0.5)
+        >>> input = torch.arange(-3, 3)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["threshold", "value", "inplace"]
+
+    threshold: float
+    value: float
+    inplace: bool
+
+    def __init__(self, threshold: float, value: float, inplace: bool = False) -> None:
+        super().__init__()
+        self.threshold = threshold
+        self.value = value
+        self.inplace = inplace
+        # TODO: check in THNN (if inplace == True, then assert value <= threshold)
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.threshold(input, self.threshold, self.value, self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = ", inplace=True" if self.inplace else ""
+        return f"threshold={self.threshold}, value={self.value}{inplace_str}"
+
+
+class ReLU(Module):
+    r"""Applies the rectified linear unit function element-wise.
+
+    :math:`\text{ReLU}(x) = (x)^+ = \max(0, x)`
+
+    Args:
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/ReLU.png
+
+    Examples::
+
+        >>> m = nn.ReLU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+
+
+      An implementation of CReLU - https://arxiv.org/abs/1603.05201
+
+        >>> m = nn.ReLU()
+        >>> input = torch.randn(2).unsqueeze(0)
+        >>> output = torch.cat((m(input), m(-input)))
+    """
+
+    __constants__ = ["inplace"]
+    inplace: bool
+
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__()
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.relu(input, inplace=self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = "inplace=True" if self.inplace else ""
+        return inplace_str
+
+
+class RReLU(Module):
+    r"""Applies the randomized leaky rectified linear unit function, element-wise.
+
+    Method described in the paper:
+    `Empirical Evaluation of Rectified Activations in Convolutional Network <https://arxiv.org/abs/1505.00853>`_.
+
+    The function is defined as:
+
+    .. math::
+        \text{RReLU}(x) =
+        \begin{cases}
+            x & \text{if } x \geq 0 \\
+            ax & \text{ otherwise }
+        \end{cases}
+
+    where :math:`a` is randomly sampled from uniform distribution
+    :math:`\mathcal{U}(\text{lower}, \text{upper})` during training while during
+    evaluation :math:`a` is fixed with :math:`a = \frac{\text{lower} + \text{upper}}{2}`.
+
+    Args:
+        lower: lower bound of the uniform distribution. Default: :math:`\frac{1}{8}`
+        upper: upper bound of the uniform distribution. Default: :math:`\frac{1}{3}`
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/RReLU.png
+
+    Examples::
+
+        >>> m = nn.RReLU(0.1, 0.3)
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+
+    """
+
+    __constants__ = ["lower", "upper", "inplace"]
+
+    lower: float
+    upper: float
+    inplace: bool
+
+    def __init__(
+        self, lower: float = 1.0 / 8, upper: float = 1.0 / 3, inplace: bool = False
+    ) -> None:
+        super().__init__()
+        self.lower = lower
+        self.upper = upper
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.rrelu(input, self.lower, self.upper, self.training, self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = ", inplace=True" if self.inplace else ""
+        return f"lower={self.lower}, upper={self.upper}{inplace_str}"
+
+
+class Hardtanh(Module):
+    r"""Applies the HardTanh function element-wise.
+
+    HardTanh is defined as:
+
+    .. math::
+        \text{HardTanh}(x) = \begin{cases}
+            \text{max\_val} & \text{ if } x > \text{ max\_val } \\
+            \text{min\_val} & \text{ if } x < \text{ min\_val } \\
+            x & \text{ otherwise } \\
+        \end{cases}
+
+    Args:
+        min_val: minimum value of the linear region range. Default: -1
+        max_val: maximum value of the linear region range. Default: 1
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Keyword arguments :attr:`min_value` and :attr:`max_value`
+    have been deprecated in favor of :attr:`min_val` and :attr:`max_val`.
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Hardtanh.png
+
+    Examples::
+
+        >>> m = nn.Hardtanh(-2, 2)
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["min_val", "max_val", "inplace"]
+
+    min_val: float
+    max_val: float
+    inplace: bool
+
+    def __init__(
+        self,
+        min_val: float = -1.0,
+        max_val: float = 1.0,
+        inplace: bool = False,
+        min_value: float | None = None,
+        max_value: float | None = None,
+    ) -> None:
+        super().__init__()
+        if min_value is not None:
+            warnings.warn(
+                "keyword argument `min_value` is deprecated and rename to `min_val`",
+                FutureWarning,
+                stacklevel=2,
+            )
+            min_val = min_value
+        if max_value is not None:
+            warnings.warn(
+                "keyword argument `max_value` is deprecated and rename to `max_val`",
+                FutureWarning,
+                stacklevel=2,
+            )
+            max_val = max_value
+
+        self.min_val = min_val
+        self.max_val = max_val
+        self.inplace = inplace
+        if self.max_val <= self.min_val:
+            raise AssertionError(
+                f"max_val ({self.max_val}) must be greater than min_val ({self.min_val})"
+            )
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.hardtanh(input, self.min_val, self.max_val, self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = ", inplace=True" if self.inplace else ""
+        return f"min_val={self.min_val}, max_val={self.max_val}{inplace_str}"
+
+
+class ReLU6(Hardtanh):
+    r"""Applies the ReLU6 function element-wise.
+
+    .. math::
+        \text{ReLU6}(x) = \min(\max(0,x), 6)
+
+    Args:
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/ReLU6.png
+
+    Examples::
+
+        >>> m = nn.ReLU6()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__(0.0, 6.0, inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = "inplace=True" if self.inplace else ""
+        return inplace_str
+
+
+class Sigmoid(Module):
+    r"""Applies the Sigmoid function element-wise.
+
+    .. math::
+        \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
+
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Sigmoid.png
+
+    Examples::
+
+        >>> m = nn.Sigmoid()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return torch.sigmoid(input)
+
+
+class Hardsigmoid(Module):
+    r"""Applies the Hardsigmoid function element-wise.
+
+    Hardsigmoid is defined as:
+
+    .. math::
+        \text{Hardsigmoid}(x) = \begin{cases}
+            0 & \text{if~} x \le -3, \\
+            1 & \text{if~} x \ge +3, \\
+            x / 6 + 1 / 2 & \text{otherwise}
+        \end{cases}
+
+    Args:
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Hardsigmoid.png
+
+    Examples::
+
+        >>> m = nn.Hardsigmoid()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["inplace"]
+
+    inplace: bool
+
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__()
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.hardsigmoid(input, self.inplace)
+
+
+class Tanh(Module):
+    r"""Applies the Hyperbolic Tangent (Tanh) function element-wise.
+
+    Tanh is defined as:
+
+    .. math::
+        \text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)} {\exp(x) + \exp(-x)}
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Tanh.png
+
+    Examples::
+
+        >>> m = nn.Tanh()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return torch.tanh(input)
+
+
+class SiLU(Module):
+    r"""Applies the Sigmoid Linear Unit (SiLU) function, element-wise.
+
+    The SiLU function is also known as the swish function.
+
+    .. math::
+        \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}
+
+    .. note::
+        See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_
+        where the SiLU (Sigmoid Linear Unit) was originally coined, and see
+        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation
+        in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:
+        a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_
+        where the SiLU was experimented with later.
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/SiLU.png
+
+    Examples::
+
+        >>> m = nn.SiLU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["inplace"]
+    inplace: bool
+
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__()
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.silu(input, inplace=self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = "inplace=True" if self.inplace else ""
+        return inplace_str
+
+
+class Mish(Module):
+    r"""Applies the Mish function, element-wise.
+
+    Mish: A Self Regularized Non-Monotonic Neural Activation Function.
+
+    .. math::
+        \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
+
+    .. note::
+        See `Mish: A Self Regularized Non-Monotonic Neural Activation Function <https://arxiv.org/abs/1908.08681>`_
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Mish.png
+
+    Examples::
+
+        >>> m = nn.Mish()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["inplace"]
+    inplace: bool
+
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__()
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.mish(input, inplace=self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = "inplace=True" if self.inplace else ""
+        return inplace_str
+
+
+class Hardswish(Module):
+    r"""Applies the Hardswish function, element-wise.
+
+    Method described in the paper: `Searching for MobileNetV3 <https://arxiv.org/abs/1905.02244>`_.
+
+    Hardswish is defined as:
+
+    .. math::
+        \text{Hardswish}(x) = \begin{cases}
+            0 & \text{if~} x \le -3, \\
+            x & \text{if~} x \ge +3, \\
+            x \cdot (x + 3) /6 & \text{otherwise}
+        \end{cases}
+
+    Args:
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Hardswish.png
+
+    Examples::
+
+        >>> m = nn.Hardswish()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["inplace"]
+
+    inplace: bool
+
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__()
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.hardswish(input, self.inplace)
+
+
+class ELU(Module):
+    r"""Applies the Exponential Linear Unit (ELU) function, element-wise.
+
+    Method described in the paper: `Fast and Accurate Deep Network Learning by Exponential Linear
+    Units (ELUs) <https://arxiv.org/abs/1511.07289>`__.
+
+    ELU is defined as:
+
+    .. math::
+        \text{ELU}(x) = \begin{cases}
+        x, & \text{ if } x > 0\\
+        \alpha * (\exp(x) - 1), & \text{ if } x \leq 0
+        \end{cases}
+
+    Args:
+        alpha: the :math:`\alpha` value for the ELU formulation. Default: 1.0
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/ELU.png
+
+    Examples::
+
+        >>> m = nn.ELU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["alpha", "inplace"]
+    alpha: float
+    inplace: bool
+
+    def __init__(self, alpha: float = 1.0, inplace: bool = False) -> None:
+        super().__init__()
+        self.alpha = alpha
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.elu(input, self.alpha, self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = ", inplace=True" if self.inplace else ""
+        return f"alpha={self.alpha}{inplace_str}"
+
+
+class CELU(Module):
+    r"""Applies the CELU function element-wise.
+
+    .. math::
+        \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))
+
+    More details can be found in the paper `Continuously Differentiable Exponential Linear Units`_ .
+
+    Args:
+        alpha: the :math:`\alpha` value for the CELU formulation. Default: 1.0
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/CELU.png
+
+    Examples::
+
+        >>> m = nn.CELU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+
+    .. _`Continuously Differentiable Exponential Linear Units`:
+        https://arxiv.org/abs/1704.07483
+    """
+
+    __constants__ = ["alpha", "inplace"]
+    alpha: float
+    inplace: bool
+
+    def __init__(self, alpha: float = 1.0, inplace: bool = False) -> None:
+        super().__init__()
+        self.alpha = alpha
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.celu(input, self.alpha, self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = ", inplace=True" if self.inplace else ""
+        return f"alpha={self.alpha}{inplace_str}"
+
+
+class SELU(Module):
+    r"""Applies the SELU function element-wise.
+
+    .. math::
+        \text{SELU}(x) = \text{scale} * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))
+
+    with :math:`\alpha = 1.6732632423543772848170429916717` and
+    :math:`\text{scale} = 1.0507009873554804934193349852946`.
+
+    .. warning::
+        When using ``kaiming_normal`` or ``kaiming_normal_`` for initialisation,
+        ``nonlinearity='linear'`` should be used instead of ``nonlinearity='selu'``
+        in order to get `Self-Normalizing Neural Networks`_.
+        See :func:`torch.nn.init.calculate_gain` for more information.
+
+    More details can be found in the paper `Self-Normalizing Neural Networks`_ .
+
+    Args:
+        inplace (bool, optional): can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/SELU.png
+
+    Examples::
+
+        >>> m = nn.SELU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+
+    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
+    """
+
+    __constants__ = ["inplace"]
+    inplace: bool
+
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__()
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.selu(input, self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = "inplace=True" if self.inplace else ""
+        return inplace_str
+
+
+class GLU(Module):
+    r"""Applies the gated linear unit function.
+
+    :math:`{GLU}(a, b)= a \otimes \sigma(b)` where :math:`a` is the first half
+    of the input matrices and :math:`b` is the second half.
+
+    Args:
+        dim (int): the dimension on which to split the input. Default: -1
+
+    Shape:
+        - Input: :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional
+          dimensions
+        - Output: :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2`
+
+    Examples::
+
+        >>> m = nn.GLU()
+        >>> input = torch.randn(4, 2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["dim"]
+    dim: int
+
+    def __init__(self, dim: int = -1) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.glu(input, self.dim)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"dim={self.dim}"
+
+
+class GELU(Module):
+    r"""Applies the Gaussian Error Linear Units function.
+
+    .. math:: \text{GELU}(x) = x * \Phi(x)
+
+    where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.
+
+    When the approximate argument is 'tanh', Gelu is estimated with:
+
+    .. math:: \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt{2 / \pi} * (x + 0.044715 * x^3)))
+
+    Args:
+        approximate (str, optional): the gelu approximation algorithm to use:
+            ``'none'`` | ``'tanh'``. Default: ``'none'``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/GELU.png
+
+    Examples::
+
+        >>> m = nn.GELU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["approximate"]
+    approximate: str
+
+    def __init__(self, approximate: str = "none") -> None:
+        super().__init__()
+        self.approximate = approximate
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.gelu(input, approximate=self.approximate)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"approximate={repr(self.approximate)}"
+
+
+class Hardshrink(Module):
+    r"""Applies the Hard Shrinkage (Hardshrink) function element-wise.
+
+    Hardshrink is defined as:
+
+    .. math::
+        \text{HardShrink}(x) =
+        \begin{cases}
+        x, & \text{ if } x > \lambda \\
+        x, & \text{ if } x < -\lambda \\
+        0, & \text{ otherwise }
+        \end{cases}
+
+    Args:
+        lambd: the :math:`\lambda` value for the Hardshrink formulation. Default: 0.5
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Hardshrink.png
+
+    Examples::
+
+        >>> m = nn.Hardshrink()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["lambd"]
+    lambd: float
+
+    def __init__(self, lambd: float = 0.5) -> None:
+        super().__init__()
+        self.lambd = lambd
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Run forward pass.
+        """
+        return F.hardshrink(input, self.lambd)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"{self.lambd}"
+
+
+class LeakyReLU(Module):
+    r"""Applies the LeakyReLU function element-wise.
+
+    .. math::
+        \text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)
+
+
+    or
+
+    .. math::
+        \text{LeakyReLU}(x) =
+        \begin{cases}
+        x, & \text{ if } x \geq 0 \\
+        \text{negative\_slope} \times x, & \text{ otherwise }
+        \end{cases}
+
+    Args:
+        negative_slope: Controls the angle of the negative slope (which is used for
+          negative input values). Default: 1e-2
+        inplace: can optionally do the operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)` where `*` means, any number of additional
+          dimensions
+        - Output: :math:`(*)`, same shape as the input
+
+    .. image:: ../scripts/activation_images/LeakyReLU.png
+
+    Examples::
+
+        >>> m = nn.LeakyReLU(0.1)
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["inplace", "negative_slope"]
+    inplace: bool
+    negative_slope: float
+
+    def __init__(self, negative_slope: float = 1e-2, inplace: bool = False) -> None:
+        super().__init__()
+        self.negative_slope = negative_slope
+        self.inplace = inplace
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Run forward pass.
+        """
+        return F.leaky_relu(input, self.negative_slope, self.inplace)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        inplace_str = ", inplace=True" if self.inplace else ""
+        return f"negative_slope={self.negative_slope}{inplace_str}"
+
+
+class LogSigmoid(Module):
+    r"""Applies the Logsigmoid function element-wise.
+
+    .. math::
+        \text{LogSigmoid}(x) = \log\left(\frac{ 1 }{ 1 + \exp(-x)}\right)
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/LogSigmoid.png
+
+    Examples::
+
+        >>> m = nn.LogSigmoid()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Run forward pass.
+        """
+        return F.logsigmoid(input)
+
+
+class Softplus(Module):
+    r"""Applies the Softplus function element-wise.
+
+    .. math::
+        \text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))
+
+    SoftPlus is a smooth approximation to the ReLU function and can be used
+    to constrain the output of a machine to always be positive.
+
+    For numerical stability the implementation reverts to the linear function
+    when :math:`input \times \beta > threshold`.
+
+    Args:
+        beta: the :math:`\beta` value for the Softplus formulation. Default: 1
+        threshold: values above this revert to a linear function. Default: 20
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Softplus.png
+
+    Examples::
+
+        >>> m = nn.Softplus()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["beta", "threshold"]
+    beta: float
+    threshold: float
+
+    def __init__(self, beta: float = 1.0, threshold: float = 20.0) -> None:
+        super().__init__()
+        self.beta = beta
+        self.threshold = threshold
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Run forward pass.
+        """
+        return F.softplus(input, self.beta, self.threshold)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"beta={self.beta}, threshold={self.threshold}"
+
+
+class Softshrink(Module):
+    r"""Applies the soft shrinkage function element-wise.
+
+    .. math::
+        \text{SoftShrinkage}(x) =
+        \begin{cases}
+        x - \lambda, & \text{ if } x > \lambda \\
+        x + \lambda, & \text{ if } x < -\lambda \\
+        0, & \text{ otherwise }
+        \end{cases}
+
+    Args:
+        lambd: the :math:`\lambda` (must be no less than zero) value for the Softshrink formulation. Default: 0.5
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Softshrink.png
+
+    Examples::
+
+        >>> m = nn.Softshrink()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["lambd"]
+    lambd: float
+
+    def __init__(self, lambd: float = 0.5) -> None:
+        super().__init__()
+        self.lambd = lambd
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Run forward pass.
+        """
+        return F.softshrink(input, self.lambd)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return str(self.lambd)
+
+
+def _check_arg_device(x: torch.Tensor | None) -> bool:
+    if x is not None:
+        return x.device.type in [
+            "cpu",
+            "cuda",
+            torch.utils.backend_registration._privateuse1_backend_name,
+        ]
+    return True
+
+
+def _arg_requires_grad(x: torch.Tensor | None) -> bool:
+    if x is not None:
+        return x.requires_grad
+    return False
+
+
+def _is_make_fx_tracing():
+    if not torch.jit.is_scripting():
+        torch_dispatch_mode_stack = (
+            torch.utils._python_dispatch._get_current_dispatch_mode_stack()
+        )
+        # this can be triggered when dynamo inlining the module too.
+        return (
+            any(
+                type(x) is torch.fx.experimental.proxy_tensor.ProxyTorchDispatchMode
+                for x in torch_dispatch_mode_stack
+            )
+            or torch.compiler.is_exporting()
+        )
+    else:
+        return False
+
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information from different representation subspaces.
+
+    This MultiheadAttention layer implements the original architecture described
+    in the `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_ paper. The
+    intent of this layer is as a reference implementation for foundational understanding
+    and thus it contains only limited features relative to newer architectures.
+    Given the fast pace of innovation in transformer-like architectures, we recommend
+    exploring this `tutorial <https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html>`_
+    to build efficient layers from building blocks in core or using higher
+    level libraries from the `PyTorch Ecosystem <https://landscape.pytorch.org/>`_.
+
+    Multi-Head Attention is defined as:
+
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1,\dots,\text{head}_h)W^O
+
+    where :math:`\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+
+    ``nn.MultiheadAttention`` will use the optimized implementations of
+    ``scaled_dot_product_attention()`` when possible.
+
+    In addition to support for the new ``scaled_dot_product_attention()``
+    function, for speeding up Inference, MHA will use
+    fastpath inference with support for Nested Tensors, iff:
+
+    - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor).
+    - inputs are batched (3D) with ``batch_first==True``
+    - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``
+    - training is disabled (using ``.eval()``)
+    - ``add_bias_kv`` is ``False``
+    - ``add_zero_attn`` is ``False``
+    - ``kdim`` and ``vdim`` are equal to ``embed_dim``
+    - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``
+      nor ``attn_mask`` is passed
+    - autocast is disabled
+
+    If the optimized inference fastpath implementation is in use, a
+    `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for
+    ``query``/``key``/``value`` to represent padding more efficiently than using a
+    padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_
+    will be returned, and an additional speedup proportional to the fraction of the input
+    that is padding can be expected.
+
+    Args:
+        embed_dim: Total dimension of the model.
+        num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split
+            across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).
+        dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).
+        bias: If specified, adds bias to input / output projection layers. Default: ``True``.
+        add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.
+        add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.
+            Default: ``False``.
+        kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
+        vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+
+    Examples::
+
+        >>> # xdoctest: +SKIP
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+
+    .. _`FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness`:
+         https://arxiv.org/abs/2205.14135
+
+    """
+
+    __constants__ = ["batch_first"]
+    bias_k: torch.Tensor | None
+    bias_v: torch.Tensor | None
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        kdim=None,
+        vdim=None,
+        batch_first=False,
+        device=None,
+        dtype=None,
+    ) -> None:
+        if embed_dim <= 0 or num_heads <= 0:
+            raise ValueError(
+                f"embed_dim and num_heads must be greater than 0,"
+                f" got embed_dim={embed_dim} and num_heads={num_heads} instead"
+            )
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.batch_first = batch_first
+        self.head_dim = embed_dim // num_heads
+        if self.head_dim * num_heads != self.embed_dim:
+            raise AssertionError("embed_dim must be divisible by num_heads")
+
+        if not self._qkv_same_embed_dim:
+            self.q_proj_weight = Parameter(
+                torch.empty((embed_dim, embed_dim), **factory_kwargs)
+            )
+            self.k_proj_weight = Parameter(
+                torch.empty((embed_dim, self.kdim), **factory_kwargs)
+            )
+            self.v_proj_weight = Parameter(
+                torch.empty((embed_dim, self.vdim), **factory_kwargs)
+            )
+            self.register_parameter("in_proj_weight", None)
+        else:
+            self.in_proj_weight = Parameter(
+                torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)
+            )
+            self.register_parameter("q_proj_weight", None)
+            self.register_parameter("k_proj_weight", None)
+            self.register_parameter("v_proj_weight", None)
+
+        if bias:
+            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim, **factory_kwargs))
+        else:
+            self.register_parameter("in_proj_bias", None)
+        self.out_proj = NonDynamicallyQuantizableLinear(
+            embed_dim, embed_dim, bias=bias, **factory_kwargs
+        )
+
+        if add_bias_kv:
+            self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
+            self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        if self._qkv_same_embed_dim:
+            xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.0)
+            constant_(self.out_proj.bias, 0.0)
+        if self.bias_k is not None:
+            xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            xavier_normal_(self.bias_v)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if "_qkv_same_embed_dim" not in state:
+            state["_qkv_same_embed_dim"] = True
+
+        super().__setstate__(state)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        key_padding_mask: Tensor | None = None,
+        need_weights: bool = True,
+        attn_mask: Tensor | None = None,
+        average_attn_weights: bool = True,
+        is_causal: bool = False,
+    ) -> tuple[Tensor, Tensor | None]:
+        r"""Compute attention outputs using query, key, and value embeddings.
+
+            Supports optional parameters for padding, masks and attention weights.
+
+        Args:
+            query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``
+                or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,
+                :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.
+                Queries are compared against key-value pairs to produce the output.
+                See "Attention Is All You Need" for more details.
+            key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``
+                or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,
+                :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.
+                See "Attention Is All You Need" for more details.
+            value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when
+                ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source
+                sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.
+                See "Attention Is All You Need" for more details.
+            key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``
+                to ignore for the purpose of attention (i.e. treat as "padding"). For unbatched `query`, shape should be :math:`(S)`.
+                Binary and float masks are supported.
+                For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for
+                the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.
+            need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.
+                Set ``need_weights=False`` to use the optimized ``scaled_dot_product_attention``
+                and achieve the best performance for MHA.
+                Default: ``True``.
+            attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
+                :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
+                :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be
+                broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
+                Binary and float masks are supported. For a binary mask, a ``True`` value indicates that the
+                corresponding position is not allowed to attend. For a float mask, the mask values will be added to
+                the attention weight.
+                If both attn_mask and key_padding_mask are supplied, their types should match.
+            average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
+                heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
+                effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)
+            is_causal: If specified, applies a causal mask as attention mask.
+                Default: ``False``.
+                Warning:
+                ``is_causal`` provides a hint that ``attn_mask`` is the
+                causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+
+        Outputs:
+            - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,
+              :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,
+              where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the
+              embedding dimension ``embed_dim``.
+            - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,
+              returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+              :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+              :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+              head of shape :math:`(\text{num\_heads}, L, S)` when input is unbatched or :math:`(N, \text{num\_heads}, L, S)`.
+
+            .. note::
+                `batch_first` argument is ignored for unbatched inputs.
+        """  # noqa: B950
+        why_not_fast_path = ""
+        if (
+            (attn_mask is not None and torch.is_floating_point(attn_mask))
+            or (key_padding_mask is not None)
+            and torch.is_floating_point(key_padding_mask)
+        ):
+            why_not_fast_path = "floating-point masks are not supported for fast path."
+
+        is_batched = query.dim() == 3
+
+        key_padding_mask = F._canonical_mask(
+            mask=key_padding_mask,
+            mask_name="key_padding_mask",
+            other_type=F._none_or_dtype(attn_mask),
+            other_name="attn_mask",
+            target_type=query.dtype,
+        )
+
+        attn_mask = F._canonical_mask(
+            mask=attn_mask,
+            mask_name="attn_mask",
+            other_type=None,
+            other_name="",
+            target_type=query.dtype,
+            check_other=False,
+        )
+
+        is_fastpath_enabled = torch.backends.mha.get_fastpath_enabled()
+
+        if not is_fastpath_enabled:
+            why_not_fast_path = "torch.backends.mha.get_fastpath_enabled() was not True"
+        elif not is_batched:
+            why_not_fast_path = (
+                f"input not batched; expected query.dim() of 3 but got {query.dim()}"
+            )
+        elif query is not key or key is not value:
+            # When lifting this restriction, don't forget to either
+            # enforce that the dtypes all match or test cases where
+            # they don't!
+            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
+        elif self.in_proj_bias is not None and query.dtype != self.in_proj_bias.dtype:
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
+        elif self.in_proj_weight is None:
+            why_not_fast_path = "in_proj_weight was None"
+        elif query.dtype != self.in_proj_weight.dtype:
+            # this case will fail anyway, but at least they'll get a useful error message.
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
+        elif self.training:
+            why_not_fast_path = "training is enabled"
+        elif (self.num_heads % 2) != 0:
+            why_not_fast_path = "self.num_heads is not even"
+        elif not self.batch_first:
+            why_not_fast_path = "batch_first was not True"
+        elif self.bias_k is not None:
+            why_not_fast_path = "self.bias_k was not None"
+        elif self.bias_v is not None:
+            why_not_fast_path = "self.bias_v was not None"
+        elif self.add_zero_attn:
+            why_not_fast_path = "add_zero_attn was enabled"
+        elif not self._qkv_same_embed_dim:
+            why_not_fast_path = "_qkv_same_embed_dim was not True"
+        elif query.is_nested and (
+            key_padding_mask is not None or attn_mask is not None
+        ):
+            why_not_fast_path = (
+                "supplying both src_key_padding_mask and src_mask at the same time \
+                                 is not supported with NestedTensor input"
+            )
+        elif torch.is_autocast_enabled():
+            why_not_fast_path = "autocast is enabled"
+
+        fast_path_blocked_by_tracing = False
+        if not why_not_fast_path:
+            tensor_args = (
+                query,
+                key,
+                value,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.out_proj.weight,
+                self.out_proj.bias,
+            )
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_fast_path = "some Tensor argument has_torch_function"
+            elif not all(_check_arg_device(x) for x in tensor_args):
+                why_not_fast_path = (
+                    "some Tensor argument's device is neither one of "
+                    f"cpu, cuda or {torch.utils.backend_registration._privateuse1_backend_name}"
+                )
+            elif torch.is_grad_enabled() and any(
+                _arg_requires_grad(x) for x in tensor_args
+            ):
+                why_not_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+            elif _is_make_fx_tracing():
+                why_not_fast_path = "we are running make_fx tracing"
+                fast_path_blocked_by_tracing = True
+            if not why_not_fast_path:
+                merged_mask, mask_type = self.merge_masks(
+                    attn_mask, key_padding_mask, query
+                )
+
+                if self.in_proj_bias is not None and self.in_proj_weight is not None:
+                    return torch._native_multi_head_attention(
+                        query,
+                        key,
+                        value,
+                        self.embed_dim,
+                        self.num_heads,
+                        self.in_proj_weight,
+                        self.in_proj_bias,
+                        self.out_proj.weight,
+                        self.out_proj.bias,
+                        merged_mask,
+                        need_weights,
+                        average_attn_weights,
+                        mask_type,
+                    )
+
+        any_nested = query.is_nested or key.is_nested or value.is_nested
+        if any_nested:
+            raise AssertionError(
+                "MultiheadAttention does not support NestedTensor outside of its fast path. "
+                + f"The fast path was not hit because {why_not_fast_path}"
+            )
+
+        if self.batch_first and is_batched:
+            # make sure that the transpose op does not affect the "is" property
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = (x.transpose(1, 0) for x in (query, key))
+                    value = key
+            else:
+                query, key, value = (x.transpose(1, 0) for x in (query, key, value))
+
+        if not self._qkv_same_embed_dim:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight,
+                k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight,
+                average_attn_weights=average_attn_weights,
+                is_causal=is_causal,
+            )
+        else:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                average_attn_weights=average_attn_weights,
+                is_causal=is_causal,
+            )
+        if self.batch_first and is_batched:
+            attn_output = attn_output.transpose(1, 0)
+            if fast_path_blocked_by_tracing:
+                # Keep the traced slowpath layout aligned with eager fastpath.
+                attn_output = attn_output.contiguous()
+            return attn_output, attn_output_weights
+        else:
+            return attn_output, attn_output_weights
+
+    def merge_masks(
+        self,
+        attn_mask: Tensor | None,
+        key_padding_mask: Tensor | None,
+        query: Tensor,
+    ) -> tuple[Tensor | None, int | None]:
+        r"""Determine mask type and combine masks if necessary.
+
+        If only one mask is provided, that mask
+        and the corresponding mask type will be returned. If both masks are provided, they will be both
+        expanded to shape ``(batch_size, num_heads, seq_len, seq_len)``, combined with logical ``or``
+        and mask type 2 will be returned
+        Args:
+            attn_mask: attention mask of shape ``(seq_len, seq_len)``, mask type 0
+            key_padding_mask: padding mask of shape ``(batch_size, seq_len)``, mask type 1
+            query: query embeddings of shape ``(batch_size, seq_len, embed_dim)``
+        Returns:
+            merged_mask: merged mask
+            mask_type: merged mask type (0, 1, or 2)
+        """
+        mask_type: int | None = None
+        merged_mask: Tensor | None = None
+
+        if key_padding_mask is not None:
+            mask_type = 1
+            merged_mask = key_padding_mask
+
+        if attn_mask is not None:
+            # In this branch query can't be a nested tensor, so it has a shape
+            batch_size, seq_len, _ = query.shape
+            mask_type = 2
+
+            # Always expands attn_mask to 4D
+            if attn_mask.dim() == 3:
+                attn_mask_expanded = attn_mask.view(batch_size, -1, seq_len, seq_len)
+            else:  # attn_mask.dim() == 2:
+                attn_mask_expanded = attn_mask.view(1, 1, seq_len, seq_len).expand(
+                    batch_size, self.num_heads, -1, -1
+                )
+            merged_mask = attn_mask_expanded
+
+            if key_padding_mask is not None:
+                key_padding_mask_expanded = key_padding_mask.view(
+                    batch_size, 1, 1, seq_len
+                ).expand(-1, self.num_heads, -1, -1)
+                merged_mask = attn_mask_expanded + key_padding_mask_expanded
+
+        # no attn_mask and no key_padding_mask, returns None, None
+        return merged_mask, mask_type
+
+
+class PReLU(Module):
+    r"""Applies the element-wise PReLU function.
+
+    .. math::
+        \text{PReLU}(x) = \max(0,x) + a * \min(0,x)
+
+    or
+
+    .. math::
+        \text{PReLU}(x) =
+        \begin{cases}
+        x, & \text{ if } x \ge 0 \\
+        ax, & \text{ otherwise }
+        \end{cases}
+
+    Here :math:`a` is a learnable parameter. When called without arguments, `nn.PReLU()` uses a single
+    parameter :math:`a` across all input channels. If called with `nn.PReLU(nChannels)`,
+    a separate :math:`a` is used for each input channel.
+
+
+    .. note::
+        weight decay should not be used when learning :math:`a` for good performance.
+
+    .. note::
+        Channel dim is the 2nd dim of input. When input has dims < 2, then there is
+        no channel dim and the number of channels = 1.
+
+    Args:
+        num_parameters (int): number of :math:`a` to learn.
+            Although it takes an int as input, there is only two values are legitimate:
+            1, or the number of channels at input. Default: 1
+        init (float): the initial value of :math:`a`. Default: 0.25
+
+    Shape:
+        - Input: :math:`( *)` where `*` means, any number of additional
+          dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    Attributes:
+        weight (Tensor): the learnable weights of shape (:attr:`num_parameters`).
+
+    .. image:: ../scripts/activation_images/PReLU.png
+
+    Examples::
+
+        >>> m = nn.PReLU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["num_parameters"]
+    num_parameters: int
+
+    def __init__(
+        self, num_parameters: int = 1, init: float = 0.25, device=None, dtype=None
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.num_parameters = num_parameters
+        super().__init__()
+        self.init = init
+        self.weight = Parameter(torch.empty(num_parameters, **factory_kwargs))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in ``__init__``.
+        """
+        torch.nn.init.constant_(self.weight, self.init)
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.prelu(input, self.weight)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"num_parameters={self.num_parameters}"
+
+
+class Softsign(Module):
+    r"""Applies the element-wise Softsign function.
+
+    .. math::
+        \text{SoftSign}(x) = \frac{x}{ 1 + |x|}
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Softsign.png
+
+    Examples::
+
+        >>> m = nn.Softsign()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.softsign(input)
+
+
+class Tanhshrink(Module):
+    r"""Applies the element-wise Tanhshrink function.
+
+    .. math::
+        \text{Tanhshrink}(x) = x - \tanh(x)
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    .. image:: ../scripts/activation_images/Tanhshrink.png
+
+    Examples::
+
+        >>> m = nn.Tanhshrink()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.tanhshrink(input)
+
+
+class Softmin(Module):
+    r"""Applies the Softmin function to an n-dimensional input Tensor.
+
+    Rescales them so that the elements of the n-dimensional output Tensor
+    lie in the range `[0, 1]` and sum to 1.
+
+    Softmin is defined as:
+
+    .. math::
+        \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)}
+
+    Shape:
+        - Input: :math:`(*)` where `*` means, any number of additional
+          dimensions
+        - Output: :math:`(*)`, same shape as the input
+
+    Args:
+        dim (int): A dimension along which Softmin will be computed (so every slice
+            along dim will sum to 1).
+
+    Returns:
+        a Tensor of the same dimension and shape as the input, with
+        values in the range [0, 1]
+
+    Examples::
+
+        >>> m = nn.Softmin(dim=1)
+        >>> input = torch.randn(2, 3)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["dim"]
+    dim: int | None
+
+    def __init__(self, dim: int | None = None) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        if not hasattr(self, "dim"):
+            self.dim = None
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.softmin(input, self.dim, _stacklevel=5)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"dim={self.dim}"
+
+
+class Softmax(Module):
+    r"""Applies the Softmax function to an n-dimensional input Tensor.
+
+    Rescales them so that the elements of the n-dimensional output Tensor
+    lie in the range [0,1] and sum to 1.
+
+    Softmax is defined as:
+
+    .. math::
+        \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
+
+    When the input Tensor is a sparse tensor then the unspecified
+    values are treated as ``-inf``.
+
+    Shape:
+        - Input: :math:`(*)` where `*` means, any number of additional
+          dimensions
+        - Output: :math:`(*)`, same shape as the input
+
+    Returns:
+        a Tensor of the same dimension and shape as the input with
+        values in the range [0, 1]
+
+    Args:
+        dim (int): A dimension along which Softmax will be computed (so every slice
+            along dim will sum to 1).
+
+    .. note::
+        This module doesn't work directly with NLLLoss,
+        which expects the Log to be computed between the Softmax and itself.
+        Use `LogSoftmax` instead (it's faster and has better numerical properties).
+
+    Examples::
+
+        >>> m = nn.Softmax(dim=1)
+        >>> input = torch.randn(2, 3)
+        >>> output = m(input)
+
+    """
+
+    __constants__ = ["dim"]
+    dim: int | None
+
+    def __init__(self, dim: int | None = None) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        if not hasattr(self, "dim"):
+            self.dim = None
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.softmax(input, self.dim, _stacklevel=5)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"dim={self.dim}"
+
+
+class Softmax2d(Module):
+    r"""Applies SoftMax over features to each spatial location.
+
+    When given an image of ``Channels x Height x Width``, it will
+    apply `Softmax` to each location :math:`(Channels, h_i, w_j)`
+
+    Shape:
+        - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`.
+        - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)
+
+    Returns:
+        a Tensor of the same dimension and shape as the input with
+        values in the range [0, 1]
+
+    Examples::
+
+        >>> m = nn.Softmax2d()
+        >>> # you softmax over the 2nd dimension
+        >>> input = torch.randn(2, 3, 12, 13)
+        >>> output = m(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        if input.dim() not in (3, 4):
+            raise ValueError(
+                f"Softmax2d: expected input to be 3D or 4D, got {input.dim()}D instead"
+            )
+        return F.softmax(input, -3, _stacklevel=5)
+
+
+class LogSoftmax(Module):
+    r"""Applies the :math:`\log(\text{Softmax}(x))` function to an n-dimensional input Tensor.
+
+    The LogSoftmax formulation can be simplified as:
+
+    .. math::
+        \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)
+
+    Shape:
+        - Input: :math:`(*)` where `*` means, any number of additional
+          dimensions
+        - Output: :math:`(*)`, same shape as the input
+
+    Args:
+        dim (int): A dimension along which LogSoftmax will be computed.
+
+    Returns:
+        a Tensor of the same dimension and shape as the input with
+        values in the range [-inf, 0)
+
+    Examples::
+
+        >>> m = nn.LogSoftmax(dim=1)
+        >>> input = torch.randn(2, 3)
+        >>> output = m(input)
+    """
+
+    __constants__ = ["dim"]
+    dim: int | None
+
+    def __init__(self, dim: int | None = None) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        if not hasattr(self, "dim"):
+            self.dim = None
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.log_softmax(input, self.dim, _stacklevel=5)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"dim={self.dim}"
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/adaptive.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/adaptive.py
new file mode 100644
index 0000000000000000000000000000000000000000..4267ed9993bff1ff69d57028308f4a3121ef2050
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/adaptive.py
@@ -0,0 +1,339 @@
+# mypy: allow-untyped-defs
+
+import itertools
+from collections import namedtuple
+from collections.abc import Sequence
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from .container import ModuleList, Sequential
+from .linear import Linear
+from .module import Module
+
+
+__all__ = ["AdaptiveLogSoftmaxWithLoss"]
+
+_ASMoutput = namedtuple("_ASMoutput", ["output", "loss"])
+
+
+class AdaptiveLogSoftmaxWithLoss(Module):
+    (
+        """Efficient softmax approximation.
+
+    As described in
+    `Efficient softmax approximation for GPUs by Edouard Grave, Armand Joulin,
+    Moustapha Ciss\u00e9, David Grangier, and Herv\u00e9 J\u00e9gou
+    <https://arxiv.org/abs/1609.04309>`__.
+"""
+        r"""
+    Adaptive softmax is an approximate strategy for training models with large
+    output spaces. It is most effective when the label distribution is highly
+    imbalanced, for example in natural language modelling, where the word
+    frequency distribution approximately follows the `Zipf's law`_.
+
+    Adaptive softmax partitions the labels into several clusters, according to
+    their frequency. These clusters may contain different number of targets
+    each.
+    Additionally, clusters containing less frequent labels assign lower
+    dimensional embeddings to those labels, which speeds up the computation.
+    For each minibatch, only clusters for which at least one target is
+    present are evaluated.
+
+    The idea is that the clusters which are accessed frequently
+    (like the first one, containing most frequent labels), should also be cheap
+    to compute -- that is, contain a small number of assigned labels.
+
+    We highly recommend taking a look at the original paper for more details.
+
+    * :attr:`cutoffs` should be an ordered Sequence of integers sorted
+      in the increasing order.
+      It controls number of clusters and the partitioning of targets into
+      clusters. For example setting ``cutoffs = [10, 100, 1000]``
+      means that first `10` targets will be assigned
+      to the 'head' of the adaptive softmax, targets `11, 12, ..., 100` will be
+      assigned to the first cluster, and targets `101, 102, ..., 1000` will be
+      assigned to the second cluster, while targets
+      `1001, 1002, ..., n_classes - 1` will be assigned
+      to the last, third cluster.
+
+    * :attr:`div_value` is used to compute the size of each additional cluster,
+      which is given as
+      :math:`\left\lfloor\frac{\texttt{in\_features}}{\texttt{div\_value}^{idx}}\right\rfloor`,
+      where :math:`idx` is the cluster index (with clusters
+      for less frequent words having larger indices,
+      and indices starting from :math:`1`).
+
+    * :attr:`head_bias` if set to True, adds a bias term to the 'head' of the
+      adaptive softmax. See paper for details. Set to False in the official
+      implementation.
+
+    .. warning::
+        Labels passed as inputs to this module should be sorted according to
+        their frequency. This means that the most frequent label should be
+        represented by the index `0`, and the least frequent
+        label should be represented by the index `n_classes - 1`.
+
+    .. note::
+        This module returns a ``NamedTuple`` with ``output``
+        and ``loss`` fields. See further documentation for details.
+
+    .. note::
+        To compute log-probabilities for all classes, the ``log_prob``
+        method can be used.
+
+    Args:
+        in_features (int): Number of features in the input tensor
+        n_classes (int): Number of classes in the dataset
+        cutoffs (Sequence): Cutoffs used to assign targets to their buckets
+        div_value (float, optional): value used as an exponent to compute sizes
+            of the clusters. Default: 4.0
+        head_bias (bool, optional): If ``True``, adds a bias term to the 'head' of the
+            adaptive softmax. Default: ``False``
+
+    Returns:
+        ``NamedTuple`` with ``output`` and ``loss`` fields:
+            * **output** is a Tensor of size ``N`` containing computed target
+              log probabilities for each example
+            * **loss** is a Scalar representing the computed negative
+              log likelihood loss
+
+    Shape:
+        - input: :math:`(N, \texttt{in\_features})` or :math:`(\texttt{in\_features})`
+        - target: :math:`(N)` or :math:`()` where each value satisfies :math:`0 <= \texttt{target[i]} <= \texttt{n\_classes}`
+        - output1: :math:`(N)` or :math:`()`
+        - output2: ``Scalar``
+
+    .. _Zipf's law: https://en.wikipedia.org/wiki/Zipf%27s_law
+    """
+    )
+
+    in_features: int
+    n_classes: int
+    cutoffs: list[int]
+    div_value: float
+    head_bias: bool
+    head: Linear
+    tail: ModuleList
+
+    def __init__(
+        self,
+        in_features: int,
+        n_classes: int,
+        cutoffs: Sequence[int],
+        div_value: float = 4.0,
+        head_bias: bool = False,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        cutoffs = list(cutoffs)
+
+        if len(cutoffs) == 0:
+            raise ValueError("cutoffs should be a sequence of length larger than 0")
+
+        if (
+            (cutoffs != sorted(cutoffs))
+            or (min(cutoffs) <= 0)
+            or (max(cutoffs) > (n_classes - 1))
+            or (len(set(cutoffs)) != len(cutoffs))
+            or any(int(c) != c for c in cutoffs)
+        ):
+            raise ValueError(
+                "cutoffs should be a sequence of unique, positive "
+                "integers sorted in an increasing order, where "
+                "each value is between 1 and n_classes-1"
+            )
+
+        self.in_features = in_features
+        self.n_classes = n_classes
+        self.cutoffs = cutoffs + [n_classes]
+        self.div_value = div_value
+        self.head_bias = head_bias
+
+        self.shortlist_size = self.cutoffs[0]
+        self.n_clusters = len(self.cutoffs) - 1
+        self.head_size = self.shortlist_size + self.n_clusters
+
+        self.head = Linear(
+            self.in_features, self.head_size, bias=self.head_bias, **factory_kwargs
+        )
+        self.tail = ModuleList()
+
+        for i in range(self.n_clusters):
+            hsz = int(self.in_features // (self.div_value ** (i + 1)))
+            osz = self.cutoffs[i + 1] - self.cutoffs[i]
+
+            projection = Sequential(
+                Linear(self.in_features, hsz, bias=False, **factory_kwargs),
+                Linear(hsz, osz, bias=False, **factory_kwargs),
+            )
+
+            self.tail.append(projection)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in ``__init__``.
+        """
+        self.head.reset_parameters()
+        for i2h, h2o in self.tail:  # type: ignore[misc]
+            i2h.reset_parameters()  # type: ignore[has-type]
+            h2o.reset_parameters()  # type: ignore[has-type]
+
+    def forward(self, input_: Tensor, target_: Tensor) -> _ASMoutput:
+        """
+        Runs the forward pass.
+        """
+        targ_dim = target_.dim()
+
+        if targ_dim == 1:
+            if input_.size(0) != target_.size(0):
+                raise RuntimeError(
+                    "Input and target should have the same size in the batch dimension."
+                )
+            if input_.dim() != 2:
+                raise RuntimeError(
+                    "1D target tensor expects 2D input tensors, "
+                    "but found inputs with size",
+                    input_.size(),
+                )
+        elif targ_dim == 0:
+            if input_.dim() != 1:
+                raise RuntimeError(
+                    "0D target tensor expects 1D input tensors, "
+                    "but found inputs with size",
+                    input_.size(),
+                )
+        else:
+            raise RuntimeError(
+                "0D or 1D target tensor expected, multi-target not supported"
+            )
+
+        is_batched = targ_dim > 0
+        input = input_ if is_batched else input_.unsqueeze(0)
+        target = target_ if is_batched else target_.unsqueeze(0)
+
+        used_rows = 0
+        batch_size = target.size(0)
+
+        output = input.new_zeros(batch_size)
+        gather_inds = target.new_empty(batch_size)
+
+        cutoff_values = [0] + self.cutoffs
+        for i in range(len(cutoff_values) - 1):
+            low_idx = cutoff_values[i]
+            high_idx = cutoff_values[i + 1]
+
+            target_mask = (target >= low_idx) & (target < high_idx)
+            row_indices = target_mask.nonzero().squeeze()
+
+            if row_indices.numel() == 0:
+                continue
+
+            if i == 0:
+                gather_inds.index_copy_(0, row_indices, target[target_mask])
+
+            else:
+                relative_target = target[target_mask] - low_idx
+                input_subset = input.index_select(0, row_indices)
+
+                cluster_output = self.tail[i - 1](input_subset)
+                cluster_index = self.shortlist_size + i - 1
+
+                gather_inds.index_fill_(0, row_indices, cluster_index)
+                cluster_logprob = F.log_softmax(cluster_output, dim=1)
+                local_logprob = cluster_logprob.gather(1, relative_target.unsqueeze(1))
+                output.index_copy_(0, row_indices, local_logprob.squeeze(1))
+
+            used_rows += row_indices.numel()
+
+        if used_rows != batch_size:
+            raise RuntimeError(
+                f"Target values should be in [0, {self.n_classes - 1}], "
+                f"but values in range [{target.min().item()}, {target.max().item()}] "
+                "were found. "
+            )
+
+        head_output = self.head(input)
+        head_logprob = F.log_softmax(head_output, dim=1)
+        output += head_logprob.gather(1, gather_inds.unsqueeze(1)).squeeze()
+        loss = (-output).mean()
+
+        if not is_batched:
+            output = output.squeeze(0)
+
+        return _ASMoutput(output, loss)
+
+    def _get_full_log_prob(self, input, head_output):
+        """Given input tensor, and output of ``self.head``, compute the log of the full distribution."""
+        out = input.new_empty((head_output.size(0), self.n_classes))
+        head_logprob = F.log_softmax(head_output, dim=1)
+
+        out[:, : self.shortlist_size] = head_logprob[:, : self.shortlist_size]
+
+        for i, (start_idx, stop_idx) in enumerate(itertools.pairwise(self.cutoffs)):
+            cluster_output = self.tail[i](input)
+            cluster_logprob = F.log_softmax(cluster_output, dim=1)
+            output_logprob = cluster_logprob + head_logprob[
+                :, self.shortlist_size + i
+            ].unsqueeze(1)
+
+            out[:, start_idx:stop_idx] = output_logprob
+
+        return out
+
+    def log_prob(self, input: Tensor) -> Tensor:
+        r"""Compute log probabilities for all :math:`\texttt{n\_classes}`.
+
+        Args:
+            input (Tensor): a minibatch of examples
+
+        Returns:
+            log-probabilities of for each class :math:`c`
+            in range :math:`0 <= c <= \texttt{n\_classes}`, where :math:`\texttt{n\_classes}` is a
+            parameter passed to ``AdaptiveLogSoftmaxWithLoss`` constructor.
+
+        Shape:
+            - Input: :math:`(N, \texttt{in\_features})`
+            - Output: :math:`(N, \texttt{n\_classes})`
+
+        """
+        head_output = self.head(input)
+        return self._get_full_log_prob(input, head_output)
+
+    def predict(self, input: Tensor) -> Tensor:
+        r"""Return the class with the highest probability for each example in the input minibatch.
+
+        This is equivalent to ``self.log_prob(input).argmax(dim=1)``, but is more efficient in some cases.
+
+        Args:
+            input (Tensor): a minibatch of examples
+
+        Returns:
+            output (Tensor): a class with the highest probability for each example
+
+        Shape:
+            - Input: :math:`(N, \texttt{in\_features})`
+            - Output: :math:`(N)`
+        """
+        head_output = self.head(input)
+        output = torch.argmax(head_output, dim=1)
+        not_in_shortlist = output >= self.shortlist_size
+        all_in_shortlist = not (not_in_shortlist.any())
+
+        if all_in_shortlist:
+            return output
+
+        elif not_in_shortlist.all():
+            log_prob = self._get_full_log_prob(input, head_output)
+            return torch.argmax(log_prob, dim=1)
+
+        else:
+            log_prob = self._get_full_log_prob(
+                input[not_in_shortlist], head_output[not_in_shortlist]
+            )
+            output[not_in_shortlist] = torch.argmax(log_prob, dim=1)
+            return output
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/batchnorm.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/batchnorm.py
new file mode 100644
index 0000000000000000000000000000000000000000..e56a747e141904ee926991ea3aa07934578a0f22
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/batchnorm.py
@@ -0,0 +1,951 @@
+# mypy: allow-untyped-defs
+from typing import Any
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F, init
+from torch.nn.parameter import Parameter, UninitializedBuffer, UninitializedParameter
+
+from ._functions import SyncBatchNorm as sync_batch_norm
+from .lazy import LazyModuleMixin
+from .module import Module
+
+
+__all__ = [
+    "BatchNorm1d",
+    "LazyBatchNorm1d",
+    "BatchNorm2d",
+    "LazyBatchNorm2d",
+    "BatchNorm3d",
+    "LazyBatchNorm3d",
+    "SyncBatchNorm",
+]
+
+
+class _NormBase(Module):
+    """Common base of _InstanceNorm and _BatchNorm."""
+
+    _version = 2
+    __constants__ = ["track_running_stats", "momentum", "eps", "num_features", "affine"]
+    num_features: int
+    eps: float
+    momentum: float | None
+    affine: bool
+    track_running_stats: bool
+    # WARNING: weight and bias purposely not defined here.
+    # See https://github.com/pytorch/pytorch/issues/39670
+
+    def __init__(
+        self,
+        num_features: int,
+        eps: float = 1e-5,
+        momentum: float | None = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+        device=None,
+        dtype=None,
+        *,
+        bias: bool = True,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+        if self.affine:
+            self.weight = Parameter(torch.empty(num_features, **factory_kwargs))
+            if bias:
+                self.bias = Parameter(torch.empty(num_features, **factory_kwargs))
+            else:
+                self.register_parameter("bias", None)
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+        if self.track_running_stats:
+            self.register_buffer(
+                "running_mean", torch.zeros(num_features, **factory_kwargs)
+            )
+            self.register_buffer(
+                "running_var", torch.ones(num_features, **factory_kwargs)
+            )
+            self.running_mean: Tensor | None
+            self.running_var: Tensor | None
+            self.register_buffer(
+                "num_batches_tracked",
+                torch.tensor(
+                    0,
+                    dtype=torch.long,
+                    # pyrefly: ignore [bad-argument-type]
+                    **{k: v for k, v in factory_kwargs.items() if k != "dtype"},
+                ),
+            )
+            self.num_batches_tracked: Tensor | None
+        else:
+            self.register_buffer("running_mean", None)
+            self.register_buffer("running_var", None)
+            self.register_buffer("num_batches_tracked", None)
+        self.reset_parameters()
+
+    def reset_running_stats(self) -> None:
+        if self.track_running_stats:
+            # running_mean/running_var/num_batches... are registered at runtime depending
+            # if self.track_running_stats is on
+            self.running_mean.zero_()  # type: ignore[union-attr]
+            self.running_var.fill_(1)  # type: ignore[union-attr]
+            self.num_batches_tracked.zero_()  # type: ignore[union-attr,operator]
+
+    def reset_parameters(self) -> None:
+        self.reset_running_stats()
+        if self.affine:
+            init.ones_(self.weight)
+            if self.bias is not None:
+                init.zeros_(self.bias)
+
+    def _check_input_dim(self, input):
+        raise NotImplementedError
+
+    def extra_repr(self):
+        return (
+            "{num_features}, eps={eps}, momentum={momentum}, affine={affine}, "
+            "bias={use_bias}, track_running_stats={track_running_stats}".format(
+                **self.__dict__, use_bias=self.bias is not None
+            )
+        )
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ) -> None:
+        version = local_metadata.get("version", None)
+
+        if (version is None or version < 2) and self.track_running_stats:
+            # at version 2: added num_batches_tracked buffer
+            #               this should have a default value of 0
+            num_batches_tracked_key = prefix + "num_batches_tracked"
+            if num_batches_tracked_key not in state_dict:
+                state_dict[num_batches_tracked_key] = (
+                    self.num_batches_tracked
+                    if self.num_batches_tracked is not None
+                    and self.num_batches_tracked.device != torch.device("meta")
+                    else torch.tensor(0, dtype=torch.long)
+                )
+
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
+
+class _BatchNorm(_NormBase):
+    def __init__(
+        self,
+        num_features: int,
+        eps: float = 1e-5,
+        momentum: float | None = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+        device=None,
+        dtype=None,
+        *,
+        bias: bool = True,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(
+            num_features,
+            eps,
+            momentum,
+            affine,
+            track_running_stats,
+            **factory_kwargs,
+            bias=bias,
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        self._check_input_dim(input)
+
+        # exponential_average_factor is set to self.momentum
+        # (when it is available) only so that it gets updated
+        # in ONNX graph when this node is exported to ONNX.
+        if self.momentum is None:
+            exponential_average_factor = 0.0
+        else:
+            exponential_average_factor = self.momentum
+
+        if self.training and self.track_running_stats:
+            # TODO: if statement only here to tell the jit to skip emitting this when it is None
+            if self.num_batches_tracked is not None:  # type: ignore[has-type]
+                self.num_batches_tracked.add_(1)  # type: ignore[has-type]
+                if self.momentum is None:  # use cumulative moving average
+                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
+                else:  # use exponential moving average
+                    exponential_average_factor = self.momentum
+
+        r"""
+        Decide whether the mini-batch stats should be used for normalization rather than the buffers.
+        Mini-batch stats are used in training mode, and in eval mode when buffers are None.
+        """
+        if self.training:
+            bn_training = True
+        else:
+            bn_training = (self.running_mean is None) and (self.running_var is None)
+
+        r"""
+        Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
+        passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
+        used for normalization (i.e. in eval mode when buffers are not None).
+        """
+        return F.batch_norm(
+            input,
+            # If buffers are not to be tracked, ensure that they won't be updated
+            (
+                self.running_mean
+                if not self.training or self.track_running_stats
+                else None
+            ),
+            self.running_var if not self.training or self.track_running_stats else None,
+            self.weight,
+            self.bias,
+            bn_training,
+            exponential_average_factor,
+            self.eps,
+        )
+
+
+class _LazyNormBase(LazyModuleMixin, _NormBase):
+    weight: UninitializedParameter  # type: ignore[assignment]
+    bias: UninitializedParameter  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        eps=1e-5,
+        momentum=0.1,
+        affine=True,
+        track_running_stats=True,
+        device=None,
+        dtype=None,
+        *,
+        bias=True,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(
+            # affine, bias and track_running_stats are hardcoded to False to
+            # avoid creating tensors that will soon be overwritten.
+            0,
+            eps,
+            momentum,
+            False,
+            False,
+            **factory_kwargs,
+            bias=False,
+        )
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+        if self.affine:
+            # pyrefly: ignore [unexpected-keyword]
+            self.weight = UninitializedParameter(**factory_kwargs)
+            if bias:
+                # pyrefly: ignore  # bad-argument-type
+                self.bias = UninitializedParameter(**factory_kwargs)
+        if self.track_running_stats:
+            # pyrefly: ignore [unexpected-keyword]
+            self.running_mean = UninitializedBuffer(**factory_kwargs)
+            # pyrefly: ignore [unexpected-keyword]
+            self.running_var = UninitializedBuffer(**factory_kwargs)
+            self.num_batches_tracked = torch.tensor(
+                0,
+                dtype=torch.long,
+                # pyrefly: ignore [bad-argument-type]
+                **{k: v for k, v in factory_kwargs.items() if k != "dtype"},
+            )
+
+    def reset_parameters(self) -> None:
+        # pyrefly: ignore [bad-argument-type]
+        if not self.has_uninitialized_params() and self.num_features != 0:
+            super().reset_parameters()
+
+    def initialize_parameters(self, input) -> None:  # type: ignore[override]
+        # pyrefly: ignore [bad-argument-type]
+        if self.has_uninitialized_params():
+            self.num_features = input.shape[1]
+            if self.affine:
+                if not isinstance(self.weight, UninitializedParameter):
+                    raise AssertionError(
+                        "self.weight must be an UninitializedParameter"
+                    )
+                self.weight.materialize((self.num_features,))
+                if self.bias is not None:
+                    if not isinstance(self.bias, UninitializedParameter):
+                        raise AssertionError(
+                            "self.bias must be an UninitializedParameter"
+                        )
+                    self.bias.materialize((self.num_features,))
+            if self.track_running_stats:
+                self.running_mean.materialize(  # type:ignore[union-attr]
+                    (self.num_features,)
+                )
+                self.running_var.materialize(  # type:ignore[union-attr]
+                    (self.num_features,)
+                )
+            self.reset_parameters()
+
+
+class BatchNorm1d(_BatchNorm):
+    r"""Applies Batch Normalization over a 2D or 3D input.
+
+    Method described in the paper
+    `Batch Normalization: Accelerating Deep Network Training by Reducing
+    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the number of features or channels of the input). By default, the
+    elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0.
+    At train time in the forward pass, the variance is calculated via the biased estimator,
+    equivalent to ``torch.var(input, correction=0)``. However, the value stored in the
+    moving average of the variance is calculated via the unbiased  estimator, equivalent to
+    ``torch.var(input, correction=1)``.
+
+    Also by default, during training this layer keeps running estimates of its
+    computed mean and variance, which are then used for normalization during
+    evaluation. The running estimates are kept with a default :attr:`momentum`
+    of 0.1.
+
+    If :attr:`track_running_stats` is set to ``False``, this layer then does not
+    keep running estimates, and batch statistics are instead used during
+    evaluation time as well.
+
+    .. note::
+        This :attr:`momentum` argument is different from one used in optimizer
+        classes and the conventional notion of momentum. Mathematically, the
+        update rule for running statistics here is
+        :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
+        where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
+        new observed value.
+
+    Because the Batch Normalization is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal Batch Normalization.
+
+    Args:
+        num_features: number of features or channels :math:`C` of the input
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics, and initializes statistics
+            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
+            When these buffers are ``None``, this module always uses batch statistics.
+            in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`, where :math:`N` is the batch size,
+          :math:`C` is the number of features or channels, and :math:`L` is the sequence length
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+
+    Examples::
+
+        >>> # With Learnable Parameters
+        >>> m = nn.BatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = nn.BatchNorm1d(100, affine=False)
+        >>> input = torch.randn(20, 100)
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
+
+
+class LazyBatchNorm1d(_LazyNormBase, _BatchNorm):
+    r"""A :class:`torch.nn.BatchNorm1d` module with lazy initialization.
+
+    Lazy initialization based on the ``num_features`` argument of the :class:`BatchNorm1d` that is inferred
+    from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight`, `bias`,
+    `running_mean` and `running_var`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics, and initializes statistics
+            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
+            When these buffers are ``None``, this module always uses batch statistics.
+            in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+    """
+
+    cls_to_become = BatchNorm1d  # type: ignore[assignment]
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
+
+
+class BatchNorm2d(_BatchNorm):
+    r"""Applies Batch Normalization over a 4D input.
+
+    4D is a mini-batch of 2D inputs
+    with additional channel dimension. Method described in the paper
+    `Batch Normalization: Accelerating Deep Network Training by Reducing
+    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set
+    to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the
+    standard-deviation is calculated via the biased estimator, equivalent to
+    ``torch.var(input, correction=0)``. However, the value stored in the moving average of the
+    standard-deviation is calculated via the unbiased  estimator, equivalent to
+    ``torch.var(input, correction=1)``.
+
+    Also by default, during training this layer keeps running estimates of its
+    computed mean and variance, which are then used for normalization during
+    evaluation. The running estimates are kept with a default :attr:`momentum`
+    of 0.1.
+
+    If :attr:`track_running_stats` is set to ``False``, this layer then does not
+    keep running estimates, and batch statistics are instead used during
+    evaluation time as well.
+
+    .. note::
+        This :attr:`momentum` argument is different from one used in optimizer
+        classes and the conventional notion of momentum. Mathematically, the
+        update rule for running statistics here is
+        :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
+        where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
+        new observed value.
+
+    Because the Batch Normalization is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial Batch Normalization.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, H, W)`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics, and initializes statistics
+            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
+            When these buffers are ``None``, this module always uses batch statistics.
+            in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+
+    Examples::
+
+        >>> # With Learnable Parameters
+        >>> m = nn.BatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = nn.BatchNorm2d(100, affine=False)
+        >>> input = torch.randn(20, 100, 35, 45)
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 4:
+            raise ValueError(f"expected 4D input (got {input.dim()}D input)")
+
+
+class LazyBatchNorm2d(_LazyNormBase, _BatchNorm):
+    r"""A :class:`torch.nn.BatchNorm2d` module with lazy initialization.
+
+    Lazy initialization is done for the ``num_features`` argument of the :class:`BatchNorm2d` that is inferred
+    from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight`, `bias`,
+    `running_mean` and `running_var`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics, and initializes statistics
+            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
+            When these buffers are ``None``, this module always uses batch statistics.
+            in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+    """
+
+    cls_to_become = BatchNorm2d  # type: ignore[assignment]
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 4:
+            raise ValueError(f"expected 4D input (got {input.dim()}D input)")
+
+
+class BatchNorm3d(_BatchNorm):
+    r"""Applies Batch Normalization over a 5D input.
+
+    5D is a mini-batch of 3D inputs with additional channel dimension as described in the paper
+    `Batch Normalization: Accelerating Deep Network Training by Reducing
+    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set
+    to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the
+    standard-deviation is calculated via the biased estimator, equivalent to
+    ``torch.var(input, correction=0)``. However, the value stored in the moving average of the
+    standard-deviation is calculated via the unbiased  estimator, equivalent to
+    ``torch.var(input, correction=1)``.
+
+    Also by default, during training this layer keeps running estimates of its
+    computed mean and variance, which are then used for normalization during
+    evaluation. The running estimates are kept with a default :attr:`momentum`
+    of 0.1.
+
+    If :attr:`track_running_stats` is set to ``False``, this layer then does not
+    keep running estimates, and batch statistics are instead used during
+    evaluation time as well.
+
+    .. note::
+        This :attr:`momentum` argument is different from one used in optimizer
+        classes and the conventional notion of momentum. Mathematically, the
+        update rule for running statistics here is
+        :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
+        where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
+        new observed value.
+
+    Because the Batch Normalization is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric Batch Normalization
+    or Spatio-temporal Batch Normalization.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, D, H, W)`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics, and initializes statistics
+            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
+            When these buffers are ``None``, this module always uses batch statistics.
+            in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+
+    Examples::
+
+        >>> # With Learnable Parameters
+        >>> m = nn.BatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = nn.BatchNorm3d(100, affine=False)
+        >>> input = torch.randn(20, 100, 35, 45, 10)
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 5:
+            raise ValueError(f"expected 5D input (got {input.dim()}D input)")
+
+
+class LazyBatchNorm3d(_LazyNormBase, _BatchNorm):
+    r"""A :class:`torch.nn.BatchNorm3d` module with lazy initialization.
+
+    Lazy initialization is done for the ``num_features`` argument of the :class:`BatchNorm3d` that is inferred
+    from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight`, `bias`,
+    `running_mean` and `running_var`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics, and initializes statistics
+            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
+            When these buffers are ``None``, this module always uses batch statistics.
+            in both training and eval modes. Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+    """
+
+    cls_to_become = BatchNorm3d  # type: ignore[assignment]
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 5:
+            raise ValueError(f"expected 5D input (got {input.dim()}D input)")
+
+
+class SyncBatchNorm(_BatchNorm):
+    r"""Applies Batch Normalization over a N-Dimensional input.
+
+    The N-D input is a mini-batch of [N-2]D inputs with additional channel dimension) as described in the paper
+    `Batch Normalization: Accelerating Deep Network Training by Reducing
+    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension over all
+    mini-batches of the same process groups. :math:`\gamma` and :math:`\beta`
+    are learnable parameter vectors of size `C` (where `C` is the input size).
+    By default, the elements of :math:`\gamma` are sampled from
+    :math:`\mathcal{U}(0, 1)` and the elements of :math:`\beta` are set to 0.
+    The standard-deviation is calculated via the biased estimator, equivalent to
+    `torch.var(input, correction=0)`.
+
+    Also by default, during training this layer keeps running estimates of its
+    computed mean and variance, which are then used for normalization during
+    evaluation. The running estimates are kept with a default :attr:`momentum`
+    of 0.1.
+
+    If :attr:`track_running_stats` is set to ``False``, this layer then does not
+    keep running estimates, and batch statistics are instead used during
+    evaluation time as well.
+
+    .. note::
+        This :attr:`momentum` argument is different from one used in optimizer
+        classes and the conventional notion of momentum. Mathematically, the
+        update rule for running statistics here is
+        :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
+        where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
+        new observed value.
+
+    Because the Batch Normalization is done for each channel in the ``C`` dimension, computing
+    statistics on ``(N, +)`` slices, it's common terminology to call this Volumetric Batch
+    Normalization or Spatio-temporal Batch Normalization.
+
+    Currently :class:`SyncBatchNorm` only supports
+    :class:`~torch.nn.DistributedDataParallel` (DDP) with single GPU per process. Use
+    :meth:`torch.nn.SyncBatchNorm.convert_sync_batchnorm()` to convert
+    :attr:`BatchNorm*D` layer to :class:`SyncBatchNorm` before wrapping
+    Network with DDP.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, +)`
+        eps: a value added to the denominator for numerical stability.
+            Default: ``1e-5``
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics, and initializes statistics
+            buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
+            When these buffers are ``None``, this module always uses batch statistics.
+            in both training and eval modes. Default: ``True``
+        process_group: synchronization of stats happen within each process group
+            individually. Default behavior is synchronization across the whole
+            world
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, +)`
+        - Output: :math:`(N, C, +)` (same shape as input)
+
+    .. note::
+        Synchronization of batchnorm statistics occurs only while training, i.e.
+        synchronization is disabled when ``model.eval()`` is set or if
+        ``self.training`` is otherwise ``False``.
+
+    Examples::
+
+        >>> # xdoctest: +SKIP
+        >>> # With Learnable Parameters
+        >>> m = nn.SyncBatchNorm(100)
+        >>> # creating process group (optional)
+        >>> # ranks is a list of int identifying rank ids.
+        >>> ranks = list(range(8))
+        >>> r1, r2 = ranks[:4], ranks[4:]
+        >>> # Note: every rank calls into new_group for every
+        >>> # process group created, even if that rank is not
+        >>> # part of the group.
+        >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
+        >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
+        >>> # Without Learnable Parameters
+        >>> m = nn.BatchNorm3d(100, affine=False, process_group=process_group)
+        >>> input = torch.randn(20, 100, 35, 45, 10)
+        >>> output = m(input)
+
+        >>> # network is nn.BatchNorm layer
+        >>> sync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network, process_group)
+        >>> # only single gpu per process is currently supported
+        >>> ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel(
+        >>>                         sync_bn_network,
+        >>>                         device_ids=[args.local_rank],
+        >>>                         output_device=args.local_rank)
+    """
+
+    def __init__(
+        self,
+        num_features: int,
+        eps: float = 1e-5,
+        momentum: float | None = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+        process_group: Any | None = None,
+        device=None,
+        dtype=None,
+        *,
+        bias: bool = True,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(
+            num_features,
+            eps,
+            momentum,
+            affine,
+            track_running_stats,
+            **factory_kwargs,
+            bias=bias,
+        )
+        self.process_group = process_group
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() < 2:
+            raise ValueError(f"expected at least 2D input (got {input.dim()}D input)")
+
+    def _check_non_zero_input_channels(self, input) -> None:
+        if input.size(1) == 0:
+            raise ValueError(
+                "SyncBatchNorm number of input channels should be non-zero"
+            )
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        self._check_input_dim(input)
+        self._check_non_zero_input_channels(input)
+
+        # exponential_average_factor is set to self.momentum
+        # (when it is available) only so that it gets updated
+        # in ONNX graph when this node is exported to ONNX.
+        if self.momentum is None:
+            exponential_average_factor = 0.0
+        else:
+            exponential_average_factor = self.momentum
+
+        if self.training and self.track_running_stats:
+            if self.num_batches_tracked is None:
+                raise AssertionError("num_batches_tracked must not be None")
+            self.num_batches_tracked.add_(1)
+            if self.momentum is None:  # use cumulative moving average
+                exponential_average_factor = 1.0 / self.num_batches_tracked.item()
+            else:  # use exponential moving average
+                exponential_average_factor = self.momentum
+
+        r"""
+        Decide whether the mini-batch stats should be used for normalization rather than the buffers.
+        Mini-batch stats are used in training mode, and in eval mode when buffers are None.
+        """
+        if self.training:
+            bn_training = True
+        else:
+            bn_training = (self.running_mean is None) and (self.running_var is None)
+
+        r"""
+        Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
+        passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
+        used for normalization (i.e. in eval mode when buffers are not None).
+        """
+        # If buffers are not to be tracked, ensure that they won't be updated
+        running_mean = (
+            self.running_mean if not self.training or self.track_running_stats else None
+        )
+        running_var = (
+            self.running_var if not self.training or self.track_running_stats else None
+        )
+
+        # Don't sync batchnorm stats in inference mode (model.eval()).
+        need_sync = (
+            bn_training
+            and self.training
+            and torch.distributed.is_available()
+            and torch.distributed.is_initialized()
+        )
+        if need_sync:
+            # currently only GPU/PrivateUse1 input is supported
+            if input.device.type not in [
+                "cuda",
+                "hpu",
+                "xpu",
+                torch._C._get_privateuse1_backend_name(),
+            ]:
+                raise ValueError(
+                    "SyncBatchNorm expected input tensor to be on GPU or XPU or "
+                    f"{torch._C._get_privateuse1_backend_name()}"
+                )
+
+            process_group = torch.distributed.group.WORLD
+            if self.process_group:
+                process_group = self.process_group
+            world_size = torch.distributed.get_world_size(process_group)
+            need_sync = world_size > 1
+
+        # fallback to framework BN when synchronization is not necessary
+        if not need_sync:
+            return F.batch_norm(
+                input,
+                running_mean,
+                running_var,
+                self.weight,
+                self.bias,
+                bn_training,
+                exponential_average_factor,
+                self.eps,
+            )
+        else:
+            if not bn_training:
+                raise AssertionError("bn_training must be True")
+            return sync_batch_norm.apply(
+                input,
+                self.weight,
+                self.bias,
+                running_mean,
+                running_var,
+                self.eps,
+                exponential_average_factor,
+                process_group,  # type: ignore[possibly-undefined]
+                world_size,  # type: ignore[possibly-undefined]
+            )
+
+    @classmethod
+    def convert_sync_batchnorm(cls, module, process_group=None):
+        r"""Converts all :attr:`BatchNorm*D` layers in the model to :class:`torch.nn.SyncBatchNorm` layers.
+
+        Args:
+            module (nn.Module): module containing one or more :attr:`BatchNorm*D` layers
+            process_group (optional): process group to scope synchronization,
+                default is the whole world
+
+        Returns:
+            The original :attr:`module` with the converted :class:`torch.nn.SyncBatchNorm`
+            layers. If the original :attr:`module` is a :attr:`BatchNorm*D` layer,
+            a new :class:`torch.nn.SyncBatchNorm` layer object will be returned
+            instead.
+
+        Example::
+
+            >>> # Network with nn.BatchNorm layer
+            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+            >>> module = torch.nn.Sequential(
+            >>>            torch.nn.Linear(20, 100),
+            >>>            torch.nn.BatchNorm1d(100),
+            >>>          ).cuda()
+            >>> # creating process group (optional)
+            >>> # ranks is a list of int identifying rank ids.
+            >>> ranks = list(range(8))
+            >>> r1, r2 = ranks[:4], ranks[4:]
+            >>> # Note: every rank calls into new_group for every
+            >>> # process group created, even if that rank is not
+            >>> # part of the group.
+            >>> # xdoctest: +SKIP("distributed")
+            >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
+            >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
+            >>> sync_bn_module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)
+
+        """
+        module_output = module
+        if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
+            module_output = torch.nn.SyncBatchNorm(
+                module.num_features,
+                module.eps,
+                module.momentum,
+                module.affine,
+                module.track_running_stats,
+                process_group,
+                bias=module.bias is not None,
+            )
+            if module.affine:
+                with torch.no_grad():
+                    module_output.weight = module.weight
+                    module_output.bias = module.bias
+            module_output.running_mean = module.running_mean
+            module_output.running_var = module.running_var
+            module_output.num_batches_tracked = module.num_batches_tracked
+            module_output.training = module.training
+            if hasattr(module, "qconfig"):
+                module_output.qconfig = module.qconfig
+        for name, child in module.named_children():
+            module_output.add_module(
+                name, cls.convert_sync_batchnorm(child, process_group)
+            )
+        del module
+        return module_output
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/channelshuffle.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/channelshuffle.py
new file mode 100644
index 0000000000000000000000000000000000000000..34a48f04f853dd1c458b035635728a122e9cc4d3
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/channelshuffle.py
@@ -0,0 +1,62 @@
+import torch.nn.functional as F
+from torch import Tensor
+
+from .module import Module
+
+
+__all__ = ["ChannelShuffle"]
+
+
+class ChannelShuffle(Module):
+    r"""Divides and rearranges the channels in a tensor.
+
+    This operation divides the channels in a tensor of shape :math:`(N, C, *)`
+    into g groups as :math:`(N, \frac{C}{g}, g, *)` and shuffles them,
+    while retaining the original tensor shape in the final output.
+
+    Args:
+        groups (int): number of groups to divide channels in.
+
+    Examples::
+
+        >>> channel_shuffle = nn.ChannelShuffle(2)
+        >>> input = torch.arange(1, 17, dtype=torch.float32).view(1, 4, 2, 2)
+        >>> input
+        tensor([[[[ 1.,  2.],
+                  [ 3.,  4.]],
+                 [[ 5.,  6.],
+                  [ 7.,  8.]],
+                 [[ 9., 10.],
+                  [11., 12.]],
+                 [[13., 14.],
+                  [15., 16.]]]])
+        >>> output = channel_shuffle(input)
+        >>> output
+        tensor([[[[ 1.,  2.],
+                  [ 3.,  4.]],
+                 [[ 9., 10.],
+                  [11., 12.]],
+                 [[ 5.,  6.],
+                  [ 7.,  8.]],
+                 [[13., 14.],
+                  [15., 16.]]]])
+    """
+
+    __constants__ = ["groups"]
+    groups: int
+
+    def __init__(self, groups: int) -> None:
+        super().__init__()
+        self.groups = groups
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.channel_shuffle(input, self.groups)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"groups={self.groups}"
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/container.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/container.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1a1d7cc990d1512bf5ca7343e3255518e901374
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/container.py
@@ -0,0 +1,1039 @@
+# mypy: allow-untyped-defs
+from __future__ import annotations
+
+import operator
+from collections import abc as container_abcs, OrderedDict
+from itertools import chain, islice
+from typing import Any, overload, TYPE_CHECKING, TypeVar
+from typing_extensions import deprecated, Self
+
+import torch
+from torch._jit_internal import _copy_to_script_wrapper
+from torch.nn.parameter import Parameter
+
+from .module import Module
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator, Mapping
+
+
+__all__ = [
+    "Container",
+    "Sequential",
+    "ModuleList",
+    "ModuleDict",
+    "ParameterList",
+    "ParameterDict",
+]
+
+T = TypeVar("T", bound=Module)
+_V = TypeVar("_V")
+
+
+# Copied from torch.nn.modules.module, required for a custom __repr__ for ModuleList
+def _addindent(s_, numSpaces):
+    s = s_.split("\n")
+    # don't do anything for single-line stuff
+    if len(s) == 1:
+        return s_
+    first = s.pop(0)
+    s = [(numSpaces * " ") + line for line in s]
+    s = "\n".join(s)
+    s = first + "\n" + s
+    return s
+
+
+@deprecated(
+    "`nn.Container` is deprecated. "
+    "All of it's functionality is now implemented in `nn.Module`. Subclass that instead.",
+    category=FutureWarning,
+)
+class Container(Module):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__()
+        for key, value in kwargs.items():
+            self.add_module(key, value)
+
+
+class Sequential(Module):
+    r"""A sequential container.
+
+    Modules will be added to it in the order they are passed in the
+    constructor. Alternatively, an ``OrderedDict`` of modules can be
+    passed in. The ``forward()`` method of ``Sequential`` accepts any
+    input and forwards it to the first module it contains. It then
+    "chains" outputs to inputs sequentially for each subsequent module,
+    finally returning the output of the last module.
+
+    The value a ``Sequential`` provides over manually calling a sequence
+    of modules is that it allows treating the whole container as a
+    single module, such that performing a transformation on the
+    ``Sequential`` applies to each of the modules it stores (which are
+    each a registered submodule of the ``Sequential``).
+
+    What's the difference between a ``Sequential`` and a
+    :class:`torch.nn.ModuleList`? A ``ModuleList`` is exactly what it
+    sounds like--a list for storing ``Module`` s! On the other hand,
+    the layers in a ``Sequential`` are connected in a cascading way.
+
+    Example::
+
+        # Using Sequential to create a small model. When `model` is run,
+        # input will first be passed to `Conv2d(1,20,5)`. The output of
+        # `Conv2d(1,20,5)` will be used as the input to the first
+        # `ReLU`; the output of the first `ReLU` will become the input
+        # for `Conv2d(20,64,5)`. Finally, the output of
+        # `Conv2d(20,64,5)` will be used as input to the second `ReLU`
+        model = nn.Sequential(
+            nn.Conv2d(1, 20, 5), nn.ReLU(), nn.Conv2d(20, 64, 5), nn.ReLU()
+        )
+
+        # Using Sequential with OrderedDict. This is functionally the
+        # same as the above code
+        model = nn.Sequential(
+            OrderedDict(
+                [
+                    ("conv1", nn.Conv2d(1, 20, 5)),
+                    ("relu1", nn.ReLU()),
+                    ("conv2", nn.Conv2d(20, 64, 5)),
+                    ("relu2", nn.ReLU()),
+                ]
+            )
+        )
+    """
+
+    _modules: dict[str, Module]  # type: ignore[assignment]
+
+    @overload
+    def __init__(self, *args: Module) -> None: ...
+
+    @overload
+    # pyrefly: ignore [inconsistent-overload]
+    def __init__(self, arg: OrderedDict[str, Module]) -> None: ...
+
+    def __init__(self, *args):
+        super().__init__()
+        if len(args) == 1 and isinstance(args[0], OrderedDict):
+            for key, module in args[0].items():
+                self.add_module(key, module)
+        else:
+            for idx, module in enumerate(args):
+                self.add_module(str(idx), module)
+
+    def _get_item_by_idx(self, iterator: Iterable[_V], idx: int) -> _V:
+        """Get the idx-th item of the iterator."""
+        size = len(self)
+        idx = operator.index(idx)
+        if not -size <= idx < size:
+            raise IndexError(f"index {idx} is out of range")
+        idx %= size
+        return next(islice(iterator, idx, None))
+
+    @_copy_to_script_wrapper
+    def __getitem__(self, idx: slice | int) -> Sequential | Module:
+        if isinstance(idx, slice):
+            return self.__class__(OrderedDict(list(self._modules.items())[idx]))
+        else:
+            return self._get_item_by_idx(self._modules.values(), idx)
+
+    def __setitem__(self, idx: int, module: Module) -> None:
+        key: str = self._get_item_by_idx(self._modules.keys(), idx)
+        return setattr(self, key, module)
+
+    def __delitem__(self, idx: slice | int) -> None:
+        if isinstance(idx, slice):
+            for key in list(self._modules.keys())[idx]:
+                delattr(self, key)
+        else:
+            key = self._get_item_by_idx(self._modules.keys(), idx)
+            delattr(self, key)
+        # To preserve numbering
+        str_indices = [str(i) for i in range(len(self._modules))]
+        self._modules = OrderedDict(
+            zip(str_indices, self._modules.values(), strict=True)
+        )
+
+    @_copy_to_script_wrapper
+    def __len__(self) -> int:
+        return len(self._modules)
+
+    def __add__(self, other) -> Sequential:
+        if isinstance(other, Sequential):
+            ret = Sequential()
+            for layer in self:
+                ret.append(layer)
+            for layer in other:
+                ret.append(layer)
+            return ret
+        else:
+            raise ValueError(
+                "add operator supports only objects "
+                f"of Sequential class, but {str(type(other))} is given."
+            )
+
+    def pop(self, key: int | slice) -> Module:
+        """
+        Pop ``key`` from self.
+        """
+        v = self[key]
+        del self[key]
+        return v
+
+    def __iadd__(self, other) -> Self:
+        if isinstance(other, Sequential):
+            offset = len(self)
+            for i, module in enumerate(other):
+                self.add_module(str(i + offset), module)
+            return self
+        else:
+            raise ValueError(
+                "add operator supports only objects "
+                f"of Sequential class, but {str(type(other))} is given."
+            )
+
+    def __mul__(self, other: int) -> Sequential:
+        if not isinstance(other, int):
+            raise TypeError(
+                f"unsupported operand type(s) for *: {type(self)} and {type(other)}"
+            )
+        elif other <= 0:
+            raise ValueError(
+                f"Non-positive multiplication factor {other} for {type(self)}"
+            )
+        else:
+            combined = Sequential()
+            offset = 0
+            for _ in range(other):
+                for module in self:
+                    combined.add_module(str(offset), module)
+                    offset += 1
+            return combined
+
+    def __rmul__(self, other: int) -> Sequential:
+        return self.__mul__(other)
+
+    def __imul__(self, other: int) -> Self:
+        if not isinstance(other, int):
+            raise TypeError(
+                f"unsupported operand type(s) for *: {type(self)} and {type(other)}"
+            )
+        elif other <= 0:
+            raise ValueError(
+                f"Non-positive multiplication factor {other} for {type(self)}"
+            )
+        else:
+            len_original = len(self)
+            offset = len(self)
+            for _ in range(other - 1):
+                for i in range(len_original):
+                    self.add_module(str(i + offset), self._modules[str(i)])
+                offset += len_original
+            return self
+
+    @_copy_to_script_wrapper
+    def __dir__(self) -> list[str]:
+        keys = super().__dir__()
+        keys = [key for key in keys if not key.isdigit()]
+        return keys
+
+    @_copy_to_script_wrapper
+    def __iter__(self) -> Iterator[Module]:
+        return iter(self._modules.values())
+
+    # NB: We can't really type check this function as the type of input
+    # may change dynamically (as is tested in
+    # TestScript.test_sequential_intermediary_types).  Cannot annotate
+    # with Any as TorchScript expects a more precise type
+    def forward(self, input):
+        """
+        Runs the forward pass.
+        """
+        for module in self:
+            input = module(input)
+        return input
+
+    def append(self, module: Module) -> Self:
+        r"""Append a given module to the end.
+
+        Args:
+            module (nn.Module): module to append
+
+        Example::
+
+            >>> import torch.nn as nn
+            >>> n = nn.Sequential(nn.Linear(1, 2), nn.Linear(2, 3))
+            >>> n.append(nn.Linear(3, 4))
+            Sequential(
+                (0): Linear(in_features=1, out_features=2, bias=True)
+                (1): Linear(in_features=2, out_features=3, bias=True)
+                (2): Linear(in_features=3, out_features=4, bias=True)
+            )
+
+        """
+        self.add_module(str(len(self)), module)
+        return self
+
+    def insert(self, index: int, module: Module) -> Self:
+        """
+        Inserts a module into the Sequential container at the specified index.
+
+        Args:
+            index (int): The index to insert the module.
+            module (Module): The module to be inserted.
+
+        Example::
+
+            >>> import torch.nn as nn
+            >>> n = nn.Sequential(nn.Linear(1, 2), nn.Linear(2, 3))
+            >>> n.insert(0, nn.Linear(3, 4))
+            Sequential(
+                (0): Linear(in_features=3, out_features=4, bias=True)
+                (1): Linear(in_features=1, out_features=2, bias=True)
+                (2): Linear(in_features=2, out_features=3, bias=True)
+            )
+
+        """
+        if not isinstance(module, Module):
+            raise AssertionError(f"module should be of type: {Module}")
+        n = len(self._modules)
+        if not (-n <= index <= n):
+            raise IndexError(f"Index out of range: {index}")
+        if index < 0:
+            index += n
+        for i in range(n, index, -1):
+            self._modules[str(i)] = self._modules[str(i - 1)]
+        self._modules[str(index)] = module
+        return self
+
+    def extend(self, sequential: Iterable[Module]) -> Self:
+        """
+        Extends the current Sequential container with layers from another Sequential container.
+
+        Args:
+            sequential (Sequential): A Sequential container whose layers will be added to the current container.
+
+        Example::
+
+            >>> import torch.nn as nn
+            >>> n = nn.Sequential(nn.Linear(1, 2), nn.Linear(2, 3))
+            >>> other = nn.Sequential(nn.Linear(3, 4), nn.Linear(4, 5))
+            >>> n.extend(other) # or `n + other`
+            Sequential(
+                (0): Linear(in_features=1, out_features=2, bias=True)
+                (1): Linear(in_features=2, out_features=3, bias=True)
+                (2): Linear(in_features=3, out_features=4, bias=True)
+                (3): Linear(in_features=4, out_features=5, bias=True)
+            )
+
+        """
+        for layer in sequential:
+            self.append(layer)
+        return self
+
+
+class ModuleList(Module):
+    r"""Holds submodules in a list.
+
+    :class:`~torch.nn.ModuleList` can be indexed like a regular Python list, but
+    modules it contains are properly registered, and will be visible by all
+    :class:`~torch.nn.Module` methods.
+
+    Args:
+        modules (iterable, optional): an iterable of modules to add
+
+    Example::
+
+        class MyModule(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])
+
+            def forward(self, x):
+                # ModuleList can act as an iterable, or be indexed using ints
+                for i, l in enumerate(self.linears):
+                    x = self.linears[i // 2](x) + l(x)
+                return x
+    """
+
+    _modules: dict[str, Module]  # type: ignore[assignment]
+
+    def __init__(self, modules: Iterable[Module] | None = None) -> None:
+        super().__init__()
+        if modules is not None:
+            self += modules
+
+    def _get_abs_string_index(self, idx):
+        """Get the absolute index for the list of modules."""
+        idx = operator.index(idx)
+        if not (-len(self) <= idx < len(self)):
+            raise IndexError(f"index {idx} is out of range")
+        if idx < 0:
+            idx += len(self)
+        return str(idx)
+
+    @overload
+    def __getitem__(self, idx: slice) -> ModuleList: ...
+
+    @overload
+    def __getitem__(self, idx: int) -> Module: ...
+
+    @_copy_to_script_wrapper
+    def __getitem__(self, idx: int | slice) -> Module | ModuleList:
+        if isinstance(idx, slice):
+            return self.__class__(list(self._modules.values())[idx])
+        else:
+            return self._modules[self._get_abs_string_index(idx)]
+
+    def __setitem__(self, idx: int, module: Module) -> None:
+        idx = self._get_abs_string_index(idx)
+        return setattr(self, str(idx), module)
+
+    def __delitem__(self, idx: int | slice) -> None:
+        if isinstance(idx, slice):
+            for k in range(len(self._modules))[idx]:
+                delattr(self, str(k))
+        else:
+            delattr(self, self._get_abs_string_index(idx))
+        # To preserve numbering, self._modules is being reconstructed with modules after deletion
+        str_indices = [str(i) for i in range(len(self._modules))]
+        self._modules = OrderedDict(
+            zip(str_indices, self._modules.values(), strict=True)
+        )
+
+    @_copy_to_script_wrapper
+    def __len__(self) -> int:
+        return len(self._modules)
+
+    @_copy_to_script_wrapper
+    def __iter__(self) -> Iterator[Module]:
+        return iter(self._modules.values())
+
+    def __iadd__(self, modules: Iterable[Module]) -> Self:
+        return self.extend(modules)
+
+    def __add__(self, other: Iterable[Module]) -> ModuleList:
+        combined = ModuleList()
+        for i, module in enumerate(chain(self, other)):
+            combined.add_module(str(i), module)
+        return combined
+
+    def __repr__(self) -> str:
+        """Return a custom repr for ModuleList that compresses repeated module representations."""
+        list_of_reprs = [repr(item) for item in self]
+        if len(list_of_reprs) == 0:
+            return self._get_name() + "()"
+
+        start_end_indices = [[0, 0]]
+        repeated_blocks = [list_of_reprs[0]]
+        for i, r in enumerate(list_of_reprs[1:], 1):
+            if r == repeated_blocks[-1]:
+                start_end_indices[-1][1] += 1
+                continue
+
+            start_end_indices.append([i, i])
+            repeated_blocks.append(r)
+
+        lines = []
+        main_str = self._get_name() + "("
+        for (start_id, end_id), b in zip(
+            start_end_indices, repeated_blocks, strict=True
+        ):
+            local_repr = f"({start_id}): {b}"  # default repr
+
+            if start_id != end_id:
+                n = end_id - start_id + 1
+                local_repr = f"({start_id}-{end_id}): {n} x {b}"
+
+            local_repr = _addindent(local_repr, 2)
+            lines.append(local_repr)
+
+        main_str += "\n  " + "\n  ".join(lines) + "\n"
+        main_str += ")"
+        return main_str
+
+    @_copy_to_script_wrapper
+    def __dir__(self) -> list[str]:
+        keys = super().__dir__()
+        keys = [key for key in keys if not key.isdigit()]
+        return keys
+
+    def insert(self, index: int, module: Module) -> None:
+        r"""Insert a given module before a given index in the list.
+
+        Args:
+            index (int): index to insert.
+            module (nn.Module): module to insert
+        """
+        for i in range(len(self._modules), index, -1):
+            self._modules[str(i)] = self._modules[str(i - 1)]
+        self._modules[str(index)] = module
+
+    def append(self, module: Module) -> Self:
+        r"""Append a given module to the end of the list.
+
+        Args:
+            module (nn.Module): module to append
+        """
+        self.add_module(str(len(self)), module)
+        return self
+
+    def pop(self, key: int | slice) -> Module:
+        v = self[key]
+        del self[key]
+        return v
+
+    def extend(self, modules: Iterable[Module]) -> Self:
+        r"""Append modules from a Python iterable to the end of the list.
+
+        Args:
+            modules (iterable): iterable of modules to append
+        """
+        if not isinstance(modules, container_abcs.Iterable):
+            raise TypeError(
+                "ModuleList.extend should be called with an "
+                "iterable, but got " + type(modules).__name__
+            )
+        offset = len(self)
+        for i, module in enumerate(modules):
+            self.add_module(str(offset + i), module)
+        return self
+
+    # remove forward altogether to fallback on Module's _forward_unimplemented
+
+
+class ModuleDict(Module):
+    r"""Holds submodules in a dictionary.
+
+    :class:`~torch.nn.ModuleDict` can be indexed like a regular Python dictionary,
+    but modules it contains are properly registered, and will be visible by all
+    :class:`~torch.nn.Module` methods.
+
+    :class:`~torch.nn.ModuleDict` is an **ordered** dictionary that respects
+
+    * the order of insertion, and
+
+    * in :meth:`~torch.nn.ModuleDict.update`, the order of the merged
+      ``OrderedDict``, ``dict`` or another
+      :class:`~torch.nn.ModuleDict` (the argument to
+      :meth:`~torch.nn.ModuleDict.update`).
+
+    Note that :meth:`~torch.nn.ModuleDict.update` with other unordered mapping
+    types does not preserve the order of the merged mapping.
+
+    Args:
+        modules (iterable, optional): a mapping (dictionary) of (string: module)
+            or an iterable of key-value pairs of type (string, module)
+
+    Example::
+
+        class MyModule(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.choices = nn.ModuleDict(
+                    {"conv": nn.Conv2d(10, 10, 3), "pool": nn.MaxPool2d(3)}
+                )
+                self.activations = nn.ModuleDict(
+                    [["lrelu", nn.LeakyReLU()], ["prelu", nn.PReLU()]]
+                )
+
+            def forward(self, x, choice, act):
+                x = self.choices[choice](x)
+                x = self.activations[act](x)
+                return x
+    """
+
+    _modules: dict[str, Module]  # type: ignore[assignment]
+
+    def __init__(self, modules: Mapping[str, Module] | None = None) -> None:
+        super().__init__()
+        if modules is not None:
+            self.update(modules)
+
+    @_copy_to_script_wrapper
+    def __getitem__(self, key: str) -> Module:
+        return self._modules[key]
+
+    def __setitem__(self, key: str, module: Module) -> None:
+        self.add_module(key, module)
+
+    def __delitem__(self, key: str) -> None:
+        del self._modules[key]
+
+    @_copy_to_script_wrapper
+    def __len__(self) -> int:
+        return len(self._modules)
+
+    @_copy_to_script_wrapper
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._modules)
+
+    @_copy_to_script_wrapper
+    def __contains__(self, key: str) -> bool:
+        return key in self._modules
+
+    def clear(self) -> None:
+        """Remove all items from the ModuleDict."""
+        self._modules.clear()
+
+    def pop(self, key: str) -> Module:
+        r"""Remove key from the ModuleDict and return its module.
+
+        Args:
+            key (str): key to pop from the ModuleDict
+        """
+        v = self[key]
+        del self[key]
+        return v
+
+    @_copy_to_script_wrapper
+    def keys(self) -> container_abcs.KeysView[str]:
+        r"""Return an iterable of the ModuleDict keys."""
+        return self._modules.keys()
+
+    @_copy_to_script_wrapper
+    def items(self) -> container_abcs.ItemsView[str, Module]:
+        r"""Return an iterable of the ModuleDict key/value pairs."""
+        return self._modules.items()
+
+    @_copy_to_script_wrapper
+    def values(self) -> container_abcs.ValuesView[Module]:
+        r"""Return an iterable of the ModuleDict values."""
+        return self._modules.values()
+
+    def update(self, modules: Mapping[str, Module]) -> None:
+        r"""Update the :class:`~torch.nn.ModuleDict` with key-value pairs from a mapping, overwriting existing keys.
+
+        .. note::
+            If :attr:`modules` is an ``OrderedDict``, a :class:`~torch.nn.ModuleDict`, or
+            an iterable of key-value pairs, the order of new elements in it is preserved.
+
+        Args:
+            modules (iterable): a mapping (dictionary) from string to :class:`~torch.nn.Module`,
+                or an iterable of key-value pairs of type (string, :class:`~torch.nn.Module`)
+        """
+        if not isinstance(modules, container_abcs.Iterable):
+            raise TypeError(
+                "ModuleDict.update should be called with an "
+                "iterable of key/value pairs, but got " + type(modules).__name__
+            )
+
+        if isinstance(modules, (OrderedDict, ModuleDict, container_abcs.Mapping)):
+            for key, module in modules.items():
+                self[key] = module
+        else:
+            # modules here can be a list with two items
+            for j, m in enumerate(modules):
+                if not isinstance(m, container_abcs.Iterable):
+                    raise TypeError(
+                        "ModuleDict update sequence element "
+                        "#" + str(j) + " should be Iterable; is" + type(m).__name__
+                    )
+                # pyrefly: ignore [bad-argument-type]
+                if not len(m) == 2:
+                    raise ValueError(
+                        "ModuleDict update sequence element "
+                        # pyrefly: ignore [bad-argument-type]
+                        "#" + str(j) + " has length " + str(len(m)) + "; 2 is required"
+                    )
+                # modules can be Mapping (what it's typed at), or a list: [(name1, module1), (name2, module2)]
+                # that's too cumbersome to type correctly with overloads, so we add an ignore here
+                self[m[0]] = m[1]  # type: ignore[assignment]
+
+    # remove forward altogether to fallback on Module's _forward_unimplemented
+
+
+class ParameterList(Module):
+    r"""Holds parameters in a list.
+
+    :class:`~torch.nn.ParameterList` can be used like a regular Python
+    list, but Tensors that are :class:`~torch.nn.Parameter` are properly registered,
+    and will be visible by all :class:`~torch.nn.Module` methods.
+
+    Note that the constructor, assigning an element of the list, the
+    :meth:`~torch.nn.ParameterList.append` method and the :meth:`~torch.nn.ParameterList.extend`
+    method will convert any :class:`~torch.Tensor` into :class:`~torch.nn.Parameter`.
+
+    Args:
+        parameters (iterable, optional): an iterable of elements to add to the list.
+
+    Example::
+
+        class MyModule(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.params = nn.ParameterList(
+                    [nn.Parameter(torch.randn(10, 10)) for i in range(10)]
+                )
+
+            def forward(self, x):
+                # ParameterList can act as an iterable, or be indexed using ints
+                for i, p in enumerate(self.params):
+                    x = self.params[i // 2].mm(x) + p.mm(x)
+                return x
+    """
+
+    def __init__(self, values: Iterable[Any] | None = None) -> None:
+        super().__init__()
+        self._size = 0
+        if values is not None:
+            self += values
+
+    def _get_abs_string_index(self, idx):
+        """Get the absolute index for the list of modules."""
+        idx = operator.index(idx)
+        if not (-len(self) <= idx < len(self)):
+            raise IndexError(f"index {idx} is out of range")
+        if idx < 0:
+            idx += len(self)
+        return str(idx)
+
+    @overload
+    def __getitem__(self, idx: int) -> Any: ...
+
+    @overload
+    # pyrefly: ignore [inconsistent-overload]
+    def __getitem__(self: T, idx: slice) -> T: ...
+
+    def __getitem__(self, idx):
+        if isinstance(idx, slice):
+            start, stop, step = idx.indices(len(self))
+            out = self.__class__()
+            for i in range(start, stop, step):
+                out.append(self[i])
+            return out
+        else:
+            idx = self._get_abs_string_index(idx)
+            return getattr(self, str(idx))
+
+    def __setitem__(self, idx: int, param: Any) -> None:
+        # Note that all other function that add an entry to the list part of
+        # the ParameterList end up here. So this is the only place where we need
+        # to wrap things into Parameter if needed.
+        # Objects added via setattr() are not in the list part and thus won't
+        # call into this function.
+        idx = self._get_abs_string_index(idx)
+        if isinstance(param, torch.Tensor) and not isinstance(param, Parameter):
+            param = Parameter(param)
+        return setattr(self, str(idx), param)
+
+    def __len__(self) -> int:
+        return self._size
+
+    def __iter__(self) -> Iterator[Any]:
+        return iter(self[i] for i in range(len(self)))
+
+    def __iadd__(self, parameters: Iterable[Any]) -> Self:
+        return self.extend(parameters)
+
+    def __dir__(self) -> list[str]:
+        keys = super().__dir__()
+        keys = [key for key in keys if not key.isdigit()]
+        return keys
+
+    def append(self, value: Any) -> Self:
+        """Append a given value at the end of the list.
+
+        Args:
+            value (Any): value to append
+        """
+        new_idx = len(self)
+        self._size += 1
+        self[new_idx] = value
+        return self
+
+    def extend(self, values: Iterable[Any]) -> Self:
+        """Append values from a Python iterable to the end of the list.
+
+        Args:
+            values (iterable): iterable of values to append
+        """
+        # Tensor is an iterable but we never want to unpack it here
+        if not isinstance(values, container_abcs.Iterable) or isinstance(
+            values, torch.Tensor
+        ):
+            raise TypeError(
+                "ParameterList.extend should be called with an "
+                "iterable, but got " + type(values).__name__
+            )
+        for value in values:
+            self.append(value)
+        return self
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        child_lines = []
+        for k, p in enumerate(self):
+            if isinstance(p, torch.Tensor):
+                size_str = "x".join(str(size) for size in p.size())
+                if p.device.type in ["cuda", torch._C._get_privateuse1_backend_name()]:
+                    device_str = f" ({p.device})"
+                else:
+                    device_str = ""
+                parastr = "{} containing: [{} of size {}{}]".format(
+                    "Parameter" if isinstance(p, Parameter) else "Tensor",
+                    p.dtype,
+                    size_str,
+                    device_str,
+                )
+                child_lines.append("  (" + str(k) + "): " + parastr)
+            else:
+                child_lines.append(
+                    "  (" + str(k) + "): Object of type: " + type(p).__name__
+                )
+
+        tmpstr = "\n".join(child_lines)
+        return tmpstr
+
+    def __call__(self, *args, **kwargs):
+        raise RuntimeError("ParameterList should not be called.")
+
+
+class ParameterDict(Module):
+    r"""Holds parameters in a dictionary.
+
+    ParameterDict can be indexed like a regular Python dictionary, but Parameters it
+    contains are properly registered, and will be visible by all Module methods.
+    Other objects are treated as would be done by a regular Python dictionary
+
+    :class:`~torch.nn.ParameterDict` is an **ordered** dictionary.
+    :meth:`~torch.nn.ParameterDict.update` with other unordered mapping
+    types (e.g., Python's plain ``dict``) does not preserve the order of the
+    merged mapping. On the other hand, ``OrderedDict`` or another :class:`~torch.nn.ParameterDict`
+    will preserve their ordering.
+
+    Note that the constructor, assigning an element of the dictionary and the
+    :meth:`~torch.nn.ParameterDict.update` method will convert any :class:`~torch.Tensor` into
+    :class:`~torch.nn.Parameter`.
+
+    Args:
+        values (iterable, optional): a mapping (dictionary) of
+            (string : Any) or an iterable of key-value pairs
+            of type (string, Any)
+
+    Example::
+
+        class MyModule(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.params = nn.ParameterDict(
+                    {
+                        "left": nn.Parameter(torch.randn(5, 10)),
+                        "right": nn.Parameter(torch.randn(5, 10)),
+                    }
+                )
+
+            def forward(self, x, choice):
+                x = self.params[choice].mm(x)
+                return x
+    """
+
+    def __init__(self, parameters: Any = None) -> None:
+        super().__init__()
+        self._keys: dict[str, None] = {}
+        if parameters is not None:
+            self.update(parameters)
+
+    def _key_to_attr(self, key: str) -> str:
+        if not isinstance(key, str):
+            raise TypeError(
+                "Index given to ParameterDict cannot be used as a key as it is "
+                f"not a string (type is '{type(key).__name__}'). Open an issue on "
+                "github if you need non-string keys."
+            )
+        else:
+            # Use the key as-is so that `.named_parameters()` returns the right thing
+            return key
+
+    def __getitem__(self, key: str) -> Any:
+        attr = self._key_to_attr(key)
+        return getattr(self, attr)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        # Note that all other function that add an entry to the dictionary part of
+        # the ParameterDict end up here. So this is the only place where we need
+        # to wrap things into Parameter if needed.
+        # Objects added via setattr() are not in the dictionary part and thus won't
+        # call into this function.
+        self._keys[key] = None
+        attr = self._key_to_attr(key)
+        if isinstance(value, torch.Tensor) and not isinstance(value, Parameter):
+            value = Parameter(value)
+        setattr(self, attr, value)
+
+    def __delitem__(self, key: str) -> None:
+        del self._keys[key]
+        attr = self._key_to_attr(key)
+        delattr(self, attr)
+
+    def __len__(self) -> int:
+        return len(self._keys)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._keys)
+
+    def __reversed__(self) -> Iterator[str]:
+        return reversed(self._keys)
+
+    def copy(self) -> ParameterDict:
+        """Return a copy of this :class:`~torch.nn.ParameterDict` instance."""
+        # We have to use an OrderedDict because the ParameterDict constructor
+        # behaves differently on plain dict vs OrderedDict
+        return ParameterDict(OrderedDict((k, self[k]) for k in self._keys))
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._keys
+
+    def setdefault(self, key: str, default: Any | None = None) -> Any:
+        """Set the default for a key in the Parameterdict.
+
+        If key is in the ParameterDict, return its value.
+        If not, insert `key` with a parameter `default` and return `default`.
+        `default` defaults to `None`.
+
+        Args:
+            key (str): key to set default for
+            default (Any): the parameter set to the key
+        """
+        if key not in self:
+            self[key] = default
+        return self[key]
+
+    def clear(self) -> None:
+        """Remove all items from the ParameterDict."""
+        for k in self._keys.copy():
+            del self[k]
+
+    def pop(self, key: str) -> Any:
+        r"""Remove key from the ParameterDict and return its parameter.
+
+        Args:
+            key (str): key to pop from the ParameterDict
+        """
+        v = self[key]
+        del self[key]
+        return v
+
+    def popitem(self) -> tuple[str, Any]:
+        """Remove and return the last inserted `(key, parameter)` pair from the ParameterDict."""
+        k, _ = self._keys.popitem()
+        # We need the key in the _keys to be able to access/del
+        self._keys[k] = None
+        val = self[k]
+        del self[k]
+        return k, val
+
+    def get(self, key: str, default: Any | None = None) -> Any:
+        r"""Return the parameter associated with key if present. Otherwise return default if provided, None if not.
+
+        Args:
+            key (str): key to get from the ParameterDict
+            default (Parameter, optional): value to return if key not present
+        """
+        return self[key] if key in self else default  # noqa: SIM401
+
+    def fromkeys(
+        self, keys: Iterable[str], default: Any | None = None
+    ) -> ParameterDict:
+        r"""Return a new ParameterDict with the keys provided.
+
+        Args:
+            keys (iterable, string): keys to make the new ParameterDict from
+            default (Parameter, optional): value to set for all keys
+        """
+        return ParameterDict((k, default) for k in keys)
+
+    def keys(self) -> container_abcs.KeysView[str]:
+        r"""Return an iterable of the ParameterDict keys."""
+        return self._keys.keys()
+
+    def items(self) -> Iterable[tuple[str, Any]]:
+        r"""Return an iterable of the ParameterDict key/value pairs."""
+        return ((k, self[k]) for k in self._keys)
+
+    def values(self) -> Iterable[Any]:
+        r"""Return an iterable of the ParameterDict values."""
+        return (self[k] for k in self._keys)
+
+    def update(self, parameters: Mapping[str, Any] | ParameterDict) -> None:
+        r"""Update the :class:`~torch.nn.ParameterDict` with key-value pairs from ``parameters``, overwriting existing keys.
+
+        .. note::
+            If :attr:`parameters` is an ``OrderedDict``, a :class:`~torch.nn.ParameterDict`, or
+            an iterable of key-value pairs, the order of new elements in it is preserved.
+
+        Args:
+            parameters (iterable): a mapping (dictionary) from string to
+                :class:`~torch.nn.Parameter`, or an iterable of
+                key-value pairs of type (string, :class:`~torch.nn.Parameter`)
+        """
+        if not isinstance(parameters, container_abcs.Iterable):
+            raise TypeError(
+                "ParametersDict.update should be called with an "
+                "iterable of key/value pairs, but got " + type(parameters).__name__
+            )
+
+        if isinstance(parameters, (OrderedDict, ParameterDict)):
+            for key, parameter in parameters.items():
+                self[key] = parameter
+        elif isinstance(parameters, container_abcs.Mapping):
+            for key, parameter in sorted(parameters.items()):
+                self[key] = parameter
+        else:
+            for j, p in enumerate(parameters):
+                if not isinstance(p, container_abcs.Iterable):
+                    raise TypeError(
+                        "ParameterDict update sequence element "
+                        "#" + str(j) + " should be Iterable; is" + type(p).__name__
+                    )
+                # pyrefly: ignore [bad-argument-type]
+                if not len(p) == 2:
+                    raise ValueError(
+                        "ParameterDict update sequence element "
+                        # pyrefly: ignore [bad-argument-type]
+                        "#" + str(j) + " has length " + str(len(p)) + "; 2 is required"
+                    )
+                # parameters as length-2 list too cumbersome to type, see ModuleDict.update comment
+                self[p[0]] = p[1]  # type: ignore[assignment]
+
+    def extra_repr(self) -> str:
+        child_lines = []
+        for k, p in self.items():
+            if isinstance(p, torch.Tensor):
+                size_str = "x".join(str(size) for size in p.size())
+                if p.device.type in ["cuda", torch._C._get_privateuse1_backend_name()]:
+                    device_str = f" ({p.device})"
+                else:
+                    device_str = ""
+                parastr = "{} containing: [{} of size {}{}]".format(
+                    "Parameter" if isinstance(p, Parameter) else "Tensor",
+                    torch.typename(p),
+                    size_str,
+                    device_str,
+                )
+                child_lines.append("  (" + str(k) + "): " + parastr)
+            else:
+                child_lines.append(
+                    "  (" + str(k) + "): Object of type: " + type(p).__name__
+                )
+        tmpstr = "\n".join(child_lines)
+        return tmpstr
+
+    def __call__(self, input):
+        raise RuntimeError("ParameterDict should not be called.")
+
+    def __or__(self, other: ParameterDict) -> ParameterDict:
+        copy = self.copy()
+        copy.update(other)
+        return copy
+
+    def __ror__(self, other: ParameterDict) -> ParameterDict:
+        copy = other.copy()
+        copy.update(self)
+        return copy
+
+    def __ior__(self, other: ParameterDict) -> Self:
+        self.update(other)
+        return self
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/conv.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/conv.py
new file mode 100644
index 0000000000000000000000000000000000000000..bcc39e3e24f05015de75a68226a3499e54f69e11
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/conv.py
@@ -0,0 +1,1921 @@
+# mypy: allow-untyped-defs
+import math
+from typing import Literal
+from typing_extensions import deprecated
+
+import torch
+from torch import Tensor
+from torch._torch_docs import reproducibility_notes
+from torch.nn import functional as F, init
+from torch.nn.common_types import _size_1_t, _size_2_t, _size_3_t
+from torch.nn.parameter import Parameter, UninitializedParameter
+
+from .lazy import LazyModuleMixin
+from .module import Module
+from .utils import _pair, _reverse_repeat_tuple, _single, _triple
+
+
+__all__ = [
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "ConvTranspose1d",
+    "ConvTranspose2d",
+    "ConvTranspose3d",
+    "LazyConv1d",
+    "LazyConv2d",
+    "LazyConv3d",
+    "LazyConvTranspose1d",
+    "LazyConvTranspose2d",
+    "LazyConvTranspose3d",
+]
+
+convolution_notes = {
+    "groups_note": r"""* :attr:`groups` controls the connections between inputs and outputs.
+      :attr:`in_channels` and :attr:`out_channels` must both be divisible by
+      :attr:`groups`. For example,
+
+        * At groups=1, all inputs are convolved to all outputs.
+        * At groups=2, the operation becomes equivalent to having two conv
+          layers side by side, each seeing half the input channels
+          and producing half the output channels, and both subsequently
+          concatenated.
+        * At groups= :attr:`in_channels`, each input channel is convolved with
+          its own set of filters (of size
+          :math:`\frac{\text{out\_channels}}{\text{in\_channels}}`).""",
+    "depthwise_separable_note": r"""When `groups == in_channels` and `out_channels == K * in_channels`,
+        where `K` is a positive integer, this operation is also known as a "depthwise convolution".
+
+        In other words, for an input of size :math:`(N, C_{in}, L_{in})`,
+        a depthwise convolution with a depthwise multiplier `K` can be performed with the arguments
+        :math:`(C_\text{in}=C_\text{in}, C_\text{out}=C_\text{in} \times \text{K}, ..., \text{groups}=C_\text{in})`.""",
+}  # noqa: B950
+
+
+class _ConvNd(Module):
+    __constants__ = [
+        "stride",
+        "padding",
+        "dilation",
+        "groups",
+        "padding_mode",
+        "output_padding",
+        "in_channels",
+        "out_channels",
+        "kernel_size",
+    ]
+    __annotations__ = {"bias": torch.Tensor | None}
+
+    def _conv_forward(  # type: ignore[empty-body]
+        self, input: Tensor, weight: Tensor, bias: Tensor | None
+    ) -> Tensor: ...
+
+    in_channels: int
+    _reversed_padding_repeated_twice: list[int]
+    out_channels: int
+    kernel_size: tuple[int, ...]
+    stride: tuple[int, ...]
+    padding: str | tuple[int, ...]
+    dilation: tuple[int, ...]
+    transposed: bool
+    output_padding: tuple[int, ...]
+    groups: int
+    padding_mode: Literal["zeros", "reflect", "replicate", "circular"]
+    weight: Tensor
+    bias: Tensor | None
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: tuple[int, ...],
+        stride: tuple[int, ...],
+        padding: str | tuple[int, ...],
+        dilation: tuple[int, ...],
+        transposed: bool,
+        output_padding: tuple[int, ...],
+        groups: int,
+        bias: bool,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"],
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        if groups <= 0:
+            raise ValueError("groups must be a positive integer")
+        if in_channels % groups != 0:
+            raise ValueError("in_channels must be divisible by groups")
+        if out_channels % groups != 0:
+            raise ValueError("out_channels must be divisible by groups")
+        valid_padding_strings = {"same", "valid"}
+        if isinstance(padding, str):
+            if padding not in valid_padding_strings:
+                raise ValueError(
+                    f"Invalid padding string {padding!r}, should be one of {valid_padding_strings}"
+                )
+            if padding == "same" and any(s != 1 for s in stride):
+                raise ValueError(
+                    "padding='same' is not supported for strided convolutions"
+                )
+
+        valid_padding_modes = {"zeros", "reflect", "replicate", "circular"}
+        if padding_mode not in valid_padding_modes:
+            raise ValueError(
+                f"padding_mode must be one of {valid_padding_modes}, but got padding_mode='{padding_mode}'"
+            )
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.transposed = transposed
+        self.output_padding = output_padding
+        self.groups = groups
+        self.padding_mode = padding_mode
+        # `_reversed_padding_repeated_twice` is the padding to be passed to
+        # `F.pad` if needed (e.g., for non-zero padding types that are
+        # implemented as two ops: padding + conv). `F.pad` accepts paddings in
+        # reverse order than the dimension.
+        if isinstance(self.padding, str):
+            self._reversed_padding_repeated_twice = [0, 0] * len(kernel_size)
+            if padding == "same":
+                for d, k, i in zip(
+                    dilation,
+                    kernel_size,
+                    range(len(kernel_size) - 1, -1, -1),
+                    strict=False,
+                ):
+                    total_padding = d * (k - 1)
+                    left_pad = total_padding // 2
+                    self._reversed_padding_repeated_twice[2 * i] = left_pad
+                    self._reversed_padding_repeated_twice[2 * i + 1] = (
+                        total_padding - left_pad
+                    )
+        else:
+            self._reversed_padding_repeated_twice = _reverse_repeat_tuple(
+                self.padding, 2
+            )
+
+        if transposed:
+            self.weight = Parameter(
+                torch.empty(
+                    (in_channels, out_channels // groups, *kernel_size),
+                    **factory_kwargs,
+                )
+            )
+        else:
+            self.weight = Parameter(
+                torch.empty(
+                    (out_channels, in_channels // groups, *kernel_size),
+                    **factory_kwargs,
+                )
+            )
+        if bias:
+            self.bias = Parameter(torch.empty(out_channels, **factory_kwargs))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        # Use torch.no_grad() to allow in-place modification of weights/bias
+        with torch.no_grad():
+            # Check if the weight tensor is in channels_last_3d format to ensure numerical consistency
+            if not self.weight.is_contiguous():
+                # Use an empty contiguous buffer to optimize memory and ensure consistent initialization
+                temp_weight = torch.empty_like(
+                    self.weight, memory_format=torch.contiguous_format
+                )
+                init.kaiming_uniform_(temp_weight, a=math.sqrt(5))
+                # Copy the initialized values back to the original tensor
+                self.weight.copy_(temp_weight)
+            else:
+                # Standard initialization for default contiguous memory format
+                init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+            if self.bias is not None:
+                # Bias initialization remains independent of weight memory format
+                fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+                bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+                init.uniform_(self.bias, -bound, bound)
+
+    def extra_repr(self):
+        s = "{in_channels}, {out_channels}, kernel_size={kernel_size}, stride={stride}"
+        if self.padding != (0,) * len(self.padding):
+            s += ", padding={padding}"
+        if self.dilation != (1,) * len(self.dilation):
+            s += ", dilation={dilation}"
+        if self.output_padding != (0,) * len(self.output_padding):
+            s += ", output_padding={output_padding}"
+        if self.groups != 1:
+            s += ", groups={groups}"
+        if self.bias is None:
+            s += ", bias=False"
+        if self.padding_mode != "zeros":
+            s += ", padding_mode={padding_mode}"
+        return s.format(**self.__dict__)
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        if not hasattr(self, "padding_mode"):
+            self.padding_mode = "zeros"
+
+
+class Conv1d(_ConvNd):
+    __doc__ = (
+        r"""Applies a 1D convolution over an input signal composed of several input
+    planes.
+
+    In the simplest case, the output value of the layer with input size
+    :math:`(N, C_{\text{in}}, L)` and output :math:`(N, C_{\text{out}}, L_{\text{out}})` can be
+    precisely described as:
+
+    .. math::
+        \text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) +
+        \sum_{k = 0}^{C_{in} - 1} \text{weight}(C_{\text{out}_j}, k)
+        \star \text{input}(N_i, k)
+
+    where :math:`\star` is the valid `cross-correlation`_ operator,
+    :math:`N` is a batch size, :math:`C` denotes a number of channels,
+    :math:`L` is a length of signal sequence.
+    """
+        + r"""
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    * :attr:`stride` controls the stride for the cross-correlation, a single
+      number or a one-element tuple.
+
+    * :attr:`padding` controls the amount of padding applied to the input. It
+      can be either a string {{'valid', 'same'}} or a tuple of ints giving the
+      amount of implicit padding applied on both sides.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also
+      known as the \u00e0 trous algorithm. It is harder to describe, but this `link`_
+      has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+    {groups_note}
+
+    Note:
+        {depthwise_separable_note}
+    Note:
+        {cudnn_reproducibility_note}
+
+    Note:
+        ``padding='valid'`` is the same as no padding. ``padding='same'`` pads
+        the input so the output has the shape as the input. However, this mode
+        doesn't support any stride values other than 1.
+
+    Note:
+        This module supports complex data types i.e. ``complex32, complex64, complex128``.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to both sides of
+            the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel
+            elements. Default: 1
+        groups (int, optional): Number of blocked connections from input
+            channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the
+            output. Default: ``True``
+        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,
+            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
+
+    """.format(**reproducibility_notes, **convolution_notes)
+        + r"""
+
+    Shape:
+        - Input: :math:`(N, C_{in}, L_{in})` or :math:`(C_{in}, L_{in})`
+        - Output: :math:`(N, C_{out}, L_{out})` or :math:`(C_{out}, L_{out})`, where
+
+          .. math::
+              L_{out} = \left\lfloor\frac{L_{in} + 2 \times \text{padding} - \text{dilation}
+                        \times (\text{kernel\_size} - 1) - 1}{\text{stride}} + 1\right\rfloor
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape
+            :math:`(\text{out\_channels},
+            \frac{\text{in\_channels}}{\text{groups}}, \text{kernel\_size})`.
+            The values of these weights are sampled from
+            :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+            :math:`k = \frac{groups}{C_\text{in} * \text{kernel\_size}}`
+        bias (Tensor):   the learnable bias of the module of shape
+            (out_channels). If :attr:`bias` is ``True``, then the values of these weights are
+            sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+            :math:`k = \frac{groups}{C_\text{in} * \text{kernel\_size}}`
+
+    Examples::
+
+        >>> m = nn.Conv1d(16, 33, 3, stride=2)
+        >>> input = torch.randn(20, 16, 50)
+        >>> output = m(input)
+
+    .. _cross-correlation:
+        https://en.wikipedia.org/wiki/Cross-correlation
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+    """
+    )
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_1_t,
+        stride: _size_1_t = 1,
+        padding: str | _size_1_t = 0,
+        dilation: _size_1_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # we create new variables below to make mypy happy since kernel_size has
+        # type Union[int, Tuple[int]] and kernel_size_ has type Tuple[int]
+        kernel_size_ = _single(kernel_size)
+        stride_ = _single(stride)
+        padding_ = padding if isinstance(padding, str) else _single(padding)
+        dilation_ = _single(dilation)
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size_,
+            stride_,
+            padding_,
+            dilation_,
+            False,
+            _single(0),
+            groups,
+            bias,
+            padding_mode,
+            **factory_kwargs,
+        )
+
+    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Tensor | None):
+        if self.padding_mode != "zeros":
+            return F.conv1d(
+                F.pad(
+                    input, self._reversed_padding_repeated_twice, mode=self.padding_mode
+                ),
+                weight,
+                bias,
+                self.stride,
+                _single(0),
+                self.dilation,
+                self.groups,
+            )
+
+        return F.conv1d(
+            input, weight, bias, self.stride, self.padding, self.dilation, self.groups
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        return self._conv_forward(input, self.weight, self.bias)
+
+
+class Conv2d(_ConvNd):
+    __doc__ = (
+        r"""Applies a 2D convolution over an input signal composed of several input
+    planes.
+
+    In the simplest case, the output value of the layer with input size
+    :math:`(N, C_{\text{in}}, H, W)` and output :math:`(N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}})`
+    can be precisely described as:
+
+    .. math::
+        \text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) +
+        \sum_{k = 0}^{C_{\text{in}} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k)
+
+
+    where :math:`\star` is the valid 2D `cross-correlation`_ operator,
+    :math:`N` is a batch size, :math:`C_{\text{in}}` and :math:`C_{\text{out}}` correspond to
+    :attr:`in_channels` and :attr:`out_channels` respectively,
+    :math:`H` and :math:`W` are the input height and width in pixels.
+    See the Shape section below for how :math:`H_{\text{out}}` and :math:`W_{\text{out}}`
+    are derived from :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, and :attr:`dilation`.
+    """
+        + r"""
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    * :attr:`stride` controls the stride for the cross-correlation, a single
+      number or a tuple.
+
+    * :attr:`padding` controls the amount of padding applied to the input. It
+      can be either a string {{'valid', 'same'}} or an int / a tuple of ints giving the
+      amount of implicit padding applied on both sides.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also
+      known as the \u00e0 trous algorithm. It is harder to describe, but this `link`_
+      has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+
+    {groups_note}
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:
+
+        - a single ``int`` -- in which case the same value is used for the height and width dimension
+        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,
+          and the second `int` for the width dimension
+
+    Note:
+        {depthwise_separable_note}
+
+    Note:
+        {cudnn_reproducibility_note}
+
+    Note:
+        ``padding='valid'`` is the same as no padding. ``padding='same'`` pads
+        the input so the output has the shape as the input. However, this mode
+        doesn't support any stride values other than 1.
+
+    Note:
+        This module supports complex data types i.e. ``complex32, complex64, complex128``.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to all four sides of
+            the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input
+            channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the
+            output. Default: ``True``
+        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,
+            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
+    """.format(**reproducibility_notes, **convolution_notes)
+        + r"""
+
+    Shape:
+        - Input: :math:`(N, C_{in}, H_{in}, W_{in})` or :math:`(C_{in}, H_{in}, W_{in})`
+        - Output: :math:`(N, C_{out}, H_{out}, W_{out})` or :math:`(C_{out}, H_{out}, W_{out})`, where
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in}  + 2 \times \text{padding}[0] - \text{dilation}[0]
+                        \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in}  + 2 \times \text{padding}[1] - \text{dilation}[1]
+                        \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape
+            :math:`(\text{out\_channels}, \frac{\text{in\_channels}}{\text{groups}},`
+            :math:`\text{kernel\_size[0]}, \text{kernel\_size[1]})`.
+            The values of these weights are sampled from
+            :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+            :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
+        bias (Tensor):   the learnable bias of the module of shape
+            (out_channels). If :attr:`bias` is ``True``,
+            then the values of these weights are
+            sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+            :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
+
+    Examples:
+
+        >>> # With square kernels and equal stride
+        >>> m = nn.Conv2d(16, 33, 3, stride=2)
+        >>> # non-square kernels and unequal stride and with padding
+        >>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
+        >>> # non-square kernels and unequal stride and with padding and dilation
+        >>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))
+        >>> input = torch.randn(20, 16, 50, 100)
+        >>> output = m(input)
+
+    .. _cross-correlation:
+        https://en.wikipedia.org/wiki/Cross-correlation
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+    """
+    )
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_2_t,
+        stride: _size_2_t = 1,
+        padding: str | _size_2_t = 0,
+        dilation: _size_2_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        kernel_size_ = _pair(kernel_size)
+        stride_ = _pair(stride)
+        padding_ = padding if isinstance(padding, str) else _pair(padding)
+        dilation_ = _pair(dilation)
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size_,
+            stride_,
+            padding_,
+            dilation_,
+            False,
+            _pair(0),
+            groups,
+            bias,
+            padding_mode,
+            **factory_kwargs,
+        )
+
+    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Tensor | None):
+        if self.padding_mode != "zeros":
+            return F.conv2d(
+                F.pad(
+                    input, self._reversed_padding_repeated_twice, mode=self.padding_mode
+                ),
+                weight,
+                bias,
+                self.stride,
+                _pair(0),
+                self.dilation,
+                self.groups,
+            )
+
+        return F.conv2d(
+            input, weight, bias, self.stride, self.padding, self.dilation, self.groups
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        return self._conv_forward(input, self.weight, self.bias)
+
+
+class Conv3d(_ConvNd):
+    __doc__ = (
+        r"""Applies a 3D convolution over an input signal composed of several input
+    planes.
+
+    In the simplest case, the output value of the layer with input size :math:`(N, C_{in}, D, H, W)`
+    and output :math:`(N, C_{out}, D_{out}, H_{out}, W_{out})` can be precisely described as:
+
+    .. math::
+        out(N_i, C_{out_j}) = bias(C_{out_j}) +
+                                \sum_{k = 0}^{C_{in} - 1} weight(C_{out_j}, k) \star input(N_i, k)
+
+    where :math:`\star` is the valid 3D `cross-correlation`_ operator
+    """
+        + r"""
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    * :attr:`stride` controls the stride for the cross-correlation.
+
+    * :attr:`padding` controls the amount of padding applied to the input. It
+      can be either a string {{'valid', 'same'}} or a tuple of ints giving the
+      amount of implicit padding applied on both sides.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
+      It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+
+    {groups_note}
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:
+
+        - a single ``int`` -- in which case the same value is used for the depth, height and width dimension
+        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,
+          the second `int` for the height dimension and the third `int` for the width dimension
+
+    Note:
+        {depthwise_separable_note}
+
+    Note:
+        {cudnn_reproducibility_note}
+
+    Note:
+        ``padding='valid'`` is the same as no padding. ``padding='same'`` pads
+        the input so the output has the shape as the input. However, this mode
+        doesn't support any stride values other than 1.
+
+    Note:
+        This module supports complex data types i.e. ``complex32, complex64, complex128``.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to all six sides of
+            the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
+        padding_mode (str, optional): ``'zeros'``, ``'reflect'``, ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
+    """.format(**reproducibility_notes, **convolution_notes)
+        + r"""
+
+    Shape:
+        - Input: :math:`(N, C_{in}, D_{in}, H_{in}, W_{in})` or :math:`(C_{in}, D_{in}, H_{in}, W_{in})`
+        - Output: :math:`(N, C_{out}, D_{out}, H_{out}, W_{out})` or :math:`(C_{out}, D_{out}, H_{out}, W_{out})`,
+          where
+
+          .. math::
+              D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] - \text{dilation}[0]
+                    \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] - \text{dilation}[1]
+                    \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] - \text{dilation}[2]
+                    \times (\text{kernel\_size}[2] - 1) - 1}{\text{stride}[2]} + 1\right\rfloor
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape
+                         :math:`(\text{out\_channels}, \frac{\text{in\_channels}}{\text{groups}},`
+                         :math:`\text{kernel\_size[0]}, \text{kernel\_size[1]}, \text{kernel\_size[2]})`.
+                         The values of these weights are sampled from
+                         :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{2}\text{kernel\_size}[i]}`
+        bias (Tensor):   the learnable bias of the module of shape (out_channels). If :attr:`bias` is ``True``,
+                         then the values of these weights are
+                         sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{2}\text{kernel\_size}[i]}`
+
+    Examples::
+
+        >>> # With square kernels and equal stride
+        >>> m = nn.Conv3d(16, 33, 3, stride=2)
+        >>> # non-square kernels and unequal stride and with padding
+        >>> m = nn.Conv3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(4, 2, 0))
+        >>> input = torch.randn(20, 16, 10, 50, 100)
+        >>> output = m(input)
+
+    .. _cross-correlation:
+        https://en.wikipedia.org/wiki/Cross-correlation
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+    """
+    )
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_3_t,
+        stride: _size_3_t = 1,
+        padding: str | _size_3_t = 0,
+        dilation: _size_3_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        kernel_size_ = _triple(kernel_size)
+        stride_ = _triple(stride)
+        padding_ = padding if isinstance(padding, str) else _triple(padding)
+        dilation_ = _triple(dilation)
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size_,
+            stride_,
+            padding_,
+            dilation_,
+            False,
+            _triple(0),
+            groups,
+            bias,
+            padding_mode,
+            **factory_kwargs,
+        )
+
+    def _conv_forward(self, input: Tensor, weight: Tensor, bias: Tensor | None):
+        if self.padding_mode != "zeros":
+            return F.conv3d(
+                F.pad(
+                    input, self._reversed_padding_repeated_twice, mode=self.padding_mode
+                ),
+                weight,
+                bias,
+                self.stride,
+                _triple(0),
+                self.dilation,
+                self.groups,
+            )
+
+        return F.conv3d(
+            input, weight, bias, self.stride, self.padding, self.dilation, self.groups
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        return self._conv_forward(input, self.weight, self.bias)
+
+
+class _ConvTransposeNd(_ConvNd):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        transposed,
+        output_padding,
+        groups,
+        bias,
+        padding_mode,
+        device=None,
+        dtype=None,
+    ) -> None:
+        if padding_mode != "zeros":
+            raise ValueError(
+                f'Only "zeros" padding mode is supported for {self.__class__.__name__}'
+            )
+
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            transposed,
+            output_padding,
+            groups,
+            bias,
+            padding_mode,
+            **factory_kwargs,
+        )
+
+    # dilation being an optional parameter is for backwards
+    # compatibility
+    def _output_padding(
+        self,
+        input: Tensor,
+        output_size: list[int] | None,
+        stride: list[int],
+        padding: list[int],
+        kernel_size: list[int],
+        num_spatial_dims: int,
+        dilation: list[int] | None = None,
+    ) -> list[int]:
+        if output_size is None:
+            ret = _single(self.output_padding)  # converting to list if was not already
+        else:
+            has_batch_dim = input.dim() == num_spatial_dims + 2
+            num_non_spatial_dims = 2 if has_batch_dim else 1
+            if len(output_size) == num_non_spatial_dims + num_spatial_dims:
+                output_size = output_size[num_non_spatial_dims:]
+            if len(output_size) != num_spatial_dims:
+                raise ValueError(
+                    f"ConvTranspose{num_spatial_dims}D: for {input.dim()}D input, output_size must have {num_spatial_dims} "
+                    f"or {num_non_spatial_dims + num_spatial_dims} elements (got {len(output_size)})"
+                )
+
+            min_sizes = torch.jit.annotate(list[int], [])
+            max_sizes = torch.jit.annotate(list[int], [])
+            for d in range(num_spatial_dims):
+                dim_size = (
+                    (input.size(d + num_non_spatial_dims) - 1) * stride[d]
+                    - 2 * padding[d]
+                    + (dilation[d] if dilation is not None else 1)
+                    * (kernel_size[d] - 1)
+                    + 1
+                )
+                min_sizes.append(dim_size)
+                max_sizes.append(min_sizes[d] + stride[d] - 1)
+
+            for i in range(len(output_size)):
+                size = output_size[i]
+                min_size = min_sizes[i]
+                max_size = max_sizes[i]
+                if size < min_size or size > max_size:
+                    raise ValueError(
+                        f"requested an output size of {output_size}, but valid sizes range "
+                        f"from {min_sizes} to {max_sizes} (for an input of {input.size()[2:]})"
+                    )
+
+            res = torch.jit.annotate(list[int], [])
+            for d in range(num_spatial_dims):
+                res.append(output_size[d] - min_sizes[d])
+
+            ret = res
+        return ret
+
+
+class ConvTranspose1d(_ConvTransposeNd):
+    __doc__ = (
+        r"""Applies a 1D transposed convolution operator over an input image
+    composed of several input planes.
+
+    This module can be seen as the gradient of Conv1d with respect to its input.
+    It is also known as a fractionally-strided convolution or
+    a deconvolution (although it is not an actual deconvolution operation as it does
+    not compute a true inverse of convolution). For more information, see the visualizations
+    `here`_ and the `Deconvolutional Networks`_ paper.
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    * :attr:`stride` controls the stride for the cross-correlation.
+
+    * :attr:`padding` controls the amount of implicit zero padding on both
+      sides for ``dilation * (kernel_size - 1) - padding`` number of points. See note
+      below for details.
+
+    * :attr:`output_padding` controls the additional size added to one side
+      of the output shape. See note below for details.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
+      It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+    {groups_note}
+
+    Note:
+        The :attr:`padding` argument effectively adds ``dilation * (kernel_size - 1) - padding``
+        amount of zero padding to both sizes of the input. This is set so that
+        when a :class:`~torch.nn.Conv1d` and a :class:`~torch.nn.ConvTranspose1d`
+        are initialized with same parameters, they are inverses of each other in
+        regard to the input and output shapes. However, when ``stride > 1``,
+        :class:`~torch.nn.Conv1d` maps multiple input shapes to the same output
+        shape. :attr:`output_padding` is provided to resolve this ambiguity by
+        effectively increasing the calculated output shape on one side. Note
+        that :attr:`output_padding` is only used to find output shape, but does
+        not actually add zero-padding to output.
+
+    Note:
+        In some circumstances when using the CUDA backend with CuDNN, this operator
+        may select a nondeterministic algorithm to increase performance. If this is
+        undesirable, you can try to make the operation deterministic (potentially at
+        a performance cost) by setting ``torch.backends.cudnn.deterministic =
+        True``.
+        Please see the notes on :doc:`/notes/randomness` for background.
+
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding
+            will be added to both sides of the input. Default: 0
+        output_padding (int or tuple, optional): Additional size added to one side
+            of the output shape. Default: 0
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+    """.format(**reproducibility_notes, **convolution_notes)
+        + r"""
+
+    Shape:
+        - Input: :math:`(N, C_{in}, L_{in})` or :math:`(C_{in}, L_{in})`
+        - Output: :math:`(N, C_{out}, L_{out})` or :math:`(C_{out}, L_{out})`, where
+
+          .. math::
+              L_{out} = (L_{in} - 1) \times \text{stride} - 2 \times \text{padding} + \text{dilation}
+                        \times (\text{kernel\_size} - 1) + \text{output\_padding} + 1
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape
+                         :math:`(\text{in\_channels}, \frac{\text{out\_channels}}{\text{groups}},`
+                         :math:`\text{kernel\_size})`.
+                         The values of these weights are sampled from
+                         :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{out} * \text{kernel\_size}}`
+        bias (Tensor):   the learnable bias of the module of shape (out_channels).
+                         If :attr:`bias` is ``True``, then the values of these weights are
+                         sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{out} * \text{kernel\_size}}`
+
+    Examples::
+
+        >>> # With square kernels and equal stride
+        >>> m = nn.ConvTranspose1d(16, 33, 3, stride=2)
+        >>> input = torch.randn(20, 16, 50)
+        >>> output = m(input)
+        >>> # exact output size can be also specified as an argument
+        >>> input = torch.randn(1, 16, 12)
+        >>> downsample = nn.Conv1d(16, 16, 3, stride=2, padding=1)
+        >>> upsample = nn.ConvTranspose1d(16, 16, 3, stride=2, padding=1)
+        >>> h = downsample(input)
+        >>> h.size()
+        torch.Size([1, 16, 6])
+        >>> output = upsample(h, output_size=input.size())
+        >>> output.size()
+        torch.Size([1, 16, 12])
+
+    .. _`here`:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+
+    .. _`Deconvolutional Networks`:
+        https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf
+    """
+    )
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_1_t,
+        stride: _size_1_t = 1,
+        padding: _size_1_t = 0,
+        output_padding: _size_1_t = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: _size_1_t = 1,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        kernel_size = _single(kernel_size)
+        stride = _single(stride)
+        padding = _single(padding)
+        dilation = _single(dilation)
+        output_padding = _single(output_padding)
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            True,
+            output_padding,
+            groups,
+            bias,
+            padding_mode,
+            **factory_kwargs,
+        )
+
+    def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor:
+        if self.padding_mode != "zeros":
+            raise ValueError(
+                "Only `zeros` padding mode is supported for ConvTranspose1d"
+            )
+
+        if not isinstance(self.padding, tuple):
+            raise AssertionError("self.padding must be a tuple")
+        # One cannot replace List by Tuple or Sequence in "_output_padding" because
+        # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.
+        num_spatial_dims = 1
+        output_padding = self._output_padding(
+            input,
+            output_size,
+            self.stride,  # type: ignore[arg-type]
+            self.padding,  # type: ignore[arg-type]
+            self.kernel_size,  # type: ignore[arg-type]
+            num_spatial_dims,
+            self.dilation,  # type: ignore[arg-type]
+        )
+        return F.conv_transpose1d(
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            output_padding,
+            self.groups,
+            self.dilation,
+        )
+
+
+class ConvTranspose2d(_ConvTransposeNd):
+    __doc__ = (
+        r"""Applies a 2D transposed convolution operator over an input image
+    composed of several input planes.
+
+    This module can be seen as the gradient of Conv2d with respect to its input.
+    It is also known as a fractionally-strided convolution or
+    a deconvolution (although it is not an actual deconvolution operation as it does
+    not compute a true inverse of convolution). For more information, see the visualizations
+    `here`_ and the `Deconvolutional Networks`_ paper.
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    * :attr:`stride` controls the stride for the cross-correlation. When stride > 1, ConvTranspose2d inserts zeros between input
+      elements along the spatial dimensions before applying the convolution kernel. This zero-insertion operation is the standard
+      behavior of transposed convolutions, which can increase the spatial resolution and is equivalent to a learnable
+      upsampling operation.
+
+    * :attr:`padding` controls the amount of implicit zero padding on both
+      sides for ``dilation * (kernel_size - 1) - padding`` number of points. See note
+      below for details.
+
+    * :attr:`output_padding` controls the additional size added to one side
+      of the output shape. See note below for details.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
+      It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+    {groups_note}
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`output_padding`
+    can either be:
+
+        - a single ``int`` -- in which case the same value is used for the height and width dimensions
+        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,
+          and the second `int` for the width dimension
+
+    Note:
+        The :attr:`padding` argument effectively adds ``dilation * (kernel_size - 1) - padding``
+        amount of zero padding to both sizes of the input. This is set so that
+        when a :class:`~torch.nn.Conv2d` and a :class:`~torch.nn.ConvTranspose2d`
+        are initialized with same parameters, they are inverses of each other in
+        regard to the input and output shapes. However, when ``stride > 1``,
+        :class:`~torch.nn.Conv2d` maps multiple input shapes to the same output
+        shape. :attr:`output_padding` is provided to resolve this ambiguity by
+        effectively increasing the calculated output shape on one side. Note
+        that :attr:`output_padding` is only used to find output shape, but does
+        not actually add zero-padding to output.
+
+    Note:
+        {cudnn_reproducibility_note}
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding
+            will be added to both sides of each dimension in the input. Default: 0
+        output_padding (int or tuple, optional): Additional size added to one side
+            of each dimension in the output shape. Default: 0
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+    """.format(**reproducibility_notes, **convolution_notes)
+        + r"""
+
+    Shape:
+        - Input: :math:`(N, C_{in}, H_{in}, W_{in})` or :math:`(C_{in}, H_{in}, W_{in})`
+        - Output: :math:`(N, C_{out}, H_{out}, W_{out})` or :math:`(C_{out}, H_{out}, W_{out})`, where
+
+        .. math::
+              H_{out} = (H_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] + \text{dilation}[0]
+                        \times (\text{kernel\_size}[0] - 1) + \text{output\_padding}[0] + 1
+        .. math::
+              W_{out} = (W_{in} - 1) \times \text{stride}[1] - 2 \times \text{padding}[1] + \text{dilation}[1]
+                        \times (\text{kernel\_size}[1] - 1) + \text{output\_padding}[1] + 1
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape
+                         :math:`(\text{in\_channels}, \frac{\text{out\_channels}}{\text{groups}},`
+                         :math:`\text{kernel\_size[0]}, \text{kernel\_size[1]})`.
+                         The values of these weights are sampled from
+                         :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{out} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
+        bias (Tensor):   the learnable bias of the module of shape (out_channels)
+                         If :attr:`bias` is ``True``, then the values of these weights are
+                         sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{out} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
+
+    Examples::
+
+        >>> # With square kernels and equal stride
+        >>> m = nn.ConvTranspose2d(16, 33, 3, stride=2)
+        >>> # non-square kernels and unequal stride and with padding
+        >>> m = nn.ConvTranspose2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
+        >>> input = torch.randn(20, 16, 50, 100)
+        >>> output = m(input)
+        >>> # exact output size can be also specified as an argument
+        >>> input = torch.randn(1, 16, 12, 12)
+        >>> downsample = nn.Conv2d(16, 16, 3, stride=2, padding=1)
+        >>> upsample = nn.ConvTranspose2d(16, 16, 3, stride=2, padding=1)
+        >>> h = downsample(input)
+        >>> h.size()
+        torch.Size([1, 16, 6, 6])
+        >>> output = upsample(h, output_size=input.size())
+        >>> output.size()
+        torch.Size([1, 16, 12, 12])
+
+    .. _`here`:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+
+    .. _`Deconvolutional Networks`:
+        https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf
+    """
+    )
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_2_t,
+        stride: _size_2_t = 1,
+        padding: _size_2_t = 0,
+        output_padding: _size_2_t = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: _size_2_t = 1,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        kernel_size = _pair(kernel_size)
+        stride = _pair(stride)
+        padding = _pair(padding)
+        dilation = _pair(dilation)
+        output_padding = _pair(output_padding)
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            True,
+            output_padding,
+            groups,
+            bias,
+            padding_mode,
+            **factory_kwargs,
+        )
+
+    def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor:
+        """
+        Performs the forward pass.
+
+        Attributes:
+            input (Tensor): The input tensor.
+            output_size (list[int], optional): A list of integers representing
+                the size of the output tensor. Default is None.
+        """
+        if self.padding_mode != "zeros":
+            raise ValueError(
+                "Only `zeros` padding mode is supported for ConvTranspose2d"
+            )
+
+        if not isinstance(self.padding, tuple):
+            raise AssertionError("self.padding must be a tuple")
+        # One cannot replace List by Tuple or Sequence in "_output_padding" because
+        # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.
+        num_spatial_dims = 2
+        output_padding = self._output_padding(
+            input,
+            output_size,
+            self.stride,  # type: ignore[arg-type]
+            self.padding,  # type: ignore[arg-type]
+            self.kernel_size,  # type: ignore[arg-type]
+            num_spatial_dims,
+            self.dilation,  # type: ignore[arg-type]
+        )
+
+        return F.conv_transpose2d(
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            output_padding,
+            self.groups,
+            self.dilation,
+        )
+
+
+class ConvTranspose3d(_ConvTransposeNd):
+    __doc__ = (
+        r"""Applies a 3D transposed convolution operator over an input image composed of several input
+    planes.
+    The transposed convolution operator multiplies each input value element-wise by a learnable kernel,
+    and sums over the outputs from all input feature planes.
+
+    This module can be seen as the gradient of Conv3d with respect to its input.
+    It is also known as a fractionally-strided convolution or
+    a deconvolution (although it is not an actual deconvolution operation as it does
+    not compute a true inverse of convolution). For more information, see the visualizations
+    `here`_ and the `Deconvolutional Networks`_ paper.
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    * :attr:`stride` controls the stride for the cross-correlation.
+
+    * :attr:`padding` controls the amount of implicit zero padding on both
+      sides for ``dilation * (kernel_size - 1) - padding`` number of points. See note
+      below for details.
+
+    * :attr:`output_padding` controls the additional size added to one side
+      of the output shape. See note below for details.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
+      It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+    {groups_note}
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`output_padding`
+    can either be:
+
+        - a single ``int`` -- in which case the same value is used for the depth, height and width dimensions
+        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,
+          the second `int` for the height dimension and the third `int` for the width dimension
+
+    Note:
+        The :attr:`padding` argument effectively adds ``dilation * (kernel_size - 1) - padding``
+        amount of zero padding to both sizes of the input. This is set so that
+        when a :class:`~torch.nn.Conv3d` and a :class:`~torch.nn.ConvTranspose3d`
+        are initialized with same parameters, they are inverses of each other in
+        regard to the input and output shapes. However, when ``stride > 1``,
+        :class:`~torch.nn.Conv3d` maps multiple input shapes to the same output
+        shape. :attr:`output_padding` is provided to resolve this ambiguity by
+        effectively increasing the calculated output shape on one side. Note
+        that :attr:`output_padding` is only used to find output shape, but does
+        not actually add zero-padding to output.
+
+    Note:
+        {cudnn_reproducibility_note}
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding
+            will be added to both sides of each dimension in the input. Default: 0
+        output_padding (int or tuple, optional): Additional size added to one side
+            of each dimension in the output shape. Default: 0
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+    """.format(**reproducibility_notes, **convolution_notes)
+        + r"""
+
+    Shape:
+        - Input: :math:`(N, C_{in}, D_{in}, H_{in}, W_{in})` or :math:`(C_{in}, D_{in}, H_{in}, W_{in})`
+        - Output: :math:`(N, C_{out}, D_{out}, H_{out}, W_{out})` or
+          :math:`(C_{out}, D_{out}, H_{out}, W_{out})`, where
+
+        .. math::
+              D_{out} = (D_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] + \text{dilation}[0]
+                        \times (\text{kernel\_size}[0] - 1) + \text{output\_padding}[0] + 1
+        .. math::
+              H_{out} = (H_{in} - 1) \times \text{stride}[1] - 2 \times \text{padding}[1] + \text{dilation}[1]
+                        \times (\text{kernel\_size}[1] - 1) + \text{output\_padding}[1] + 1
+        .. math::
+              W_{out} = (W_{in} - 1) \times \text{stride}[2] - 2 \times \text{padding}[2] + \text{dilation}[2]
+                        \times (\text{kernel\_size}[2] - 1) + \text{output\_padding}[2] + 1
+
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape
+                         :math:`(\text{in\_channels}, \frac{\text{out\_channels}}{\text{groups}},`
+                         :math:`\text{kernel\_size[0]}, \text{kernel\_size[1]}, \text{kernel\_size[2]})`.
+                         The values of these weights are sampled from
+                         :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{out} * \prod_{i=0}^{2}\text{kernel\_size}[i]}`
+        bias (Tensor):   the learnable bias of the module of shape (out_channels)
+                         If :attr:`bias` is ``True``, then the values of these weights are
+                         sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                         :math:`k = \frac{groups}{C_\text{out} * \prod_{i=0}^{2}\text{kernel\_size}[i]}`
+
+    Examples::
+
+        >>> # With square kernels and equal stride
+        >>> m = nn.ConvTranspose3d(16, 33, 3, stride=2)
+        >>> # non-square kernels and unequal stride and with padding
+        >>> m = nn.ConvTranspose3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(0, 4, 2))
+        >>> input = torch.randn(20, 16, 10, 50, 100)
+        >>> output = m(input)
+
+    .. _`here`:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+
+    .. _`Deconvolutional Networks`:
+        https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf
+    """
+    )
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_3_t,
+        stride: _size_3_t = 1,
+        padding: _size_3_t = 0,
+        output_padding: _size_3_t = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: _size_3_t = 1,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        kernel_size = _triple(kernel_size)
+        stride = _triple(stride)
+        padding = _triple(padding)
+        dilation = _triple(dilation)
+        output_padding = _triple(output_padding)
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            True,
+            output_padding,
+            groups,
+            bias,
+            padding_mode,
+            **factory_kwargs,
+        )
+
+    def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor:
+        if self.padding_mode != "zeros":
+            raise ValueError(
+                "Only `zeros` padding mode is supported for ConvTranspose3d"
+            )
+
+        if not isinstance(self.padding, tuple):
+            raise AssertionError("self.padding must be a tuple")
+        # One cannot replace List by Tuple or Sequence in "_output_padding" because
+        # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.
+        num_spatial_dims = 3
+        output_padding = self._output_padding(
+            input,
+            output_size,
+            self.stride,  # type: ignore[arg-type]
+            self.padding,  # type: ignore[arg-type]
+            self.kernel_size,  # type: ignore[arg-type]
+            num_spatial_dims,
+            self.dilation,  # type: ignore[arg-type]
+        )
+
+        return F.conv_transpose3d(
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            output_padding,
+            self.groups,
+            self.dilation,
+        )
+
+
+# TODO: Deprecate and remove the following alias `_ConvTransposeMixin`.
+#
+# `_ConvTransposeMixin` was a mixin that was removed.  It is meant to be used
+# with `_ConvNd` to construct actual module classes that implements conv
+# transpose ops:
+#
+#   class MyConvTranspose(_ConvNd, _ConvTransposeMixin):
+#       ...
+#
+# In PyTorch, it has been replaced by `_ConvTransposeNd`, which is a proper
+# subclass of `_ConvNd`.  However, some user code in the wild still (incorrectly)
+# use the internal class `_ConvTransposeMixin`.  Hence, we provide this alias
+# for BC, because it is cheap and easy for us to do so, even though that
+# `_ConvTransposeNd` is really not a mixin anymore (but multiple inheritance as
+# above would still work).
+class _ConvTransposeMixin(_ConvTransposeNd):
+    @deprecated(
+        "`_ConvTransposeMixin` is a deprecated internal class. "
+        "Please consider using public APIs.",
+        category=FutureWarning,
+    )
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+
+# TODO: Conv2dLocal
+# TODO: Conv2dMap
+# TODO: ConvTranspose2dMap
+
+
+class _LazyConvXdMixin(LazyModuleMixin):
+    groups: int
+    transposed: bool
+    in_channels: int
+    out_channels: int
+    kernel_size: tuple[int, ...]
+    weight: UninitializedParameter
+    bias: UninitializedParameter
+
+    def reset_parameters(self) -> None:
+        # has_uninitialized_params is defined in parent class and it is using a protocol on self
+        if not self.has_uninitialized_params() and self.in_channels != 0:  # type: ignore[misc]
+            # "type:ignore[..]" is required because mypy thinks that "reset_parameters" is undefined
+            # in super class. Turns out that it is defined in _ConvND which is inherited by any class
+            # that also inherits _LazyConvXdMixin
+            super().reset_parameters()  # type: ignore[misc]
+
+    # Signature of "initialize_parameters" is incompatible with the definition in supertype LazyModuleMixin
+    def initialize_parameters(self, input: Tensor, *args, **kwargs) -> None:  # type: ignore[override]
+        # defined by parent class but using a protocol
+        if self.has_uninitialized_params():  # type: ignore[misc]
+            self.in_channels = self._get_in_channels(input)
+            if self.in_channels % self.groups != 0:
+                raise ValueError("in_channels must be divisible by groups")
+            if not isinstance(self.weight, UninitializedParameter):
+                raise AssertionError("self.weight must be an UninitializedParameter")
+            if self.transposed:
+                self.weight.materialize(
+                    (
+                        self.in_channels,
+                        self.out_channels // self.groups,
+                        *self.kernel_size,
+                    )
+                )
+            else:
+                self.weight.materialize(
+                    (
+                        self.out_channels,
+                        self.in_channels // self.groups,
+                        *self.kernel_size,
+                    )
+                )
+            if self.bias is not None:
+                if not isinstance(self.bias, UninitializedParameter):
+                    raise AssertionError("self.bias must be an UninitializedParameter")
+                self.bias.materialize((self.out_channels,))
+            self.reset_parameters()
+
+    # Function to extract in_channels from first input.
+    def _get_in_channels(self, input: Tensor) -> int:
+        num_spatial_dims = self._get_num_spatial_dims()
+        num_dims_no_batch = num_spatial_dims + 1  # +1 for channels dim
+        num_dims_batch = num_dims_no_batch + 1
+        if input.dim() not in (num_dims_no_batch, num_dims_batch):
+            raise RuntimeError(
+                f"Expected {num_dims_no_batch}D (unbatched) or {num_dims_batch}D (batched) input "
+                f"to {self.__class__.__name__}, but "
+                f"got input of size: {input.shape}"
+            )
+        return input.shape[1] if input.dim() == num_dims_batch else input.shape[0]
+
+    # Function to return the number of spatial dims expected for inputs to the module.
+    # This is expected to be implemented by subclasses.
+    def _get_num_spatial_dims(self) -> int:
+        raise NotImplementedError
+
+
+# LazyConv1d defines weight as a Tensor but derived class defines it as UninitializeParameter
+class LazyConv1d(_LazyConvXdMixin, Conv1d):  # type: ignore[misc]
+    r"""A :class:`torch.nn.Conv1d` module with lazy initialization of the ``in_channels`` argument.
+
+    The ``in_channels`` argument of the :class:`Conv1d` is inferred from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight` and `bias`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of
+            the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel
+            elements. Default: 1
+        groups (int, optional): Number of blocked connections from input
+            channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the
+            output. Default: ``True``
+        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,
+            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
+
+    .. seealso:: :class:`torch.nn.Conv1d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`
+    """
+
+    # super class define this variable as None. "type: ignore[..] is required
+    # since we are redefining the variable.
+    cls_to_become = Conv1d  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        out_channels: int,
+        kernel_size: _size_1_t,
+        stride: _size_1_t = 1,
+        padding: _size_1_t = 0,
+        dilation: _size_1_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(
+            0,
+            0,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            # bias is hardcoded to False to avoid creating tensor
+            # that will soon be overwritten.
+            False,
+            padding_mode,
+            **factory_kwargs,
+        )
+        # pyrefly: ignore [bad-override, unexpected-keyword]
+        self.weight = UninitializedParameter(**factory_kwargs)
+        self.out_channels = out_channels
+        if bias:
+            # pyrefly: ignore [bad-override, unexpected-keyword]
+            self.bias = UninitializedParameter(**factory_kwargs)
+
+    def _get_num_spatial_dims(self) -> int:
+        return 1
+
+
+# LazyConv2d defines weight as a Tensor but derived class defines it as UninitializeParameter
+class LazyConv2d(_LazyConvXdMixin, Conv2d):  # type: ignore[misc]
+    r"""A :class:`torch.nn.Conv2d` module with lazy initialization of the ``in_channels`` argument.
+
+    The ``in_channels`` argument of the :class:`Conv2d` that is inferred from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight` and `bias`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of
+            the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel
+            elements. Default: 1
+        groups (int, optional): Number of blocked connections from input
+            channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the
+            output. Default: ``True``
+        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,
+            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
+
+    .. seealso:: :class:`torch.nn.Conv2d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`
+    """
+
+    # super class define this variable as None. "type: ignore[..] is required
+    # since we are redefining the variable.
+    cls_to_become = Conv2d  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        out_channels: int,
+        kernel_size: _size_2_t,
+        stride: _size_2_t = 1,
+        padding: _size_2_t = 0,
+        dilation: _size_2_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(
+            0,
+            0,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            # bias is hardcoded to False to avoid creating tensor
+            # that will soon be overwritten.
+            False,
+            padding_mode,
+            **factory_kwargs,
+        )
+        # pyrefly: ignore [bad-override, unexpected-keyword]
+        self.weight = UninitializedParameter(**factory_kwargs)
+        self.out_channels = out_channels
+        if bias:
+            # pyrefly: ignore [bad-override, unexpected-keyword]
+            self.bias = UninitializedParameter(**factory_kwargs)
+
+    def _get_num_spatial_dims(self) -> int:
+        return 2
+
+
+# LazyConv3d defines weight as a Tensor but derived class defines it as UninitializeParameter
+class LazyConv3d(_LazyConvXdMixin, Conv3d):  # type: ignore[misc]
+    r"""A :class:`torch.nn.Conv3d` module with lazy initialization of the ``in_channels`` argument.
+
+    The ``in_channels`` argument of the :class:`Conv3d` that is inferred from
+    the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight` and `bias`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of
+            the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel
+            elements. Default: 1
+        groups (int, optional): Number of blocked connections from input
+            channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the
+            output. Default: ``True``
+        padding_mode (str, optional): ``'zeros'``, ``'reflect'``,
+            ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
+
+    .. seealso:: :class:`torch.nn.Conv3d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`
+    """
+
+    # super class define this variable as None. "type: ignore[..] is required
+    # since we are redefining the variable.
+    cls_to_become = Conv3d  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        out_channels: int,
+        kernel_size: _size_3_t,
+        stride: _size_3_t = 1,
+        padding: _size_3_t = 0,
+        dilation: _size_3_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(
+            0,
+            0,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            # bias is hardcoded to False to avoid creating tensor
+            # that will soon be overwritten.
+            False,
+            padding_mode,
+            **factory_kwargs,
+        )
+        # pyrefly: ignore [bad-override, unexpected-keyword]
+        self.weight = UninitializedParameter(**factory_kwargs)
+        self.out_channels = out_channels
+        if bias:
+            # pyrefly: ignore [bad-override, unexpected-keyword]
+            self.bias = UninitializedParameter(**factory_kwargs)
+
+    def _get_num_spatial_dims(self) -> int:
+        return 3
+
+
+# LazyConvTranspose1d defines weight as a Tensor but derived class defines it as UninitializeParameter
+class LazyConvTranspose1d(_LazyConvXdMixin, ConvTranspose1d):  # type: ignore[misc]
+    r"""A :class:`torch.nn.ConvTranspose1d` module with lazy initialization of the ``in_channels`` argument.
+
+    The ``in_channels`` argument of the :class:`ConvTranspose1d` that is inferred from
+    the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight` and `bias`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding
+            will be added to both sides of the input. Default: 0
+        output_padding (int or tuple, optional): Additional size added to one side
+            of the output shape. Default: 0
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+
+    .. seealso:: :class:`torch.nn.ConvTranspose1d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`
+    """
+
+    # super class define this variable as None. "type: ignore[..] is required
+    # since we are redefining the variable.
+    cls_to_become = ConvTranspose1d  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        out_channels: int,
+        kernel_size: _size_1_t,
+        stride: _size_1_t = 1,
+        padding: _size_1_t = 0,
+        output_padding: _size_1_t = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: _size_1_t = 1,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(
+            0,
+            0,
+            kernel_size,
+            stride,
+            padding,
+            output_padding,
+            groups,
+            # bias is hardcoded to False to avoid creating tensor
+            # that will soon be overwritten.
+            False,
+            dilation,
+            padding_mode,
+            **factory_kwargs,
+        )
+        # pyrefly: ignore [bad-override, unexpected-keyword]
+        self.weight = UninitializedParameter(**factory_kwargs)
+        self.out_channels = out_channels
+        if bias:
+            # pyrefly: ignore [bad-override, unexpected-keyword]
+            self.bias = UninitializedParameter(**factory_kwargs)
+
+    def _get_num_spatial_dims(self) -> int:
+        return 1
+
+
+# LazyConvTranspose2d defines weight as a Tensor but derived class defines it as UninitializeParameter
+class LazyConvTranspose2d(_LazyConvXdMixin, ConvTranspose2d):  # type: ignore[misc]
+    r"""A :class:`torch.nn.ConvTranspose2d` module with lazy initialization of the ``in_channels`` argument.
+
+    The ``in_channels`` argument of the :class:`ConvTranspose2d` is inferred from
+    the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight` and `bias`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding
+            will be added to both sides of each dimension in the input. Default: 0
+        output_padding (int or tuple, optional): Additional size added to one side
+            of each dimension in the output shape. Default: 0
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+
+    .. seealso:: :class:`torch.nn.ConvTranspose2d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`
+    """
+
+    # super class define this variable as None. "type: ignore[..] is required
+    # since we are redefining the variable.
+    cls_to_become = ConvTranspose2d  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        out_channels: int,
+        kernel_size: _size_2_t,
+        stride: _size_2_t = 1,
+        padding: _size_2_t = 0,
+        output_padding: _size_2_t = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: int = 1,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(
+            0,
+            0,
+            kernel_size,
+            stride,
+            padding,
+            output_padding,
+            groups,
+            # bias is hardcoded to False to avoid creating tensor
+            # that will soon be overwritten.
+            False,
+            dilation,
+            padding_mode,
+            **factory_kwargs,
+        )
+        # pyrefly: ignore [bad-override, unexpected-keyword]
+        self.weight = UninitializedParameter(**factory_kwargs)
+        self.out_channels = out_channels
+        if bias:
+            # pyrefly: ignore [bad-override, unexpected-keyword]
+            self.bias = UninitializedParameter(**factory_kwargs)
+
+    def _get_num_spatial_dims(self) -> int:
+        return 2
+
+
+# LazyConvTranspose3d defines weight as a Tensor but derived class defines it as UninitializeParameter
+class LazyConvTranspose3d(_LazyConvXdMixin, ConvTranspose3d):  # type: ignore[misc]
+    r"""A :class:`torch.nn.ConvTranspose3d` module with lazy initialization of the ``in_channels`` argument.
+
+    The ``in_channels`` argument of the :class:`ConvTranspose3d` is inferred from
+    the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight` and `bias`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): ``dilation * (kernel_size - 1) - padding`` zero-padding
+            will be added to both sides of each dimension in the input. Default: 0
+        output_padding (int or tuple, optional): Additional size added to one side
+            of each dimension in the output shape. Default: 0
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
+        dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
+
+    .. seealso:: :class:`torch.nn.ConvTranspose3d` and :class:`torch.nn.modules.lazy.LazyModuleMixin`
+    """
+
+    # super class define this variable as None. "type: ignore[..] is required
+    # since we are redefining the variable.
+    cls_to_become = ConvTranspose3d  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        out_channels: int,
+        kernel_size: _size_3_t,
+        stride: _size_3_t = 1,
+        padding: _size_3_t = 0,
+        output_padding: _size_3_t = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: _size_3_t = 1,
+        padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(
+            0,
+            0,
+            kernel_size,
+            stride,
+            padding,
+            output_padding,
+            groups,
+            # bias is hardcoded to False to avoid creating tensor
+            # that will soon be overwritten.
+            False,
+            dilation,
+            padding_mode,
+            **factory_kwargs,
+        )
+        # pyrefly: ignore [bad-override, unexpected-keyword]
+        self.weight = UninitializedParameter(**factory_kwargs)
+        self.out_channels = out_channels
+        if bias:
+            # pyrefly: ignore [bad-override, unexpected-keyword]
+            self.bias = UninitializedParameter(**factory_kwargs)
+
+    def _get_num_spatial_dims(self) -> int:
+        return 3
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/distance.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/distance.py
new file mode 100644
index 0000000000000000000000000000000000000000..27ab92fef5eb4a6da80d97d8559a204a6956ac4d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/distance.py
@@ -0,0 +1,100 @@
+import torch.nn.functional as F
+from torch import Tensor
+
+from .module import Module
+
+
+__all__ = ["PairwiseDistance", "CosineSimilarity"]
+
+
+class PairwiseDistance(Module):
+    r"""
+    Computes the pairwise distance between input vectors, or between columns of input matrices.
+
+    Distances are computed using ``p``-norm, with constant ``eps`` added to avoid division by zero
+    if ``p`` is negative, i.e.:
+
+    .. math ::
+        \mathrm{dist}\left(x, y\right) = \left\Vert x-y + \epsilon e \right\Vert_p,
+
+    where :math:`e` is the vector of ones and the ``p``-norm is given by.
+
+    .. math ::
+        \Vert x \Vert _p = \left( \sum_{i=1}^n  \vert x_i \vert ^ p \right) ^ {1/p}.
+
+    Args:
+        p (real, optional): the norm degree. Can be negative. Default: 2
+        eps (float, optional): Small value to avoid division by zero.
+            Default: 1e-6
+        keepdim (bool, optional): Determines whether or not to keep the vector dimension.
+            Default: False
+    Shape:
+        - Input1: :math:`(N, D)` or :math:`(D)` where `N = batch dimension` and `D = vector dimension`
+        - Input2: :math:`(N, D)` or :math:`(D)`, same shape as the Input1
+        - Output: :math:`(N)` or :math:`()` based on input dimension.
+          If :attr:`keepdim` is ``True``, then :math:`(N, 1)` or :math:`(1)` based on input dimension.
+
+    Examples:
+        >>> pdist = nn.PairwiseDistance(p=2)
+        >>> input1 = torch.randn(100, 128)
+        >>> input2 = torch.randn(100, 128)
+        >>> output = pdist(input1, input2)
+    """
+
+    __constants__ = ["norm", "eps", "keepdim"]
+    norm: float
+    eps: float
+    keepdim: bool
+
+    def __init__(
+        self, p: float = 2.0, eps: float = 1e-6, keepdim: bool = False
+    ) -> None:
+        super().__init__()
+        self.norm = p
+        self.eps = eps
+        self.keepdim = keepdim
+
+    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.pairwise_distance(x1, x2, self.norm, self.eps, self.keepdim)
+
+
+class CosineSimilarity(Module):
+    r"""Returns cosine similarity between :math:`x_1` and :math:`x_2`, computed along `dim`.
+
+    .. math ::
+        \text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)}.
+
+    Args:
+        dim (int, optional): Dimension where cosine similarity is computed. Default: 1
+        eps (float, optional): Small value to avoid division by zero.
+            Default: 1e-8
+    Shape:
+        - Input1: :math:`(\ast_1, D, \ast_2)` where D is at position `dim`
+        - Input2: :math:`(\ast_1, D, \ast_2)`, same number of dimensions as x1, matching x1 size at dimension `dim`,
+          and broadcastable with x1 at other dimensions.
+        - Output: :math:`(\ast_1, \ast_2)`
+
+    Examples:
+        >>> input1 = torch.randn(100, 128)
+        >>> input2 = torch.randn(100, 128)
+        >>> cos = nn.CosineSimilarity(dim=1, eps=1e-6)
+        >>> output = cos(input1, input2)
+    """
+
+    __constants__ = ["dim", "eps"]
+    dim: int
+    eps: float
+
+    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.cosine_similarity(x1, x2, self.dim, self.eps)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/dropout.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/dropout.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee3de5d61dc0b56d6f708a242611bfc5b2850288
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/dropout.py
@@ -0,0 +1,323 @@
+import torch.nn.functional as F
+from torch import Tensor
+
+from .module import Module
+
+
+__all__ = [
+    "Dropout",
+    "Dropout1d",
+    "Dropout2d",
+    "Dropout3d",
+    "AlphaDropout",
+    "FeatureAlphaDropout",
+]
+
+
+class _DropoutNd(Module):
+    __constants__ = ["p", "inplace"]
+    p: float
+    inplace: bool
+
+    def __init__(self, p: float = 0.5, inplace: bool = False) -> None:
+        super().__init__()
+        if p < 0 or p > 1:
+            raise ValueError(
+                f"dropout probability has to be between 0 and 1, but got {p}"
+            )
+        self.p = p
+        self.inplace = inplace
+
+    def extra_repr(self) -> str:
+        return f"p={self.p}, inplace={self.inplace}"
+
+
+class Dropout(_DropoutNd):
+    r"""During training, randomly zeroes some of the elements of the input tensor with probability :attr:`p`.
+
+    The zeroed elements are chosen independently for each forward call and are sampled from a Bernoulli distribution.
+
+    Each channel will be zeroed out independently on every forward call.
+
+    This has proven to be an effective technique for regularization and
+    preventing the co-adaptation of neurons as described in the paper
+    `Improving neural networks by preventing co-adaptation of feature
+    detectors`_ .
+
+    Furthermore, the outputs are scaled by a factor of :math:`\frac{1}{1-p}` during
+    training. This means that during evaluation the module simply computes an
+    identity function.
+
+    Args:
+        p: probability of an element to be zeroed. Default: 0.5
+        inplace: If set to ``True``, will do this operation in-place. Default: ``False``
+
+    Shape:
+        - Input: :math:`(*)`. Input can be of any shape
+        - Output: :math:`(*)`. Output is of the same shape as input
+
+    Examples::
+
+        >>> m = nn.Dropout(p=0.2)
+        >>> input = torch.randn(20, 16)
+        >>> output = m(input)
+
+    .. _Improving neural networks by preventing co-adaptation of feature
+        detectors: https://arxiv.org/abs/1207.0580
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.dropout(input, self.p, self.training, self.inplace)
+
+
+class Dropout1d(_DropoutNd):
+    r"""Randomly zero out entire channels.
+
+    A channel is a 1D feature map,
+    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the
+    batched input is a 1D tensor :math:`\text{input}[i, j]`.
+
+    Each channel will be zeroed out independently on every forward call with
+    probability :attr:`p` using samples from a Bernoulli distribution.
+
+    Usually the input comes from :class:`nn.Conv1d` modules.
+
+    As described in the paper
+    `Efficient Object Localization Using Convolutional Networks`_ ,
+    if adjacent pixels within feature maps are strongly correlated
+    (as is normally the case in early convolution layers) then i.i.d. dropout
+    will not regularize the activations and will otherwise just result
+    in an effective learning rate decrease.
+
+    In this case, :func:`nn.Dropout1d` will help promote independence between
+    feature maps and should be used instead.
+
+    Args:
+        p (float, optional): probability of an element to be zero-ed.
+        inplace (bool, optional): If set to ``True``, will do this operation
+            in-place
+
+    Shape:
+        - Input: :math:`(N, C, L)` or :math:`(C, L)`.
+        - Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input).
+
+    Examples::
+
+        >>> m = nn.Dropout1d(p=0.2)
+        >>> input = torch.randn(20, 16, 32)
+        >>> output = m(input)
+
+    .. _Efficient Object Localization Using Convolutional Networks:
+       https://arxiv.org/abs/1411.4280
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.dropout1d(input, self.p, self.training, self.inplace)
+
+
+class Dropout2d(_DropoutNd):
+    r"""Randomly zero out entire channels.
+
+    A channel is a 2D feature map,
+    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the
+    batched input is a 2D tensor :math:`\text{input}[i, j]`.
+
+    Each channel will be zeroed out independently on every forward call with
+    probability :attr:`p` using samples from a Bernoulli distribution.
+
+    Usually the input comes from :class:`nn.Conv2d` modules.
+
+    As described in the paper
+    `Efficient Object Localization Using Convolutional Networks`_ ,
+    if adjacent pixels within feature maps are strongly correlated
+    (as is normally the case in early convolution layers) then i.i.d. dropout
+    will not regularize the activations and will otherwise just result
+    in an effective learning rate decrease.
+
+    In this case, :func:`nn.Dropout2d` will help promote independence between
+    feature maps and should be used instead.
+
+    Args:
+        p (float, optional): probability of an element to be zero-ed.
+        inplace (bool, optional): If set to ``True``, will do this operation
+            in-place
+
+    .. warning ::
+        Due to historical reasons, this class will perform 1D channel-wise dropout
+        for 3D inputs (as done by :class:`nn.Dropout1d`). Thus, it currently does NOT
+        support inputs without a batch dimension of shape :math:`(C, H, W)`. This
+        behavior will change in a future release to interpret 3D inputs as no-batch-dim
+        inputs. To maintain the old behavior, switch to :class:`nn.Dropout1d`.
+
+    Shape:
+        - Input: :math:`(N, C, H, W)` or :math:`(N, C, L)`.
+        - Output: :math:`(N, C, H, W)` or :math:`(N, C, L)` (same shape as input).
+
+    Examples::
+
+        >>> m = nn.Dropout2d(p=0.2)
+        >>> input = torch.randn(20, 16, 32, 32)
+        >>> output = m(input)
+
+    .. _Efficient Object Localization Using Convolutional Networks:
+       https://arxiv.org/abs/1411.4280
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.dropout2d(input, self.p, self.training, self.inplace)
+
+
+class Dropout3d(_DropoutNd):
+    r"""Randomly zero out entire channels.
+
+    A channel is a 3D feature map,
+    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the
+    batched input is a 3D tensor :math:`\text{input}[i, j]`.
+
+    Each channel will be zeroed out independently on every forward call with
+    probability :attr:`p` using samples from a Bernoulli distribution.
+
+    Usually the input comes from :class:`nn.Conv3d` modules.
+
+    As described in the paper
+    `Efficient Object Localization Using Convolutional Networks`_ ,
+    if adjacent pixels within feature maps are strongly correlated
+    (as is normally the case in early convolution layers) then i.i.d. dropout
+    will not regularize the activations and will otherwise just result
+    in an effective learning rate decrease.
+
+    In this case, :func:`nn.Dropout3d` will help promote independence between
+    feature maps and should be used instead.
+
+    Args:
+        p (float, optional): probability of an element to be zeroed.
+        inplace (bool, optional): If set to ``True``, will do this operation
+            in-place
+
+    Shape:
+        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`.
+        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input).
+
+    Examples::
+
+        >>> m = nn.Dropout3d(p=0.2)
+        >>> input = torch.randn(20, 16, 4, 32, 32)
+        >>> output = m(input)
+
+    .. _Efficient Object Localization Using Convolutional Networks:
+       https://arxiv.org/abs/1411.4280
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.dropout3d(input, self.p, self.training, self.inplace)
+
+
+class AlphaDropout(_DropoutNd):
+    r"""Applies Alpha Dropout over the input.
+
+    Alpha Dropout is a type of Dropout that maintains the self-normalizing
+    property.
+    For an input with zero mean and unit standard deviation, the output of
+    Alpha Dropout maintains the original mean and standard deviation of the
+    input.
+    Alpha Dropout goes hand-in-hand with SELU activation function, which ensures
+    that the outputs have zero mean and unit standard deviation.
+
+    During training, it randomly masks some of the elements of the input
+    tensor with probability *p* using samples from a bernoulli distribution.
+    The elements to masked are randomized on every forward call, and scaled
+    and shifted to maintain zero mean and unit standard deviation.
+
+    During evaluation the module simply computes an identity function.
+
+    More details can be found in the paper `Self-Normalizing Neural Networks`_ .
+
+    Args:
+        p (float): probability of an element to be dropped. Default: 0.5
+        inplace (bool, optional): If set to ``True``, will do this operation
+            in-place
+
+    Shape:
+        - Input: :math:`(*)`. Input can be of any shape
+        - Output: :math:`(*)`. Output is of the same shape as input
+
+    Examples::
+
+        >>> m = nn.AlphaDropout(p=0.2)
+        >>> input = torch.randn(20, 16)
+        >>> output = m(input)
+
+    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.alpha_dropout(input, self.p, self.training)
+
+
+class FeatureAlphaDropout(_DropoutNd):
+    r"""Randomly masks out entire channels.
+
+    A channel is a feature map,
+    e.g. the :math:`j`-th channel of the :math:`i`-th sample in the batch input
+    is a tensor :math:`\text{input}[i, j]` of the input tensor). Instead of
+    setting activations to zero, as in regular Dropout, the activations are set
+    to the negative saturation value of the SELU activation function. More details
+    can be found in the paper `Self-Normalizing Neural Networks`_ .
+
+    Each element will be masked independently for each sample on every forward
+    call with probability :attr:`p` using samples from a Bernoulli distribution.
+    The elements to be masked are randomized on every forward call, and scaled
+    and shifted to maintain zero mean and unit variance.
+
+    Usually the input comes from :class:`nn.AlphaDropout` modules.
+
+    As described in the paper
+    `Efficient Object Localization Using Convolutional Networks`_ ,
+    if adjacent pixels within feature maps are strongly correlated
+    (as is normally the case in early convolution layers) then i.i.d. dropout
+    will not regularize the activations and will otherwise just result
+    in an effective learning rate decrease.
+
+    In this case, :func:`nn.AlphaDropout` will help promote independence between
+    feature maps and should be used instead.
+
+    Args:
+        p (float, optional): probability of an element to be zeroed. Default: 0.5
+        inplace (bool, optional): If set to ``True``, will do this operation
+            in-place
+
+    Shape:
+        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`.
+        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input).
+
+    Examples::
+
+        >>> m = nn.FeatureAlphaDropout(p=0.2)
+        >>> input = torch.randn(20, 16, 4, 32, 32)
+        >>> output = m(input)
+
+    .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
+    .. _Efficient Object Localization Using Convolutional Networks:
+       https://arxiv.org/abs/1411.4280
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.feature_alpha_dropout(input, self.p, self.training)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/flatten.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/flatten.py
new file mode 100644
index 0000000000000000000000000000000000000000..146a1890d422475712c9d62d0ff841530282d30e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/flatten.py
@@ -0,0 +1,167 @@
+# mypy: allow-untyped-defs
+
+from torch import Tensor
+from torch.types import _size
+
+from .module import Module
+
+
+__all__ = ["Flatten", "Unflatten"]
+
+
+class Flatten(Module):
+    r"""
+    Flattens a contiguous range of dims into a tensor.
+
+    For use with :class:`~nn.Sequential`, see :meth:`torch.flatten` for details.
+
+    Shape:
+        - Input: :math:`(*, S_{\text{start}},..., S_{i}, ..., S_{\text{end}}, *)`,'
+          where :math:`S_{i}` is the size at dimension :math:`i` and :math:`*` means any
+          number of dimensions including none.
+        - Output: :math:`(*, \prod_{i=\text{start}}^{\text{end}} S_{i}, *)`.
+
+    Args:
+        start_dim: first dim to flatten (default = 1).
+        end_dim: last dim to flatten (default = -1).
+
+    Examples::
+        >>> input = torch.randn(32, 1, 5, 5)
+        >>> # With default parameters
+        >>> m = nn.Flatten()
+        >>> output = m(input)
+        >>> output.size()
+        torch.Size([32, 25])
+        >>> # With non-default parameters
+        >>> m = nn.Flatten(0, 2)
+        >>> output = m(input)
+        >>> output.size()
+        torch.Size([160, 5])
+    """
+
+    __constants__ = ["start_dim", "end_dim"]
+    start_dim: int
+    end_dim: int
+
+    def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
+        super().__init__()
+        self.start_dim = start_dim
+        self.end_dim = end_dim
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return input.flatten(self.start_dim, self.end_dim)
+
+    def extra_repr(self) -> str:
+        """
+        Returns the extra representation of the module.
+        """
+        return f"start_dim={self.start_dim}, end_dim={self.end_dim}"
+
+
+class Unflatten(Module):
+    r"""
+    Unflattens a tensor dim expanding it to a desired shape. For use with :class:`~nn.Sequential`.
+
+    * :attr:`dim` specifies the dimension of the input tensor to be unflattened, and it can
+      be either `int` or `str` when `Tensor` or `NamedTensor` is used, respectively.
+
+    * :attr:`unflattened_size` is the new shape of the unflattened dimension of the tensor and it can be
+      a `tuple` of ints or a `list` of ints or `torch.Size` for `Tensor` input;  a `NamedShape`
+      (tuple of `(name, size)` tuples) for `NamedTensor` input.
+
+    Shape:
+        - Input: :math:`(*, S_{\text{dim}}, *)`, where :math:`S_{\text{dim}}` is the size at
+          dimension :attr:`dim` and :math:`*` means any number of dimensions including none.
+        - Output: :math:`(*, U_1, ..., U_n, *)`, where :math:`U` = :attr:`unflattened_size` and
+          :math:`\prod_{i=1}^n U_i = S_{\text{dim}}`.
+
+    Args:
+        dim (Union[int, str]): Dimension to be unflattened
+        unflattened_size (Union[torch.Size, Tuple, List, NamedShape]): New shape of the unflattened dimension
+
+    Examples:
+        >>> input = torch.randn(2, 50)
+        >>> # With tuple of ints
+        >>> m = nn.Sequential(
+        >>>     nn.Linear(50, 50),
+        >>>     nn.Unflatten(1, (2, 5, 5))
+        >>> )
+        >>> output = m(input)
+        >>> output.size()
+        torch.Size([2, 2, 5, 5])
+        >>> # With torch.Size
+        >>> m = nn.Sequential(
+        >>>     nn.Linear(50, 50),
+        >>>     nn.Unflatten(1, torch.Size([2, 5, 5]))
+        >>> )
+        >>> output = m(input)
+        >>> output.size()
+        torch.Size([2, 2, 5, 5])
+        >>> # With namedshape (tuple of tuples)
+        >>> input = torch.randn(2, 50, names=("N", "features"))
+        >>> unflatten = nn.Unflatten("features", (("C", 2), ("H", 5), ("W", 5)))
+        >>> output = unflatten(input)
+        >>> output.size()
+        torch.Size([2, 2, 5, 5])
+    """
+
+    NamedShape = tuple[tuple[str, int]]
+
+    __constants__ = ["dim", "unflattened_size"]
+    dim: int | str
+    unflattened_size: _size | NamedShape
+
+    def __init__(self, dim: int | str, unflattened_size: _size | NamedShape) -> None:
+        super().__init__()
+
+        if isinstance(dim, int):
+            self._require_tuple_int(unflattened_size)
+        elif isinstance(dim, str):
+            self._require_tuple_tuple(unflattened_size)
+        else:
+            raise TypeError("invalid argument type for dim parameter")
+
+        self.dim = dim
+        self.unflattened_size = unflattened_size
+
+    def _require_tuple_tuple(self, input) -> None:
+        if isinstance(input, tuple):
+            for idx, elem in enumerate(input):
+                if not isinstance(elem, tuple):
+                    raise TypeError(
+                        "unflattened_size must be tuple of tuples, "
+                        + f"but found element of type {type(elem).__name__} at pos {idx}"
+                    )
+            return
+        raise TypeError(
+            "unflattened_size must be a tuple of tuples, "
+            + f"but found type {type(input).__name__}"
+        )
+
+    def _require_tuple_int(self, input) -> None:
+        if isinstance(input, (tuple, list)):
+            for idx, elem in enumerate(input):
+                if not isinstance(elem, int):
+                    raise TypeError(
+                        "unflattened_size must be tuple of ints, "
+                        + f"but found element of type {type(elem).__name__} at pos {idx}"
+                    )
+            return
+        raise TypeError(
+            f"unflattened_size must be a tuple of ints, but found type {type(input).__name__}"
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return input.unflatten(self.dim, self.unflattened_size)
+
+    def extra_repr(self) -> str:
+        """
+        Returns the extra representation of the module.
+        """
+        return f"dim={self.dim}, unflattened_size={self.unflattened_size}"
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/fold.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/fold.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab1a58882c852370141e1e1dd911278334b425d8
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/fold.py
@@ -0,0 +1,335 @@
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.common_types import _size_any_t
+
+from .module import Module
+
+
+__all__ = ["Fold", "Unfold"]
+
+
+class Fold(Module):
+    (
+        r"""Combines an array of sliding local blocks into a large containing tensor.
+
+    Consider a batched :attr:`input` tensor containing sliding local blocks,
+    e.g., patches of images, of shape :math:`(N, C \times  \prod(\text{kernel\_size}), L)`,
+    where :math:`N` is batch dimension, :math:`C \times \prod(\text{kernel\_size})`
+    is the number of values within a block (a block has :math:`\prod(\text{kernel\_size})`
+    spatial locations each containing a :math:`C`-channeled vector), and
+    :math:`L` is the total number of blocks. (This is exactly the
+    same specification as the output shape of :class:`~torch.nn.Unfold`.) This
+    operation combines these local blocks into the large :attr:`output` tensor
+    of shape :math:`(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)`
+    by summing the overlapping values. Similar to :class:`~torch.nn.Unfold`, the
+    arguments must satisfy
+
+    .. math::
+        L = \prod_d \left\lfloor\frac{\text{output\_size}[d] + 2 \times \text{padding}[d] %
+            - \text{dilation}[d] \times (\text{kernel\_size}[d] - 1) - 1}{\text{stride}[d]} + 1\right\rfloor,
+
+    where :math:`d` is over all spatial dimensions.
+
+    * :attr:`output_size` describes the spatial shape of the large containing
+      tensor of the sliding local blocks. It is useful to resolve the ambiguity
+      when multiple input shapes map to same number of sliding blocks, e.g.,
+      with ``stride > 0``.
+
+    The :attr:`padding`, :attr:`stride` and :attr:`dilation` arguments specify
+    how the sliding blocks are retrieved.
+
+    * :attr:`stride` controls the stride for the sliding blocks.
+
+    * :attr:`padding` controls the amount of implicit zero-paddings on both
+      sides for :attr:`padding` number of points for each dimension before
+      reshaping.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
+      It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+    Args:
+        output_size (int or tuple): the shape of the spatial dimensions of the
+                                    output (i.e., ``output.sizes()[2:]``)
+        kernel_size (int or tuple): the size of the sliding blocks
+        dilation (int or tuple, optional): a parameter that controls the
+                                           stride of elements within the
+                                           neighborhood. Default: 1
+        padding (int or tuple, optional): implicit zero padding to be added on
+                                          both sides of input. Default: 0
+        stride (int or tuple): the stride of the sliding blocks in the input
+                               spatial dimensions. Default: 1
+
+    * If :attr:`output_size`, :attr:`kernel_size`, :attr:`dilation`,
+      :attr:`padding` or :attr:`stride` is an int or a tuple of length 1 then
+      their values will be replicated across all spatial dimensions.
+
+    * For the case of two output spatial dimensions this operation is sometimes
+      called ``col2im``.
+
+    .. note::
+        :class:`~torch.nn.Fold` calculates each combined value in the resulting
+        large tensor by summing all values from all containing blocks.
+        :class:`~torch.nn.Unfold` extracts the values in the local blocks by
+        copying from the large tensor. So, if the blocks overlap, they are not
+        inverses of each other.
+
+        In general, folding and unfolding operations are related as
+        follows. Consider :class:`~torch.nn.Fold` and
+        :class:`~torch.nn.Unfold` instances created with the same
+        parameters:
+
+        >>> fold_params = dict(kernel_size=..., dilation=..., padding=..., stride=...)
+        >>> fold = nn.Fold(output_size=..., **fold_params)
+        >>> unfold = nn.Unfold(**fold_params)
+
+        Then for any (supported) ``input`` tensor the following
+        equality holds:
+
+        ::
+
+            fold(unfold(input)) == divisor * input
+
+        where ``divisor`` is a tensor that depends only on the shape
+        and dtype of the ``input``:
+
+        >>> # xdoctest: +SKIP
+        >>> input_ones = torch.ones(input.shape, dtype=input.dtype)
+        >>> divisor = fold(unfold(input_ones))
+
+        When the ``divisor`` tensor contains no zero elements, then
+        ``fold`` and ``unfold`` operations are inverses of each
+        other (up to constant divisor).
+
+    .. warning::
+        Currently, only unbatched (3D) or batched (4D) image-like output tensors are supported.
+
+    Shape:
+        - Input: :math:`(N, C \times \prod(\text{kernel\_size}), L)` or :math:`(C \times \prod(\text{kernel\_size}), L)`
+        - Output: :math:`(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)`
+          or :math:`(C, \text{output\_size}[0], \text{output\_size}[1], \dots)` as described above
+
+    Examples::
+
+        >>> fold = nn.Fold(output_size=(4, 5), kernel_size=(2, 2))
+        >>> input = torch.randn(1, 3 * 2 * 2, 12)
+        >>> output = fold(input)
+        >>> output.size()
+        torch.Size([1, 3, 4, 5])
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+
+    """
+    )
+
+    __constants__ = ["output_size", "kernel_size", "dilation", "padding", "stride"]
+    output_size: _size_any_t
+    kernel_size: _size_any_t
+    dilation: _size_any_t
+    padding: _size_any_t
+    stride: _size_any_t
+
+    def __init__(
+        self,
+        output_size: _size_any_t,
+        kernel_size: _size_any_t,
+        dilation: _size_any_t = 1,
+        padding: _size_any_t = 0,
+        stride: _size_any_t = 1,
+    ) -> None:
+        super().__init__()
+        self.output_size = output_size
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.padding = padding
+        self.stride = stride
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.fold(
+            input,
+            self.output_size,
+            self.kernel_size,
+            self.dilation,
+            self.padding,
+            self.stride,
+        )
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return (
+            "output_size={output_size}, kernel_size={kernel_size}, "
+            "dilation={dilation}, padding={padding}, stride={stride}".format(
+                **self.__dict__
+            )
+        )
+
+
+class Unfold(Module):
+    (
+        r"""Extracts sliding local blocks from a batched input tensor.
+
+    Consider a batched :attr:`input` tensor of shape :math:`(N, C, *)`,
+    where :math:`N` is the batch dimension, :math:`C` is the channel dimension,
+    and :math:`*` represent arbitrary spatial dimensions. This operation flattens
+    each sliding :attr:`kernel_size`-sized block within the spatial dimensions
+    of :attr:`input` into a column (i.e., last dimension) of a 3-D :attr:`output`
+    tensor of shape :math:`(N, C \times \prod(\text{kernel\_size}), L)`, where
+    :math:`C \times \prod(\text{kernel\_size})` is the total number of values
+    within each block (a block has :math:`\prod(\text{kernel\_size})` spatial
+    locations each containing a :math:`C`-channeled vector), and :math:`L` is
+    the total number of such blocks:
+
+    .. math::
+        L = \prod_d \left\lfloor\frac{\text{spatial\_size}[d] + 2 \times \text{padding}[d] %
+            - \text{dilation}[d] \times (\text{kernel\_size}[d] - 1) - 1}{\text{stride}[d]} + 1\right\rfloor,
+
+    where :math:`\text{spatial\_size}` is formed by the spatial dimensions
+    of :attr:`input` (:math:`*` above), and :math:`d` is over all spatial
+    dimensions.
+
+    Therefore, indexing :attr:`output` at the last dimension (column dimension)
+    gives all values within a certain block.
+
+    The :attr:`padding`, :attr:`stride` and :attr:`dilation` arguments specify
+    how the sliding blocks are retrieved.
+
+    * :attr:`stride` controls the stride for the sliding blocks.
+
+    * :attr:`padding` controls the amount of implicit zero-paddings on both
+      sides for :attr:`padding` number of points for each dimension before
+      reshaping.
+"""
+        """
+    * :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
+      It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
+"""
+        r"""
+    Args:
+        kernel_size (int or tuple): the size of the sliding blocks
+        dilation (int or tuple, optional): a parameter that controls the
+                                           stride of elements within the
+                                           neighborhood. Default: 1
+        padding (int or tuple, optional): implicit zero padding to be added on
+                                          both sides of input. Default: 0
+        stride (int or tuple, optional): the stride of the sliding blocks in the input
+                                         spatial dimensions. Default: 1
+
+    * If :attr:`kernel_size`, :attr:`dilation`, :attr:`padding` or
+      :attr:`stride` is an int or a tuple of length 1, their values will be
+      replicated across all spatial dimensions.
+
+    * For the case of two input spatial dimensions this operation is sometimes
+      called ``im2col``.
+
+    .. note::
+        :class:`~torch.nn.Fold` calculates each combined value in the resulting
+        large tensor by summing all values from all containing blocks.
+        :class:`~torch.nn.Unfold` extracts the values in the local blocks by
+        copying from the large tensor. So, if the blocks overlap, they are not
+        inverses of each other.
+
+        In general, folding and unfolding operations are related as
+        follows. Consider :class:`~torch.nn.Fold` and
+        :class:`~torch.nn.Unfold` instances created with the same
+        parameters:
+
+        >>> fold_params = dict(kernel_size=..., dilation=..., padding=..., stride=...)
+        >>> fold = nn.Fold(output_size=..., **fold_params)
+        >>> unfold = nn.Unfold(**fold_params)
+
+        Then for any (supported) ``input`` tensor the following
+        equality holds:
+
+        ::
+
+            fold(unfold(input)) == divisor * input
+
+        where ``divisor`` is a tensor that depends only on the shape
+        and dtype of the ``input``:
+
+        >>> # xdoctest: +SKIP
+        >>> input_ones = torch.ones(input.shape, dtype=input.dtype)
+        >>> divisor = fold(unfold(input_ones))
+
+        When the ``divisor`` tensor contains no zero elements, then
+        ``fold`` and ``unfold`` operations are inverses of each
+        other (up to constant divisor).
+
+    .. warning::
+        Currently, only 4-D input tensors (batched image-like tensors) are
+        supported.
+
+    Shape:
+        - Input: :math:`(N, C, *)`
+        - Output: :math:`(N, C \times \prod(\text{kernel\_size}), L)` as described above
+
+    Examples::
+
+        >>> unfold = nn.Unfold(kernel_size=(2, 3))
+        >>> input = torch.randn(2, 5, 3, 4)
+        >>> output = unfold(input)
+        >>> # each patch contains 30 values (2x3=6 vectors, each of 5 channels)
+        >>> # 4 blocks (2x3 kernels) in total in the 3x4 input
+        >>> output.size()
+        torch.Size([2, 30, 4])
+
+        >>> # xdoctest: +IGNORE_WANT
+        >>> # Convolution is equivalent with Unfold + Matrix Multiplication + Fold (or view to output shape)
+        >>> inp = torch.randn(1, 3, 10, 12)
+        >>> w = torch.randn(2, 3, 4, 5)
+        >>> inp_unf = torch.nn.functional.unfold(inp, (4, 5))
+        >>> out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
+        >>> out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))
+        >>> # or equivalently (and avoiding a copy),
+        >>> # out = out_unf.view(1, 2, 7, 8)
+        >>> (torch.nn.functional.conv2d(inp, w) - out).abs().max()
+        tensor(1.9073e-06)
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+
+    """
+    )
+
+    __constants__ = ["kernel_size", "dilation", "padding", "stride"]
+    kernel_size: _size_any_t
+    dilation: _size_any_t
+    padding: _size_any_t
+    stride: _size_any_t
+
+    def __init__(
+        self,
+        kernel_size: _size_any_t,
+        dilation: _size_any_t = 1,
+        padding: _size_any_t = 0,
+        stride: _size_any_t = 1,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.padding = padding
+        self.stride = stride
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.unfold(
+            input, self.kernel_size, self.dilation, self.padding, self.stride
+        )
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return (
+            "kernel_size={kernel_size}, dilation={dilation}, padding={padding},"
+            " stride={stride}".format(**self.__dict__)
+        )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/instancenorm.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/instancenorm.py
new file mode 100644
index 0000000000000000000000000000000000000000..76a343afbd0dd5fed3241b6a20a0f445be8ecc12
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/instancenorm.py
@@ -0,0 +1,492 @@
+# mypy: allow-untyped-defs
+
+import warnings
+
+import torch.nn.functional as F
+from torch import Tensor
+
+from .batchnorm import _LazyNormBase, _NormBase
+
+
+__all__ = [
+    "InstanceNorm1d",
+    "InstanceNorm2d",
+    "InstanceNorm3d",
+    "LazyInstanceNorm1d",
+    "LazyInstanceNorm2d",
+    "LazyInstanceNorm3d",
+]
+
+
+class _InstanceNorm(_NormBase):
+    def __init__(
+        self,
+        num_features: int,
+        eps: float = 1e-5,
+        momentum: float = 0.1,
+        affine: bool = False,
+        track_running_stats: bool = False,
+        device=None,
+        dtype=None,
+        *,
+        bias: bool = True,  # for backward compatibility
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(
+            num_features,
+            eps,
+            momentum,
+            affine,
+            track_running_stats,
+            **factory_kwargs,
+            bias=bias,
+        )
+
+    def _check_input_dim(self, input):
+        raise NotImplementedError
+
+    def _get_no_batch_dim(self):
+        raise NotImplementedError
+
+    def _handle_no_batch_input(self, input):
+        return self._apply_instance_norm(input.unsqueeze(0)).squeeze(0)
+
+    def _apply_instance_norm(self, input):
+        return F.instance_norm(
+            input,
+            self.running_mean,
+            self.running_var,
+            self.weight,
+            self.bias,
+            self.training or not self.track_running_stats,
+            self.momentum if self.momentum is not None else 0.0,
+            self.eps,
+        )
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ) -> None:
+        version = local_metadata.get("version", None)
+        # at version 1: removed running_mean and running_var when
+        # track_running_stats=False (default)
+        if version is None and not self.track_running_stats:
+            running_stats_keys = []
+            for name in ("running_mean", "running_var"):
+                key = prefix + name
+                if key in state_dict:
+                    running_stats_keys.append(key)
+            if len(running_stats_keys) > 0:
+                error_msgs.append(
+                    "Unexpected running stats buffer(s) {names} for {klass} "
+                    "with track_running_stats=False. If state_dict is a "
+                    "checkpoint saved before 0.4.0, this may be expected "
+                    "because {klass} does not track running stats by default "
+                    "since 0.4.0. Please remove these keys from state_dict. If "
+                    "the running stats are actually needed, instead set "
+                    "track_running_stats=True in {klass} to enable them. See "
+                    "the documentation of {klass} for details.".format(
+                        names=" and ".join(f'"{k}"' for k in running_stats_keys),
+                        klass=self.__class__.__name__,
+                    )
+                )
+                for key in running_stats_keys:
+                    state_dict.pop(key)
+
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        self._check_input_dim(input)
+
+        feature_dim = input.dim() - self._get_no_batch_dim()
+        if input.size(feature_dim) != self.num_features:
+            if self.affine:
+                raise ValueError(
+                    f"expected input's size at dim={feature_dim} to match num_features"
+                    f" ({self.num_features}), but got: {input.size(feature_dim)}."
+                )
+            else:
+                warnings.warn(
+                    f"input's size at dim={feature_dim} does not match num_features. "
+                    "You can silence this warning by not passing in num_features, "
+                    "which is not used because affine=False",
+                    stacklevel=2,
+                )
+
+        if input.dim() == self._get_no_batch_dim():
+            return self._handle_no_batch_input(input)
+
+        return self._apply_instance_norm(input)
+
+
+class InstanceNorm1d(_InstanceNorm):
+    r"""Applies Instance Normalization.
+
+    This operation applies Instance Normalization
+    over a 2D (unbatched) or 3D (batched) input as described in the paper
+    `Instance Normalization: The Missing Ingredient for Fast Stylization
+    <https://arxiv.org/abs/1607.08022>`__.
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension separately
+    for each object in a mini-batch. :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the number of features or channels of the input) if :attr:`affine` is ``True``.
+    The variance is calculated via the biased estimator, equivalent to
+    `torch.var(input, correction=0)`.
+
+    By default, this layer uses instance statistics computed from input data in
+    both training and evaluation modes.
+
+    If :attr:`track_running_stats` is set to ``True``, during training this
+    layer keeps running estimates of its computed mean and variance, which are
+    then used for normalization during evaluation. The running estimates are
+    kept with a default :attr:`momentum` of 0.1.
+
+    .. note::
+        This :attr:`momentum` argument is different from one used in optimizer
+        classes and the conventional notion of momentum. Mathematically, the
+        update rule for running statistics here is
+        :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
+        where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
+        new observed value.
+
+    .. note::
+        :class:`InstanceNorm1d` and :class:`LayerNorm` are very similar, but
+        have some subtle differences. :class:`InstanceNorm1d` is applied
+        on each channel of channeled data like multidimensional time series, but
+        :class:`LayerNorm` is usually applied on entire sample and often in NLP
+        tasks. Additionally, :class:`LayerNorm` applies elementwise affine
+        transform, while :class:`InstanceNorm1d` usually don't apply affine
+        transform.
+
+    Args:
+        num_features: number of features or channels :math:`C` of the input
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        momentum: the value used for the running_mean and running_var computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters, initialized the same way as done for batch normalization.
+            Default: ``False``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, L)` or :math:`(C, L)`
+        - Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input)
+
+    Examples::
+
+        >>> # Without Learnable Parameters
+        >>> m = nn.InstanceNorm1d(100)
+        >>> # With Learnable Parameters
+        >>> m = nn.InstanceNorm1d(100, affine=True)
+        >>> input = torch.randn(20, 100, 40)
+        >>> output = m(input)
+    """
+
+    def _get_no_batch_dim(self) -> int:
+        return 2
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() not in (2, 3):
+            raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
+
+
+class LazyInstanceNorm1d(_LazyNormBase, _InstanceNorm):
+    r"""A :class:`torch.nn.InstanceNorm1d` module with lazy initialization of the ``num_features`` argument.
+
+    The ``num_features`` argument of the :class:`InstanceNorm1d` is inferred from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight`, `bias`, `running_mean` and `running_var`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, L)` or :math:`(C, L)`
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        momentum: the value used for the running_mean and running_var computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters, initialized the same way as done for batch normalization.
+            Default: ``False``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, L)` or :math:`(C, L)`
+        - Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input)
+    """
+
+    cls_to_become = InstanceNorm1d  # type: ignore[assignment]
+
+    def _get_no_batch_dim(self) -> int:
+        return 2
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() not in (2, 3):
+            raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
+
+
+class InstanceNorm2d(_InstanceNorm):
+    r"""Applies Instance Normalization.
+
+    This operation applies Instance Normalization
+    over a 4D input (a mini-batch of 2D inputs
+    with additional channel dimension) as described in the paper
+    `Instance Normalization: The Missing Ingredient for Fast Stylization
+    <https://arxiv.org/abs/1607.08022>`__.
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension separately
+    for each object in a mini-batch. :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the input size) if :attr:`affine` is ``True``.
+    The standard-deviation is calculated via the biased estimator, equivalent to
+    `torch.var(input, correction=0)`.
+
+    By default, this layer uses instance statistics computed from input data in
+    both training and evaluation modes.
+
+    If :attr:`track_running_stats` is set to ``True``, during training this
+    layer keeps running estimates of its computed mean and variance, which are
+    then used for normalization during evaluation. The running estimates are
+    kept with a default :attr:`momentum` of 0.1.
+
+    .. note::
+        This :attr:`momentum` argument is different from one used in optimizer
+        classes and the conventional notion of momentum. Mathematically, the
+        update rule for running statistics here is
+        :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
+        where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
+        new observed value.
+
+    .. note::
+        :class:`InstanceNorm2d` and :class:`LayerNorm` are very similar, but
+        have some subtle differences. :class:`InstanceNorm2d` is applied
+        on each channel of channeled data like RGB images, but
+        :class:`LayerNorm` is usually applied on entire sample and often in NLP
+        tasks. Additionally, :class:`LayerNorm` applies elementwise affine
+        transform, while :class:`InstanceNorm2d` usually don't apply affine
+        transform.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, H, W)` or :math:`(C, H, W)`
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        momentum: the value used for the running_mean and running_var computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters, initialized the same way as done for batch normalization.
+            Default: ``False``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`
+        - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)
+
+    Examples::
+
+        >>> # Without Learnable Parameters
+        >>> m = nn.InstanceNorm2d(100)
+        >>> # With Learnable Parameters
+        >>> m = nn.InstanceNorm2d(100, affine=True)
+        >>> input = torch.randn(20, 100, 35, 45)
+        >>> output = m(input)
+    """
+
+    def _get_no_batch_dim(self) -> int:
+        return 3
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() not in (3, 4):
+            raise ValueError(f"expected 3D or 4D input (got {input.dim()}D input)")
+
+
+class LazyInstanceNorm2d(_LazyNormBase, _InstanceNorm):
+    r"""A :class:`torch.nn.InstanceNorm2d` module with lazy initialization of the ``num_features`` argument.
+
+    The ``num_features`` argument of the :class:`InstanceNorm2d` is inferred from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight`, `bias`,
+    `running_mean` and `running_var`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, H, W)` or :math:`(C, H, W)`
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        momentum: the value used for the running_mean and running_var computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters, initialized the same way as done for batch normalization.
+            Default: ``False``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`
+        - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)
+    """
+
+    cls_to_become = InstanceNorm2d  # type: ignore[assignment]
+
+    def _get_no_batch_dim(self) -> int:
+        return 3
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() not in (3, 4):
+            raise ValueError(f"expected 3D or 4D input (got {input.dim()}D input)")
+
+
+class InstanceNorm3d(_InstanceNorm):
+    r"""Applies Instance Normalization.
+
+    This operation applies Instance Normalization
+    over a 5D input (a mini-batch of 3D inputs with additional channel dimension) as described in the paper
+    `Instance Normalization: The Missing Ingredient for Fast Stylization
+    <https://arxiv.org/abs/1607.08022>`__.
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension separately
+    for each object in a mini-batch. :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size C (where C is the input size) if :attr:`affine` is ``True``.
+    The standard-deviation is calculated via the biased estimator, equivalent to
+    `torch.var(input, correction=0)`.
+
+    By default, this layer uses instance statistics computed from input data in
+    both training and evaluation modes.
+
+    If :attr:`track_running_stats` is set to ``True``, during training this
+    layer keeps running estimates of its computed mean and variance, which are
+    then used for normalization during evaluation. The running estimates are
+    kept with a default :attr:`momentum` of 0.1.
+
+    .. note::
+        This :attr:`momentum` argument is different from one used in optimizer
+        classes and the conventional notion of momentum. Mathematically, the
+        update rule for running statistics here is
+        :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
+        where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
+        new observed value.
+
+    .. note::
+        :class:`InstanceNorm3d` and :class:`LayerNorm` are very similar, but
+        have some subtle differences. :class:`InstanceNorm3d` is applied
+        on each channel of channeled data like 3D models with RGB color, but
+        :class:`LayerNorm` is usually applied on entire sample and often in NLP
+        tasks. Additionally, :class:`LayerNorm` applies elementwise affine
+        transform, while :class:`InstanceNorm3d` usually don't apply affine
+        transform.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        momentum: the value used for the running_mean and running_var computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters, initialized the same way as done for batch normalization.
+            Default: ``False``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input)
+
+    Examples::
+
+        >>> # Without Learnable Parameters
+        >>> m = nn.InstanceNorm3d(100)
+        >>> # With Learnable Parameters
+        >>> m = nn.InstanceNorm3d(100, affine=True)
+        >>> input = torch.randn(20, 100, 35, 45, 10)
+        >>> output = m(input)
+    """
+
+    def _get_no_batch_dim(self) -> int:
+        return 4
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() not in (4, 5):
+            raise ValueError(f"expected 4D or 5D input (got {input.dim()}D input)")
+
+
+class LazyInstanceNorm3d(_LazyNormBase, _InstanceNorm):
+    r"""A :class:`torch.nn.InstanceNorm3d` module with lazy initialization of the ``num_features`` argument.
+
+    The ``num_features`` argument of the :class:`InstanceNorm3d` is inferred from the ``input.size(1)``.
+    The attributes that will be lazily initialized are `weight`, `bias`,
+    `running_mean` and `running_var`.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        momentum: the value used for the running_mean and running_var computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters, initialized the same way as done for batch normalization.
+            Default: ``False``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``False``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input)
+    """
+
+    cls_to_become = InstanceNorm3d  # type: ignore[assignment]
+
+    def _get_no_batch_dim(self) -> int:
+        return 4
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() not in (4, 5):
+            raise ValueError(f"expected 4D or 5D input (got {input.dim()}D input)")
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/lazy.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/lazy.py
new file mode 100644
index 0000000000000000000000000000000000000000..72d90d1c10364ea380b1f27069dd69dda6ec80cc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/lazy.py
@@ -0,0 +1,278 @@
+# mypy: allow-untyped-defs
+import itertools
+from typing import Any, Protocol
+
+import torch
+from torch.nn.parameter import is_lazy
+
+
+__all__ = ["LazyModuleMixin"]
+
+
+class _LazyProtocol(Protocol):
+    """This class is used to avoid errors with mypy checks for the attributes in a mixin.
+
+    https://mypy.readthedocs.io/en/latest/more_types.html#mixin-classes
+    """
+
+    def _register_load_state_dict_pre_hook(self, hook): ...
+
+    def register_forward_pre_hook(self, hook, *, prepend=False, with_kwargs=False): ...
+
+    def _lazy_load_hook(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ): ...
+
+    def _get_name(self): ...
+
+    def _infer_parameters(self, module, input): ...
+
+    @property
+    def _parameters(self): ...
+
+    @property
+    def _buffers(self): ...
+
+    @property
+    def _non_persistent_buffers_set(self): ...
+
+    @property
+    def _load_hook(self): ...
+
+    @property
+    def _initialize_hook(self): ...
+
+
+class LazyModuleMixin:
+    r"""A mixin for modules that lazily initialize parameters, also known as "lazy modules".
+
+    .. warning:
+        Lazy modules are an experimental new feature under active development,
+        and their API is likely to change.
+
+    Modules that lazily initialize parameters, or "lazy modules",
+    derive the shapes of their parameters from the first input(s)
+    to their forward method. Until that first forward they contain
+    :class:`torch.nn.UninitializedParameter` s that should not be accessed
+    or used, and afterward they contain regular :class:`torch.nn.Parameter` s.
+    Lazy modules are convenient since they don't require computing some
+    module arguments, like the :attr:`in_features` argument of a
+    typical :class:`torch.nn.Linear`.
+
+    After construction, networks with lazy modules should first
+    be converted to the desired dtype and placed on the expected device.
+    This is because lazy modules only perform shape inference so the usual dtype
+    and device placement behavior applies.
+    The lazy modules should then perform "dry runs" to initialize all the components in the module.
+    These "dry runs" send inputs of the correct size, dtype, and device through
+    the network and to each one of its lazy modules. After this the network can be used as usual.
+
+    >>> # xdoctest: +SKIP
+    >>> class LazyMLP(torch.nn.Module):
+    ...     def __init__(self) -> None:
+    ...         super().__init__()
+    ...         self.fc1 = torch.nn.LazyLinear(10)
+    ...         self.relu1 = torch.nn.ReLU()
+    ...         self.fc2 = torch.nn.LazyLinear(1)
+    ...         self.relu2 = torch.nn.ReLU()
+    ...
+    ...     def forward(self, input):
+    ...         x = self.relu1(self.fc1(input))
+    ...         y = self.relu2(self.fc2(x))
+    ...         return y
+    >>> # constructs a network with lazy modules
+    >>> lazy_mlp = LazyMLP()
+    >>> # transforms the network's device and dtype
+    >>> # NOTE: these transforms can and should be applied after construction and before any 'dry runs'
+    >>> lazy_mlp = lazy_mlp.cuda()
+    >>> lazy_mlp
+    LazyMLP( (fc1): LazyLinear(in_features=0, out_features=10, bias=True)
+      (relu1): ReLU()
+      (fc2): LazyLinear(in_features=0, out_features=1, bias=True)
+      (relu2): ReLU()
+    )
+    >>> # performs a dry run to initialize the network's lazy modules
+    >>> lazy_mlp(torch.ones(10, 10).cuda())
+    >>> # after initialization, LazyLinear modules become regular Linear modules
+    >>> lazy_mlp
+    LazyMLP(
+      (fc1): Linear(in_features=10, out_features=10, bias=True)
+      (relu1): ReLU()
+      (fc2): Linear(in_features=10, out_features=1, bias=True)
+      (relu2): ReLU()
+    )
+    >>> # attaches an optimizer, since parameters can now be used as usual
+    >>> optim = torch.optim.SGD(lazy_mlp.parameters(), lr=0.01)
+
+    A final caveat when using lazy modules is that the order of initialization of a network's
+    parameters may change, since the lazy modules are always initialized after other modules.
+    For example, if the LazyMLP class defined above had a :class:`torch.nn.LazyLinear` module
+    first and then a regular :class:`torch.nn.Linear` second, the second module would be
+    initialized on construction and the first module would be initialized during the first dry run.
+    This can cause the parameters of a network using lazy modules to be initialized differently
+    than the parameters of a network without lazy modules as the order of parameter initializations,
+    which often depends on a stateful random number generator, is different.
+    Check :doc:`/notes/randomness` for more details.
+
+    Lazy modules can be serialized with a state dict like other modules. For example:
+
+    >>> lazy_mlp = LazyMLP()
+    >>> # The state dict shows the uninitialized parameters
+    >>> lazy_mlp.state_dict()
+    OrderedDict({'fc1.weight': <UninitializedParameter>,
+                 'fc1.bias': <UninitializedParameter>,
+                 'fc2.weight': <UninitializedParameter>,
+                 'fc2.bias': <UninitializedParameter>})
+
+    Lazy modules can load regular :class:`torch.nn.Parameter` s (i.e. you can serialize/deserialize
+    initialized LazyModules and they will remain initialized)
+
+
+    >>> full_mlp = LazyMLP()
+    >>> # Dry run to initialize another module
+    >>> full_mlp.forward(torch.ones(10, 1))
+    >>> # Load an initialized state into a lazy module
+    >>> lazy_mlp.load_state_dict(full_mlp.state_dict())
+    >>> # The state dict now holds valid values
+    >>> lazy_mlp.state_dict()
+    OrderedDict([('fc1.weight',
+                  tensor([[-0.3837],
+                          [ 0.0907],
+                          [ 0.6708],
+                          [-0.5223],
+                          [-0.9028],
+                          [ 0.2851],
+                          [-0.4537],
+                          [ 0.6813],
+                          [ 0.5766],
+                          [-0.8678]])),
+                 ('fc1.bias',
+                  tensor([-1.8832e+25,  4.5636e-41, -1.8832e+25,  4.5636e-41, -6.1598e-30,
+                           4.5637e-41, -1.8788e+22,  4.5636e-41, -2.0042e-31,  4.5637e-41])),
+                 ('fc2.weight',
+                  tensor([[ 0.1320,  0.2938,  0.0679,  0.2793,  0.1088, -0.1795, -0.2301,  0.2807,
+                            0.2479,  0.1091]])),
+                 ('fc2.bias', tensor([0.0019]))])
+
+    Note, however, that the loaded parameters will not be replaced when doing a "dry run" if they are initialized
+    when the state is loaded. This prevents using initialized modules in different contexts.
+    """
+
+    # modules inheriting from this will change their __class__ to the specified
+    # one after they are fully initialized
+    cls_to_become: type[Any] | None = None
+
+    def __init__(self: _LazyProtocol, *args, **kwargs):
+        # Mypy doesn't like this super call in a mixin
+        super().__init__(*args, **kwargs)  # type: ignore[misc]
+        # pyrefly: ignore [read-only]
+        self._load_hook = self._register_load_state_dict_pre_hook(self._lazy_load_hook)
+        # pyrefly: ignore [read-only]
+        self._initialize_hook = self.register_forward_pre_hook(
+            self._infer_parameters, with_kwargs=True
+        )
+
+    def _save_to_state_dict(self: _LazyProtocol, destination, prefix, keep_vars):
+        # This should be ideally implemented as a hook,
+        # but we should override `detach` in the UninitializedParameter to return itself
+        # which is not clean
+        for name, param in self._parameters.items():
+            if param is not None:
+                if not (is_lazy(param) or keep_vars):
+                    param = param.detach()
+                destination[prefix + name] = param
+        for name, buf in self._buffers.items():
+            if buf is not None and name not in self._non_persistent_buffers_set:
+                if not (is_lazy(buf) or keep_vars):
+                    buf = buf.detach()
+                destination[prefix + name] = buf
+
+    def _lazy_load_hook(
+        self: _LazyProtocol,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        """load_state_dict pre-hook function for lazy buffers and parameters.
+
+        The purpose of this hook is to adjust the current state and/or
+        ``state_dict`` being loaded so that a module instance serialized in
+        both un/initialized state can be deserialized onto both un/initialized
+        module instance.
+        See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``
+        for the details of the hook specification.
+        """
+        for name, param in itertools.chain(
+            self._parameters.items(), self._buffers.items()
+        ):
+            key = prefix + name
+            if key in state_dict and param is not None:
+                input_param = state_dict[key]
+                if is_lazy(param):
+                    # The current parameter is not initialized but the one being loaded one is
+                    # create a new parameter based on the uninitialized one
+                    if not is_lazy(input_param):
+                        with torch.no_grad():
+                            param.materialize(input_param.shape)
+
+    def initialize_parameters(self: _LazyProtocol, *args, **kwargs):
+        r"""Initialize parameters according to the input batch properties.
+
+        This adds an interface to isolate parameter initialization from the
+        forward pass when doing parameter shape inference.
+        """
+        raise NotImplementedError(
+            f"initialize_parameters is not implemented for {self.__class__.__name__}"
+        )
+
+    def has_uninitialized_params(self: _LazyProtocol):
+        r"""Check if a module has parameters that are not initialized."""
+        # This is to avoid the JIT to track this parameter and force
+        # custom modules __setstate__ to add it
+        params = self._parameters.values()
+        buffers = self._buffers.values()
+        for param in itertools.chain(params, buffers):
+            if is_lazy(param):
+                return True
+        return False
+
+    # torchrec tests the code consistency with the following code
+    # fmt: off
+    def _infer_parameters(self: _LazyProtocol, module, args, kwargs=None):
+        r"""Infers the size and initializes the parameters according to the provided input batch.
+
+        Given a module that contains parameters that were declared inferable
+        using :class:`torch.nn.parameter.ParameterMode.Infer`, runs a forward pass
+        in the complete module using the provided input to initialize all the parameters
+        as needed.
+        The module is set into evaluation mode before running the forward pass in order
+        to avoid saving statistics or calculating gradients
+        """
+        kwargs = kwargs if kwargs else {}
+        module.initialize_parameters(*args, **kwargs)
+        if module.has_uninitialized_params():
+            raise RuntimeError(f'module {self._get_name()} has not been fully initialized')
+        module._initialize_hook.remove()
+        module._load_hook.remove()
+        delattr(module, '_initialize_hook')
+        delattr(module, '_load_hook')
+        if module.cls_to_become is not None:
+            module.__class__ = module.cls_to_become
+    # fmt: on
+
+    def _replicate_for_data_parallel(self: _LazyProtocol):
+        raise RuntimeError(
+            "Modules with uninitialized parameters can't be used with `DataParallel`. "
+            "Run a dummy forward pass to correctly initialize the modules"
+        )
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/linear.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..30d693b3278b2ef9012f4cf817e785d554af7c7a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/linear.py
@@ -0,0 +1,338 @@
+# mypy: allow-untyped-defs
+import math
+from typing import Any
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F, init
+from torch.nn.parameter import Parameter, UninitializedParameter
+
+from .lazy import LazyModuleMixin
+from .module import Module
+
+
+__all__ = [
+    "Bilinear",
+    "Identity",
+    "LazyLinear",
+    "Linear",
+]
+
+
+class Identity(Module):
+    r"""A placeholder identity operator that is argument-insensitive.
+
+    Args:
+        args: any argument (unused)
+        kwargs: any keyword argument (unused)
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Output: :math:`(*)`, same shape as the input.
+
+    Examples::
+
+        >>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False)
+        >>> input = torch.randn(128, 20)
+        >>> output = m(input)
+        >>> print(output.size())
+        torch.Size([128, 20])
+
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__()
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return input
+
+
+class Linear(Module):
+    r"""Applies an affine linear transformation to the incoming data: :math:`y = xA^T + b`.
+
+    This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    Args:
+        in_features: size of each input sample
+        out_features: size of each output sample
+        bias: If set to ``False``, the layer will not learn an additive bias.
+            Default: ``True``
+
+    Shape:
+        - Input: :math:`(*, H_\text{in})` where :math:`*` means any number of
+          dimensions including none and :math:`H_\text{in} = \text{in\_features}`.
+        - Output: :math:`(*, H_\text{out})` where all but the last dimension
+          are the same shape as the input and :math:`H_\text{out} = \text{out\_features}`.
+
+    Attributes:
+        weight: the learnable weights of the module of shape
+            :math:`(\text{out\_features}, \text{in\_features})`. The values are
+            initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
+            :math:`k = \frac{1}{\text{in\_features}}`
+        bias:   the learnable bias of the module of shape :math:`(\text{out\_features})`.
+                If :attr:`bias` is ``True``, the values are initialized from
+                :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                :math:`k = \frac{1}{\text{in\_features}}`
+
+    Examples::
+
+        >>> m = nn.Linear(20, 30)
+        >>> input = torch.randn(128, 20)
+        >>> output = m(input)
+        >>> print(output.size())
+        torch.Size([128, 30])
+    """
+
+    __constants__ = ["in_features", "out_features"]
+    in_features: int
+    out_features: int
+    weight: Tensor
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(
+            torch.empty((out_features, in_features), **factory_kwargs)
+        )
+        if bias:
+            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
+        else:
+            self.register_parameter("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in ``__init__``.
+        """
+        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
+        # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see
+        # https://github.com/pytorch/pytorch/issues/57109
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.linear(input, self.weight, self.bias)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}"
+
+
+# This class exists solely to avoid triggering an obscure error when scripting
+# an improperly quantized attention layer. See this issue for details:
+# https://github.com/pytorch/pytorch/issues/58969
+# TODO: fail fast on quantization API usage error, then remove this class
+# and replace uses of it with plain Linear
+class NonDynamicallyQuantizableLinear(Linear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__(
+            in_features, out_features, bias=bias, device=device, dtype=dtype
+        )
+
+
+class Bilinear(Module):
+    r"""Applies a bilinear transformation to the incoming data: :math:`y = x_1^T A x_2 + b`.
+
+    Args:
+        in1_features: size of each first input sample, must be > 0
+        in2_features: size of each second input sample, must be > 0
+        out_features: size of each output sample, must be > 0
+        bias: If set to ``False``, the layer will not learn an additive bias.
+            Default: ``True``
+
+    Shape:
+        - Input1: :math:`(*, H_\text{in1})` where :math:`H_\text{in1}=\text{in1\_features}` and
+          :math:`*` means any number of additional dimensions including none. All but the last dimension
+          of the inputs should be the same.
+        - Input2: :math:`(*, H_\text{in2})` where :math:`H_\text{in2}=\text{in2\_features}`.
+        - Output: :math:`(*, H_\text{out})` where :math:`H_\text{out}=\text{out\_features}`
+          and all but the last dimension are the same shape as the input.
+
+    Attributes:
+        weight: the learnable weights of the module of shape
+            :math:`(\text{out\_features}, \text{in1\_features}, \text{in2\_features})`.
+            The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
+            :math:`k = \frac{1}{\text{in1\_features}}`
+        bias:   the learnable bias of the module of shape :math:`(\text{out\_features})`.
+                If :attr:`bias` is ``True``, the values are initialized from
+                :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
+                :math:`k = \frac{1}{\text{in1\_features}}`
+
+    Examples::
+
+        >>> m = nn.Bilinear(20, 30, 40)
+        >>> input1 = torch.randn(128, 20)
+        >>> input2 = torch.randn(128, 30)
+        >>> output = m(input1, input2)
+        >>> print(output.size())
+        torch.Size([128, 40])
+    """
+
+    __constants__ = ["in1_features", "in2_features", "out_features"]
+    in1_features: int
+    in2_features: int
+    out_features: int
+    weight: Tensor
+
+    def __init__(
+        self,
+        in1_features: int,
+        in2_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.in1_features = in1_features
+        self.in2_features = in2_features
+        self.out_features = out_features
+        self.weight = Parameter(
+            torch.empty((out_features, in1_features, in2_features), **factory_kwargs)
+        )
+
+        if bias:
+            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
+        else:
+            self.register_parameter("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in ``__init__``.
+        """
+        if self.in1_features <= 0:
+            raise ValueError(
+                f"in1_features must be > 0, but got (in1_features={self.in1_features})"
+            )
+        bound = 1 / math.sqrt(self.weight.size(1))
+        init.uniform_(self.weight, -bound, bound)
+        if self.bias is not None:
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input1: Tensor, input2: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.bilinear(input1, input2, self.weight, self.bias)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return (
+            f"in1_features={self.in1_features}, in2_features={self.in2_features}, "
+            f"out_features={self.out_features}, bias={self.bias is not None}"
+        )
+
+
+class LazyLinear(LazyModuleMixin, Linear):
+    r"""A :class:`torch.nn.Linear` module where `in_features` is inferred.
+
+    In this module, the `weight` and `bias` are of :class:`torch.nn.UninitializedParameter`
+    class. They will be initialized after the first call to ``forward`` is done and the
+    module will become a regular :class:`torch.nn.Linear` module. The ``in_features`` argument
+    of the :class:`Linear` is inferred from the ``input.shape[-1]``.
+
+    Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
+    on lazy modules and their limitations.
+
+    Args:
+        out_features: size of each output sample
+        bias: If set to ``False``, the layer will not learn an additive bias.
+            Default: ``True``
+
+    Attributes:
+        weight: the learnable weights of the module of shape
+            :math:`(\text{out\_features}, \text{in\_features})`. The values are
+            initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
+            :math:`k = \frac{1}{\text{in\_features}}`
+        bias:   the learnable bias of the module of shape :math:`(\text{out\_features})`.
+                If :attr:`bias` is ``True``, the values are initialized from
+                :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
+                :math:`k = \frac{1}{\text{in\_features}}`
+
+
+    """
+
+    cls_to_become = Linear  # type: ignore[assignment]
+    # pyrefly: ignore [bad-override]
+    weight: UninitializedParameter
+    bias: UninitializedParameter  # type: ignore[assignment]
+
+    def __init__(
+        self, out_features: int, bias: bool = True, device=None, dtype=None
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # bias is hardcoded to False to avoid creating tensor
+        # that will soon be overwritten.
+        # pyrefly: ignore [bad-argument-type]
+        super().__init__(0, 0, False)
+        # pyrefly: ignore [unexpected-keyword]
+        self.weight = UninitializedParameter(**factory_kwargs)
+        self.out_features = out_features
+        if bias:
+            # pyrefly: ignore [unexpected-keyword]
+            self.bias = UninitializedParameter(**factory_kwargs)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in ``__init__``.
+        """
+        # pyrefly: ignore [bad-argument-type]
+        if not self.has_uninitialized_params() and self.in_features != 0:
+            super().reset_parameters()
+
+    def initialize_parameters(self, input) -> None:  # type: ignore[override]
+        """
+        Infers ``in_features`` based on ``input`` and initializes parameters.
+        """
+        # pyrefly: ignore [bad-argument-type]
+        if self.has_uninitialized_params():
+            with torch.no_grad():
+                self.in_features = input.shape[-1]
+                self.weight.materialize((self.out_features, self.in_features))
+                if self.bias is not None:
+                    self.bias.materialize((self.out_features,))
+                self.reset_parameters()
+        if self.in_features == 0:
+            if input.shape[-1] != self.weight.shape[-1]:
+                raise AssertionError(
+                    f"The in_features inferred from input: {input.shape[-1]} "
+                    f"is not equal to in_features from self.weight: "
+                    f"{self.weight.shape[-1]}"
+                )
+            self.in_features = input.shape[-1]
+
+
+# TODO: PartialLinear - maybe in sparse?
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/loss.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..9af79927c1b31d7c4a51fe9b3f00cf6b3bc8a1c4
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/loss.py
@@ -0,0 +1,2123 @@
+# mypy: allow-untyped-defs
+from collections.abc import Callable
+from typing_extensions import deprecated
+
+from torch import Tensor
+from torch.nn import _reduction as _Reduction, functional as F
+
+from .distance import PairwiseDistance
+from .module import Module
+
+
+__all__ = [
+    "L1Loss",
+    "NLLLoss",
+    "NLLLoss2d",
+    "PoissonNLLLoss",
+    "GaussianNLLLoss",
+    "KLDivLoss",
+    "MSELoss",
+    "BCELoss",
+    "BCEWithLogitsLoss",
+    "HingeEmbeddingLoss",
+    "MultiLabelMarginLoss",
+    "SmoothL1Loss",
+    "HuberLoss",
+    "SoftMarginLoss",
+    "CrossEntropyLoss",
+    "MultiLabelSoftMarginLoss",
+    "CosineEmbeddingLoss",
+    "MarginRankingLoss",
+    "MultiMarginLoss",
+    "TripletMarginLoss",
+    "TripletMarginWithDistanceLoss",
+    "CTCLoss",
+]
+
+
+class _Loss(Module):
+    reduction: str
+
+    def __init__(self, size_average=None, reduce=None, reduction: str = "mean") -> None:
+        super().__init__()
+        if size_average is not None or reduce is not None:
+            self.reduction: str = _Reduction.legacy_get_string(size_average, reduce)
+        else:
+            self.reduction = reduction
+
+
+class _WeightedLoss(_Loss):
+    def __init__(
+        self,
+        weight: Tensor | None = None,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.register_buffer("weight", weight)
+        self.weight: Tensor | None
+
+
+class L1Loss(_Loss):
+    r"""Creates a criterion that measures the mean absolute error (MAE) between each element in
+    the input :math:`x` and target :math:`y`.
+
+    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = \left| x_n - y_n \right|,
+
+    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``
+    (default ``'mean'``), then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total
+    of :math:`N` elements each.
+
+    The sum operation still operates over all the elements, and divides by :math:`N`.
+
+    The division by :math:`N` can be avoided if one sets ``reduction = 'sum'``.
+
+    Supports real-valued and complex-valued inputs.
+
+    Args:
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then
+          :math:`(*)`, same shape as the input.
+
+    Examples:
+
+        >>> loss = nn.L1Loss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5)
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["reduction"]
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.l1_loss(input, target, reduction=self.reduction)
+
+
+class NLLLoss(_WeightedLoss):
+    r"""The negative log likelihood loss. It is useful to train a classification
+    problem with `C` classes.
+
+    If provided, the optional argument :attr:`weight` should be a 1D Tensor assigning
+    weight to each of the classes. This is particularly useful when you have an
+    unbalanced training set.
+
+    The `input` given through a forward call is expected to contain
+    log-probabilities of each class. `input` has to be a Tensor of size either
+    :math:`(minibatch, C)` or :math:`(minibatch, C, d_1, d_2, ..., d_K)`
+    with :math:`K \geq 1` for the `K`-dimensional case. The latter is useful for
+    higher dimension inputs, such as computing NLL loss per-pixel for 2D images.
+
+    Obtaining log-probabilities in a neural network is easily achieved by
+    adding a  `LogSoftmax`  layer in the last layer of your network.
+    You may use `CrossEntropyLoss` instead, if you prefer not to add an extra
+    layer.
+
+    The `target` that this loss expects should be a class index in the range :math:`[0, C-1]`
+    where `C = number of classes`; if `ignore_index` is specified, this loss also accepts
+    this class index (this index may not necessarily be in the class range).
+
+    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \\
+        l_n = - w_{y_n} x_{n,y_n}, \\
+        w_{c} = \text{weight}[c] \cdot \mathbb{1}\{c \not= \text{ignore\_index}\},
+
+    where :math:`x` is the input, :math:`y` is the target, :math:`w` is the weight, and
+    :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``
+    (default ``'mean'``), then
+
+    .. math::
+        \ell(x, y) = \begin{cases}
+            \sum_{n=1}^N \frac{1}{\sum_{n=1}^N w_{y_n}} l_n, &
+            \text{if reduction} = \text{`mean';}\\
+            \sum_{n=1}^N l_n,  &
+            \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    Args:
+        weight (Tensor, optional): a manual rescaling weight given to each
+            class. If given, it has to be a Tensor of size `C`. Otherwise, it is
+            treated as if having all ones.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``None``
+        ignore_index (int, optional): Specifies a target value that is ignored
+            and does not contribute to the input gradient. When
+            :attr:`size_average` is ``True``, the loss is averaged over
+            non-ignored targets.
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``None``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will
+            be applied, ``'mean'``: the weighted mean of the output is taken,
+            ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in
+            the meantime, specifying either of those two args will override
+            :attr:`reduction`. Default: ``'mean'``
+
+    Shape::
+        - Input: :math:`(N, C)` or :math:`(C)`, where `C = number of classes`, `N = batch size`, or
+          :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \geq 1`
+          in the case of `K`-dimensional loss.
+        - Target: :math:`(N)` or :math:`()`, where each value is
+          :math:`0 \leq \text{targets}[i] \leq C-1`, or
+          :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \geq 1` in the case of
+          K-dimensional loss.
+        - Output: If :attr:`reduction` is ``'none'``, shape :math:`(N)` or
+          :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \geq 1` in the case of K-dimensional loss.
+          Otherwise, scalar.
+
+    Examples:
+
+        >>> log_softmax = nn.LogSoftmax(dim=1)
+        >>> loss_fn = nn.NLLLoss()
+        >>> # input to NLLLoss is of size N x C = 3 x 5
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> # each element in target must have 0 <= value < C
+        >>> target = torch.tensor([1, 0, 4])
+        >>> loss = loss_fn(log_softmax(input), target)
+        >>> loss.backward()
+        >>>
+        >>>
+        >>> # 2D loss example (used, for example, with image inputs)
+        >>> N, C = 5, 4
+        >>> loss_fn = nn.NLLLoss()
+        >>> data = torch.randn(N, 16, 10, 10)
+        >>> conv = nn.Conv2d(16, C, (3, 3))
+        >>> log_softmax = nn.LogSoftmax(dim=1)
+        >>> # output of conv forward is of shape [N, C, 8, 8]
+        >>> output = log_softmax(conv(data))
+        >>> # each element in target must have 0 <= value < C
+        >>> target = torch.empty(N, 8, 8, dtype=torch.long).random_(0, C)
+        >>> # input to NLLLoss is of size N x C x height (8) x width (8)
+        >>> loss = loss_fn(output, target)
+        >>> loss.backward()
+    """
+
+    __constants__ = ["ignore_index", "reduction"]
+    ignore_index: int
+
+    def __init__(
+        self,
+        weight: Tensor | None = None,
+        size_average=None,
+        ignore_index: int = -100,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(weight, size_average, reduce, reduction)
+        self.ignore_index = ignore_index
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.nll_loss(
+            input,
+            target,
+            weight=self.weight,
+            ignore_index=self.ignore_index,
+            reduction=self.reduction,
+        )
+
+
+@deprecated(
+    "`NLLLoss2d` has been deprecated. "
+    "Please use `NLLLoss` instead as a drop-in replacement and see "
+    "https://pytorch.org/docs/main/nn.html#torch.nn.NLLLoss for more details.",
+    category=FutureWarning,
+)
+class NLLLoss2d(NLLLoss):
+    def __init__(
+        self,
+        weight: Tensor | None = None,
+        size_average=None,
+        ignore_index: int = -100,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(weight, size_average, ignore_index, reduce, reduction)
+
+
+class PoissonNLLLoss(_Loss):
+    r"""Negative log likelihood loss with Poisson distribution of target.
+
+    The loss can be described as:
+
+    .. math::
+        \text{target} \sim \mathrm{Poisson}(\text{input})
+
+        \text{loss}(\text{input}, \text{target}) = \text{input} - \text{target} * \log(\text{input})
+                                    + \log(\text{target!})
+
+    The last term can be omitted or approximated with Stirling formula. The
+    approximation is used for target values more than 1. For targets less or
+    equal to 1 zeros are added to the loss.
+
+    Args:
+        log_input (bool, optional): if ``True`` the loss is computed as
+            :math:`\exp(\text{input}) - \text{target}*\text{input}`, if ``False`` the loss is
+            :math:`\text{input} - \text{target}*\log(\text{input}+\text{eps})`.
+        full (bool, optional): whether to compute full loss, i. e. to add the
+            Stirling approximation term
+
+            .. math::
+                \text{target}*\log(\text{target}) - \text{target} + 0.5 * \log(2\pi\text{target}).
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        eps (float, optional): Small value to avoid evaluation of :math:`\log(0)` when
+            :attr:`log_input = False`. Default: 1e-8
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Examples:
+
+        >>> loss = nn.PoissonNLLLoss()
+        >>> log_input = torch.randn(5, 2, requires_grad=True)
+        >>> target = torch.randn(5, 2)
+        >>> output = loss(log_input, target)
+        >>> output.backward()
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar by default. If :attr:`reduction` is ``'none'``, then :math:`(*)`,
+          the same shape as the input.
+    """
+
+    __constants__ = ["log_input", "full", "eps", "reduction"]
+    log_input: bool
+    full: bool
+    eps: float
+
+    def __init__(
+        self,
+        log_input: bool = True,
+        full: bool = False,
+        size_average=None,
+        eps: float = 1e-8,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.log_input = log_input
+        self.full = full
+        self.eps = eps
+
+    def forward(self, log_input: Tensor, target: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.poisson_nll_loss(
+            log_input,
+            target,
+            log_input=self.log_input,
+            full=self.full,
+            eps=self.eps,
+            reduction=self.reduction,
+        )
+
+
+class GaussianNLLLoss(_Loss):
+    r"""Gaussian negative log likelihood loss.
+
+    The targets are treated as samples from Gaussian distributions with
+    expectations and variances predicted by the neural network. For a
+    ``target`` tensor modelled as having Gaussian distribution with a tensor
+    of expectations ``input`` and a tensor of positive variances ``var`` the loss is:
+
+    .. math::
+        \text{loss} = \frac{1}{2}\left(\log\left(\text{max}\left(\text{var},
+        \ \text{eps}\right)\right) + \frac{\left(\text{input} - \text{target}\right)^2}
+        {\text{max}\left(\text{var}, \ \text{eps}\right)}\right) + \text{const.}
+
+    where :attr:`eps` is used for stability. By default, the constant term of
+    the loss function is omitted unless :attr:`full` is ``True``. If ``var`` is not the same
+    size as ``input`` (due to a homoscedastic assumption), it must either have a final dimension
+    of 1 or have one fewer dimension (with all other sizes being the same) for correct broadcasting.
+
+    Args:
+        full (bool, optional): include the constant term in the loss
+            calculation. Default: ``False``.
+        eps (float, optional): value used to clamp ``var`` (see note below), for
+            stability. Default: 1e-6.
+        reduction (str, optional): specifies the reduction to apply to the
+            output:``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction
+            will be applied, ``'mean'``: the output is the average of all batch
+            member losses, ``'sum'``: the output is the sum of all batch member
+            losses. Default: ``'mean'``.
+
+    Shape:
+        - Input: :math:`(N, *)` or :math:`(*)` where :math:`*` means any number of additional
+          dimensions
+        - Target: :math:`(N, *)` or :math:`(*)`, same shape as the input, or same shape as the input
+          but with one dimension equal to 1 (to allow for broadcasting)
+        - Var: :math:`(N, *)` or :math:`(*)`, same shape as the input, or same shape as the input but
+          with one dimension equal to 1, or same shape as the input but with one fewer
+          dimension (to allow for broadcasting), or a scalar value
+        - Output: scalar if :attr:`reduction` is ``'mean'`` (default) or
+          ``'sum'``. If :attr:`reduction` is ``'none'``, then :math:`(N, *)`, same
+          shape as the input
+
+    Examples:
+        >>> loss = nn.GaussianNLLLoss()
+        >>> input = torch.randn(5, 2, requires_grad=True)
+        >>> target = torch.randn(5, 2)
+        >>> var = torch.ones(5, 2, requires_grad=True)  # heteroscedastic
+        >>> output = loss(input, target, var)
+        >>> output.backward()
+
+        >>> loss = nn.GaussianNLLLoss()
+        >>> input = torch.randn(5, 2, requires_grad=True)
+        >>> target = torch.randn(5, 2)
+        >>> var = torch.ones(5, 1, requires_grad=True)  # homoscedastic
+        >>> output = loss(input, target, var)
+        >>> output.backward()
+
+    Note:
+        The clamping of ``var`` is ignored with respect to autograd, and so the
+        gradients are unaffected by it.
+
+    Reference:
+        Nix, D. A. and Weigend, A. S., "Estimating the mean and variance of the
+        target probability distribution", Proceedings of 1994 IEEE International
+        Conference on Neural Networks (ICNN'94), Orlando, FL, USA, 1994, pp. 55-60
+        vol.1, doi: 10.1109/ICNN.1994.374138.
+    """
+
+    __constants__ = ["full", "eps", "reduction"]
+    full: bool
+    eps: float
+
+    def __init__(
+        self, *, full: bool = False, eps: float = 1e-6, reduction: str = "mean"
+    ) -> None:
+        super().__init__(None, None, reduction)
+        self.full = full
+        self.eps = eps
+
+    def forward(self, input: Tensor, target: Tensor, var: Tensor | float) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.gaussian_nll_loss(
+            input, target, var, full=self.full, eps=self.eps, reduction=self.reduction
+        )
+
+
+class KLDivLoss(_Loss):
+    r"""The Kullback-Leibler divergence loss.
+
+    For tensors of the same shape :math:`y_{\text{pred}},\ y_{\text{true}}`,
+    where :math:`y_{\text{pred}}` is the :attr:`input` and :math:`y_{\text{true}}` is the
+    :attr:`target`, we define the **pointwise KL-divergence** as
+
+    .. math::
+
+        L(y_{\text{pred}},\ y_{\text{true}})
+            = y_{\text{true}} \cdot \log \frac{y_{\text{true}}}{y_{\text{pred}}}
+            = y_{\text{true}} \cdot (\log y_{\text{true}} - \log y_{\text{pred}})
+
+    To avoid underflow issues when computing this quantity, this loss expects the argument
+    :attr:`input` in the log-space. The argument :attr:`target` may also be provided in the
+    log-space if :attr:`log_target`\ `= True`.
+
+    To summarise, this function is roughly equivalent to computing
+
+    .. code-block:: python
+
+        if not log_target:  # default
+            loss_pointwise = target * (target.log() - input)
+        else:
+            loss_pointwise = target.exp() * (target - input)
+
+    and then reducing this result depending on the argument :attr:`reduction` as
+
+    .. code-block:: python
+
+        if reduction == "mean":  # default
+            loss = loss_pointwise.mean()
+        elif reduction == "batchmean":  # mathematically correct
+            loss = loss_pointwise.sum() / input.size(0)
+        elif reduction == "sum":
+            loss = loss_pointwise.sum()
+        else:  # reduction == "none"
+            loss = loss_pointwise
+
+    .. note::
+        As all the other losses in PyTorch, this function expects the first argument,
+        :attr:`input`, to be the output of the model (e.g. the neural network)
+        and the second, :attr:`target`, to be the observations in the dataset.
+        This differs from the standard mathematical notation :math:`KL(P\ ||\ Q)` where
+        :math:`P` denotes the distribution of the observations and :math:`Q` denotes the model.
+
+    .. warning::
+        :attr:`reduction`\ `= "mean"` doesn't return the true KL divergence value, please use
+        :attr:`reduction`\ `= "batchmean"` which aligns with the mathematical definition.
+
+    Args:
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to `False`, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is `False`. Default: `True`
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is `False`, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: `True`
+        reduction (str, optional): Specifies the reduction to apply to the output. Default: `"mean"`
+        log_target (bool, optional): Specifies whether `target` is the log space. Default: `False`
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar by default. If :attr:`reduction` is `'none'`, then :math:`(*)`,
+          same shape as the input.
+
+    Examples:
+        >>> kl_loss = nn.KLDivLoss(reduction="batchmean")
+        >>> # input should be a distribution in the log space
+        >>> input = F.log_softmax(torch.randn(3, 5, requires_grad=True), dim=1)
+        >>> # Sample a batch of distributions. Usually this would come from the dataset
+        >>> target = F.softmax(torch.rand(3, 5), dim=1)
+        >>> output = kl_loss(input, target)
+        >>>
+        >>> kl_loss = nn.KLDivLoss(reduction="batchmean", log_target=True)
+        >>> log_target = F.log_softmax(torch.rand(3, 5), dim=1)
+        >>> output = kl_loss(input, log_target)
+    """
+
+    __constants__ = ["reduction"]
+
+    def __init__(
+        self,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+        log_target: bool = False,
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.log_target = log_target
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.kl_div(
+            input, target, reduction=self.reduction, log_target=self.log_target
+        )
+
+
+class MSELoss(_Loss):
+    r"""Creates a criterion that measures the mean squared error (squared L2 norm) between
+    each element in the input :math:`x` and target :math:`y`.
+
+    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = \left( x_n - y_n \right)^2,
+
+    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``
+    (default ``'mean'``), then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), &  \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  &  \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    :math:`x` and :math:`y` are tensors of arbitrary shapes with a total
+    of :math:`N` elements each.
+
+    The mean operation still operates over all the elements, and divides by :math:`N`.
+
+    The division by :math:`N` can be avoided if one sets ``reduction = 'sum'``.
+
+    Args:
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+
+    Examples:
+
+        >>> loss = nn.MSELoss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5)
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["reduction"]
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.mse_loss(input, target, reduction=self.reduction)
+
+
+class BCELoss(_WeightedLoss):
+    r"""Creates a criterion that measures the Binary Cross Entropy between the target and
+    the input probabilities:
+
+    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = - w_n \left[ y_n \cdot \log x_n + (1 - y_n) \cdot \log (1 - x_n) \right],
+
+    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``
+    (default ``'mean'``), then
+
+    .. math::
+        \ell(x, y) = \begin{cases}
+            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    This is used for measuring the error of a reconstruction in for example
+    an auto-encoder. Note that the targets :math:`y` should be numbers
+    between 0 and 1.
+
+    Notice that if :math:`x_n` is either 0 or 1, one of the log terms would be
+    mathematically undefined in the above loss equation. PyTorch chooses to set
+    :math:`\log (0) = -\infty`, since :math:`\lim_{x\to 0} \log (x) = -\infty`.
+    However, an infinite term in the loss equation is not desirable for several reasons.
+
+    For one, if either :math:`y_n = 0` or :math:`(1 - y_n) = 0`, then we would be
+    multiplying 0 with infinity. Secondly, if we have an infinite loss value, then
+    we would also have an infinite term in our gradient, since
+    :math:`\lim_{x\to 0} \frac{d}{dx} \log (x) = \infty`.
+    This would make BCELoss's backward method nonlinear with respect to :math:`x_n`,
+    and using it for things like linear regression would not be straight-forward.
+
+    Our solution is that BCELoss clamps its log function outputs to be greater than
+    or equal to -100. This way, we can always have a finite loss value and a linear
+    backward method.
+
+
+    Args:
+        weight (Tensor, optional): a manual rescaling weight given to the loss
+            of each batch element. If given, has to be a Tensor of size `nbatch`.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same
+          shape as input.
+
+    Examples:
+
+        >>> m = nn.Sigmoid()
+        >>> loss = nn.BCELoss()
+        >>> input = torch.randn(3, 2, requires_grad=True)
+        >>> target = torch.rand(3, 2, requires_grad=False)
+        >>> output = loss(m(input), target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["reduction"]
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.binary_cross_entropy(
+            input, target, weight=self.weight, reduction=self.reduction
+        )
+
+
+class BCEWithLogitsLoss(_Loss):
+    r"""This loss combines a `Sigmoid` layer and the `BCELoss` in one single
+    class. This version is more numerically stable than using a plain `Sigmoid`
+    followed by a `BCELoss` as, by combining the operations into one layer,
+    we take advantage of the log-sum-exp trick for numerical stability.
+
+    The unreduced (i.e. with :attr:`reduction` set to ``'none'``) loss can be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_n = - w_n \left[ y_n \cdot \log \sigma(x_n)
+        + (1 - y_n) \cdot \log (1 - \sigma(x_n)) \right],
+
+    where :math:`N` is the batch size. If :attr:`reduction` is not ``'none'``
+    (default ``'mean'``), then
+
+    .. math::
+        \ell(x, y) = \begin{cases}
+            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    This is used for measuring the error of a reconstruction in for example
+    an auto-encoder. Note that the targets `t[i]` should be numbers
+    between 0 and 1.
+
+    It's possible to trade off recall and precision by adding weights to positive examples.
+    In the case of multi-label classification the loss can be described as:
+
+    .. math::
+        \ell_c(x, y) = L_c = \{l_{1,c},\dots,l_{N,c}\}^\top, \quad
+        l_{n,c} = - w_{n,c} \left[ p_c y_{n,c} \cdot \log \sigma(x_{n,c})
+        + (1 - y_{n,c}) \cdot \log (1 - \sigma(x_{n,c})) \right],
+
+    where :math:`c` is the class number (:math:`c > 1` for multi-label binary classification,
+    :math:`c = 1` for single-label binary classification),
+    :math:`n` is the number of the sample in the batch and
+    :math:`p_c` is the weight of the positive answer for the class :math:`c`.
+
+    :math:`p_c > 1` increases the recall, :math:`p_c < 1` increases the precision.
+
+    For example, if a dataset contains 100 positive and 300 negative examples of a single class,
+    then ``pos_weight`` for the class should be equal to :math:`\frac{300}{100}=3`.
+    The loss would act as if the dataset contains :math:`3\times 100=300` positive examples.
+
+    Examples:
+
+        >>> target = torch.ones([10, 64], dtype=torch.float32)  # 64 classes, batch size = 10
+        >>> output = torch.full([10, 64], 1.5)  # A prediction (logit)
+        >>> pos_weight = torch.ones([64])  # All weights are equal to 1
+        >>> criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
+        >>> criterion(output, target)  # -log(sigmoid(1.5))
+        tensor(0.20...)
+
+    In the above example, the ``pos_weight`` tensor's elements correspond to the 64 distinct classes
+    in a multi-label binary classification scenario. Each element in ``pos_weight`` is designed to adjust the
+    loss function based on the imbalance between negative and positive samples for the respective class.
+    This approach is useful in datasets with varying levels of class imbalance, ensuring that the loss
+    calculation accurately accounts for the distribution in each class.
+
+    Args:
+        weight (Tensor, optional): a manual rescaling weight given to the loss
+            of each batch element. The dimension of weight supports :ref:`broadcasting to a common shape <broadcasting-semantics>`
+            with respect to the output (and target) shape.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+        pos_weight (Tensor, optional): a weight of positive examples to be broadcasted with target.
+            Must be a tensor with equal size along the class dimension to the number of classes.
+            Pay close attention to PyTorch's broadcasting semantics in order to achieve the desired
+            operations. For a target of size [B, C, H, W] (where B is batch size) pos_weight of
+            size [B, C, H, W] will apply different pos_weights to each element of the batch or
+            [C, H, W] the same pos_weights across the batch. To apply the same positive weight
+            along all spatial dimensions for a 2D multi-class target [C, H, W] use: [C, 1, 1].
+            Default: ``None``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same
+          shape as input.
+
+    Examples:
+
+        >>> loss = nn.BCEWithLogitsLoss()
+        >>> input = torch.randn(3, requires_grad=True)
+        >>> target = torch.empty(3).random_(2)
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    def __init__(
+        self,
+        weight: Tensor | None = None,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+        pos_weight: Tensor | None = None,
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.register_buffer("weight", weight)
+        self.register_buffer("pos_weight", pos_weight)
+        self.weight: Tensor | None
+        self.pos_weight: Tensor | None
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.binary_cross_entropy_with_logits(
+            input,
+            target,
+            self.weight,
+            pos_weight=self.pos_weight,
+            reduction=self.reduction,
+        )
+
+
+class HingeEmbeddingLoss(_Loss):
+    r"""Measures the loss given an input tensor :math:`x` and a labels tensor :math:`y`
+    (containing 1 or -1).
+    This is usually used for measuring whether two inputs are similar or
+    dissimilar, e.g. using the L1 pairwise distance as :math:`x`, and is typically
+    used for learning nonlinear embeddings or semi-supervised learning.
+
+    The loss function for :math:`n`-th sample in the mini-batch is
+
+    .. math::
+        l_n = \begin{cases}
+            x_n, & \text{if}\; y_n = 1,\\
+            \max \{0, margin - x_n\}, & \text{if}\; y_n = -1,
+        \end{cases}
+
+    and the total loss functions is
+
+    .. math::
+        \ell(x, y) = \begin{cases}
+            \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  & \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    where :math:`L = \{l_1,\dots,l_N\}^\top`.
+
+    Args:
+        margin (float, optional): Has a default value of `1`.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(*)` where :math:`*` means, any number of dimensions. The sum operation
+          operates over all the elements.
+        - Target: :math:`(*)`, same shape as the input
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then same shape as the input
+
+    Examples:
+
+        >>> loss = nn.HingeEmbeddingLoss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5).sign()
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["margin", "reduction"]
+    margin: float
+
+    def __init__(
+        self,
+        margin: float = 1.0,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.margin = margin
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.hinge_embedding_loss(
+            input, target, margin=self.margin, reduction=self.reduction
+        )
+
+
+class MultiLabelMarginLoss(_Loss):
+    r"""Creates a criterion that optimizes a multi-class multi-classification
+    hinge loss (margin-based loss) between input :math:`x` (a 2D mini-batch `Tensor`)
+    and output :math:`y` (which is a 2D `Tensor` of target class indices).
+    For each sample in the mini-batch:
+
+    .. math::
+        \text{loss}(x, y) = \sum_{ij}\frac{\max(0, 1 - (x[y[j]] - x[i]))}{\text{x.size}(0)}
+
+    where :math:`x \in \left\{0, \; \cdots , \; \text{x.size}(0) - 1\right\}`, \
+    :math:`y \in \left\{0, \; \cdots , \; \text{y.size}(0) - 1\right\}`, \
+    :math:`0 \leq y[j] \leq \text{x.size}(0)-1`, \
+    and :math:`i \neq y[j]` for all :math:`i` and :math:`j`.
+
+    :math:`y` and :math:`x` must have the same size.
+
+    The criterion only considers a contiguous block of non-negative targets that
+    starts at the front.
+
+    This allows for different samples to have variable amounts of target classes.
+
+    Args:
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(C)` or :math:`(N, C)` where `N` is the batch size and `C`
+          is the number of classes.
+        - Target: :math:`(C)` or :math:`(N, C)`, label targets padded by -1 ensuring same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(N)`.
+
+    Examples:
+
+        >>> loss = nn.MultiLabelMarginLoss()
+        >>> x = torch.FloatTensor([[0.1, 0.2, 0.4, 0.8]])
+        >>> # for target y, only consider labels 3 and 0, not after label -1
+        >>> y = torch.LongTensor([[3, 0, -1, 1]])
+        >>> # 0.25 * ((1-(0.1-0.2)) + (1-(0.1-0.4)) + (1-(0.8-0.2)) + (1-(0.8-0.4)))
+        >>> loss(x, y)
+        tensor(0.85...)
+
+    """
+
+    __constants__ = ["reduction"]
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.multilabel_margin_loss(input, target, reduction=self.reduction)
+
+
+class SmoothL1Loss(_Loss):
+    r"""Creates a criterion that uses a squared term if the absolute
+    element-wise error falls below beta and an L1 term otherwise.
+    It is less sensitive to outliers than :class:`torch.nn.MSELoss` and in some cases
+    prevents exploding gradients (e.g. see the paper `Fast R-CNN`_ by Ross Girshick).
+
+    For a batch of size :math:`N`, the unreduced loss can be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1, ..., l_N\}^T
+
+    with
+
+    .. math::
+        l_n = \begin{cases}
+        0.5 (x_n - y_n)^2 / beta, & \text{if } |x_n - y_n| < beta \\
+        |x_n - y_n| - 0.5 * beta, & \text{otherwise }
+        \end{cases}
+
+    If `reduction` is not `none`, then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), &  \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  &  \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    .. note::
+        Smooth L1 loss can be seen as exactly :class:`L1Loss`, but with the :math:`|x - y| < beta`
+        portion replaced with a quadratic function such that its slope is 1 at :math:`|x - y| = beta`.
+        The quadratic segment smooths the L1 loss near :math:`|x - y| = 0`.
+
+    .. note::
+        Smooth L1 loss is closely related to :class:`HuberLoss`, being
+        equivalent to :math:`huber(x, y) / beta` (note that Smooth L1's beta hyper-parameter is
+        also known as delta for Huber). This leads to the following differences:
+
+        * As beta -> 0, Smooth L1 loss converges to :class:`L1Loss`, while :class:`HuberLoss`
+          converges to a constant 0 loss. When beta is 0, Smooth L1 loss is equivalent to L1 loss.
+        * As beta -> :math:`+\infty`, Smooth L1 loss converges to a constant 0 loss, while
+          :class:`HuberLoss` converges to :class:`MSELoss`.
+        * For Smooth L1 loss, as beta varies, the L1 segment of the loss has a constant slope of 1.
+          For :class:`HuberLoss`, the slope of the L1 segment is beta.
+
+    .. _`Fast R-CNN`: https://arxiv.org/abs/1504.08083
+
+    Args:
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+        beta (float, optional): Specifies the threshold at which to change between L1 and L2 loss.
+            The value must be non-negative. Default: 1.0
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same shape as the input.
+
+    Examples:
+
+        >>> loss = nn.SmoothL1Loss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5)
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["reduction"]
+
+    def __init__(
+        self, size_average=None, reduce=None, reduction: str = "mean", beta: float = 1.0
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.beta = beta
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.smooth_l1_loss(input, target, reduction=self.reduction, beta=self.beta)
+
+
+class HuberLoss(_Loss):
+    r"""Creates a criterion that uses a squared term if the absolute
+    element-wise error falls below delta and a delta-scaled L1 term otherwise.
+    This loss combines advantages of both :class:`L1Loss` and :class:`MSELoss`; the
+    delta-scaled L1 region makes the loss less sensitive to outliers than :class:`MSELoss`,
+    while the L2 region provides smoothness over :class:`L1Loss` near 0. See
+    `Huber loss <https://en.wikipedia.org/wiki/Huber_loss>`_ for more information.
+
+    For a batch of size :math:`N`, the unreduced loss can be described as:
+
+    .. math::
+        \ell(x, y) = L = \{l_1, ..., l_N\}^T
+
+    with
+
+    .. math::
+        l_n = \begin{cases}
+        0.5 (x_n - y_n)^2, & \text{if } |x_n - y_n| < delta \\
+        delta * (|x_n - y_n| - 0.5 * delta), & \text{otherwise }
+        \end{cases}
+
+    If `reduction` is not `none`, then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), &  \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  &  \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    .. note::
+        When delta is set to 1, this loss is equivalent to :class:`SmoothL1Loss`.
+        In general, this loss differs from :class:`SmoothL1Loss` by a factor of delta (AKA beta
+        in Smooth L1).
+        See :class:`SmoothL1Loss` for additional discussion on the differences in behavior
+        between the two losses.
+
+    Args:
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Default: ``'mean'``
+        delta (float, optional): Specifies the threshold at which to change between delta-scaled L1 and L2 loss.
+            The value must be positive.  Default: 1.0
+
+    Shape:
+        - Input: :math:`(*)` where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same shape as the input.
+
+    Examples:
+
+        >>> loss = nn.HuberLoss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5)
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["reduction", "delta"]
+
+    def __init__(self, reduction: str = "mean", delta: float = 1.0) -> None:
+        super().__init__(reduction=reduction)
+        self.delta = delta
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.huber_loss(input, target, reduction=self.reduction, delta=self.delta)
+
+
+class SoftMarginLoss(_Loss):
+    r"""Creates a criterion that optimizes a two-class classification
+    logistic loss between input tensor :math:`x` and target tensor :math:`y`
+    (containing 1 or -1).
+
+    .. math::
+        \text{loss}(x, y) = \sum_i \frac{\log(1 + \exp(-y[i]*x[i]))}{\text{x.nelement}()}
+
+    Args:
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
+        - Target: :math:`(*)`, same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(*)`, same
+          shape as input.
+
+    Examples:
+
+        >>> loss = nn.SoftMarginLoss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5).sign()
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["reduction"]
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.soft_margin_loss(input, target, reduction=self.reduction)
+
+
+class CrossEntropyLoss(_WeightedLoss):
+    r"""This criterion computes the cross entropy loss between input logits
+    and target.
+
+    It is useful when training a classification problem with `C` classes.
+    If provided, the optional argument :attr:`weight` should be a 1D `Tensor`
+    assigning weight to each of the classes.
+    This is particularly useful when you have an unbalanced training set.
+
+    The `input` is expected to contain the unnormalized logits for each class (which do `not` need
+    to be positive or sum to 1, in general).
+    `input` has to be a Tensor of size :math:`(C)` for unbatched input,
+    :math:`(minibatch, C)` or :math:`(minibatch, C, d_1, d_2, ..., d_K)` with :math:`K \geq 1` for the
+    `K`-dimensional case. The last being useful for higher dimension inputs, such
+    as computing cross entropy loss per-pixel for 2D images.
+
+    The `target` that this criterion expects should contain either:
+
+    - Class indices in the range :math:`[0, C)` where :math:`C` is the number of classes; if
+      `ignore_index` is specified, this loss also accepts this class index (this index
+      may not necessarily be in the class range). The unreduced (i.e. with :attr:`reduction`
+      set to ``'none'``) loss for this case can be described as:
+
+      .. math::
+          \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+          l_n = - w_{y_n} \log \frac{\exp(x_{n,y_n})}{\sum_{c=1}^C \exp(x_{n,c})}
+          \cdot \mathbb{1}\{y_n \not= \text{ignore\_index}\}
+
+      where :math:`x` is the input, :math:`y` is the target, :math:`w` is the weight,
+      :math:`C` is the number of classes, and :math:`N` spans the minibatch dimension as well as
+      :math:`d_1, ..., d_k` for the `K`-dimensional case. If
+      :attr:`reduction` is not ``'none'`` (default ``'mean'``), then
+
+      .. math::
+          \ell(x, y) = \begin{cases}
+              \sum_{n=1}^N \frac{1}{\sum_{n=1}^N w_{y_n} \cdot \mathbb{1}\{y_n \not= \text{ignore\_index}\}} l_n, &
+               \text{if reduction} = \text{`mean';}\\
+                \sum_{n=1}^N l_n,  &
+                \text{if reduction} = \text{`sum'.}
+            \end{cases}
+
+      Note that this case is equivalent to applying :class:`~torch.nn.LogSoftmax`
+      on an input, followed by :class:`~torch.nn.NLLLoss`.
+
+    - Probabilities for each class; useful when labels beyond a single class per minibatch item
+      are required, such as for blended labels, label smoothing, etc. The unreduced (i.e. with
+      :attr:`reduction` set to ``'none'``) loss for this case can be described as:
+
+      .. math::
+          \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad
+          l_n = - \sum_{c=1}^C w_c \log \frac{\exp(x_{n,c})}{\sum_{i=1}^C \exp(x_{n,i})} y_{n,c}
+
+      where :math:`x` is the input, :math:`y` is the target, :math:`w` is the weight,
+      :math:`C` is the number of classes, and :math:`N` spans the minibatch dimension as well as
+      :math:`d_1, ..., d_k` for the `K`-dimensional case. If
+      :attr:`reduction` is not ``'none'`` (default ``'mean'``), then
+
+      .. math::
+          \ell(x, y) = \begin{cases}
+              \frac{\sum_{n=1}^N l_n}{N}, &
+               \text{if reduction} = \text{`mean';}\\
+                \sum_{n=1}^N l_n,  &
+                \text{if reduction} = \text{`sum'.}
+            \end{cases}
+
+    .. note::
+        The performance of this criterion is generally better when `target` contains class
+        indices, as this allows for optimized computation. Consider providing `target` as
+        class probabilities only when a single class label per minibatch item is too restrictive.
+
+    Args:
+        weight (Tensor, optional): a manual rescaling weight given to each class.
+            If given, has to be a Tensor of size `C`.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        ignore_index (int, optional): Specifies a target value that is ignored
+            and does not contribute to the input gradient. When :attr:`size_average` is
+            ``True``, the loss is averaged over non-ignored targets. Note that
+            :attr:`ignore_index` is only applicable when the target contains class indices.
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will
+            be applied, ``'mean'``: the weighted mean of the output is taken,
+            ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in
+            the meantime, specifying either of those two args will override
+            :attr:`reduction`. Default: ``'mean'``
+        label_smoothing (float, optional): A float in [0.0, 1.0]. Specifies the amount
+            of smoothing when computing the loss, where 0.0 means no smoothing. The targets
+            become a mixture of the original ground truth and a uniform distribution as described in
+            `Rethinking the Inception Architecture for Computer Vision <https://arxiv.org/abs/1512.00567>`__. Default: :math:`0.0`.
+
+    Shape:
+        - Input: Shape :math:`(C)`, :math:`(N, C)` or :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \geq 1`
+          in the case of `K`-dimensional loss.
+        - Target: If containing class indices, shape :math:`()`, :math:`(N)` or :math:`(N, d_1, d_2, ..., d_K)` with
+          :math:`K \geq 1` in the case of K-dimensional loss where each value should be between :math:`[0, C)`. The
+          target data type is required to be long when using class indices. If containing class probabilities, the
+          target must be the same shape input, and each value should be between :math:`[0, 1]`. This means the target
+          data type is required to be float when using class probabilities. Note that PyTorch does not strictly enforce
+          probability constraints on the class probabilities and that it is the user's responsibility to ensure
+          ``target`` contains valid probability distributions (see below examples section for more details).
+        - Output: If reduction is 'none', shape :math:`()`, :math:`(N)` or :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \geq 1`
+          in the case of K-dimensional loss, depending on the shape of the input. Otherwise, scalar.
+
+
+        where:
+
+        .. math::
+            \begin{aligned}
+                C ={} & \text{number of classes} \\
+                N ={} & \text{batch size} \\
+            \end{aligned}
+
+    Examples:
+
+        >>> # Example of target with class indices
+        >>> loss = nn.CrossEntropyLoss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.empty(3, dtype=torch.long).random_(5)
+        >>> output = loss(input, target)
+        >>> output.backward()
+        >>>
+        >>> # Example of target with class probabilities
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5).softmax(dim=1)
+        >>> output = loss(input, target)
+        >>> output.backward()
+
+    .. note::
+        When ``target`` contains class probabilities, it should consist of soft labels—that is,
+        each ``target`` entry should represent a probability distribution over the possible classes for a given data sample,
+        with individual probabilities between ``[0,1]`` and the total distribution summing to 1.
+        This is why the :func:`softmax()` function is applied to the ``target`` in the class probabilities example above.
+
+        PyTorch does not validate whether the values provided in ``target`` lie in the range ``[0,1]``
+        or whether the distribution of each data sample sums to ``1``.
+        No warning will be raised and it is the user's responsibility
+        to ensure that ``target`` contains valid probability distributions.
+        Providing arbitrary values may yield misleading loss values and unstable gradients during training.
+
+    Examples:
+        >>> # xdoctest: +SKIP
+        >>> # Example of target with incorrectly specified class probabilities
+        >>> loss = nn.CrossEntropyLoss()
+        >>> torch.manual_seed(283)
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.randn(3, 5)
+        >>> # Provided target class probabilities are not in range [0,1]
+        >>> target
+        tensor([[ 0.7105,  0.4446,  2.0297,  0.2671, -0.6075],
+                [-1.0496, -0.2753, -0.3586,  0.9270,  1.0027],
+                [ 0.7551,  0.1003,  1.3468, -0.3581, -0.9569]])
+        >>> # Provided target class probabilities do not sum to 1
+        >>> target.sum(axis=1)
+        tensor([2.8444, 0.2462, 0.8873])
+        >>> # No error message and possible misleading loss value
+        >>> loss(input, target).item()
+        4.6379876136779785
+        >>>
+        >>> # Example of target with correctly specified class probabilities
+        >>> # Use .softmax() to ensure true probability distribution
+        >>> target_new = target.softmax(dim=1)
+        >>> # New target class probabilities all in range [0,1]
+        >>> target_new
+        tensor([[0.1559, 0.1195, 0.5830, 0.1000, 0.0417],
+                [0.0496, 0.1075, 0.0990, 0.3579, 0.3860],
+                [0.2607, 0.1355, 0.4711, 0.0856, 0.0471]])
+        >>> # New target class probabilities sum to 1
+        >>> target_new.sum(axis=1)
+        tensor([1.0000, 1.0000, 1.0000])
+        >>> loss(input, target_new).item()
+        2.55349063873291
+    """
+
+    __constants__ = ["ignore_index", "reduction", "label_smoothing"]
+    ignore_index: int
+    label_smoothing: float
+
+    def __init__(
+        self,
+        weight: Tensor | None = None,
+        size_average=None,
+        ignore_index: int = -100,
+        reduce=None,
+        reduction: str = "mean",
+        label_smoothing: float = 0.0,
+    ) -> None:
+        super().__init__(weight, size_average, reduce, reduction)
+        self.ignore_index = ignore_index
+        self.label_smoothing = label_smoothing
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.cross_entropy(
+            input,
+            target,
+            weight=self.weight,
+            ignore_index=self.ignore_index,
+            reduction=self.reduction,
+            label_smoothing=self.label_smoothing,
+        )
+
+
+class MultiLabelSoftMarginLoss(_WeightedLoss):
+    r"""Creates a criterion that optimizes a multi-label one-versus-all
+    loss based on max-entropy, between input :math:`x` and target :math:`y` of size
+    :math:`(N, C)`.
+    For each sample in the minibatch:
+
+    .. math::
+        loss(x, y) = - \frac{1}{C} * \sum_i y[i] * \log((1 + \exp(-x[i]))^{-1})
+                         + (1-y[i]) * \log\left(\frac{\exp(-x[i])}{(1 + \exp(-x[i]))}\right)
+
+    where :math:`i \in \left\{0, \; \cdots , \; \text{x.nElement}() - 1\right\}`,
+    :math:`y[i] \in \left\{0, \; 1\right\}`.
+
+    Args:
+        weight (Tensor, optional): a manual rescaling weight given to each
+            class. If given, it has to be a Tensor of size `C`. Otherwise, it is
+            treated as if having all ones.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(N, C)` where `N` is the batch size and `C` is the number of classes.
+        - Target: :math:`(N, C)`, label targets must have the same shape as the input.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then :math:`(N)`.
+
+    Examples:
+
+        >>> loss = nn.MultiLabelSoftMarginLoss()
+        >>> input = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.empty(3, 5).random_(2)
+        >>> output = loss(input, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["reduction"]
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.multilabel_soft_margin_loss(
+            input, target, weight=self.weight, reduction=self.reduction
+        )
+
+
+class CosineEmbeddingLoss(_Loss):
+    r"""Creates a criterion that measures the loss given input tensors
+    :math:`x_1`, :math:`x_2` and a `Tensor` label :math:`y` with values 1 or -1.
+    Use (:math:`y=1`) to maximize the cosine similarity of two inputs, and (:math:`y=-1`) otherwise.
+    This is typically used for learning nonlinear
+    embeddings or semi-supervised learning.
+
+    The loss function for each sample is:
+
+    .. math::
+        \text{loss}(x, y) =
+        \begin{cases}
+        1 - \cos(x_1, x_2), & \text{if } y = 1 \\
+        \max(0, \cos(x_1, x_2) - \text{margin}), & \text{if } y = -1
+        \end{cases}
+
+    Args:
+        margin (float, optional): Should be a number from :math:`-1` to :math:`1`,
+            :math:`0` to :math:`0.5` is suggested. If :attr:`margin` is missing, the
+            default value is :math:`0`.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input1: :math:`(N, D)` or :math:`(D)`, where `N` is the batch size and `D` is the embedding dimension.
+        - Input2: :math:`(N, D)` or :math:`(D)`, same shape as Input1.
+        - Target: :math:`(N)` or :math:`()`.
+        - Output: If :attr:`reduction` is ``'none'``, then :math:`(N)`, otherwise scalar.
+
+    Examples:
+
+        >>> loss = nn.CosineEmbeddingLoss()
+        >>> input1 = torch.randn(3, 5, requires_grad=True)
+        >>> input2 = torch.randn(3, 5, requires_grad=True)
+        >>> target = torch.ones(3)
+        >>> output = loss(input1, input2, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["margin", "reduction"]
+    margin: float
+
+    def __init__(
+        self,
+        margin: float = 0.0,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.margin = margin
+
+    def forward(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.cosine_embedding_loss(
+            input1, input2, target, margin=self.margin, reduction=self.reduction
+        )
+
+
+class MarginRankingLoss(_Loss):
+    r"""Creates a criterion that measures the loss given
+    inputs :math:`x1`, :math:`x2`, two 1D mini-batch or 0D `Tensors`,
+    and a label 1D mini-batch or 0D `Tensor` :math:`y` (containing 1 or -1).
+
+    If :math:`y = 1` then it assumed the first input should be ranked higher
+    (have a larger value) than the second input, and vice-versa for :math:`y = -1`.
+
+    The loss function for each pair of samples in the mini-batch is:
+
+    .. math::
+        \text{loss}(x1, x2, y) = \max(0, -y * (x1 - x2) + \text{margin})
+
+    Args:
+        margin (float, optional): Has a default value of :math:`0`.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input1: :math:`(N)` or :math:`()` where `N` is the batch size.
+        - Input2: :math:`(N)` or :math:`()`, same shape as the Input1.
+        - Target: :math:`(N)` or :math:`()`, same shape as the inputs.
+        - Output: scalar. If :attr:`reduction` is ``'none'`` and Input size is not :math:`()`, then :math:`(N)`.
+
+    Examples:
+
+        >>> loss = nn.MarginRankingLoss()
+        >>> input1 = torch.randn(3, requires_grad=True)
+        >>> input2 = torch.randn(3, requires_grad=True)
+        >>> target = torch.randn(3).sign()
+        >>> output = loss(input1, input2, target)
+        >>> output.backward()
+    """
+
+    __constants__ = ["margin", "reduction"]
+    margin: float
+
+    def __init__(
+        self,
+        margin: float = 0.0,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        self.margin = margin
+
+    def forward(self, input1: Tensor, input2: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.margin_ranking_loss(
+            input1, input2, target, margin=self.margin, reduction=self.reduction
+        )
+
+
+class MultiMarginLoss(_WeightedLoss):
+    r"""Creates a criterion that optimizes a multi-class classification hinge
+    loss (margin-based loss) between input :math:`x` (a 2D mini-batch `Tensor`) and
+    output :math:`y` (which is a 1D tensor of target class indices,
+    :math:`0 \leq y \leq \text{x.size}(1)-1`):
+
+    For each mini-batch sample, the loss in terms of the 1D input :math:`x` and scalar
+    output :math:`y` is:
+
+    .. math::
+        \text{loss}(x, y) = \frac{\sum_i \max(0, \text{margin} - x[y] + x[i])^p}{\text{x.size}(0)}
+
+    where :math:`i \in \left\{0, \; \cdots , \; \text{x.size}(0) - 1\right\}`
+    and :math:`i \neq y`.
+
+    Optionally, you can give non-equal weighting on the classes by passing
+    a 1D :attr:`weight` tensor into the constructor.
+
+    The loss function then becomes:
+
+    .. math::
+        \text{loss}(x, y) = \frac{\sum_i w[y] * \max(0, \text{margin} - x[y] + x[i])^p}{\text{x.size}(0)}
+
+    Args:
+        p (int, optional): Has a default value of :math:`1`. :math:`1` and :math:`2`
+            are the only supported values.
+        margin (float, optional): Has a default value of :math:`1`.
+        weight (Tensor, optional): a manual rescaling weight given to each
+            class. If given, it has to be a Tensor of size `C`. Otherwise, it is
+            treated as if having all ones.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(N, C)` or :math:`(C)`, where :math:`N` is the batch size and :math:`C` is the number of classes.
+        - Target: :math:`(N)` or :math:`()`, where each value is :math:`0 \leq \text{targets}[i] \leq C-1`.
+        - Output: scalar. If :attr:`reduction` is ``'none'``, then same shape as the target.
+
+    Examples:
+
+        >>> loss = nn.MultiMarginLoss()
+        >>> x = torch.tensor([[0.1, 0.2, 0.4, 0.8]])
+        >>> y = torch.tensor([3])
+        >>> # 0.25 * ((1-(0.8-0.1)) + (1-(0.8-0.2)) + (1-(0.8-0.4)))
+        >>> loss(x, y)
+        tensor(0.32...)
+    """
+
+    __constants__ = ["p", "margin", "reduction"]
+    margin: float
+    p: int
+
+    def __init__(
+        self,
+        p: int = 1,
+        margin: float = 1.0,
+        weight: Tensor | None = None,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(weight, size_average, reduce, reduction)
+        if p != 1 and p != 2:
+            raise ValueError("only p == 1 and p == 2 supported")
+        if weight is not None and weight.dim() != 1:
+            raise ValueError(
+                f"MultiMarginLoss: expected weight to be None or 1D tensor, got {weight.dim()}D instead"
+            )
+        self.p = p
+        self.margin = margin
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.multi_margin_loss(
+            input,
+            target,
+            p=self.p,
+            margin=self.margin,
+            weight=self.weight,
+            reduction=self.reduction,
+        )
+
+
+class TripletMarginLoss(_Loss):
+    r"""Creates a criterion that measures the triplet loss given an input
+    tensors :math:`x1`, :math:`x2`, :math:`x3` and a margin with a value greater than :math:`0`.
+    This is used for measuring a relative similarity between samples. A triplet
+    is composed by `a`, `p` and `n` (i.e., `anchor`, `positive examples` and `negative
+    examples` respectively). The shapes of all input tensors should be
+    :math:`(N, D)`.
+
+    The distance swap is described in detail in the paper `Learning shallow
+    convolutional feature descriptors with triplet losses`_ by
+    V. Balntas, E. Riba et al.
+
+    The loss function for each sample in the mini-batch is:
+
+    .. math::
+        L(a, p, n) = \max \{d(a_i, p_i) - d(a_i, n_i) + {\rm margin}, 0\}
+
+
+    where
+
+    .. math::
+        d(x_i, y_i) = \left\lVert {\bf x}_i - {\bf y}_i \right\rVert_p
+
+    The norm is calculated using the specified p value and a small constant :math:`\varepsilon` is
+    added for numerical stability.
+
+    See also :class:`~torch.nn.TripletMarginWithDistanceLoss`, which computes the
+    triplet margin loss for input tensors using a custom distance function.
+
+    Args:
+        margin (float, optional): Default: :math:`1`.
+        p (int, optional): The norm degree for pairwise distance. Default: :math:`2`.
+        eps (float, optional): Small constant for numerical stability. Default: :math:`1e-6`.
+        swap (bool, optional): The distance swap is described in detail in the paper
+            `Learning shallow convolutional feature descriptors with triplet losses` by
+            V. Balntas, E. Riba et al. Default: ``False``.
+        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,
+            the losses are averaged over each loss element in the batch. Note that for
+            some losses, there are multiple elements per sample. If the field :attr:`size_average`
+            is set to ``False``, the losses are instead summed for each minibatch. Ignored
+            when :attr:`reduce` is ``False``. Default: ``True``
+        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the
+            losses are averaged or summed over observations for each minibatch depending
+            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per
+            batch element instead and ignores :attr:`size_average`. Default: ``True``
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Note: :attr:`size_average`
+            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
+            specifying either of those two args will override :attr:`reduction`. Default: ``'mean'``
+
+    Shape:
+        - Input: :math:`(N, D)` or :math:`(D)` where :math:`D` is the vector dimension.
+        - Output: A Tensor of shape :math:`(N)` if :attr:`reduction` is ``'none'`` and
+          input shape is :math:`(N, D)`; a scalar otherwise.
+
+    Examples:
+
+    >>> triplet_loss = nn.TripletMarginLoss(margin=1.0, p=2, eps=1e-7)
+    >>> anchor = torch.randn(100, 128, requires_grad=True)
+    >>> positive = torch.randn(100, 128, requires_grad=True)
+    >>> negative = torch.randn(100, 128, requires_grad=True)
+    >>> output = triplet_loss(anchor, positive, negative)
+    >>> output.backward()
+
+    .. _Learning shallow convolutional feature descriptors with triplet losses:
+        https://bmva-archive.org.uk/bmvc/2016/papers/paper119/index.html
+    """
+
+    __constants__ = ["margin", "p", "eps", "swap", "reduction"]
+    margin: float
+    p: float
+    eps: float
+    swap: bool
+
+    def __init__(
+        self,
+        margin: float = 1.0,
+        p: float = 2.0,
+        eps: float = 1e-6,
+        swap: bool = False,
+        size_average=None,
+        reduce=None,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(size_average, reduce, reduction)
+        if margin <= 0:
+            raise ValueError(
+                f"TripletMarginLoss: expected margin to be greater than 0, got {margin} instead"
+            )
+        self.margin = margin
+        self.p = p
+        self.eps = eps
+        self.swap = swap
+
+    def forward(self, anchor: Tensor, positive: Tensor, negative: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.triplet_margin_loss(
+            anchor,
+            positive,
+            negative,
+            margin=self.margin,
+            p=self.p,
+            eps=self.eps,
+            swap=self.swap,
+            reduction=self.reduction,
+        )
+
+
+class TripletMarginWithDistanceLoss(_Loss):
+    r"""Creates a criterion that measures the triplet loss given input
+    tensors :math:`a`, :math:`p`, and :math:`n` (representing anchor,
+    positive, and negative examples, respectively), and a nonnegative,
+    real-valued function ("distance function") used to compute the relationship
+    between the anchor and positive example ("positive distance") and the
+    anchor and negative example ("negative distance").
+
+    The unreduced loss (i.e., with :attr:`reduction` set to ``'none'``)
+    can be described as:
+
+    .. math::
+        \ell(a, p, n) = L = \{l_1,\dots,l_N\}^\top, \quad
+        l_i = \max \{d(a_i, p_i) - d(a_i, n_i) + {\rm margin}, 0\}
+
+    where :math:`N` is the batch size; :math:`d` is a nonnegative, real-valued function
+    quantifying the closeness of two tensors, referred to as the :attr:`distance_function`;
+    and :math:`margin` is a nonnegative margin representing the minimum difference
+    between the positive and negative distances that is required for the loss to
+    be 0.  The input tensors have :math:`N` elements each and can be of any shape
+    that the distance function can handle.
+
+    If :attr:`reduction` is not ``'none'``
+    (default ``'mean'``), then:
+
+    .. math::
+        \ell(x, y) =
+        \begin{cases}
+            \operatorname{mean}(L), &  \text{if reduction} = \text{`mean';}\\
+            \operatorname{sum}(L),  &  \text{if reduction} = \text{`sum'.}
+        \end{cases}
+
+    See also :class:`~torch.nn.TripletMarginLoss`, which computes the triplet
+    loss for input tensors using the :math:`l_p` distance as the distance function.
+
+    Args:
+        distance_function (Callable, optional): A nonnegative, real-valued function that
+            quantifies the closeness of two tensors. If not specified,
+            `nn.PairwiseDistance` will be used.  Default: ``None``
+        margin (float, optional): A nonnegative margin representing the minimum difference
+            between the positive and negative distances required for the loss to be 0. Larger
+            margins penalize cases where the negative examples are not distant enough from the
+            anchors, relative to the positives. Default: :math:`1`.
+        swap (bool, optional): Whether to use the distance swap described in the paper
+            `Learning shallow convolutional feature descriptors with triplet losses` by
+            V. Balntas, E. Riba et al. If True, and if the positive example is closer to the
+            negative example than the anchor is, swaps the positive example and the anchor in
+            the loss computation. Default: ``False``.
+        reduction (str, optional): Specifies the (optional) reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of
+            elements in the output, ``'sum'``: the output will be summed. Default: ``'mean'``
+
+
+    Shape:
+        - Input: :math:`(N, *)` where :math:`*` represents any number of additional dimensions
+          as supported by the distance function.
+        - Output: A Tensor of shape :math:`(N)` if :attr:`reduction` is ``'none'``, or a scalar
+          otherwise.
+
+    Examples:
+
+    >>> # Initialize embeddings
+    >>> embedding = nn.Embedding(1000, 128)
+    >>> anchor_ids = torch.randint(0, 1000, (1,))
+    >>> positive_ids = torch.randint(0, 1000, (1,))
+    >>> negative_ids = torch.randint(0, 1000, (1,))
+    >>> anchor = embedding(anchor_ids)
+    >>> positive = embedding(positive_ids)
+    >>> negative = embedding(negative_ids)
+    >>>
+    >>> # Built-in Distance Function
+    >>> triplet_loss = \
+    >>>     nn.TripletMarginWithDistanceLoss(distance_function=nn.PairwiseDistance())
+    >>> output = triplet_loss(anchor, positive, negative)
+    >>> output.backward()
+    >>>
+    >>> # Custom Distance Function
+    >>> def l_infinity(x1, x2):
+    >>>     return torch.max(torch.abs(x1 - x2), dim=1).values
+    >>>
+    >>> # xdoctest: +SKIP("FIXME: Would call backwards a second time")
+    >>> triplet_loss = (
+    >>>     nn.TripletMarginWithDistanceLoss(distance_function=l_infinity, margin=1.5))
+    >>> output = triplet_loss(anchor, positive, negative)
+    >>> output.backward()
+    >>>
+    >>> # Custom Distance Function (Lambda)
+    >>> triplet_loss = (
+    >>>     nn.TripletMarginWithDistanceLoss(
+    >>>         distance_function=lambda x, y: 1.0 - F.cosine_similarity(x, y)))
+    >>> output = triplet_loss(anchor, positive, negative)
+    >>> output.backward()
+
+    Reference:
+        V. Balntas, et al.: Learning shallow convolutional feature descriptors with triplet losses:
+        https://bmva-archive.org.uk/bmvc/2016/papers/paper119/index.html
+    """
+
+    __constants__ = ["margin", "swap", "reduction"]
+    margin: float
+    swap: bool
+
+    def __init__(
+        self,
+        *,
+        distance_function: Callable[[Tensor, Tensor], Tensor] | None = None,
+        margin: float = 1.0,
+        swap: bool = False,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__(size_average=None, reduce=None, reduction=reduction)
+        if margin <= 0:
+            raise ValueError(
+                f"TripletMarginWithDistanceLoss: expected margin to be greater than 0, got {margin} instead"
+            )
+        self.distance_function: Callable[[Tensor, Tensor], Tensor] | None = (
+            distance_function if distance_function is not None else PairwiseDistance()
+        )
+        self.margin = margin
+        self.swap = swap
+
+    def forward(self, anchor: Tensor, positive: Tensor, negative: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.triplet_margin_with_distance_loss(
+            anchor,
+            positive,
+            negative,
+            distance_function=self.distance_function,
+            margin=self.margin,
+            swap=self.swap,
+            reduction=self.reduction,
+        )
+
+
+class CTCLoss(_Loss):
+    r"""The Connectionist Temporal Classification loss.
+
+    Calculates loss between a continuous (unsegmented) time series and a target sequence. CTCLoss sums over the
+    probability of possible alignments of input to target, producing a loss value which is differentiable
+    with respect to each input node. The alignment of input to target is assumed to be "many-to-one", which
+    limits the length of the target sequence such that it must be :math:`\leq` the input length.
+
+    Args:
+        blank (int, optional): blank label. Default :math:`0`.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+            ``'mean'``: the output losses will be divided by the target lengths and
+            then the mean over the batch is taken, ``'sum'``: the output losses will be summed.
+            Default: ``'mean'``
+        zero_infinity (bool, optional):
+            Whether to zero infinite losses and the associated gradients.
+            Default: ``False``
+            Infinite losses mainly occur when the inputs are too short
+            to be aligned to the targets.
+
+    Shape:
+        - Log_probs: Tensor of size :math:`(T, N, C)` or :math:`(T, C)`,
+          where :math:`T = \text{input length}`,
+          :math:`N = \text{batch size}`, and
+          :math:`C = \text{number of classes (including blank)}`.
+          The logarithmized probabilities of the outputs (e.g. obtained with
+          :func:`torch.nn.functional.log_softmax`).
+        - Targets: Tensor of size :math:`(N, S)` or
+          :math:`(\operatorname{sum}(\text{target\_lengths}))`,
+          where :math:`N = \text{batch size}` and
+          :math:`S = \text{max target length, if shape is } (N, S)`.
+          It represents the target sequences. Each element in the target
+          sequence is a class index. And the target index cannot be blank (default=0).
+          In the :math:`(N, S)` form, targets are padded to the
+          length of the longest sequence, and stacked.
+          In the :math:`(\operatorname{sum}(\text{target\_lengths}))` form,
+          the targets are assumed to be un-padded and
+          concatenated within 1 dimension.
+        - Input_lengths: Tuple or tensor of size :math:`(N)` or :math:`()`,
+          where :math:`N = \text{batch size}`. It represents the lengths of the
+          inputs (must each be :math:`\leq T`). And the lengths are specified
+          for each sequence to achieve masking under the assumption that sequences
+          are padded to equal lengths.
+        - Target_lengths: Tuple or tensor of size :math:`(N)` or :math:`()`,
+          where :math:`N = \text{batch size}`. It represents lengths of the targets.
+          Lengths are specified for each sequence to achieve masking under the
+          assumption that sequences are padded to equal lengths. If target shape is
+          :math:`(N,S)`, target_lengths are effectively the stop index
+          :math:`s_n` for each target sequence, such that ``target_n = targets[n,0:s_n]`` for
+          each target in a batch. Lengths must each be :math:`\leq S`
+          If the targets are given as a 1d tensor that is the concatenation of individual
+          targets, the target_lengths must add up to the total length of the tensor.
+        - Output: scalar if :attr:`reduction` is ``'mean'`` (default) or
+          ``'sum'``. If :attr:`reduction` is ``'none'``, then :math:`(N)` if input is batched or
+          :math:`()` if input is unbatched, where :math:`N = \text{batch size}`.
+
+    Examples:
+
+        >>> # Target are to be padded
+        >>> T = 50  # Input sequence length
+        >>> C = 20  # Number of classes (including blank)
+        >>> N = 16  # Batch size
+        >>> S = 30  # Target sequence length of longest target in batch (padding length)
+        >>> S_min = 10  # Minimum target length, for demonstration purposes
+        >>>
+        >>> # Initialize random batch of input vectors, for *size = (T,N,C)
+        >>> input = torch.randn(T, N, C).log_softmax(2).detach().requires_grad_()
+        >>>
+        >>> # Initialize random batch of targets (0 = blank, 1:C = classes)
+        >>> target = torch.randint(low=1, high=C, size=(N, S), dtype=torch.long)
+        >>>
+        >>> input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long)
+        >>> target_lengths = torch.randint(
+        ...     low=S_min,
+        ...     high=S,
+        ...     size=(N,),
+        ...     dtype=torch.long,
+        ... )
+        >>> ctc_loss = nn.CTCLoss()
+        >>> loss = ctc_loss(input, target, input_lengths, target_lengths)
+        >>> loss.backward()
+        >>>
+        >>>
+        >>> # Target are to be un-padded
+        >>> T = 50  # Input sequence length
+        >>> C = 20  # Number of classes (including blank)
+        >>> N = 16  # Batch size
+        >>>
+        >>> # Initialize random batch of input vectors, for *size = (T,N,C)
+        >>> input = torch.randn(T, N, C).log_softmax(2).detach().requires_grad_()
+        >>> input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long)
+        >>>
+        >>> # Initialize random batch of targets (0 = blank, 1:C = classes)
+        >>> target_lengths = torch.randint(low=1, high=T, size=(N,), dtype=torch.long)
+        >>> target = torch.randint(
+        ...     low=1,
+        ...     high=C,
+        ...     size=(sum(target_lengths),),
+        ...     dtype=torch.long,
+        ... )
+        >>> ctc_loss = nn.CTCLoss()
+        >>> loss = ctc_loss(input, target, input_lengths, target_lengths)
+        >>> loss.backward()
+        >>>
+        >>>
+        >>> # Target are to be un-padded and unbatched (effectively N=1)
+        >>> T = 50  # Input sequence length
+        >>> C = 20  # Number of classes (including blank)
+        >>>
+        >>> # Initialize random batch of input vectors, for *size = (T,C)
+        >>> # xdoctest: +SKIP("FIXME: error in doctest")
+        >>> input = torch.randn(T, C).log_softmax(1).detach().requires_grad_()
+        >>> input_lengths = torch.tensor(T, dtype=torch.long)
+        >>>
+        >>> # Initialize random batch of targets (0 = blank, 1:C = classes)
+        >>> target_lengths = torch.randint(low=1, high=T, size=(), dtype=torch.long)
+        >>> target = torch.randint(
+        ...     low=1,
+        ...     high=C,
+        ...     size=(target_lengths,),
+        ...     dtype=torch.long,
+        ... )
+        >>> ctc_loss = nn.CTCLoss()
+        >>> loss = ctc_loss(input, target, input_lengths, target_lengths)
+        >>> loss.backward()
+
+    Reference:
+        A. Graves et al.: Connectionist Temporal Classification:
+        Labelling Unsegmented Sequence Data with Recurrent Neural Networks:
+        https://www.cs.toronto.edu/~graves/icml_2006.pdf
+
+    Note:
+        In order to use CuDNN, the following must be satisfied: the :attr:`targets` must be
+        in concatenated format, all :attr:`input_lengths` must be `T`.  :math:`blank=0`,
+        :attr:`target_lengths` :math:`\leq 256`, the integer arguments must be of
+        dtype :attr:`torch.int32`, and the :attr:`log_probs` itself must be of
+        dtype :attr:`torch.float32`.
+
+        The regular implementation uses the (more common in PyTorch) `torch.long` dtype.
+
+
+    Note:
+        In some circumstances when using the CUDA backend with CuDNN, this operator
+        may select a nondeterministic algorithm to increase performance. If this is
+        undesirable, you can try to make the operation deterministic (potentially at
+        a performance cost) by setting ``torch.backends.cudnn.deterministic =
+        True``.
+        Please see the notes on :doc:`/notes/randomness` for background.
+    """
+
+    __constants__ = ["blank", "reduction"]
+    blank: int
+    zero_infinity: bool
+
+    def __init__(
+        self, blank: int = 0, reduction: str = "mean", zero_infinity: bool = False
+    ) -> None:
+        super().__init__(reduction=reduction)
+        self.blank = blank
+        self.zero_infinity = zero_infinity
+
+    def forward(
+        self,
+        log_probs: Tensor,
+        targets: Tensor,
+        input_lengths: Tensor,
+        target_lengths: Tensor,
+    ) -> Tensor:
+        """Runs the forward pass."""
+        return F.ctc_loss(
+            log_probs,
+            targets,
+            input_lengths,
+            target_lengths,
+            self.blank,
+            self.reduction,
+            self.zero_infinity,
+        )
+
+
+# TODO: L1HingeEmbeddingCriterion
+# TODO: MSECriterion weight
+# TODO: ClassSimplexCriterion
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/module.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/module.py
new file mode 100644
index 0000000000000000000000000000000000000000..e004ef8d14f13cc5dbb2afb4d434bfc0fa976cc9
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/module.py
@@ -0,0 +1,3053 @@
+# mypy: allow-untyped-defs
+
+import functools
+import inspect
+import itertools
+import warnings
+import weakref
+from collections import namedtuple, OrderedDict
+from collections.abc import Callable, Iterator, Mapping
+from typing import Any, Optional, overload, TypeVar, Union
+from typing_extensions import Self
+
+import torch
+from torch import device, dtype, Tensor
+from torch._prims_common import DeviceLikeType
+from torch.nn.parameter import Buffer, Parameter
+from torch.utils._python_dispatch import is_traceable_wrapper_subclass
+from torch.utils.hooks import BackwardHook, RemovableHandle
+
+
+__all__ = [
+    "register_module_forward_pre_hook",
+    "register_module_forward_hook",
+    "register_module_full_backward_pre_hook",
+    "register_module_backward_hook",
+    "register_module_full_backward_hook",
+    "register_module_buffer_registration_hook",
+    "register_module_module_registration_hook",
+    "register_module_parameter_registration_hook",
+    "Module",
+]
+
+_grad_t = tuple[Tensor, ...] | Tensor
+# See https://mypy.readthedocs.io/en/latest/generics.html#generic-methods-and-generic-self for the use
+# of `T` to annotate `self`. Many methods of `Module` return `self` and we want those return values to be
+# the type of the subclass, not the looser type of `Module`.
+T = TypeVar("T", bound="Module")
+
+
+class _IncompatibleKeys(
+    # pyrefly: ignore [invalid-inheritance]
+    namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]),
+):
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        if not self.missing_keys and not self.unexpected_keys:
+            return "<All keys matched successfully>"
+        return super().__repr__()
+
+    __str__ = __repr__
+
+
+def _addindent(s_, numSpaces):
+    s = s_.split("\n")
+    # don't do anything for single-line stuff
+    if len(s) == 1:
+        return s_
+    first = s.pop(0)
+    # Only add indentation to non-blank lines; blank lines stay empty
+    s = [(numSpaces * " ") + line if line.strip() else "" for line in s]
+    s = "\n".join(s)
+    s = first + "\n" + s
+    return s
+
+
+r"""This tracks hooks common to all modules that are executed immediately before
+.registering the buffer/module/parameter"""
+_global_buffer_registration_hooks: dict[int, Callable] = OrderedDict()
+_global_module_registration_hooks: dict[int, Callable] = OrderedDict()
+_global_parameter_registration_hooks: dict[int, Callable] = OrderedDict()
+
+
+class _WrappedHook:
+    def __init__(self, hook: Callable, module: Optional["Module"] = None) -> None:  # noqa: UP045
+        self.hook: Callable = hook
+        functools.update_wrapper(self, hook)
+
+        self.with_module: bool = False
+
+        if module is not None:
+            self.module: weakref.ReferenceType[Module] = weakref.ref(module)
+            self.with_module = True
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if self.with_module:
+            module = self.module()
+            if module is None:
+                raise RuntimeError("You are trying to call the hook of a dead Module!")
+            return self.hook(module, *args, **kwargs)
+        return self.hook(*args, **kwargs)
+
+    def __getstate__(self) -> dict:
+        result = {"hook": self.hook, "with_module": self.with_module}
+        if self.with_module:
+            # pyrefly: ignore [bad-typed-dict-key]
+            result["module"] = self.module()
+
+        return result
+
+    def __setstate__(self, state: dict):
+        self.hook = state["hook"]
+        self.with_module = state["with_module"]
+
+        if self.with_module:
+            if state["module"] is None:
+                raise RuntimeError(
+                    "You are trying to revive the hook of a dead Module!"
+                )
+            self.module = weakref.ref(state["module"])
+
+
+r"""This tracks hooks common to all modules that are executed before/after
+calling forward and backward. This is global state used for debugging/profiling
+purposes"""
+_global_backward_pre_hooks: dict[int, Callable] = OrderedDict()
+_global_backward_hooks: dict[int, Callable] = OrderedDict()
+_global_is_full_backward_hook: bool | None = None
+_global_forward_pre_hooks: dict[int, Callable] = OrderedDict()
+_global_forward_hooks: dict[int, Callable] = OrderedDict()
+_global_forward_hooks_always_called: dict[int, bool] = OrderedDict()
+_global_forward_hooks_with_kwargs: dict[int, bool] = OrderedDict()
+
+
+def _has_any_global_hook():
+    return (
+        _global_backward_pre_hooks
+        or _global_backward_hooks
+        or _global_forward_pre_hooks
+        or _global_forward_hooks
+        or _global_forward_hooks_always_called
+        or _global_forward_hooks_with_kwargs
+    )
+
+
+_EXTRA_STATE_KEY_SUFFIX = "_extra_state"
+
+
+def register_module_buffer_registration_hook(
+    hook: Callable[..., None],
+) -> RemovableHandle:
+    r"""Register a buffer registration hook common to all modules.
+
+    .. warning ::
+
+        This adds global state to the `nn.Module` module
+
+    The hook will be called every time :func:`register_buffer` is invoked.
+    It should have the following signature::
+
+        hook(module, name, buffer) -> None or new buffer
+
+    The hook can modify the input or return a single modified value in the hook.
+
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+    """
+    handle = RemovableHandle(_global_buffer_registration_hooks)
+    _global_buffer_registration_hooks[handle.id] = hook
+    return handle
+
+
+def register_module_module_registration_hook(
+    hook: Callable[..., None],
+) -> RemovableHandle:
+    r"""Register a module registration hook common to all modules.
+
+    .. warning ::
+
+        This adds global state to the `nn.Module` module
+
+    The hook will be called every time :func:`register_module` is invoked.
+    It should have the following signature::
+
+        hook(module, name, submodule) -> None or new submodule
+
+    The hook can modify the input or return a single modified value in the hook.
+
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+    """
+    handle = RemovableHandle(_global_module_registration_hooks)
+    _global_module_registration_hooks[handle.id] = hook
+    return handle
+
+
+def register_module_parameter_registration_hook(
+    hook: Callable[..., None],
+) -> RemovableHandle:
+    r"""Register a parameter registration hook common to all modules.
+
+    .. warning ::
+
+        This adds global state to the `nn.Module` module
+
+    The hook will be called every time :func:`register_parameter` is invoked.
+    It should have the following signature::
+
+        hook(module, name, param) -> None or new parameter
+
+    The hook can modify the input or return a single modified value in the hook.
+
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+    """
+    handle = RemovableHandle(_global_parameter_registration_hooks)
+    _global_parameter_registration_hooks[handle.id] = hook
+    return handle
+
+
+def register_module_forward_pre_hook(hook: Callable[..., None]) -> RemovableHandle:
+    r"""Register a forward pre-hook common to all modules.
+
+    .. warning ::
+
+        This adds global state to the `nn.module` module
+        and it is only intended for debugging/profiling purposes.
+
+    The hook will be called every time before :func:`forward` is invoked.
+    It should have the following signature::
+
+        hook(module, input) -> None or modified input
+
+    The input contains only the positional arguments given to the module.
+    Keyword arguments won't be passed to the hooks and only to the ``forward``.
+    The hook can modify the input. User can either return a tuple or a
+    single modified value in the hook. We will wrap the value into a tuple
+    if a single value is returned(unless that value is already a tuple).
+
+    This hook has precedence over the specific module hooks registered with
+    ``register_forward_pre_hook``.
+
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+    """
+    handle = RemovableHandle(_global_forward_pre_hooks)
+    _global_forward_pre_hooks[handle.id] = hook
+    return handle
+
+
+def register_module_forward_hook(
+    hook: Callable[..., None],
+    *,
+    with_kwargs: bool = False,
+    always_call: bool = False,
+) -> RemovableHandle:
+    r"""Register a global forward hook for all the modules.
+
+    .. warning ::
+
+        This adds global state to the `nn.module` module
+        and it is only intended for debugging/profiling purposes.
+
+    The hook will be called every time after :func:`forward` has computed an output.
+    It should have the following signature::
+
+        hook(module, input, output) -> None or modified output
+
+    The input contains only the positional arguments given to the module.
+    Keyword arguments won't be passed to the hooks and only to the ``forward``.
+    You can optionally modify the output of the module by returning a new value
+    that will replace the output from the :func:`forward` function.
+
+    Parameters:
+        hook (Callable): The user defined hook to be registered.
+        always_call (bool): If ``True`` the ``hook`` will be run regardless of
+            whether an exception is raised while calling the Module.
+            Default: ``False``
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+
+    This hook will be executed before specific module hooks registered with
+    ``register_forward_hook``.
+    """
+    handle = RemovableHandle(
+        _global_forward_hooks, extra_dict=_global_forward_hooks_always_called
+    )
+    _global_forward_hooks[handle.id] = hook
+    if with_kwargs:
+        _global_forward_hooks_with_kwargs[handle.id] = True
+    if always_call:
+        _global_forward_hooks_always_called[handle.id] = True
+    return handle
+
+
+def register_module_backward_hook(
+    hook: Callable[["Module", _grad_t, _grad_t], _grad_t | None],
+) -> RemovableHandle:
+    r"""Register a backward hook common to all the modules.
+
+    This function is deprecated in favor of
+    :func:`torch.nn.modules.module.register_module_full_backward_hook`
+    and the behavior of this function will change in future versions.
+
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+
+    """
+    global _global_is_full_backward_hook
+    if _global_is_full_backward_hook is True:
+        raise RuntimeError(
+            "Cannot use both regular backward hooks and full backward hooks as a "
+            "global Module hook. Please use only one of them."
+        )
+
+    _global_is_full_backward_hook = False
+
+    handle = RemovableHandle(_global_backward_hooks)
+    _global_backward_hooks[handle.id] = hook
+    return handle
+
+
+def register_module_full_backward_pre_hook(
+    hook: Callable[["Module", _grad_t], _grad_t | None],
+) -> RemovableHandle:
+    r"""Register a backward pre-hook common to all the modules.
+
+    .. warning ::
+        This adds global state to the `nn.module` module
+        and it is only intended for debugging/profiling purposes.
+
+    Hooks registered using this function behave in the same way as those
+    registered by :meth:`torch.nn.Module.register_full_backward_pre_hook`.
+    Refer to its documentation for more details.
+
+    Hooks registered using this function will be called before hooks registered
+    using :meth:`torch.nn.Module.register_full_backward_pre_hook`.
+
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+
+    """
+    handle = RemovableHandle(_global_backward_pre_hooks)
+    _global_backward_pre_hooks[handle.id] = hook
+    return handle
+
+
+def register_module_full_backward_hook(
+    hook: Callable[["Module", _grad_t, _grad_t], _grad_t | None],
+) -> RemovableHandle:
+    r"""Register a backward hook common to all the modules.
+
+    .. warning ::
+        This adds global state to the `nn.module` module
+        and it is only intended for debugging/profiling purposes.
+
+    Hooks registered using this function behave in the same way as those
+    registered by :meth:`torch.nn.Module.register_full_backward_hook`.
+    Refer to its documentation for more details.
+
+    Hooks registered using this function will be called before hooks registered
+    using :meth:`torch.nn.Module.register_full_backward_hook`.
+
+    Returns:
+        :class:`torch.utils.hooks.RemovableHandle`:
+            a handle that can be used to remove the added hook by calling
+            ``handle.remove()``
+
+    """
+    global _global_is_full_backward_hook
+    if _global_is_full_backward_hook is False:
+        raise RuntimeError(
+            "Cannot use both regular backward hooks and full backward hooks as a "
+            "global Module hook. Please use only one of them."
+        )
+
+    _global_is_full_backward_hook = True
+
+    handle = RemovableHandle(_global_backward_hooks)
+    _global_backward_hooks[handle.id] = hook
+    return handle
+
+
+# Trick mypy into not applying contravariance rules to inputs by defining
+# forward as a value, rather than a function.  See also
+# https://github.com/python/mypy/issues/8795
+def _forward_unimplemented(self, *input: Any) -> None:
+    r"""Define the computation performed at every call.
+
+    Should be overridden by all subclasses.
+
+    .. note::
+        Although the recipe for forward pass needs to be defined within
+        this function, one should call the :class:`Module` instance afterwards
+        instead of this since the former takes care of running the
+        registered hooks while the latter silently ignores them.
+    """
+    raise NotImplementedError(
+        f'Module [{type(self).__name__}] is missing the required "forward" function'
+    )
+
+
+class Module:
+    r"""Base class for all neural network modules.
+
+    Your models should also subclass this class.
+
+    Modules can also contain other Modules, allowing them to be nested in
+    a tree structure. You can assign the submodules as regular attributes::
+
+        import torch.nn as nn
+        import torch.nn.functional as F
+
+
+        class Model(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv1 = nn.Conv2d(1, 20, 5)
+                self.conv2 = nn.Conv2d(20, 20, 5)
+
+            def forward(self, x):
+                x = F.relu(self.conv1(x))
+                return F.relu(self.conv2(x))
+
+    Submodules assigned in this way will be registered, and will also have their
+    parameters converted when you call :meth:`to`, etc.
+
+    .. note::
+        As per the example above, an ``__init__()`` call to the parent class
+        must be made before assignment on the child.
+
+    :ivar training: Boolean represents whether this module is in training or
+                    evaluation mode.
+    :vartype training: bool
+    """
+
+    dump_patches: bool = False
+
+    _version: int = 1
+    r"""This allows better BC support for :meth:`load_state_dict`. In
+    :meth:`state_dict`, the version number will be saved as in the attribute
+    `_metadata` of the returned state dict, and thus pickled. `_metadata` is a
+    dictionary with keys that follow the naming convention of state dict. See
+    ``_load_from_state_dict`` on how to use this information in loading.
+
+    If new parameters/buffers are added/removed from a module, this number shall
+    be bumped, and the module's `_load_from_state_dict` method can compare the
+    version number and do appropriate changes if the state dict is from before
+    the change."""
+
+    training: bool
+    _parameters: dict[str, Parameter | None]
+    _buffers: dict[str, Tensor | None]
+    _non_persistent_buffers_set: set[str]
+    _backward_pre_hooks: dict[int, Callable]
+    _backward_hooks: dict[int, Callable]
+    _is_full_backward_hook: bool | None
+    _forward_hooks: dict[int, Callable]
+    # Marks whether the corresponding _forward_hooks accept kwargs or not.
+    # As JIT does not support set[int], this dict is used as a set, where all
+    # hooks represented in this dict accept kwargs.
+    _forward_hooks_with_kwargs: dict[int, bool]
+    # forward hooks that should always be called even if an exception is raised
+    _forward_hooks_always_called: dict[int, bool]
+    _forward_pre_hooks: dict[int, Callable]
+    # Marks whether the corresponding _forward_hooks accept kwargs or not.
+    # As JIT does not support set[int], this dict is used as a set, where all
+    # hooks represented in this dict accept kwargs.
+    _forward_pre_hooks_with_kwargs: dict[int, bool]
+    _state_dict_hooks: dict[int, Callable]
+    _load_state_dict_pre_hooks: dict[int, Callable]
+    _state_dict_pre_hooks: dict[int, Callable]
+    _load_state_dict_post_hooks: dict[int, Callable]
+    _modules: dict[str, Optional["Module"]]  # noqa: UP045
+    call_super_init: bool = False
+    _compiled_call_impl: Callable | None = None
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize internal Module state, shared by both nn.Module and ScriptModule."""
+        torch._C._log_api_usage_once("python.nn_module")
+
+        # Backward compatibility: no args used to be allowed when call_super_init=False
+        if self.call_super_init is False and bool(kwargs):
+            raise TypeError(
+                f"{type(self).__name__}.__init__() got an unexpected keyword argument '{next(iter(kwargs))}'"
+                ""
+            )
+
+        if self.call_super_init is False and bool(args):
+            raise TypeError(
+                f"{type(self).__name__}.__init__() takes 1 positional argument but {len(args) + 1} were"
+                " given"
+            )
+
+        """
+        Calls super().__setattr__('a', a) instead of the typical self.a = a
+        to avoid Module.__setattr__ overhead. Module's __setattr__ has special
+        handling for parameters, submodules, and buffers but simply calls into
+        super().__setattr__ for all other attributes.
+        """
+        super().__setattr__("training", True)
+        super().__setattr__("_parameters", {})
+        super().__setattr__("_buffers", {})
+        super().__setattr__("_non_persistent_buffers_set", set())
+        super().__setattr__("_backward_pre_hooks", OrderedDict())
+        super().__setattr__("_backward_hooks", OrderedDict())
+        super().__setattr__("_is_full_backward_hook", None)
+        super().__setattr__("_forward_hooks", OrderedDict())
+        super().__setattr__("_forward_hooks_with_kwargs", OrderedDict())
+        super().__setattr__("_forward_hooks_always_called", OrderedDict())
+        super().__setattr__("_forward_pre_hooks", OrderedDict())
+        super().__setattr__("_forward_pre_hooks_with_kwargs", OrderedDict())
+        super().__setattr__("_state_dict_hooks", OrderedDict())
+        super().__setattr__("_state_dict_pre_hooks", OrderedDict())
+        super().__setattr__("_load_state_dict_pre_hooks", OrderedDict())
+        super().__setattr__("_load_state_dict_post_hooks", OrderedDict())
+        super().__setattr__("_modules", {})
+
+        if self.call_super_init:
+            super().__init__(*args, **kwargs)
+
+    forward: Callable[..., Any] = _forward_unimplemented
+
+    def register_buffer(
+        self, name: str, tensor: Tensor | None, persistent: bool = True
+    ) -> None:
+        r"""Add a buffer to the module.
+
+        This is typically used to register a buffer that should not be
+        considered a model parameter. For example, BatchNorm's ``running_mean``
+        is not a parameter, but is part of the module's state. Buffers, by
+        default, are persistent and will be saved alongside parameters. This
+        behavior can be changed by setting :attr:`persistent` to ``False``. The
+        only difference between a persistent buffer and a non-persistent buffer
+        is that the latter will not be a part of this module's
+        :attr:`state_dict`.
+
+        Buffers can be accessed as attributes using given names.
+
+        Args:
+            name (str): name of the buffer. The buffer can be accessed
+                from this module using the given name
+            tensor (Tensor or None): buffer to be registered. If ``None``, then operations
+                that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
+                the buffer is **not** included in the module's :attr:`state_dict`.
+            persistent (bool): whether the buffer is part of this module's
+                :attr:`state_dict`.
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> self.register_buffer('running_mean', torch.zeros(num_features))
+
+        """
+        if persistent is False and isinstance(self, torch.jit.ScriptModule):
+            raise RuntimeError("ScriptModule does not support non-persistent buffers")
+
+        if "_buffers" not in self.__dict__:
+            raise AttributeError("cannot assign buffer before Module.__init__() call")
+        elif not isinstance(name, str):
+            raise TypeError(
+                f"buffer name should be a string. Got {torch.typename(name)}"
+            )
+        elif "." in name:
+            raise KeyError('buffer name can\'t contain "."')
+        elif name == "":
+            raise KeyError('buffer name can\'t be empty string ""')
+        elif hasattr(self, name) and name not in self._buffers:
+            raise KeyError(f"attribute '{name}' already exists")
+        elif tensor is not None and not (
+            isinstance(tensor, torch.Tensor) or hasattr(tensor, "__torch_function__")
+        ):
+            raise TypeError(
+                f"cannot assign '{torch.typename(tensor)}' object to buffer '{name}' "
+                "(torch Tensor or None required)"
+            )
+        else:
+            for hook in _global_buffer_registration_hooks.values():
+                output = hook(self, name, tensor)
+                if output is not None:
+                    tensor = output
+            self._buffers[name] = tensor
+            if persistent:
+                self._non_persistent_buffers_set.discard(name)
+            else:
+                self._non_persistent_buffers_set.add(name)
+
+    def register_parameter(self, name: str, param: Parameter | None) -> None:
+        r"""Add a parameter to the module.
+
+        The parameter can be accessed as an attribute using given name.
+
+        Args:
+            name (str): name of the parameter. The parameter can be accessed
+                from this module using the given name
+            param (Parameter or None): parameter to be added to the module. If
+                ``None``, then operations that run on parameters, such as :attr:`cuda`,
+                are ignored. If ``None``, the parameter is **not** included in the
+                module's :attr:`state_dict`.
+        """
+        if "_parameters" not in self.__dict__:
+            raise AttributeError(
+                "cannot assign parameter before Module.__init__() call"
+            )
+
+        elif not isinstance(name, str):
+            raise TypeError(
+                f"parameter name should be a string. Got {torch.typename(name)}"
+            )
+        elif "." in name:
+            raise KeyError('parameter name can\'t contain "."')
+        elif name == "":
+            raise KeyError('parameter name can\'t be empty string ""')
+        elif hasattr(self, name) and name not in self._parameters:
+            raise KeyError(f"attribute '{name}' already exists")
+
+        if param is None:
+            self._parameters[name] = None
+        elif not isinstance(param, Parameter):
+            raise TypeError(
+                f"cannot assign '{torch.typename(param)}' object to parameter '{name}' "
+                "(torch.nn.Parameter or None required)"
+            )
+        elif param.grad_fn:
+            raise ValueError(
+                f"Cannot assign non-leaf Tensor to parameter '{name}'. Model "
+                f"parameters must be created explicitly. To express '{name}' "
+                "as a function of another Tensor, compute the value in "
+                "the forward() method."
+            )
+        else:
+            for hook in _global_parameter_registration_hooks.values():
+                output = hook(self, name, param)
+                if output is not None:
+                    param = output
+            self._parameters[name] = param
+
+    def add_module(self, name: str, module: Optional["Module"]) -> None:  # noqa: UP045
+        r"""Add a child module to the current module.
+
+        The module can be accessed as an attribute using the given name.
+
+        Args:
+            name (str): name of the child module. The child module can be
+                accessed from this module using the given name
+            module (Module): child module to be added to the module.
+        """
+        if not isinstance(module, Module) and module is not None:
+            raise TypeError(f"{torch.typename(module)} is not a Module subclass")
+        elif not isinstance(name, str):
+            raise TypeError(
+                f"module name should be a string. Got {torch.typename(name)}"
+            )
+        elif hasattr(self, name) and name not in self._modules:
+            raise KeyError(f"attribute '{name}' already exists")
+        elif "." in name:
+            raise KeyError(f'module name can\'t contain ".", got: {name}')
+        elif name == "":
+            raise KeyError('module name can\'t be empty string ""')
+        for hook in _global_module_registration_hooks.values():
+            output = hook(self, name, module)
+            if output is not None:
+                module = output
+        self._modules[name] = module
+
+    def register_module(self, name: str, module: Optional["Module"]) -> None:  # noqa: UP045
+        r"""Alias for :func:`add_module`."""
+        self.add_module(name, module)
+
+    def get_submodule(self, target: str) -> "Module":
+        """Return the submodule given by ``target`` if it exists, otherwise throw an error.
+
+        For example, let's say you have an ``nn.Module`` ``A`` that
+        looks like this:
+
+        .. code-block:: text
+
+            A(
+                (net_b): Module(
+                    (net_c): Module(
+                        (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
+                    )
+                    (linear): Linear(in_features=100, out_features=200, bias=True)
+                )
+            )
+
+        (The diagram shows an ``nn.Module`` ``A``. ``A`` which has a nested
+        submodule ``net_b``, which itself has two submodules ``net_c``
+        and ``linear``. ``net_c`` then has a submodule ``conv``.)
+
+        To check whether or not we have the ``linear`` submodule, we
+        would call ``get_submodule("net_b.linear")``. To check whether
+        we have the ``conv`` submodule, we would call
+        ``get_submodule("net_b.net_c.conv")``.
+
+        The runtime of ``get_submodule`` is bounded by the degree
+        of module nesting in ``target``. A query against
+        ``named_modules`` achieves the same result, but it is O(N) in
+        the number of transitive modules. So, for a simple check to see
+        if some submodule exists, ``get_submodule`` should always be
+        used.
+
+        Args:
+            target: The fully-qualified string name of the submodule
+                to look for. (See above example for how to specify a
+                fully-qualified string.)
+
+        Returns:
+            torch.nn.Module: The submodule referenced by ``target``
+
+        Raises:
+            AttributeError: If at any point along the path resulting from
+                the target string the (sub)path resolves to a non-existent
+                attribute name or an object that is not an instance of ``nn.Module``.
+        """
+        if target == "":
+            return self
+
+        atoms: list[str] = target.split(".")
+        mod: torch.nn.Module = self
+
+        for item in atoms:
+            if not hasattr(mod, item):
+                raise AttributeError(
+                    mod._get_name() + " has no attribute `" + item + "`"
+                )
+
+            mod = getattr(mod, item)
+
+            if not isinstance(mod, torch.nn.Module):
+                raise AttributeError("`" + item + "` is not an nn.Module")
+
+        return mod
+
+    def set_submodule(
+        self, target: str, module: "Module", strict: bool = False
+    ) -> None:
+        """
+        Set the submodule given by ``target`` if it exists, otherwise throw an error.
+
+        .. note::
+            If ``strict`` is set to ``False`` (default), the method will replace an existing submodule
+            or create a new submodule if the parent module exists. If ``strict`` is set to ``True``,
+            the method will only attempt to replace an existing submodule and throw an error if
+            the submodule does not exist.
+
+        For example, let's say you have an ``nn.Module`` ``A`` that
+        looks like this:
+
+        .. code-block:: text
+
+            A(
+                (net_b): Module(
+                    (net_c): Module(
+                        (conv): Conv2d(3, 3, 3)
+                    )
+                    (linear): Linear(3, 3)
+                )
+            )
+
+        (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
+        submodule ``net_b``, which itself has two submodules ``net_c``
+        and ``linear``. ``net_c`` then has a submodule ``conv``.)
+
+        To override the ``Conv2d`` with a new submodule ``Linear``, you
+        could call ``set_submodule("net_b.net_c.conv", nn.Linear(1, 1))``
+        where ``strict`` could be ``True`` or ``False``
+
+        To add a new submodule ``Conv2d`` to the existing ``net_b`` module,
+        you would call ``set_submodule("net_b.conv", nn.Conv2d(1, 1, 1))``.
+
+        In the above if you set ``strict=True`` and call
+        ``set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True)``, an AttributeError
+        will be raised because ``net_b`` does not have a submodule named ``conv``.
+
+        Args:
+            target: The fully-qualified string name of the submodule
+                to look for. (See above example for how to specify a
+                fully-qualified string.)
+            module: The module to set the submodule to.
+            strict: If ``False``, the method will replace an existing submodule
+                or create a new submodule if the parent module exists. If ``True``,
+                the method will only attempt to replace an existing submodule and throw an error
+                if the submodule doesn't already exist.
+
+        Raises:
+            ValueError: If the ``target`` string is empty or if ``module`` is not an instance of ``nn.Module``.
+            AttributeError: If at any point along the path resulting from
+                the ``target`` string the (sub)path resolves to a non-existent
+                attribute name or an object that is not an instance of ``nn.Module``.
+        """
+        if target == "":
+            raise ValueError("Cannot set the submodule without a target name!")
+
+        atoms: list[str] = target.split(".")
+        if not isinstance(module, torch.nn.Module):
+            raise ValueError(
+                "`" + "module" + f"` is not an nn.Module, found {type(module)}"
+            )
+        if len(atoms) == 1:
+            parent: torch.nn.Module = self
+        else:
+            parent_key = ".".join(atoms[:-1])
+            parent = self.get_submodule(parent_key)
+
+        if strict and not hasattr(parent, atoms[-1]):
+            raise AttributeError(
+                parent._get_name() + " has no attribute `" + atoms[-1] + "`"
+            )
+        if hasattr(parent, atoms[-1]):
+            mod = getattr(parent, atoms[-1])
+            if not isinstance(mod, torch.nn.Module):
+                raise AttributeError("`" + atoms[-1] + "` is not an nn.Module")
+        setattr(parent, atoms[-1], module)
+
+    def get_parameter(self, target: str) -> Parameter:
+        """Return the parameter given by ``target`` if it exists, otherwise throw an error.
+
+        See the docstring for ``get_submodule`` for a more detailed
+        explanation of this method's functionality as well as how to
+        correctly specify ``target``.
+
+        Args:
+            target: The fully-qualified string name of the Parameter
+                to look for. (See ``get_submodule`` for how to specify a
+                fully-qualified string.)
+
+        Returns:
+            torch.nn.Parameter: The Parameter referenced by ``target``
+
+        Raises:
+            AttributeError: If the target string references an invalid
+                path or resolves to something that is not an
+                ``nn.Parameter``
+        """
+        module_path, _, param_name = target.rpartition(".")
+
+        mod: torch.nn.Module = self.get_submodule(module_path)
+
+        if not hasattr(mod, param_name):
+            raise AttributeError(
+                mod._get_name() + " has no attribute `" + param_name + "`"
+            )
+
+        param: torch.nn.Parameter = getattr(mod, param_name)
+
+        if not isinstance(param, torch.nn.Parameter):
+            raise AttributeError("`" + param_name + "` is not an nn.Parameter")
+
+        return param
+
+    def get_buffer(self, target: str) -> Tensor:
+        """Return the buffer given by ``target`` if it exists, otherwise throw an error.
+
+        See the docstring for ``get_submodule`` for a more detailed
+        explanation of this method's functionality as well as how to
+        correctly specify ``target``.
+
+        Args:
+            target: The fully-qualified string name of the buffer
+                to look for. (See ``get_submodule`` for how to specify a
+                fully-qualified string.)
+
+        Returns:
+            torch.Tensor: The buffer referenced by ``target``
+
+        Raises:
+            AttributeError: If the target string references an invalid
+                path or resolves to something that is not a
+                buffer
+        """
+        module_path, _, buffer_name = target.rpartition(".")
+
+        mod: torch.nn.Module = self.get_submodule(module_path)
+
+        if not hasattr(mod, buffer_name):
+            raise AttributeError(
+                mod._get_name() + " has no attribute `" + buffer_name + "`"
+            )
+
+        buffer: torch.Tensor = getattr(mod, buffer_name)
+
+        if buffer_name not in mod._buffers:
+            raise AttributeError("`" + buffer_name + "` is not a buffer")
+
+        return buffer
+
+    def get_extra_state(self) -> Any:
+        """Return any extra state to include in the module's state_dict.
+
+        Implement this and a corresponding :func:`set_extra_state` for your module
+        if you need to store extra state. This function is called when building the
+        module's `state_dict()`.
+
+        Note that extra state should be picklable to ensure working serialization
+        of the state_dict. We only provide backwards compatibility guarantees
+        for serializing Tensors; other objects may break backwards compatibility if
+        their serialized pickled form changes.
+
+        Returns:
+            object: Any extra state to store in the module's state_dict
+        """
+        raise RuntimeError(
+            "Reached a code path in Module.get_extra_state() that should never be called. "
+            "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml "
+            "to report this bug."
+        )
+
+    def set_extra_state(self, state: Any) -> None:
+        """Set extra state contained in the loaded `state_dict`.
+
+        This function is called from :func:`load_state_dict` to handle any extra state
+        found within the `state_dict`. Implement this function and a corresponding
+        :func:`get_extra_state` for your module if you need to store extra state within its
+        `state_dict`.
+
+        Args:
+            state (dict): Extra state from the `state_dict`
+        """
+        raise RuntimeError(
+            "Reached a code path in Module.set_extra_state() that should never be called. "
+            "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml "
+            "to report this bug."
+        )
+
+    def _apply(self, fn, recurse=True):
+        if recurse:
+            for module in self.children():
+                module._apply(fn)
+
+        from torch._subclasses.fake_tensor import FakeTensor
+
+        def compute_should_use_set_data(tensor, tensor_applied) -> bool:
+            if torch._has_compatible_shallow_copy_type(
+                tensor, tensor_applied
+            ) and not isinstance(tensor_applied, FakeTensor):
+                # If the new tensor has compatible tensor type as the existing tensor,
+                # the current behavior is to change the tensor in-place using `.data =`,
+                # and the future behavior is to overwrite the existing tensor. However,
+                # changing the current behavior is a BC-breaking change, and we want it
+                # to happen in future releases. So for now we introduce the
+                # `torch.__future__.get_overwrite_module_params_on_conversion()`
+                # global flag to let the user control whether they want the future
+                # behavior of overwriting the existing tensor or not.
+                return not torch.__future__.get_overwrite_module_params_on_conversion()
+            else:
+                return False
+
+        should_use_swap_tensors = (
+            torch.__future__.get_swap_module_params_on_conversion()
+        )
+
+        for key, param in self._parameters.items():
+            if param is None:
+                continue
+            # Tensors stored in modules are graph leaves, and we don't want to
+            # track autograd history of `param_applied`, so we have to use
+            # `with torch.no_grad():`
+            with torch.no_grad():
+                param_applied = fn(param)
+            p_should_use_set_data = compute_should_use_set_data(param, param_applied)
+
+            # subclasses may have multiple child tensors so we need to use swap_tensors
+            p_should_use_swap_tensors = (
+                should_use_swap_tensors
+                or is_traceable_wrapper_subclass(param_applied)
+                or isinstance(param, FakeTensor)
+            )
+
+            param_grad = param.grad
+            if p_should_use_swap_tensors:
+                try:
+                    if param_grad is not None:
+                        # Accessing param.grad makes its at::Tensor's use_count 2, which will prevent swapping.
+                        # Decrement use count of the gradient by setting to None
+                        param.grad = None
+                    param_applied = torch.nn.Parameter(
+                        # pyrefly: ignore [bad-argument-type]
+                        param_applied,
+                        requires_grad=param.requires_grad,
+                    )
+                    torch.utils.swap_tensors(param, param_applied)
+                except Exception as e:
+                    if param_grad is not None:
+                        param.grad = param_grad
+                    raise RuntimeError(
+                        f"_apply(): Couldn't swap {self._get_name()}.{key}"
+                    ) from e
+                out_param = param
+            elif p_should_use_set_data:
+                param.data = param_applied
+                out_param = param
+            else:
+                if not isinstance(param, Parameter):
+                    raise AssertionError("param must be a Parameter")
+                if not param.is_leaf:
+                    raise AssertionError("param must be a leaf tensor")
+
+                out_param = Parameter(param_applied, param.requires_grad)
+                self._parameters[key] = out_param
+
+            if param_grad is not None:
+                with torch.no_grad():
+                    grad_applied = fn(param_grad)
+                g_should_use_set_data = compute_should_use_set_data(
+                    param_grad, grad_applied
+                )
+                if p_should_use_swap_tensors:
+                    grad_applied.requires_grad_(param_grad.requires_grad)
+                    try:
+                        torch.utils.swap_tensors(param_grad, grad_applied)
+                    except Exception as e:
+                        raise RuntimeError(
+                            f"_apply(): Couldn't swap {self._get_name()}.{key}.grad"
+                        ) from e
+                    out_param.grad = param_grad
+                elif g_should_use_set_data:
+                    if out_param.grad is None:
+                        raise AssertionError("out_param.grad must not be None")
+                    out_param.grad.data = grad_applied
+                else:
+                    if not param_grad.is_leaf:
+                        raise AssertionError("param_grad must be a leaf tensor")
+                    out_param.grad = grad_applied.requires_grad_(
+                        param_grad.requires_grad
+                    )
+
+        for key, buf in self._buffers.items():
+            if buf is not None:
+                self._buffers[key] = fn(buf)
+
+        return self
+
+    def apply(self, fn: Callable[["Module"], None]) -> Self:
+        r"""Apply ``fn`` recursively to every submodule (as returned by ``.children()``) as well as self.
+
+        Typical use includes initializing the parameters of a model
+        (see also :ref:`nn-init-doc`).
+
+        Args:
+            fn (:class:`Module` -> None): function to be applied to each submodule
+
+        Returns:
+            Module: self
+
+        Example::
+
+            >>> @torch.no_grad()
+            >>> def init_weights(m):
+            >>>     print(m)
+            >>>     if type(m) is nn.Linear:
+            >>>         m.weight.fill_(1.0)
+            >>>         print(m.weight)
+            >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
+            >>> net.apply(init_weights)
+            Linear(in_features=2, out_features=2, bias=True)
+            Parameter containing:
+            tensor([[1., 1.],
+                    [1., 1.]], requires_grad=True)
+            Linear(in_features=2, out_features=2, bias=True)
+            Parameter containing:
+            tensor([[1., 1.],
+                    [1., 1.]], requires_grad=True)
+            Sequential(
+              (0): Linear(in_features=2, out_features=2, bias=True)
+              (1): Linear(in_features=2, out_features=2, bias=True)
+            )
+
+        """
+        for module in self.children():
+            module.apply(fn)
+        fn(self)
+        return self
+
+    def cuda(self, device: int | device | None = None) -> Self:
+        r"""Move all model parameters and buffers to the GPU.
+
+        This also makes associated parameters and buffers different objects. So
+        it should be called before constructing the optimizer if the module will
+        live on GPU while being optimized.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Args:
+            device (int, optional): if specified, all parameters will be
+                copied to that device
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.cuda(device))
+
+    def ipu(self, device: int | device | None = None) -> Self:
+        r"""Move all model parameters and buffers to the IPU.
+
+        This also makes associated parameters and buffers different objects. So
+        it should be called before constructing the optimizer if the module will
+        live on IPU while being optimized.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Arguments:
+            device (int, optional): if specified, all parameters will be
+                copied to that device
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.ipu(device))
+
+    def xpu(self, device: int | device | None = None) -> Self:
+        r"""Move all model parameters and buffers to the XPU.
+
+        This also makes associated parameters and buffers different objects. So
+        it should be called before constructing optimizer if the module will
+        live on XPU while being optimized.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Arguments:
+            device (int, optional): if specified, all parameters will be
+                copied to that device
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.xpu(device))
+
+    def mtia(self, device: int | device | None = None) -> Self:
+        r"""Move all model parameters and buffers to the MTIA.
+
+        This also makes associated parameters and buffers different objects. So
+        it should be called before constructing the optimizer if the module will
+        live on MTIA while being optimized.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Arguments:
+            device (int, optional): if specified, all parameters will be
+                copied to that device
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.mtia(device))
+
+    def cpu(self) -> Self:
+        r"""Move all model parameters and buffers to the CPU.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.cpu())
+
+    def type(self, dst_type: dtype | str) -> Self:
+        r"""Casts all parameters and buffers to :attr:`dst_type`.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Args:
+            dst_type (type or string): the desired type
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.type(dst_type))
+
+    def float(self) -> Self:
+        r"""Casts all floating point parameters and buffers to ``float`` datatype.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.float() if t.is_floating_point() else t)
+
+    def double(self) -> Self:
+        r"""Casts all floating point parameters and buffers to ``double`` datatype.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.double() if t.is_floating_point() else t)
+
+    def half(self) -> Self:
+        r"""Casts all floating point parameters and buffers to ``half`` datatype.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.half() if t.is_floating_point() else t)
+
+    def bfloat16(self) -> Self:
+        r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Returns:
+            Module: self
+        """
+        return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
+
+    def to_empty(self, *, device: DeviceLikeType | None, recurse: bool = True) -> Self:
+        r"""Move the parameters and buffers to the specified device without copying storage.
+
+        Args:
+            device (:class:`torch.device`): The desired device of the parameters
+                and buffers in this module.
+            recurse (bool): Whether parameters and buffers of submodules should
+                be recursively moved to the specified device.
+
+        Returns:
+            Module: self
+        """
+        return self._apply(
+            lambda t: torch.empty_like(t, device=device), recurse=recurse
+        )
+
+    @overload
+    def to(
+        self,
+        device: DeviceLikeType | None = ...,
+        dtype: dtype | None = ...,
+        non_blocking: bool = ...,
+    ) -> Self: ...
+
+    @overload
+    def to(self, dtype: dtype, non_blocking: bool = ...) -> Self: ...
+
+    @overload
+    def to(self, tensor: Tensor, non_blocking: bool = ...) -> Self: ...
+
+    def to(self, *args, **kwargs):
+        r"""Move and/or cast the parameters and buffers.
+
+        This can be called as
+
+        .. function:: to(device=None, dtype=None, non_blocking=False)
+           :noindex:
+
+        .. function:: to(dtype, non_blocking=False)
+           :noindex:
+
+        .. function:: to(tensor, non_blocking=False)
+           :noindex:
+
+        .. function:: to(memory_format=torch.channels_last)
+           :noindex:
+
+        Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
+        floating point or complex :attr:`dtype`\ s. In addition, this method will
+        only cast the floating point or complex parameters and buffers to :attr:`dtype`
+        (if given). The integral parameters and buffers will be moved
+        :attr:`device`, if that is given, but with dtypes unchanged. When
+        :attr:`non_blocking` is set, it tries to convert/move asynchronously
+        with respect to the host if possible, e.g., moving CPU Tensors with
+        pinned memory to CUDA devices.
+
+        See below for examples.
+
+        .. note::
+            This method modifies the module in-place.
+
+        Args:
+            device (:class:`torch.device`): the desired device of the parameters
+                and buffers in this module
+            dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
+                the parameters and buffers in this module
+            tensor (torch.Tensor): Tensor whose dtype and device are the desired
+                dtype and device for all parameters and buffers in this module
+            memory_format (:class:`torch.memory_format`): the desired memory
+                format for 4D parameters and buffers in this module (keyword
+                only argument)
+
+        Returns:
+            Module: self
+
+        Examples::
+
+            >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+            >>> linear = nn.Linear(2, 2)
+            >>> linear.weight
+            Parameter containing:
+            tensor([[ 0.1913, -0.3420],
+                    [-0.5113, -0.2325]])
+            >>> linear.to(torch.double)
+            Linear(in_features=2, out_features=2, bias=True)
+            >>> linear.weight
+            Parameter containing:
+            tensor([[ 0.1913, -0.3420],
+                    [-0.5113, -0.2325]], dtype=torch.float64)
+            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
+            >>> gpu1 = torch.device("cuda:1")
+            >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
+            Linear(in_features=2, out_features=2, bias=True)
+            >>> linear.weight
+            Parameter containing:
+            tensor([[ 0.1914, -0.3420],
+                    [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
+            >>> cpu = torch.device("cpu")
+            >>> linear.to(cpu)
+            Linear(in_features=2, out_features=2, bias=True)
+            >>> linear.weight
+            Parameter containing:
+            tensor([[ 0.1914, -0.3420],
+                    [-0.5112, -0.2324]], dtype=torch.float16)
+
+            >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
+            >>> linear.weight
+            Parameter containing:
+            tensor([[ 0.3741+0.j,  0.2382+0.j],
+                    [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
+            >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
+            tensor([[0.6122+0.j, 0.1150+0.j],
+                    [0.6122+0.j, 0.1150+0.j],
+                    [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
+
+        """
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
+            *args,
+            **kwargs,
+        )
+
+        if dtype is not None:
+            if not (dtype.is_floating_point or dtype.is_complex):
+                raise TypeError(
+                    "nn.Module.to only accepts floating point or complex "
+                    f"dtypes, but got desired dtype={dtype}"
+                )
+            if dtype.is_complex:
+                warnings.warn(
+                    "Complex modules are a new feature under active development whose design may change, "
+                    "and some modules might not work as expected when using complex tensors as parameters or buffers. "
+                    "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml "
+                    "if a complex module does not work as expected.",
+                    stacklevel=2,
+                )
+
+        def convert(t):
+            try:
+                if convert_to_format is not None and t.dim() in (4, 5):
+                    return t.to(
+                        device,
+                        dtype if t.is_floating_point() or t.is_complex() else None,
+                        non_blocking,
+                        memory_format=convert_to_format,
+                    )
+                return t.to(
+                    device,
+                    dtype if t.is_floating_point() or t.is_complex() else None,
+                    non_blocking,
+                )
+            except NotImplementedError as e:
+                if str(e) == "Cannot copy out of meta tensor; no data!":
+                    raise NotImplementedError(
+                        f"{e} Please use torch.nn.Module.to_empty() instead of torch.nn.Module.to() "
+                        f"when moving module from meta to a different device."
+                    ) from None
+                else:
+                    raise
+
+        return self._apply(convert)
+
+    def register_full_backward_pre_hook(
+        self,
+        hook: Callable[["Module", _grad_t], _grad_t | None],
+        prepend: bool = False,
+    ) -> RemovableHandle:
+        r"""Register a backward pre-hook on the module.
+
+        The hook will be called every time the gradients for the module are computed.
+        The hook should have the following signature::
+
+            hook(module, grad_output) -> tuple[Tensor, ...], Tensor or None
+
+        The :attr:`grad_output` is a tuple. The hook should
+        not modify its arguments, but it can optionally return a new gradient with
+        respect to the output that will be used in place of :attr:`grad_output` in
+        subsequent computations. Entries in :attr:`grad_output` will be ``None`` for
+        all non-Tensor arguments.
+
+        For technical reasons, when this hook is applied to a Module, its forward function will
+        receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
+        of each Tensor returned by the Module's forward function.
+
+        .. warning ::
+            Modifying inputs inplace is not allowed when using backward hooks and
+            will raise an error.
+
+        Args:
+            hook (Callable): The user-defined hook to be registered.
+            prepend (bool): If true, the provided ``hook`` will be fired before
+                all existing ``backward_pre`` hooks on this
+                :class:`torch.nn.Module`. Otherwise, the provided
+                ``hook`` will be fired after all existing ``backward_pre`` hooks
+                on this :class:`torch.nn.Module`. Note that global
+                ``backward_pre`` hooks registered with
+                :func:`register_module_full_backward_pre_hook` will fire before
+                all hooks registered by this method.
+
+        Returns:
+            :class:`torch.utils.hooks.RemovableHandle`:
+                a handle that can be used to remove the added hook by calling
+                ``handle.remove()``
+
+        """
+        handle = RemovableHandle(self._backward_pre_hooks)
+        self._backward_pre_hooks[handle.id] = hook
+        if prepend:
+            self._backward_pre_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]
+        return handle
+
+    def register_backward_hook(
+        self, hook: Callable[["Module", _grad_t, _grad_t], _grad_t | None]
+    ) -> RemovableHandle:
+        r"""Register a backward hook on the module.
+
+        This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
+        the behavior of this function will change in future versions.
+
+        Returns:
+            :class:`torch.utils.hooks.RemovableHandle`:
+                a handle that can be used to remove the added hook by calling
+                ``handle.remove()``
+
+        """
+        if self._is_full_backward_hook is True:
+            raise RuntimeError(
+                "Cannot use both regular backward hooks and full backward hooks on a "
+                "single Module. Please use only one of them."
+            )
+
+        self._is_full_backward_hook = False
+
+        handle = RemovableHandle(self._backward_hooks)
+        self._backward_hooks[handle.id] = hook
+        return handle
+
+    def register_full_backward_hook(
+        self,
+        hook: Callable[["Module", _grad_t, _grad_t], _grad_t | None],
+        prepend: bool = False,
+    ) -> RemovableHandle:
+        r"""Register a backward hook on the module.
+
+        The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:
+
+            1. Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.
+            2. If none of the module inputs require gradients, the hook will fire when the gradients are computed
+               with respect to module outputs.
+            3. If none of the module outputs require gradients, then the hooks will not fire.
+
+        The hook should have the following signature::
+
+            hook(module, grad_input, grad_output) -> tuple(Tensor) or None
+
+        The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
+        with respect to the inputs and outputs respectively. The hook should
+        not modify its arguments, but it can optionally return a new gradient with
+        respect to the input that will be used in place of :attr:`grad_input` in
+        subsequent computations. :attr:`grad_input` will only correspond to the inputs given
+        as positional arguments and all kwarg arguments are ignored. Entries
+        in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
+        arguments.
+
+        For technical reasons, when this hook is applied to a Module, its forward function will
+        receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
+        of each Tensor returned by the Module's forward function.
+
+        .. warning ::
+            Modifying inputs or outputs inplace is not allowed when using backward hooks and
+            will raise an error.
+
+        Args:
+            hook (Callable): The user-defined hook to be registered.
+            prepend (bool): If true, the provided ``hook`` will be fired before
+                all existing ``backward`` hooks on this
+                :class:`torch.nn.Module`. Otherwise, the provided
+                ``hook`` will be fired after all existing ``backward`` hooks on
+                this :class:`torch.nn.Module`. Note that global
+                ``backward`` hooks registered with
+                :func:`register_module_full_backward_hook` will fire before
+                all hooks registered by this method.
+
+        Returns:
+            :class:`torch.utils.hooks.RemovableHandle`:
+                a handle that can be used to remove the added hook by calling
+                ``handle.remove()``
+
+        """
+        if self._is_full_backward_hook is False:
+            raise RuntimeError(
+                "Cannot use both regular backward hooks and full backward hooks on a "
+                "single Module. Please use only one of them."
+            )
+
+        self._is_full_backward_hook = True
+
+        handle = RemovableHandle(self._backward_hooks)
+        self._backward_hooks[handle.id] = hook
+        if prepend:
+            self._backward_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]
+        return handle
+
+    def _get_backward_hooks(self):
+        r"""Return the backward hooks for use in the call function.
+
+        It returns two lists, one with the full backward hooks and one with the non-full
+        backward hooks.
+        """
+        full_backward_hooks: list[Callable] = []
+        if _global_is_full_backward_hook is True:
+            full_backward_hooks += _global_backward_hooks.values()
+        if self._is_full_backward_hook is True:
+            full_backward_hooks += self._backward_hooks.values()
+
+        non_full_backward_hooks: list[Callable] = []
+        if _global_is_full_backward_hook is False:
+            non_full_backward_hooks += _global_backward_hooks.values()
+        if self._is_full_backward_hook is False:
+            non_full_backward_hooks += self._backward_hooks.values()
+
+        return full_backward_hooks, non_full_backward_hooks
+
+    def _get_backward_pre_hooks(self):
+        backward_pre_hooks: list[Callable] = []
+        backward_pre_hooks += _global_backward_pre_hooks.values()
+        backward_pre_hooks += self._backward_pre_hooks.values()
+
+        return backward_pre_hooks
+
+    def _maybe_warn_non_full_backward_hook(self, inputs, result, grad_fn) -> None:
+        if not isinstance(result, torch.Tensor):
+            if not (
+                isinstance(result, tuple)
+                and all(isinstance(r, torch.Tensor) for r in result)
+            ):
+                warnings.warn(
+                    "Using non-full backward hooks on a Module that does not return a "
+                    "single Tensor or a tuple of Tensors is deprecated and will be removed "
+                    "in future versions. This hook will be missing some of the grad_output. "
+                    "Please use register_full_backward_hook to get the documented behavior.",
+                    FutureWarning,
+                    stacklevel=2,
+                )
+                return
+        else:
+            result = (result,)
+
+        if not isinstance(inputs, torch.Tensor):
+            if not (
+                isinstance(inputs, tuple)
+                and all(isinstance(i, torch.Tensor) for i in inputs)
+            ):
+                warnings.warn(
+                    "Using non-full backward hooks on a Module that does not take as input a "
+                    "single Tensor or a tuple of Tensors is deprecated and will be removed "
+                    "in future versions. This hook will be missing some of the grad_input. "
+                    "Please use register_full_backward_hook to get the documented behavior.",
+                    FutureWarning,
+                    stacklevel=2,
+                )
+                return
+        else:
+            inputs = (inputs,)
+
+        # At this point we are sure that inputs and result are tuple of Tensors
+        out_grad_fn = {r.grad_fn for r in result if r.grad_fn is not None}
+        if len(out_grad_fn) == 0 or (
+            len(out_grad_fn) == 1 and grad_fn not in out_grad_fn
+        ):
+            warnings.warn(
+                "Using a non-full backward hook when outputs are nested in python data structure "
+                "is deprecated and will be removed in future versions. This hook will be missing "
+                "some grad_output.",
+                FutureWarning,
+                stacklevel=2,
+            )
+        elif len(out_grad_fn) > 1:
+            warnings.warn(
+                "Using a non-full backward hook when outputs are generated by different autograd Nodes "
+                "is deprecated and will be removed in future versions. This hook will be missing "
+                "some grad_output. Please use register_full_backward_hook to get the documented behavior.",
+                FutureWarning,
+                stacklevel=2,
+            )
+        else:
+            # At this point the grad_output part of the hook will most likely be correct
+            inputs_grad_fn = {i.grad_fn for i in inputs if i.grad_fn is not None}
+
+            next_functions = {n[0] for n in grad_fn.next_functions}
+
+            if inputs_grad_fn != next_functions:
+                warnings.warn(
+                    "Using a non-full backward hook when the forward contains multiple autograd Nodes "
+                    "is deprecated and will be removed in future versions. This hook will be missing "
+                    "some grad_input. Please use register_full_backward_hook to get the documented "
+                    "behavior.",
+                    FutureWarning,
+                    stacklevel=2,
+                )
+
+    def register_forward_pre_hook(
+        self,
+        hook: Callable[[T, tuple[Any, ...]], Any | None]
+        | Callable[
+            [T, tuple[Any, ...], dict[str, Any]], tuple[Any, dict[str, Any]] | None
+        ],
+        *,
+        prepend: bool = False,
+        with_kwargs: bool = False,
+    ) -> RemovableHandle:
+        r"""Register a forward pre-hook on the module.
+
+        The hook will be called every time before :func:`forward` is invoked.
+
+
+        If ``with_kwargs`` is false or not specified, the input contains only
+        the positional arguments given to the module. Keyword arguments won't be
+        passed to the hooks and only to the ``forward``. The hook can modify the
+        input. User can either return a tuple or a single modified value in the
+        hook. We will wrap the value into a tuple if a single value is returned
+        (unless that value is already a tuple). The hook should have the
+        following signature::
+
+            hook(module, args) -> None or modified input
+
+        If ``with_kwargs`` is true, the forward pre-hook will be passed the
+        kwargs given to the forward function. And if the hook modifies the
+        input, both the args and kwargs should be returned. The hook should have
+        the following signature::
+
+            hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
+
+        Args:
+            hook (Callable): The user defined hook to be registered.
+            prepend (bool): If true, the provided ``hook`` will be fired before
+                all existing ``forward_pre`` hooks on this
+                :class:`torch.nn.Module`. Otherwise, the provided
+                ``hook`` will be fired after all existing ``forward_pre`` hooks
+                on this :class:`torch.nn.Module`. Note that global
+                ``forward_pre`` hooks registered with
+                :func:`register_module_forward_pre_hook` will fire before all
+                hooks registered by this method.
+                Default: ``False``
+            with_kwargs (bool): If true, the ``hook`` will be passed the kwargs
+                given to the forward function.
+                Default: ``False``
+
+        Returns:
+            :class:`torch.utils.hooks.RemovableHandle`:
+                a handle that can be used to remove the added hook by calling
+                ``handle.remove()``
+        """
+        handle = RemovableHandle(
+            self._forward_pre_hooks, extra_dict=self._forward_pre_hooks_with_kwargs
+        )
+        self._forward_pre_hooks[handle.id] = hook
+        if with_kwargs:
+            self._forward_pre_hooks_with_kwargs[handle.id] = True
+
+        if prepend:
+            self._forward_pre_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]
+        return handle
+
+    def register_forward_hook(
+        self,
+        hook: Callable[[T, tuple[Any, ...], Any], Any | None]
+        | Callable[[T, tuple[Any, ...], dict[str, Any], Any], Any | None],
+        *,
+        prepend: bool = False,
+        with_kwargs: bool = False,
+        always_call: bool = False,
+    ) -> RemovableHandle:
+        r"""Register a forward hook on the module.
+
+        The hook will be called every time after :func:`forward` has computed an output.
+
+        If ``with_kwargs`` is ``False`` or not specified, the input contains only
+        the positional arguments given to the module. Keyword arguments won't be
+        passed to the hooks and only to the ``forward``. The hook can modify the
+        output. It can modify the input inplace but it will not have effect on
+        forward since this is called after :func:`forward` is called. The hook
+        should have the following signature::
+
+            hook(module, args, output) -> None or modified output
+
+        If ``with_kwargs`` is ``True``, the forward hook will be passed the
+        ``kwargs`` given to the forward function and be expected to return the
+        output possibly modified. The hook should have the following signature::
+
+            hook(module, args, kwargs, output) -> None or modified output
+
+        Args:
+            hook (Callable): The user defined hook to be registered.
+            prepend (bool): If ``True``, the provided ``hook`` will be fired
+                before all existing ``forward`` hooks on this
+                :class:`torch.nn.Module`. Otherwise, the provided
+                ``hook`` will be fired after all existing ``forward`` hooks on
+                this :class:`torch.nn.Module`. Note that global
+                ``forward`` hooks registered with
+                :func:`register_module_forward_hook` will fire before all hooks
+                registered by this method.
+                Default: ``False``
+            with_kwargs (bool): If ``True``, the ``hook`` will be passed the
+                kwargs given to the forward function.
+                Default: ``False``
+            always_call (bool): If ``True`` the ``hook`` will be run regardless of
+                whether an exception is raised while calling the Module.
+                Default: ``False``
+
+        Returns:
+            :class:`torch.utils.hooks.RemovableHandle`:
+                a handle that can be used to remove the added hook by calling
+                ``handle.remove()``
+        """
+        handle = RemovableHandle(
+            self._forward_hooks,
+            extra_dict=[
+                self._forward_hooks_with_kwargs,
+                self._forward_hooks_always_called,
+            ],
+        )
+        self._forward_hooks[handle.id] = hook
+        if with_kwargs:
+            self._forward_hooks_with_kwargs[handle.id] = True
+        if always_call:
+            self._forward_hooks_always_called[handle.id] = True
+        if prepend:
+            self._forward_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]
+        return handle
+
+    def _slow_forward(self, *input, **kwargs):
+        tracing_state = torch._C._get_tracing_state()
+        if not tracing_state or isinstance(self.forward, torch._C.ScriptMethod):
+            return self.forward(*input, **kwargs)
+        recording_scopes = torch.jit._trace._trace_module_map is not None
+        if recording_scopes:
+            # type ignore was added because at this point one knows that
+            # torch.jit._trace._trace_module_map is not Optional and has type Dict[Any, Any]
+            name = torch.jit._trace._trace_module_map.get(self, None)  # type: ignore[operator, union-attr]
+            if name:
+                tracing_state.push_scope(name)
+            else:
+                recording_scopes = False
+        try:
+            result = self.forward(*input, **kwargs)
+        finally:
+            if recording_scopes:
+                tracing_state.pop_scope()
+        return result
+
+    def _wrapped_call_impl(self, *args, **kwargs):
+        if self._compiled_call_impl is not None:
+            return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
+        else:
+            return self._call_impl(*args, **kwargs)
+
+    # torchrec tests the code consistency with the following code
+    # fmt: off
+    def _call_impl(self, *args, **kwargs):
+        forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.forward)
+        # If we don't have any hooks, we want to skip the rest of the logic in
+        # this function, and just call forward.
+        if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
+                or _global_backward_pre_hooks or _global_backward_hooks
+                or _global_forward_hooks or _global_forward_pre_hooks):
+            return forward_call(*args, **kwargs)
+
+        result = None
+        called_always_called_hooks = set()
+
+        def inner():
+            nonlocal result, args, kwargs
+
+            full_backward_hooks, non_full_backward_hooks = [], []
+            backward_pre_hooks = []
+            if self._backward_pre_hooks or _global_backward_pre_hooks:
+                backward_pre_hooks = self._get_backward_pre_hooks()
+
+            if self._backward_hooks or _global_backward_hooks:
+                full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()
+
+            if _global_forward_pre_hooks or self._forward_pre_hooks:
+                for hook_id, hook in (
+                    *_global_forward_pre_hooks.items(),
+                    *self._forward_pre_hooks.items(),
+                ):
+                    if hook_id in self._forward_pre_hooks_with_kwargs:
+                        args_kwargs_result = hook(self, args, kwargs)  # type: ignore[misc]
+                        if args_kwargs_result is not None:
+                            if isinstance(args_kwargs_result, tuple) and len(args_kwargs_result) == 2:
+                                args, kwargs = args_kwargs_result
+                            else:
+                                raise RuntimeError(
+                                    "forward pre-hook must return None or a tuple "
+                                    f"of (new_args, new_kwargs), but got {args_kwargs_result}."
+                                )
+                    else:
+                        args_result = hook(self, args)
+                        if args_result is not None:
+                            if not isinstance(args_result, tuple):
+                                args_result = (args_result,)
+                            args = args_result
+
+            bw_hook = None
+            if full_backward_hooks or backward_pre_hooks:
+                bw_hook = BackwardHook(self, full_backward_hooks, backward_pre_hooks)
+                args = bw_hook.setup_input_hook(args)
+
+            result = forward_call(*args, **kwargs)
+            if _global_forward_hooks or self._forward_hooks:
+                for hook_id, hook in (
+                    *_global_forward_hooks.items(),
+                    *self._forward_hooks.items(),
+                ):
+                    # mark that always called hook is run
+                    if hook_id in self._forward_hooks_always_called or hook_id in _global_forward_hooks_always_called:
+                        called_always_called_hooks.add(hook_id)
+
+                    if hook_id in self._forward_hooks_with_kwargs or hook_id in _global_forward_hooks_with_kwargs:
+                        hook_result = hook(self, args, kwargs, result)
+                    else:
+                        hook_result = hook(self, args, result)
+
+                    if hook_result is not None:
+                        result = hook_result
+
+            if bw_hook:
+                if not isinstance(result, (torch.Tensor, tuple)):
+                    warnings.warn("For backward hooks to be called,"
+                                  " module output should be a Tensor or a tuple of Tensors"
+                                  f" but received {type(result)}", stacklevel=2)
+                result = bw_hook.setup_output_hook(result)
+
+            # Handle the non-full backward hooks
+            if non_full_backward_hooks:
+                var = result
+                while not isinstance(var, torch.Tensor):
+                    if isinstance(var, dict):
+                        var = next(v for v in var.values() if isinstance(v, torch.Tensor))
+                    else:
+                        var = var[0]
+                grad_fn = var.grad_fn
+                if grad_fn is not None:
+                    for hook in non_full_backward_hooks:
+                        grad_fn.register_hook(_WrappedHook(hook, self))
+                    self._maybe_warn_non_full_backward_hook(args, result, grad_fn)
+
+            return result
+
+        # This is technically not behavior equivalent when compiling, but it's
+        # incredibly unlikely we will ever support throwing an exception in NN
+        # module, and then catching it here, and then reraising it, and then
+        # catching it again, and expecting the resulting frame to be compiled.
+        # The reraise here just gunks up our exception handling for no good
+        # reason.  Don't try to run the always called hooks in event of
+        # exception.
+        if torch.compiler.is_compiling():
+            return inner()
+
+        try:
+            return inner()
+        except Exception:
+            # run always called hooks if they have not already been run
+            # For now only forward hooks have the always_call option but perhaps
+            # this functionality should be added to full backward hooks as well.
+            for hook_id, hook in _global_forward_hooks.items():
+                if hook_id in _global_forward_hooks_always_called and hook_id not in called_always_called_hooks:  # type: ignore[possibly-undefined]
+                    try:
+                        hook_result = hook(self, args, result)  # type: ignore[possibly-undefined]
+                        if hook_result is not None:
+                            result = hook_result
+                    except Exception as e:
+                        warnings.warn("global module forward hook with ``always_call=True`` raised an exception "
+                                      f"that was silenced as another error was raised in forward: {str(e)}", stacklevel=2)
+                        continue
+
+            for hook_id, hook in self._forward_hooks.items():
+                if hook_id in self._forward_hooks_always_called and hook_id not in called_always_called_hooks:  # type: ignore[possibly-undefined]
+                    try:
+                        if hook_id in self._forward_hooks_with_kwargs:
+                            hook_result = hook(self, args, kwargs, result)  # type: ignore[possibly-undefined]
+                        else:
+                            hook_result = hook(self, args, result)  # type: ignore[possibly-undefined]
+                        if hook_result is not None:
+                            result = hook_result
+                    except Exception as e:
+                        warnings.warn("module forward hook with ``always_call=True`` raised an exception "
+                                      f"that was silenced as another error was raised in forward: {str(e)}", stacklevel=2)
+                        continue
+            # raise exception raised in try block
+            raise
+    # fmt: on
+
+    __call__: Callable[..., Any] = _wrapped_call_impl
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop("_compiled_call_impl", None)
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+        # Support loading old checkpoints that don't have the following attrs:
+        if "_forward_pre_hooks" not in self.__dict__:
+            self._forward_pre_hooks = OrderedDict()
+        if "_forward_pre_hooks_with_kwargs" not in self.__dict__:
+            self._forward_pre_hooks_with_kwargs = OrderedDict()
+        if "_forward_hooks_with_kwargs" not in self.__dict__:
+            self._forward_hooks_with_kwargs = OrderedDict()
+        if "_forward_hooks_always_called" not in self.__dict__:
+            self._forward_hooks_always_called = OrderedDict()
+        if "_state_dict_hooks" not in self.__dict__:
+            self._state_dict_hooks = OrderedDict()
+        if "_state_dict_pre_hooks" not in self.__dict__:
+            self._state_dict_pre_hooks = OrderedDict()
+        if "_load_state_dict_pre_hooks" not in self.__dict__:
+            self._load_state_dict_pre_hooks = OrderedDict()
+        if "_load_state_dict_post_hooks" not in self.__dict__:
+            self._load_state_dict_post_hooks = OrderedDict()
+        if "_non_persistent_buffers_set" not in self.__dict__:
+            self._non_persistent_buffers_set = set()
+        if "_is_full_backward_hook" not in self.__dict__:
+            self._is_full_backward_hook = None
+        if "_backward_pre_hooks" not in self.__dict__:
+            self._backward_pre_hooks = OrderedDict()
+
+    # It is crucial that the return type is not annotated as `Any`, otherwise type checking
+    # on `torch.nn.Module` and all its subclasses is largely disabled as a result. See:
+    # https://github.com/pytorch/pytorch/pull/115074
+    def __getattr__(self, name: str) -> Union[Tensor, "Module"]:
+        if "_parameters" in self.__dict__:
+            _parameters = self.__dict__["_parameters"]
+            if name in _parameters:
+                return _parameters[name]
+        if "_buffers" in self.__dict__:
+            _buffers = self.__dict__["_buffers"]
+            if name in _buffers:
+                return _buffers[name]
+        if "_modules" in self.__dict__:
+            modules = self.__dict__["_modules"]
+            if name in modules:
+                return modules[name]
+        raise AttributeError(
+            f"'{type(self).__name__}' object has no attribute '{name}'"
+        )
+
+    def __setattr__(self, name: str, value: Union[Tensor, "Module"]) -> None:
+        def remove_from(*dicts_or_sets) -> None:
+            for d in dicts_or_sets:
+                if name in d:
+                    if isinstance(d, dict):
+                        del d[name]
+                    else:
+                        d.discard(name)
+
+        params = self.__dict__.get("_parameters")
+        if isinstance(value, Parameter):
+            if params is None:
+                raise AttributeError(
+                    "cannot assign parameters before Module.__init__() call"
+                )
+            remove_from(
+                self.__dict__,
+                self._buffers,
+                self._modules,
+                self._non_persistent_buffers_set,
+            )
+            self.register_parameter(name, value)
+        elif params is not None and name in params:
+            if value is not None:
+                raise TypeError(
+                    f"cannot assign '{torch.typename(value)}' as parameter '{name}' "
+                    "(torch.nn.Parameter or None expected)"
+                )
+            self.register_parameter(name, value)
+        else:
+            modules = self.__dict__.get("_modules")
+            if isinstance(value, Module):
+                if modules is None:
+                    raise AttributeError(
+                        "cannot assign module before Module.__init__() call"
+                    )
+                remove_from(
+                    self.__dict__,
+                    self._parameters,
+                    self._buffers,
+                    self._non_persistent_buffers_set,
+                )
+                for hook in _global_module_registration_hooks.values():
+                    output = hook(self, name, value)
+                    if output is not None:
+                        value = output
+                modules[name] = value
+            elif modules is not None and name in modules:
+                if value is not None:
+                    raise TypeError(
+                        f"cannot assign '{torch.typename(value)}' as child module '{name}' "
+                        "(torch.nn.Module or None expected)"
+                    )
+                for hook in _global_module_registration_hooks.values():
+                    output = hook(self, name, value)
+                    if output is not None:
+                        value = output
+                modules[name] = value
+            else:
+                buffers = self.__dict__.get("_buffers")
+                if isinstance(value, Buffer) or buffers is not None and name in buffers:
+                    if value is not None and not (
+                        isinstance(value, torch.Tensor)
+                        or hasattr(value, "__torch_function__")
+                    ):
+                        raise TypeError(
+                            f"cannot assign '{torch.typename(value)}' as buffer '{name}' "
+                            "(torch.nn.Buffer, torch.Tensor or None expected)"
+                        )
+                    if isinstance(value, Buffer):
+                        persistent = value.persistent
+                    else:
+                        persistent = name not in self._non_persistent_buffers_set
+                    # === HACK ===
+                    # This whole block below should just be:
+                    # self.register_buffer(name, value, persistent)
+
+                    # But to support subclasses of nn.Module that (wrongfully) implement a
+                    # register_buffer() method that doesn't have the "persistent"
+                    # argument. Only pass it in if it is accepted otherwise assume
+                    # it is always true
+                    if (
+                        getattr(self.register_buffer, "__func__", None)
+                        is torch.nn.Module.register_buffer
+                    ):
+                        self.register_buffer(name, value, persistent)
+                    else:
+                        sign = inspect.signature(self.register_buffer)
+                        if "persistent" in sign.parameters:
+                            self.register_buffer(name, value, persistent)
+                        else:
+                            if not persistent:
+                                raise RuntimeError(
+                                    "Registering a non-persistent buffer "
+                                    "on a Module subclass that implements "
+                                    "register_buffer() without the persistent "
+                                    "argument is not allowed."
+                                )
+                            # Assume that the implementation without the argument has the
+                            # behavior from before the argument was added: persistent=True
+                            self.register_buffer(name, value)
+                    # === HACK END ===
+                else:
+                    super().__setattr__(name, value)
+
+    def __delattr__(self, name) -> None:
+        if name in self._parameters:
+            del self._parameters[name]
+        elif name in self._buffers:
+            del self._buffers[name]
+            self._non_persistent_buffers_set.discard(name)
+        elif name in self._modules:
+            del self._modules[name]
+        else:
+            super().__delattr__(name)
+
+    def _register_state_dict_hook(self, hook):
+        r"""Register a post-hook for the :meth:`~torch.nn.Module.state_dict` method.
+
+        It should have the following signature::
+            hook(module, state_dict, prefix, local_metadata) -> None or state_dict
+
+        The registered hooks can modify the ``state_dict`` inplace or return a new one.
+        If a new ``state_dict`` is returned, it will only be respected if it is the root
+        module that :meth:`~nn.Module.state_dict` is called from.
+        """
+        if getattr(hook, "_from_public_api", False):
+            raise RuntimeError(
+                "Cannot register the same function as the state dict post hook that was "
+                "previously registered via register_state_dict_post_hook"
+            )
+        handle = RemovableHandle(self._state_dict_hooks)
+        self._state_dict_hooks[handle.id] = hook
+        return handle
+
+    def register_state_dict_post_hook(self, hook):
+        r"""Register a post-hook for the :meth:`~torch.nn.Module.state_dict` method.
+
+        It should have the following signature::
+            hook(module, state_dict, prefix, local_metadata) -> None
+
+        The registered hooks can modify the ``state_dict`` inplace.
+        """
+        # In _register_state_dict_hook there was a bug described in
+        # https://github.com/pytorch/pytorch/issues/117437 where the return value
+        # was only respected for the root module but not child submodules.
+        # We fix this in this public version by only allowing inplace modifications on
+        # the state_dict by the hook. However, since hooks registered via both these
+        # APIs will be added to `_state_dict_hooks` and the type of `_state_dict_hooks`
+        # cannot be changed due to many dependencies on it, we mark a hook
+        # as being registered via the public API by setting `_from_public_api` on it.
+        # In the implementation of `state_dict`, if the callable does not have this
+        # flag, the old behavior of respecting the return value will be preserved
+        # for the root module, otherwise, we ensure that the hook returns None.
+        hook._from_public_api = True
+        handle = RemovableHandle(self._state_dict_hooks)
+        self._state_dict_hooks[handle.id] = hook
+        return handle
+
+    def register_state_dict_pre_hook(self, hook):
+        r"""Register a pre-hook for the :meth:`~torch.nn.Module.state_dict` method.
+
+        It should have the following signature::
+            hook(module, prefix, keep_vars) -> None
+
+        The registered hooks can be used to perform pre-processing before the ``state_dict``
+        call is made.
+        """
+        handle = RemovableHandle(self._state_dict_pre_hooks)
+        self._state_dict_pre_hooks[handle.id] = hook
+        return handle
+
+    def _save_to_state_dict(self, destination, prefix, keep_vars) -> None:
+        r"""Save module state to the `destination` dictionary.
+
+        The `destination` dictionary will contain the state
+        of the module, but not its descendants. This is called on every
+        submodule in :meth:`~torch.nn.Module.state_dict`.
+
+        In rare cases, subclasses can achieve class-specific behavior by
+        overriding this method with custom logic.
+
+        Args:
+            destination (dict): a dict where state will be stored
+            prefix (str): the prefix for parameters and buffers used in this
+                module
+        """
+        for name, param in self._parameters.items():
+            if param is not None:
+                destination[prefix + name] = param if keep_vars else param.detach()
+        for name, buf in self._buffers.items():
+            if buf is not None and name not in self._non_persistent_buffers_set:
+                destination[prefix + name] = buf if keep_vars else buf.detach()
+        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
+        if (
+            getattr(self.__class__, "get_extra_state", Module.get_extra_state)
+            is not Module.get_extra_state
+        ):
+            destination[extra_state_key] = self.get_extra_state()
+
+    # The user can pass an optional arbitrary mappable object to `state_dict`, in which case `state_dict` returns
+    # back that same object. But if they pass nothing, an `OrderedDict` is created and returned.
+    T_destination = TypeVar("T_destination", bound=dict[str, Any])
+
+    @overload
+    def state_dict(
+        self,
+        *,
+        destination: T_destination,
+        prefix: str = ...,
+        keep_vars: bool = ...,
+    ) -> T_destination: ...
+
+    @overload
+    def state_dict(
+        self,
+        *,
+        prefix: str = ...,
+        keep_vars: bool = ...,
+    ) -> dict[str, Any]: ...
+
+    # TODO: Change `*args` to `*` and remove the corresponding warning in docs when BC allows.
+    # Also remove the logic for arg parsing together.
+    def state_dict(self, *args, destination=None, prefix="", keep_vars=False):
+        r"""Return a dictionary containing references to the whole state of the module.
+
+        Both parameters and persistent buffers (e.g. running averages) are
+        included. Keys are corresponding parameter and buffer names.
+        Parameters and buffers set to ``None`` are not included.
+
+        .. note::
+            The returned object is a shallow copy. It contains references
+            to the module's parameters and buffers.
+
+        .. warning::
+            Currently ``state_dict()`` also accepts positional arguments for
+            ``destination``, ``prefix`` and ``keep_vars`` in order. However,
+            this is being deprecated and keyword arguments will be enforced in
+            future releases.
+
+        .. warning::
+            Please avoid the use of argument ``destination`` as it is not
+            designed for end-users.
+
+        Args:
+            destination (dict, optional): If provided, the state of module will
+                be updated into the dict and the same object is returned.
+                Otherwise, an ``OrderedDict`` will be created and returned.
+                Default: ``None``.
+            prefix (str, optional): a prefix added to parameter and buffer
+                names to compose the keys in state_dict. Default: ``''``.
+            keep_vars (bool, optional): by default the :class:`~torch.Tensor` s
+                returned in the state dict are detached from autograd. If it's
+                set to ``True``, detaching will not be performed.
+                Default: ``False``.
+
+        Returns:
+            dict:
+                a dictionary containing a whole state of the module
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> module.state_dict().keys()
+            ['bias', 'weight']
+
+        """
+        # TODO: Remove `args` and the parsing logic when BC allows.
+        if len(args) > 0:
+            # DeprecationWarning is ignored by default
+            warnings.warn(
+                "Positional args are being deprecated, use kwargs instead. Refer to "
+                "https://pytorch.org/docs/main/generated/torch.nn.Module.html#torch.nn.Module.state_dict"
+                " for details.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            if destination is None:
+                destination = args[0]
+            if len(args) > 1 and prefix == "":
+                prefix = args[1]
+            if len(args) > 2 and keep_vars is False:
+                keep_vars = args[2]
+
+        if destination is None:
+            destination = OrderedDict()
+            # pyrefly: ignore [missing-attribute]
+            destination._metadata = OrderedDict()
+
+        local_metadata = dict(version=self._version)
+        if hasattr(destination, "_metadata"):
+            destination._metadata[prefix[:-1]] = local_metadata
+
+        for hook in self._state_dict_pre_hooks.values():
+            hook(self, prefix, keep_vars)
+        self._save_to_state_dict(destination, prefix, keep_vars)
+        for name, module in self._modules.items():
+            if module is not None:
+                module.state_dict(
+                    destination=destination,
+                    prefix=prefix + name + ".",
+                    keep_vars=keep_vars,
+                )
+        for hook in self._state_dict_hooks.values():
+            hook_result = hook(self, destination, prefix, local_metadata)
+            if not getattr(hook, "_from_public_api", False):
+                if hook_result is not None:
+                    destination = hook_result
+            else:
+                if hook_result is not None:
+                    raise RuntimeError("state_dict post-hook must return None")
+        return destination
+
+    def _register_load_state_dict_pre_hook(self, hook, with_module=False):
+        r"""See :meth:`~torch.nn.Module.register_load_state_dict_pre_hook` for details.
+
+        A subtle difference is that if ``with_module`` is set to ``False``, then the
+        hook will not take the ``module`` as the first argument whereas
+        :meth:`~torch.nn.Module.register_load_state_dict_pre_hook` always takes the
+        ``module`` as the first argument.
+
+        Arguments:
+            hook (Callable): Callable hook that will be invoked before
+                loading the state dict.
+            with_module (bool, optional): Whether or not to pass the module
+                instance to the hook as the first parameter.
+        """
+        handle = RemovableHandle(self._load_state_dict_pre_hooks)
+        self._load_state_dict_pre_hooks[handle.id] = _WrappedHook(
+            hook, self if with_module else None
+        )
+        return handle
+
+    def register_load_state_dict_pre_hook(self, hook):
+        r"""Register a pre-hook to be run before module's :meth:`~nn.Module.load_state_dict` is called.
+
+        It should have the following signature::
+            hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None  # noqa: B950
+
+        Arguments:
+            hook (Callable): Callable hook that will be invoked before
+                loading the state dict.
+        """
+        return self._register_load_state_dict_pre_hook(hook, with_module=True)
+
+    def register_load_state_dict_post_hook(self, hook):
+        r"""Register a post-hook to be run after module's :meth:`~nn.Module.load_state_dict` is called.
+
+        It should have the following signature::
+            hook(module, incompatible_keys) -> None
+
+        The ``module`` argument is the current module that this hook is registered
+        on, and the ``incompatible_keys`` argument is a ``NamedTuple`` consisting
+        of attributes ``missing_keys`` and ``unexpected_keys``. ``missing_keys``
+        is a ``list`` of ``str`` containing the missing keys and
+        ``unexpected_keys`` is a ``list`` of ``str`` containing the unexpected keys.
+
+        The given incompatible_keys can be modified inplace if needed.
+
+        Note that the checks performed when calling :func:`load_state_dict` with
+        ``strict=True`` are affected by modifications the hook makes to
+        ``missing_keys`` or ``unexpected_keys``, as expected. Additions to either
+        set of keys will result in an error being thrown when ``strict=True``, and
+        clearing out both missing and unexpected keys will avoid an error.
+
+        Returns:
+            :class:`torch.utils.hooks.RemovableHandle`:
+                a handle that can be used to remove the added hook by calling
+                ``handle.remove()``
+        """
+        handle = RemovableHandle(self._load_state_dict_post_hooks)
+        self._load_state_dict_post_hooks[handle.id] = hook
+        return handle
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ) -> None:
+        r"""Copy parameters and buffers from :attr:`state_dict` into only this module, but not its descendants.
+
+        This is called on every submodule
+        in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this
+        module in input :attr:`state_dict` is provided as :attr:`local_metadata`.
+        For state dicts without metadata, :attr:`local_metadata` is empty.
+        Subclasses can achieve class-specific backward compatible loading using
+        the version number at `local_metadata.get("version", None)`.
+        Additionally, :attr:`local_metadata` can also contain the key
+        `assign_to_params_buffers` that indicates whether keys should be
+        assigned their corresponding tensor in the state_dict.
+
+        .. note::
+            :attr:`state_dict` is not the same object as the input
+            :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So
+            it can be modified.
+
+        Args:
+            state_dict (dict): a dict containing parameters and
+                persistent buffers.
+            prefix (str): the prefix for parameters and buffers used in this
+                module
+            local_metadata (dict): a dict containing the metadata for this module.
+                See
+            strict (bool): whether to strictly enforce that the keys in
+                :attr:`state_dict` with :attr:`prefix` match the names of
+                parameters and buffers in this module
+            missing_keys (list of str): if ``strict=True``, add missing keys to
+                this list
+            unexpected_keys (list of str): if ``strict=True``, add unexpected
+                keys to this list
+            error_msgs (list of str): error messages should be added to this
+                list, and will be reported together in
+                :meth:`~torch.nn.Module.load_state_dict`
+        """
+        for hook in self._load_state_dict_pre_hooks.values():
+            hook(
+                state_dict,
+                prefix,
+                local_metadata,
+                strict,
+                missing_keys,
+                unexpected_keys,
+                error_msgs,
+            )
+
+        persistent_buffers = {
+            k: v
+            for k, v in self._buffers.items()
+            if k not in self._non_persistent_buffers_set
+        }
+        local_name_params = itertools.chain(
+            self._parameters.items(),
+            # pyrefly: ignore [bad-argument-type]
+            persistent_buffers.items(),
+        )
+        local_state = {k: v for k, v in local_name_params if v is not None}
+        assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False)
+        use_swap_tensors = torch.__future__.get_swap_module_params_on_conversion()
+
+        for name, param in local_state.items():
+            key = prefix + name
+            if key in state_dict:
+                input_param = state_dict[key]
+                if not torch.overrides.is_tensor_like(input_param):
+                    error_msgs.append(
+                        f'While copying the parameter named "{key}", '
+                        "expected torch.Tensor or Tensor-like object from checkpoint but "
+                        f"received {type(input_param)}"
+                    )
+                    continue
+
+                # This is used to avoid copying uninitialized parameters into
+                # non-lazy modules, since they dont have the hook to do the checks
+                # in such case, it will error when accessing the .shape attribute.
+                is_param_lazy = torch.nn.parameter.is_lazy(param)
+                # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
+                if (
+                    not is_param_lazy
+                    and len(param.shape) == 0
+                    and len(input_param.shape) == 1
+                    and input_param.shape[0] == 1
+                ):
+                    input_param = input_param[0]
+
+                if not is_param_lazy and input_param.shape != param.shape:
+                    # local shape should match the one in checkpoint
+                    error_msgs.append(
+                        f"size mismatch for {key}: copying a param with shape {input_param.shape} from checkpoint, "
+                        f"the shape in current model is {param.shape}."
+                    )
+                    continue
+
+                if (
+                    param.is_meta
+                    and not input_param.is_meta
+                    and not assign_to_params_buffers
+                ):
+                    warnings.warn(
+                        f"for {key}: copying from a non-meta parameter in the checkpoint to a meta "
+                        "parameter in the current model, which is a no-op. (Did you mean to "
+                        "pass `assign=True` to assign items in the state dictionary to their "
+                        "corresponding key in the module instead of copying them in place?)",
+                        stacklevel=2,
+                    )
+
+                try:
+                    with torch.no_grad():
+                        if use_swap_tensors:
+                            new_input_param = param.module_load(
+                                input_param, assign=assign_to_params_buffers
+                            )
+                            if id(new_input_param) == id(input_param) or id(
+                                new_input_param
+                            ) == id(param):
+                                raise RuntimeError(
+                                    "module_load returned one of self or other, please .detach() "
+                                    "the result if returning one of the inputs in module_load"
+                                )
+                            if isinstance(param, torch.nn.Parameter):
+                                if not isinstance(new_input_param, torch.nn.Parameter):
+                                    new_input_param = torch.nn.Parameter(
+                                        new_input_param,
+                                        requires_grad=param.requires_grad,
+                                    )
+                                else:
+                                    new_input_param.requires_grad_(param.requires_grad)
+                            torch.utils.swap_tensors(param, new_input_param)
+                            del new_input_param
+                        elif assign_to_params_buffers:
+                            # Shape checks are already done above
+                            if isinstance(param, torch.nn.Parameter):
+                                if not isinstance(input_param, torch.nn.Parameter):
+                                    input_param = torch.nn.Parameter(
+                                        input_param, requires_grad=param.requires_grad
+                                    )
+                                else:
+                                    input_param.requires_grad_(param.requires_grad)
+                            setattr(self, name, input_param)
+                        else:
+                            param.copy_(input_param)
+                except Exception as ex:
+                    action = "swapping" if use_swap_tensors else "copying"
+                    error_msgs.append(
+                        f'While {action} the parameter named "{key}", '
+                        f"whose dimensions in the model are {param.size()} and "
+                        f"whose dimensions in the checkpoint are {input_param.size()}, "
+                        f"an exception occurred : {ex.args}."
+                    )
+            elif strict:
+                missing_keys.append(key)
+
+        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
+        if (
+            getattr(self.__class__, "set_extra_state", Module.set_extra_state)
+            is not Module.set_extra_state
+        ):
+            if extra_state_key in state_dict:
+                self.set_extra_state(state_dict[extra_state_key])
+            elif strict:
+                missing_keys.append(extra_state_key)
+        elif strict and (extra_state_key in state_dict):
+            unexpected_keys.append(extra_state_key)
+
+        if strict:
+            for key in state_dict:
+                if key.startswith(prefix) and key != extra_state_key:
+                    input_name = key[len(prefix) :].split(".", 1)
+                    # Must be Module if it have attributes
+                    if len(input_name) > 1:
+                        if input_name[0] not in self._modules:
+                            unexpected_keys.append(key)
+                    elif input_name[0] not in local_state:
+                        unexpected_keys.append(key)
+
+    def load_state_dict(
+        self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False
+    ):
+        r"""Copy parameters and buffers from :attr:`state_dict` into this module and its descendants.
+
+        If :attr:`strict` is ``True``, then
+        the keys of :attr:`state_dict` must exactly match the keys returned
+        by this module's :meth:`~torch.nn.Module.state_dict` function.
+
+        .. warning::
+            If :attr:`assign` is ``True`` the optimizer must be created after
+            the call to :attr:`load_state_dict` unless
+            :func:`~torch.__future__.get_swap_module_params_on_conversion` is ``True``.
+
+        Args:
+            state_dict (dict): a dict containing parameters and
+                persistent buffers.
+            strict (bool, optional): whether to strictly enforce that the keys
+                in :attr:`state_dict` match the keys returned by this module's
+                :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
+            assign (bool, optional): When set to ``False``, the properties of the tensors
+                in the current module are preserved whereas setting it to ``True`` preserves
+                properties of the Tensors in the state dict. The only
+                exception is the ``requires_grad`` field of :class:`~torch.nn.Parameter`
+                for which the value from the module is preserved. Default: ``False``
+
+        Returns:
+            ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
+                * ``missing_keys`` is a list of str containing any keys that are expected
+                    by this module but missing from the provided ``state_dict``.
+                * ``unexpected_keys`` is a list of str containing the keys that are not
+                    expected by this module but present in the provided ``state_dict``.
+
+        Note:
+            If a parameter or buffer is registered as ``None`` and its corresponding key
+            exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
+            ``RuntimeError``.
+        """
+        if not isinstance(state_dict, Mapping):
+            raise TypeError(
+                f"Expected state_dict to be dict-like, got {type(state_dict)}."
+            )
+
+        missing_keys: list[str] = []
+        unexpected_keys: list[str] = []
+        error_msgs: list[str] = []
+
+        # copy state_dict so _load_from_state_dict can modify it
+        metadata = getattr(state_dict, "_metadata", None)
+        state_dict = OrderedDict(state_dict)
+        if metadata is not None:
+            # mypy isn't aware that "_metadata" exists in state_dict
+            state_dict._metadata = metadata  # type: ignore[attr-defined]
+
+        def load(module, local_state_dict, prefix="") -> None:
+            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+            if assign:
+                local_metadata["assign_to_params_buffers"] = assign
+            module._load_from_state_dict(
+                local_state_dict,
+                prefix,
+                local_metadata,
+                True,
+                missing_keys,
+                unexpected_keys,
+                error_msgs,
+            )
+            for name, child in module._modules.items():
+                if child is not None:
+                    child_prefix = prefix + name + "."
+                    child_state_dict = {
+                        k: v
+                        for k, v in local_state_dict.items()
+                        if k.startswith(child_prefix)
+                    }
+                    load(child, child_state_dict, child_prefix)  # noqa: F821
+
+            # Note that the hook can modify missing_keys and unexpected_keys.
+            incompatible_keys = _IncompatibleKeys(missing_keys, unexpected_keys)
+            for hook in module._load_state_dict_post_hooks.values():
+                out = hook(module, incompatible_keys)
+                if out is not None:
+                    raise AssertionError(
+                        "Hooks registered with ``register_load_state_dict_post_hook`` are not"
+                        "expected to return new values, if incompatible_keys need to be modified,"
+                        "it should be done inplace."
+                    )
+
+        load(self, state_dict)
+        del load
+
+        if strict:
+            if len(unexpected_keys) > 0:
+                error_msgs.insert(
+                    0,
+                    "Unexpected key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in unexpected_keys)
+                    ),
+                )
+            if len(missing_keys) > 0:
+                error_msgs.insert(
+                    0,
+                    "Missing key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in missing_keys)
+                    ),
+                )
+
+        if len(error_msgs) > 0:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    self.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+        return _IncompatibleKeys(missing_keys, unexpected_keys)
+
+    def _named_members(
+        self, get_members_fn, prefix="", recurse=True, remove_duplicate: bool = True
+    ):
+        r"""Help yield various names + members of modules."""
+        memo = set()
+        modules = (
+            self.named_modules(prefix=prefix, remove_duplicate=remove_duplicate)
+            if recurse
+            else [(prefix, self)]
+        )
+        for module_prefix, module in modules:
+            members = get_members_fn(module)
+            for k, v in members:
+                if v is None or v in memo:
+                    continue
+                if remove_duplicate:
+                    memo.add(v)
+                name = module_prefix + ("." if module_prefix else "") + k
+                yield name, v
+
+    def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
+        r"""Return an iterator over module parameters.
+
+        This is typically passed to an optimizer.
+
+        Args:
+            recurse (bool): if True, then yields parameters of this module
+                and all submodules. Otherwise, yields only parameters that
+                are direct members of this module.
+
+        Yields:
+            Parameter: module parameter
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> for param in model.parameters():
+            >>>     print(type(param), param.size())
+            <class 'torch.Tensor'> (20L,)
+            <class 'torch.Tensor'> (20L, 1L, 5L, 5L)
+
+        """
+        for _name, param in self.named_parameters(recurse=recurse):
+            yield param
+
+    def named_parameters(
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[tuple[str, Parameter]]:
+        r"""Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
+
+        Args:
+            prefix (str): prefix to prepend to all parameter names.
+            recurse (bool): if True, then yields parameters of this module
+                and all submodules. Otherwise, yields only parameters that
+                are direct members of this module.
+            remove_duplicate (bool, optional): whether to remove the duplicated
+                parameters in the result. Defaults to True.
+
+        Yields:
+            (str, Parameter): Tuple containing the name and parameter
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> for name, param in self.named_parameters():
+            >>>     if name in ['bias']:
+            >>>         print(param.size())
+
+        """
+        gen = self._named_members(
+            lambda module: module._parameters.items(),
+            prefix=prefix,
+            recurse=recurse,
+            remove_duplicate=remove_duplicate,
+        )
+        yield from gen
+
+    def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
+        r"""Return an iterator over module buffers.
+
+        Args:
+            recurse (bool): if True, then yields buffers of this module
+                and all submodules. Otherwise, yields only buffers that
+                are direct members of this module.
+
+        Yields:
+            torch.Tensor: module buffer
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> for buf in model.buffers():
+            >>>     print(type(buf), buf.size())
+            <class 'torch.Tensor'> (20L,)
+            <class 'torch.Tensor'> (20L, 1L, 5L, 5L)
+
+        """
+        for _, buf in self.named_buffers(recurse=recurse):
+            yield buf
+
+    def named_buffers(
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[tuple[str, Tensor]]:
+        r"""Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
+
+        Args:
+            prefix (str): prefix to prepend to all buffer names.
+            recurse (bool, optional): if True, then yields buffers of this module
+                and all submodules. Otherwise, yields only buffers that
+                are direct members of this module. Defaults to True.
+            remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.
+
+        Yields:
+            (str, torch.Tensor): Tuple containing the name and buffer
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> for name, buf in self.named_buffers():
+            >>>     if name in ['running_var']:
+            >>>         print(buf.size())
+
+        """
+        gen = self._named_members(
+            lambda module: module._buffers.items(),
+            prefix=prefix,
+            recurse=recurse,
+            remove_duplicate=remove_duplicate,
+        )
+        yield from gen
+
+    def children(self) -> Iterator["Module"]:
+        r"""Return an iterator over immediate children modules.
+
+        Yields:
+            Module: a child module
+        """
+        for _name, module in self.named_children():
+            yield module
+
+    def named_children(self) -> Iterator[tuple[str, "Module"]]:
+        r"""Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
+
+        Yields:
+            (str, Module): Tuple containing a name and child module
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> for name, module in model.named_children():
+            >>>     if name in ['conv4', 'conv5']:
+            >>>         print(module)
+
+        """
+        memo = set()
+        for name, module in self._modules.items():
+            if module is not None and module not in memo:
+                memo.add(module)
+                yield name, module
+
+    def modules(self, remove_duplicate: bool = True) -> Iterator["Module"]:
+        r"""Return an iterator over all modules in the network.
+
+        Args:
+            remove_duplicate: whether to remove the duplicated module instances in the result
+                or not.
+
+        Yields:
+            Module: a module in the network
+
+        Note:
+            Duplicate modules are returned only once by default. In the following
+            example, ``l`` will be returned only once.
+
+        Example::
+
+            >>> l = nn.Linear(2, 2)
+            >>> net = nn.Sequential(l, l)
+            >>> for idx, m in enumerate(net.modules()):
+            ...     print(idx, '->', m)
+
+            0 -> Sequential(
+              (0): Linear(in_features=2, out_features=2, bias=True)
+              (1): Linear(in_features=2, out_features=2, bias=True)
+            )
+            1 -> Linear(in_features=2, out_features=2, bias=True)
+
+        """
+        for _, module in self.named_modules(remove_duplicate=remove_duplicate):
+            yield module
+
+    def named_modules(
+        self,
+        memo: set["Module"] | None = None,  # noqa: UP007
+        prefix: str = "",
+        remove_duplicate: bool = True,
+    ):
+        r"""Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
+
+        Args:
+            memo: a memo to store the set of modules already added to the result
+            prefix: a prefix that will be added to the name of the module
+            remove_duplicate: whether to remove the duplicated module instances in the result
+                or not
+
+        Yields:
+            (str, Module): Tuple of name and module
+
+        Note:
+            Duplicate modules are returned only once. In the following
+            example, ``l`` will be returned only once.
+
+        Example::
+
+            >>> l = nn.Linear(2, 2)
+            >>> net = nn.Sequential(l, l)
+            >>> for idx, m in enumerate(net.named_modules()):
+            ...     print(idx, '->', m)
+
+            0 -> ('', Sequential(
+              (0): Linear(in_features=2, out_features=2, bias=True)
+              (1): Linear(in_features=2, out_features=2, bias=True)
+            ))
+            1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
+
+        """
+        if memo is None:
+            memo = set()
+        if self not in memo:
+            if remove_duplicate:
+                memo.add(self)
+            yield prefix, self
+            for name, module in self._modules.items():
+                if module is None:
+                    continue
+                submodule_prefix = prefix + ("." if prefix else "") + name
+                yield from module.named_modules(
+                    memo, submodule_prefix, remove_duplicate
+                )
+
+    def train(self, mode: bool = True) -> Self:
+        r"""Set the module in training mode.
+
+        This has an effect only on certain modules. See the documentation of
+        particular modules for details of their behaviors in training/evaluation
+        mode, i.e., whether they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
+        etc.
+
+        Args:
+            mode (bool): whether to set training mode (``True``) or evaluation
+                         mode (``False``). Default: ``True``.
+
+        Returns:
+            Module: self
+        """
+        if not isinstance(mode, bool):
+            raise ValueError("training mode is expected to be boolean")
+        self.training = mode
+        for module in self.children():
+            module.train(mode)
+        return self
+
+    def eval(self) -> Self:
+        r"""Set the module in evaluation mode.
+
+        This has an effect only on certain modules. See the documentation of
+        particular modules for details of their behaviors in training/evaluation
+        mode, i.e. whether they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
+        etc.
+
+        This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.
+
+        See :ref:`locally-disable-grad-doc` for a comparison between
+        `.eval()` and several similar mechanisms that may be confused with it.
+
+        Returns:
+            Module: self
+        """
+        return self.train(False)
+
+    def requires_grad_(self, requires_grad: bool = True) -> Self:
+        r"""Change if autograd should record operations on parameters in this module.
+
+        This method sets the parameters' :attr:`requires_grad` attributes
+        in-place.
+
+        This method is helpful for freezing part of the module for finetuning
+        or training parts of a model individually (e.g., GAN training).
+
+        See :ref:`locally-disable-grad-doc` for a comparison between
+        `.requires_grad_()` and several similar mechanisms that may be confused with it.
+
+        Args:
+            requires_grad (bool): whether autograd should record operations on
+                                  parameters in this module. Default: ``True``.
+
+        Returns:
+            Module: self
+        """
+        for p in self.parameters():
+            p.requires_grad_(requires_grad)
+        return self
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        r"""Reset gradients of all model parameters.
+
+        See similar function under :class:`torch.optim.Optimizer` for more context.
+
+        Args:
+            set_to_none (bool): instead of setting to zero, set the grads to None.
+                See :meth:`torch.optim.Optimizer.zero_grad` for details.
+        """
+        if getattr(self, "_is_replica", False):
+            warnings.warn(
+                "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
+                "The parameters are copied (in a differentiable manner) from the original module. "
+                "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
+                "If you need gradients in your forward method, consider using autograd.grad instead.",
+                stacklevel=2,
+            )
+
+        for p in self.parameters():
+            if p.grad is not None:
+                if set_to_none:
+                    p.grad = None
+                else:
+                    if p.grad.grad_fn is not None:
+                        p.grad.detach_()
+                    else:
+                        p.grad.requires_grad_(False)
+                    p.grad.zero_()
+
+    def share_memory(self) -> Self:
+        r"""See :meth:`torch.Tensor.share_memory_`."""
+        return self._apply(lambda t: t.share_memory_())
+
+    def _get_name(self):
+        return self.__class__.__name__
+
+    def extra_repr(self) -> str:
+        r"""Return the extra representation of the module.
+
+        To print customized extra information, you should re-implement
+        this method in your own modules. Both single-line and multi-line
+        strings are acceptable.
+        """
+        return ""
+
+    def __repr__(self) -> str:
+        # We treat the extra repr like the sub-module, one item per line
+        extra_lines = []
+        extra_repr = self.extra_repr()
+        # empty string will be split into list ['']
+        if extra_repr:
+            extra_lines = extra_repr.split("\n")
+        child_lines = []
+        for key, module in self._modules.items():
+            mod_str = repr(module)
+            mod_str = _addindent(mod_str, 2)
+            child_lines.append("(" + key + "): " + mod_str)
+        lines = extra_lines + child_lines
+
+        main_str = self._get_name() + "("
+        if lines:
+            # simple one-liner info, which most builtin Modules will use
+            if len(extra_lines) == 1 and not child_lines:
+                main_str += extra_lines[0]
+            else:
+                main_str += "\n  " + "\n  ".join(lines) + "\n"
+
+        main_str += ")"
+        return main_str
+
+    def __dir__(self):
+        module_attrs = dir(self.__class__)
+        attrs = list(self.__dict__.keys())
+        parameters = list(self._parameters.keys())
+        modules = list(self._modules.keys())
+        buffers = list(self._buffers.keys())
+        keys = module_attrs + attrs + parameters + modules + buffers
+
+        # Eliminate attrs that are not legal Python variable names
+        keys = [key for key in keys if not key[0].isdigit()]
+
+        return sorted(keys)
+
+    def _replicate_for_data_parallel(self):
+        replica = self.__new__(type(self))
+        replica.__dict__ = self.__dict__.copy()
+
+        # replicas do not have parameters themselves, the replicas reference the original
+        # module.
+        replica._parameters = {}
+        replica._buffers = replica._buffers.copy()
+        replica._modules = replica._modules.copy()
+        replica._is_replica = True  # type: ignore[assignment]
+
+        return replica
+
+    def compile(self, *args, **kwargs) -> None:
+        """
+        Compile this Module's forward using :func:`torch.compile`.
+
+        This Module's `__call__` method is compiled and all arguments are passed as-is
+        to :func:`torch.compile`.
+
+        See :func:`torch.compile` for details on the arguments for this function.
+        """
+        self._compiled_call_impl = torch.compile(self._call_impl, *args, **kwargs)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/normalization.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/normalization.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4df21528cae2bec0424f2963a624cd462d65077
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/normalization.py
@@ -0,0 +1,441 @@
+# mypy: allow-untyped-defs
+import numbers
+
+import torch
+from torch import Size, Tensor
+from torch.nn import functional as F, init
+from torch.nn.parameter import Parameter
+
+from ._functions import CrossMapLRN2d as _cross_map_lrn2d
+from .module import Module
+
+
+__all__ = ["LocalResponseNorm", "CrossMapLRN2d", "LayerNorm", "GroupNorm", "RMSNorm"]
+
+
+class LocalResponseNorm(Module):
+    r"""Applies local response normalization over an input signal.
+
+    The input signal is composed of several input planes, where channels occupy the second dimension.
+    Applies normalization across channels.
+
+    .. math::
+        b_{c} = a_{c}\left(k + \frac{\alpha}{n}
+        \sum_{c'=\max(0, c-n/2)}^{\min(N-1,c+n/2)}a_{c'}^2\right)^{-\beta}
+
+    Args:
+        size: amount of neighbouring channels used for normalization
+        alpha: multiplicative factor. Default: 0.0001
+        beta: exponent. Default: 0.75
+        k: additive factor. Default: 1
+
+    Shape:
+        - Input: :math:`(N, C, *)`
+        - Output: :math:`(N, C, *)` (same shape as input)
+
+    Examples::
+
+        >>> lrn = nn.LocalResponseNorm(2)
+        >>> signal_2d = torch.randn(32, 5, 24, 24)
+        >>> signal_4d = torch.randn(16, 5, 7, 7, 7, 7)
+        >>> output_2d = lrn(signal_2d)
+        >>> output_4d = lrn(signal_4d)
+
+    """
+
+    __constants__ = ["size", "alpha", "beta", "k"]
+    size: int
+    alpha: float
+    beta: float
+    k: float
+
+    def __init__(
+        self, size: int, alpha: float = 1e-4, beta: float = 0.75, k: float = 1.0
+    ) -> None:
+        super().__init__()
+        self.size = size
+        self.alpha = alpha
+        self.beta = beta
+        self.k = k
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.local_response_norm(input, self.size, self.alpha, self.beta, self.k)
+
+    def extra_repr(self):
+        """
+        Return the extra representation of the module.
+        """
+        return "{size}, alpha={alpha}, beta={beta}, k={k}".format(**self.__dict__)
+
+
+class CrossMapLRN2d(Module):
+    size: int
+    alpha: float
+    beta: float
+    k: float
+
+    def __init__(
+        self, size: int, alpha: float = 1e-4, beta: float = 0.75, k: float = 1
+    ) -> None:
+        super().__init__()
+        self.size = size
+        self.alpha = alpha
+        self.beta = beta
+        self.k = k
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return _cross_map_lrn2d.apply(input, self.size, self.alpha, self.beta, self.k)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return "{size}, alpha={alpha}, beta={beta}, k={k}".format(**self.__dict__)
+
+
+_shape_t = int | list[int] | Size
+
+
+class LayerNorm(Module):
+    r"""Applies Layer Normalization over a mini-batch of inputs.
+
+    This layer implements the operation as described in
+    the paper `Layer Normalization <https://arxiv.org/abs/1607.06450>`__
+
+    .. math::
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated over the last `D` dimensions, where `D`
+    is the dimension of :attr:`normalized_shape`. For example, if :attr:`normalized_shape`
+    is ``(3, 5)`` (a 2-dimensional shape), the mean and standard-deviation are computed over
+    the last 2 dimensions of the input (i.e. ``input.mean((-2, -1))``).
+    :math:`\gamma` and :math:`\beta` are learnable affine transform parameters of
+    :attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``.
+    The variance is calculated via the biased estimator, equivalent to
+    `torch.var(input, correction=0)`.
+
+    .. note::
+        Unlike Batch Normalization and Instance Normalization, which applies
+        scalar scale and bias for each entire channel/plane with the
+        :attr:`affine` option, Layer Normalization applies per-element scale and
+        bias with :attr:`elementwise_affine`.
+
+    This layer uses statistics computed from input data in both training and
+    evaluation modes.
+
+    Args:
+        normalized_shape (int or list or torch.Size): input shape from an expected input
+            of size
+
+            .. math::
+                [* \times \text{normalized\_shape}[0] \times \text{normalized\_shape}[1]
+                    \times \ldots \times \text{normalized\_shape}[-1]]
+
+            If a single integer is used, it is treated as a singleton list, and this module will
+            normalize over the last dimension which is expected to be of that specific size.
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        elementwise_affine: a boolean value that when set to ``True``, this module
+            has learnable per-element affine parameters initialized to ones (for weights)
+            and zeros (for biases). Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`elementwise_affine` is ``True``). Default: ``True``
+
+    Attributes:
+        weight: the learnable weights of the module of shape
+            :math:`\text{normalized\_shape}` when :attr:`elementwise_affine` is set to ``True``.
+            The values are initialized to 1.
+        bias:   the learnable bias of the module of shape
+                :math:`\text{normalized\_shape}` when :attr:`elementwise_affine` is set to ``True``.
+                The values are initialized to 0.
+
+    Shape:
+        - Input: :math:`(N, *)`
+        - Output: :math:`(N, *)` (same shape as input)
+
+    Examples::
+
+        >>> # NLP Example
+        >>> batch, sentence_length, embedding_dim = 20, 5, 10
+        >>> embedding = torch.randn(batch, sentence_length, embedding_dim)
+        >>> layer_norm = nn.LayerNorm(embedding_dim)
+        >>> # Activate module
+        >>> layer_norm(embedding)
+        >>>
+        >>> # Image Example
+        >>> N, C, H, W = 20, 5, 10, 10
+        >>> input = torch.randn(N, C, H, W)
+        >>> # Normalize over the last three dimensions (i.e. the channel and spatial dimensions)
+        >>> # as shown in the image below
+        >>> layer_norm = nn.LayerNorm([C, H, W])
+        >>> output = layer_norm(input)
+
+    .. image:: ../_static/img/nn/layer_norm.jpg
+        :scale: 50 %
+
+    """
+
+    __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
+    normalized_shape: tuple[int, ...]
+    eps: float
+    elementwise_affine: bool
+
+    def __init__(
+        self,
+        normalized_shape: _shape_t,
+        eps: float = 1e-5,
+        elementwise_affine: bool = True,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+            if bias:
+                self.bias = Parameter(
+                    torch.empty(self.normalized_shape, **factory_kwargs)
+                )
+            else:
+                self.register_parameter("bias", None)
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if self.elementwise_affine:
+            init.ones_(self.weight)
+            if self.bias is not None:
+                init.zeros_(self.bias)
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+
+    def extra_repr(self) -> str:
+        return (
+            "{normalized_shape}, eps={eps}, elementwise_affine={elementwise_affine}, "
+            "bias={use_bias}".format(**self.__dict__, use_bias=self.bias is not None)
+        )
+
+
+class GroupNorm(Module):
+    r"""Applies Group Normalization over a mini-batch of inputs.
+
+    This layer implements the operation as described in
+    the paper `Group Normalization <https://arxiv.org/abs/1803.08494>`__
+
+    .. math::
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The input channels are separated into :attr:`num_groups` groups, each containing
+    ``num_channels / num_groups`` channels. :attr:`num_channels` must be divisible by
+    :attr:`num_groups`. The mean and standard-deviation are calculated
+    separately over each group. :math:`\gamma` and :math:`\beta` are learnable
+    per-channel affine transform parameter vectors of size :attr:`num_channels` if
+    :attr:`affine` is ``True``.
+    The variance is calculated via the biased estimator, equivalent to
+    `torch.var(input, correction=0)`.
+
+    This layer uses statistics computed from input data in both training and
+    evaluation modes.
+
+    Args:
+        num_groups (int): number of groups to separate the channels into
+        num_channels (int): number of channels expected in input
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        affine: a boolean value that when set to ``True``, this module
+            has learnable per-channel affine parameters initialized to ones (for weights)
+            and zeros (for biases). Default: ``True``
+        bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
+            :attr:`affine` is ``True``). Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, *)` where :math:`C=\text{num\_channels}`
+        - Output: :math:`(N, C, *)` (same shape as input)
+
+    Examples::
+
+        >>> input = torch.randn(20, 6, 10, 10)
+        >>> # Separate 6 channels into 3 groups
+        >>> m = nn.GroupNorm(3, 6)
+        >>> # Separate 6 channels into 6 groups (equivalent with InstanceNorm)
+        >>> m = nn.GroupNorm(6, 6)
+        >>> # Put all 6 channels into a single group (equivalent with LayerNorm)
+        >>> m = nn.GroupNorm(1, 6)
+        >>> # Activating the module
+        >>> output = m(input)
+    """
+
+    __constants__ = ["num_groups", "num_channels", "eps", "affine"]
+    num_groups: int
+    num_channels: int
+    eps: float
+    affine: bool
+
+    def __init__(
+        self,
+        num_groups: int,
+        num_channels: int,
+        eps: float = 1e-5,
+        affine: bool = True,
+        device=None,
+        dtype=None,
+        *,
+        bias: bool = True,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        if num_channels % num_groups != 0:
+            raise ValueError(
+                f"num_channels ({num_channels}) must be divisible by num_groups ({num_groups})"
+            )
+
+        self.num_groups = num_groups
+        self.num_channels = num_channels
+        self.eps = eps
+        self.affine = affine
+        if self.affine:
+            self.weight = Parameter(torch.empty(num_channels, **factory_kwargs))
+            if bias:
+                self.bias = Parameter(torch.empty(num_channels, **factory_kwargs))
+            else:
+                self.register_parameter("bias", None)
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if self.affine:
+            init.ones_(self.weight)
+            if self.bias is not None:
+                init.zeros_(self.bias)
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
+
+    def extra_repr(self) -> str:
+        return (
+            "{num_groups}, {num_channels}, eps={eps}, affine={affine}, "
+            "bias={use_bias}".format(**self.__dict__, use_bias=self.bias is not None)
+        )
+
+
+class RMSNorm(Module):
+    r"""Applies Root Mean Square Layer Normalization over a mini-batch of inputs.
+
+    This layer implements the operation as described in
+    the paper `Root Mean Square Layer Normalization <https://arxiv.org/pdf/1910.07467.pdf>`__
+
+    .. math::
+        y_i = \frac{x_i}{\mathrm{RMS}(x)} * \gamma_i, \quad
+        \text{where} \quad \text{RMS}(x) = \sqrt{\epsilon + \frac{1}{n} \sum_{i=1}^{n} x_i^2}
+
+    The RMS is taken over the last ``D`` dimensions, where ``D``
+    is the dimension of :attr:`normalized_shape`. For example, if :attr:`normalized_shape`
+    is ``(3, 5)`` (a 2-dimensional shape), the RMS is computed over
+    the last 2 dimensions of the input.
+
+    Args:
+        normalized_shape (int or list or torch.Size): input shape from an expected input
+            of size
+
+            .. math::
+                [* \times \text{normalized\_shape}[0] \times \text{normalized\_shape}[1]
+                    \times \ldots \times \text{normalized\_shape}[-1]]
+
+            If a single integer is used, it is treated as a singleton list, and this module will
+            normalize over the last dimension which is expected to be of that specific size.
+        eps (float, optional): a value added to the denominator for numerical stability.
+            If not specified, uses the machine epsilon of the computation (opmath) type:
+            fp16/bf16 and fp32 inputs use ``torch.finfo(torch.float32).eps``, while fp64
+            inputs use ``torch.finfo(torch.float64).eps``. Default: ``None``
+        elementwise_affine: a boolean value that when set to ``True``, this module
+            has learnable per-element affine parameters initialized to ones (for weights). Default: ``True``.
+
+    Shape:
+        - Input: :math:`(N, *)`
+        - Output: :math:`(N, *)` (same shape as input)
+
+    Examples::
+
+        >>> rms_norm = nn.RMSNorm([2, 3])
+        >>> input = torch.randn(2, 2, 3)
+        >>> rms_norm(input)
+
+    """
+
+    __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
+    normalized_shape: tuple[int, ...]
+    eps: float | None
+    elementwise_affine: bool
+
+    def __init__(
+        self,
+        normalized_shape: _shape_t,
+        eps: float | None = None,
+        elementwise_affine: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+        else:
+            self.register_parameter("weight", None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        if self.elementwise_affine:
+            init.ones_(self.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.rms_norm(x, self.normalized_shape, self.weight, self.eps)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return (
+            "{normalized_shape}, eps={eps}, "
+            "elementwise_affine={elementwise_affine}".format(**self.__dict__)
+        )
+
+
+# TODO: ContrastiveNorm2d
+# TODO: DivisiveNorm2d
+# TODO: SubtractiveNorm2d
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/padding.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/padding.py
new file mode 100644
index 0000000000000000000000000000000000000000..d5aa1e0d425548857d20b093041b190bc7f2f645
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/padding.py
@@ -0,0 +1,842 @@
+# mypy: allow-untyped-defs
+from collections.abc import Sequence
+
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.common_types import _size_2_t, _size_4_t, _size_6_t
+
+from .module import Module
+from .utils import _ntuple, _pair, _quadruple
+
+
+# TODO: grad_output size asserts in THNN
+
+__all__ = [
+    "CircularPad1d",
+    "CircularPad2d",
+    "CircularPad3d",
+    "ConstantPad1d",
+    "ConstantPad2d",
+    "ConstantPad3d",
+    "ReflectionPad1d",
+    "ReflectionPad2d",
+    "ReflectionPad3d",
+    "ReplicationPad1d",
+    "ReplicationPad2d",
+    "ReplicationPad3d",
+    "ZeroPad1d",
+    "ZeroPad2d",
+    "ZeroPad3d",
+]
+
+
+class _CircularPadNd(Module):
+    __constants__ = ["padding"]
+    padding: Sequence[int]
+
+    def _check_input_dim(self, input):
+        raise NotImplementedError
+
+    def forward(self, input: Tensor) -> Tensor:
+        self._check_input_dim(input)
+        return F.pad(input, self.padding, "circular")
+
+    def extra_repr(self) -> str:
+        return f"{self.padding}"
+
+
+class CircularPad1d(_CircularPadNd):
+    r"""Pads the input tensor using circular padding of the input boundary.
+
+    Tensor values at the beginning of the dimension are used to pad the end,
+    and values at the end are used to pad the beginning. If negative padding is
+    applied then the ends of the tensor get removed.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 2-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
+            Note that padding size should be less than or equal to the corresponding input dimension.
+
+    Shape:
+        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
+        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
+        >>> m = nn.CircularPad1d(2)
+        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
+        >>> input
+        tensor([[[0., 1., 2., 3.],
+                 [4., 5., 6., 7.]]])
+        >>> m(input)
+        tensor([[[2., 3., 0., 1., 2., 3., 0., 1.],
+                 [6., 7., 4., 5., 6., 7., 4., 5.]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.CircularPad1d((3, 1))
+        >>> m(input)
+        tensor([[[1., 2., 3., 0., 1., 2., 3., 0.],
+                 [5., 6., 7., 4., 5., 6., 7., 4.]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int]
+
+    def __init__(self, padding: _size_2_t) -> None:
+        super().__init__()
+        self.padding = _pair(padding)
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
+
+
+class CircularPad2d(_CircularPadNd):
+    r"""Pads the input tensor using circular padding of the input boundary.
+
+    Tensor values at the beginning of the dimension are used to pad the end,
+    and values at the end are used to pad the beginning. If negative padding is
+    applied then the ends of the tensor get removed.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
+            :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
+            Note that padding size should be less than or equal to the corresponding input dimension.
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> m = nn.CircularPad2d(2)
+        >>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
+        >>> input
+        tensor([[[[0., 1., 2.],
+                  [3., 4., 5.],
+                  [6., 7., 8.]]]])
+        >>> m(input)
+        tensor([[[[4., 5., 3., 4., 5., 3., 4.],
+                  [7., 8., 6., 7., 8., 6., 7.],
+                  [1., 2., 0., 1., 2., 0., 1.],
+                  [4., 5., 3., 4., 5., 3., 4.],
+                  [7., 8., 6., 7., 8., 6., 7.],
+                  [1., 2., 0., 1., 2., 0., 1.],
+                  [4., 5., 3., 4., 5., 3., 4.]]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.CircularPad2d((1, 1, 2, 0))
+        >>> m(input)
+        tensor([[[[5., 3., 4., 5., 3.],
+                  [8., 6., 7., 8., 6.],
+                  [2., 0., 1., 2., 0.],
+                  [5., 3., 4., 5., 3.],
+                  [8., 6., 7., 8., 6.]]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int]
+
+    def __init__(self, padding: _size_4_t) -> None:
+        super().__init__()
+        self.padding = _quadruple(padding)
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 3 and input.dim() != 4:
+            raise ValueError(f"expected 3D or 4D input (got {input.dim()}D input)")
+
+
+class CircularPad3d(_CircularPadNd):
+    r"""Pads the input tensor using circular padding of the input boundary.
+
+    Tensor values at the beginning of the dimension are used to pad the end,
+    and values at the end are used to pad the beginning. If negative padding is
+    applied then the ends of the tensor get removed.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 6-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
+            :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
+            :math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
+            Note that padding size should be less than or equal to the corresponding input dimension.
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,
+          where
+
+          :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> m = nn.CircularPad3d(3)
+        >>> input = torch.randn(16, 3, 8, 320, 480)
+        >>> output = m(input)
+        >>> # using different paddings for different sides
+        >>> m = nn.CircularPad3d((3, 3, 6, 6, 1, 1))
+        >>> output = m(input)
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int, int, int]
+
+    def __init__(self, padding: _size_6_t) -> None:
+        super().__init__()
+        self.padding = _ntuple(6)(padding)
+
+    def _check_input_dim(self, input) -> None:
+        if input.dim() != 4 and input.dim() != 5:
+            raise ValueError(f"expected 4D or 5D input (got {input.dim()}D input)")
+
+
+class _ConstantPadNd(Module):
+    __constants__ = ["padding", "value"]
+    value: float
+    padding: Sequence[int]
+
+    def __init__(self, value: float) -> None:
+        super().__init__()
+        self.value = value
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.pad(input, self.padding, "constant", self.value)
+
+    def extra_repr(self) -> str:
+        return f"padding={self.padding}, value={self.value}"
+
+
+class ConstantPad1d(_ConstantPadNd):
+    r"""Pads the input tensor boundaries with a constant value.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in both boundaries. If a 2-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
+
+    Shape:
+        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
+        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> m = nn.ConstantPad1d(2, 3.5)
+        >>> input = torch.randn(1, 2, 4)
+        >>> input
+        tensor([[[-1.0491, -0.7152, -0.0749,  0.8530],
+                 [-1.3287,  1.8966,  0.1466, -0.2771]]])
+        >>> m(input)
+        tensor([[[ 3.5000,  3.5000, -1.0491, -0.7152, -0.0749,  0.8530,  3.5000,
+                   3.5000],
+                 [ 3.5000,  3.5000, -1.3287,  1.8966,  0.1466, -0.2771,  3.5000,
+                   3.5000]]])
+        >>> m = nn.ConstantPad1d(2, 3.5)
+        >>> input = torch.randn(1, 2, 3)
+        >>> input
+        tensor([[[ 1.6616,  1.4523, -1.1255],
+                 [-3.6372,  0.1182, -1.8652]]])
+        >>> m(input)
+        tensor([[[ 3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000,  3.5000]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ConstantPad1d((3, 1), 3.5)
+        >>> m(input)
+        tensor([[[ 3.5000,  3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000],
+                 [ 3.5000,  3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int]
+
+    def __init__(self, padding: _size_2_t, value: float) -> None:
+        super().__init__(value)
+        self.padding = _pair(padding)
+
+
+class ConstantPad2d(_ConstantPadNd):
+    r"""Pads the input tensor boundaries with a constant value.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
+            :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> m = nn.ConstantPad2d(2, 3.5)
+        >>> input = torch.randn(1, 2, 2)
+        >>> input
+        tensor([[[ 1.6585,  0.4320],
+                 [-0.8701, -0.4649]]])
+        >>> m(input)
+        tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000,  1.6585,  0.4320,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000, -0.8701, -0.4649,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ConstantPad2d((3, 0, 2, 1), 3.5)
+        >>> m(input)
+        tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
+                 [ 3.5000,  3.5000,  3.5000,  1.6585,  0.4320],
+                 [ 3.5000,  3.5000,  3.5000, -0.8701, -0.4649],
+                 [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])
+    """
+
+    __constants__ = ["padding", "value"]
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int]
+
+    def __init__(self, padding: _size_4_t, value: float) -> None:
+        super().__init__(value)
+        self.padding = _quadruple(padding)
+
+
+class ConstantPad3d(_ConstantPadNd):
+    r"""Pads the input tensor boundaries with a constant value.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 6-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
+            :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
+            :math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or
+          :math:`(C, D_{out}, H_{out}, W_{out})`, where
+
+          :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> m = nn.ConstantPad3d(3, 3.5)
+        >>> input = torch.randn(16, 3, 10, 20, 30)
+        >>> output = m(input)
+        >>> # using different paddings for different sides
+        >>> m = nn.ConstantPad3d((3, 3, 6, 6, 0, 1), 3.5)
+        >>> output = m(input)
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int, int, int]
+
+    def __init__(self, padding: _size_6_t, value: float) -> None:
+        super().__init__(value)
+        self.padding = _ntuple(6)(padding)
+
+
+class _ReflectionPadNd(Module):
+    __constants__ = ["padding"]
+    padding: Sequence[int]
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.pad(input, self.padding, "reflect")
+
+    def extra_repr(self) -> str:
+        return f"{self.padding}"
+
+
+class ReflectionPad1d(_ReflectionPadNd):
+    r"""Pads the input tensor using the reflection of the input boundary.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 2-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
+            Note that padding size should be less than the corresponding input dimension.
+
+    Shape:
+        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
+        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> m = nn.ReflectionPad1d(2)
+        >>> # xdoctest: +IGNORE_WANT("other tests seem to modify printing styles")
+        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
+        >>> input
+        tensor([[[0., 1., 2., 3.],
+                 [4., 5., 6., 7.]]])
+        >>> m(input)
+        tensor([[[2., 1., 0., 1., 2., 3., 2., 1.],
+                 [6., 5., 4., 5., 6., 7., 6., 5.]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ReflectionPad1d((3, 1))
+        >>> m(input)
+        tensor([[[3., 2., 1., 0., 1., 2., 3., 2.],
+                 [7., 6., 5., 4., 5., 6., 7., 6.]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int]
+
+    def __init__(self, padding: _size_2_t) -> None:
+        super().__init__()
+        self.padding = _pair(padding)
+
+
+class ReflectionPad2d(_ReflectionPadNd):
+    r"""Pads the input tensor using the reflection of the input boundary.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
+            :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
+            Note that padding size should be less than the corresponding input dimension.
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})` where
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
+        >>> m = nn.ReflectionPad2d(2)
+        >>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
+        >>> input
+        tensor([[[[0., 1., 2.],
+                  [3., 4., 5.],
+                  [6., 7., 8.]]]])
+        >>> m(input)
+        tensor([[[[8., 7., 6., 7., 8., 7., 6.],
+                  [5., 4., 3., 4., 5., 4., 3.],
+                  [2., 1., 0., 1., 2., 1., 0.],
+                  [5., 4., 3., 4., 5., 4., 3.],
+                  [8., 7., 6., 7., 8., 7., 6.],
+                  [5., 4., 3., 4., 5., 4., 3.],
+                  [2., 1., 0., 1., 2., 1., 0.]]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ReflectionPad2d((1, 1, 2, 0))
+        >>> m(input)
+        tensor([[[[7., 6., 7., 8., 7.],
+                  [4., 3., 4., 5., 4.],
+                  [1., 0., 1., 2., 1.],
+                  [4., 3., 4., 5., 4.],
+                  [7., 6., 7., 8., 7.]]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int]
+
+    def __init__(self, padding: _size_4_t) -> None:
+        super().__init__()
+        self.padding = _quadruple(padding)
+
+
+class ReflectionPad3d(_ReflectionPadNd):
+    r"""Pads the input tensor using the reflection of the input boundary.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 6-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
+            :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
+            :math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
+            Note that padding size should be less than the corresponding input dimension.
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,
+          where
+
+          :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
+        >>> m = nn.ReflectionPad3d(1)
+        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 1, 2, 2, 2)
+        >>> m(input)
+        tensor([[[[[7., 6., 7., 6.],
+                   [5., 4., 5., 4.],
+                   [7., 6., 7., 6.],
+                   [5., 4., 5., 4.]],
+                  [[3., 2., 3., 2.],
+                   [1., 0., 1., 0.],
+                   [3., 2., 3., 2.],
+                   [1., 0., 1., 0.]],
+                  [[7., 6., 7., 6.],
+                   [5., 4., 5., 4.],
+                   [7., 6., 7., 6.],
+                   [5., 4., 5., 4.]],
+                  [[3., 2., 3., 2.],
+                   [1., 0., 1., 0.],
+                   [3., 2., 3., 2.],
+                   [1., 0., 1., 0.]]]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int, int, int]
+
+    def __init__(self, padding: _size_6_t) -> None:
+        super().__init__()
+        self.padding = _ntuple(6)(padding)
+
+
+class _ReplicationPadNd(Module):
+    __constants__ = ["padding"]
+    padding: Sequence[int]
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.pad(input, self.padding, "replicate")
+
+    def extra_repr(self) -> str:
+        return f"{self.padding}"
+
+
+class ReplicationPad1d(_ReplicationPadNd):
+    r"""Pads the input tensor using replication of the input boundary.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 2-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
+            Note that the output dimensions must remain positive.
+
+    Shape:
+        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
+        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
+        >>> m = nn.ReplicationPad1d(2)
+        >>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
+        >>> input
+        tensor([[[0., 1., 2., 3.],
+                 [4., 5., 6., 7.]]])
+        >>> m(input)
+        tensor([[[0., 0., 0., 1., 2., 3., 3., 3.],
+                 [4., 4., 4., 5., 6., 7., 7., 7.]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ReplicationPad1d((3, 1))
+        >>> m(input)
+        tensor([[[0., 0., 0., 0., 1., 2., 3., 3.],
+                 [4., 4., 4., 4., 5., 6., 7., 7.]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int]
+
+    def __init__(self, padding: _size_2_t) -> None:
+        super().__init__()
+        self.padding = _pair(padding)
+
+
+class ReplicationPad2d(_ReplicationPadNd):
+    r"""Pads the input tensor using replication of the input boundary.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
+            :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
+            Note that the output dimensions must remain positive.
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> m = nn.ReplicationPad2d(2)
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
+        >>> input
+        tensor([[[[0., 1., 2.],
+                  [3., 4., 5.],
+                  [6., 7., 8.]]]])
+        >>> m(input)
+        tensor([[[[0., 0., 0., 1., 2., 2., 2.],
+                  [0., 0., 0., 1., 2., 2., 2.],
+                  [0., 0., 0., 1., 2., 2., 2.],
+                  [3., 3., 3., 4., 5., 5., 5.],
+                  [6., 6., 6., 7., 8., 8., 8.],
+                  [6., 6., 6., 7., 8., 8., 8.],
+                  [6., 6., 6., 7., 8., 8., 8.]]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ReplicationPad2d((1, 1, 2, 0))
+        >>> m(input)
+        tensor([[[[0., 0., 1., 2., 2.],
+                  [0., 0., 1., 2., 2.],
+                  [0., 0., 1., 2., 2.],
+                  [3., 3., 4., 5., 5.],
+                  [6., 6., 7., 8., 8.]]]])
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int]
+
+    def __init__(self, padding: _size_4_t) -> None:
+        super().__init__()
+        self.padding = _quadruple(padding)
+
+
+class ReplicationPad3d(_ReplicationPadNd):
+    r"""Pads the input tensor using replication of the input boundary.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 6-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
+            :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
+            :math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
+            Note that the output dimensions must remain positive.
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,
+          where
+
+          :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> m = nn.ReplicationPad3d(3)
+        >>> input = torch.randn(16, 3, 8, 320, 480)
+        >>> output = m(input)
+        >>> # using different paddings for different sides
+        >>> m = nn.ReplicationPad3d((3, 3, 6, 6, 1, 1))
+        >>> output = m(input)
+    """
+
+    # pyrefly: ignore [bad-override]
+    padding: tuple[int, int, int, int, int, int]
+
+    def __init__(self, padding: _size_6_t) -> None:
+        super().__init__()
+        self.padding = _ntuple(6)(padding)
+
+
+class ZeroPad1d(ConstantPad1d):
+    r"""Pads the input tensor boundaries with zero.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in both boundaries. If a 2-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
+
+    Shape:
+        - Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
+        - Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> m = nn.ZeroPad1d(2)
+        >>> input = torch.randn(1, 2, 4)
+        >>> input
+        tensor([[[-1.0491, -0.7152, -0.0749,  0.8530],
+                 [-1.3287,  1.8966,  0.1466, -0.2771]]])
+        >>> m(input)
+        tensor([[[ 0.0000,  0.0000, -1.0491, -0.7152, -0.0749,  0.8530,  0.0000,
+                   0.0000],
+                 [ 0.0000,  0.0000, -1.3287,  1.8966,  0.1466, -0.2771,  0.0000,
+                   0.0000]]])
+        >>> m = nn.ZeroPad1d(2)
+        >>> input = torch.randn(1, 2, 3)
+        >>> input
+        tensor([[[ 1.6616,  1.4523, -1.1255],
+                 [-3.6372,  0.1182, -1.8652]]])
+        >>> m(input)
+        tensor([[[ 0.0000,  0.0000,  1.6616,  1.4523, -1.1255,  0.0000,  0.0000],
+                 [ 0.0000,  0.0000, -3.6372,  0.1182, -1.8652,  0.0000,  0.0000]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ZeroPad1d((3, 1))
+        >>> m(input)
+        tensor([[[ 0.0000,  0.0000,  0.0000,  1.6616,  1.4523, -1.1255,  0.0000],
+                 [ 0.0000,  0.0000,  0.0000, -3.6372,  0.1182, -1.8652,  0.0000]]])
+    """
+
+    padding: tuple[int, int]
+
+    def __init__(self, padding: _size_2_t) -> None:
+        super().__init__(padding, 0.0)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"{self.padding}"
+
+
+class ZeroPad2d(ConstantPad2d):
+    r"""Pads the input tensor boundaries with zero.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
+            :math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> m = nn.ZeroPad2d(2)
+        >>> input = torch.randn(1, 1, 3, 3)
+        >>> input
+        tensor([[[[-0.1678, -0.4418,  1.9466],
+                  [ 0.9604, -0.4219, -0.5241],
+                  [-0.9162, -0.5436, -0.6446]]]])
+        >>> m(input)
+        tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
+                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
+                  [ 0.0000,  0.0000, -0.1678, -0.4418,  1.9466,  0.0000,  0.0000],
+                  [ 0.0000,  0.0000,  0.9604, -0.4219, -0.5241,  0.0000,  0.0000],
+                  [ 0.0000,  0.0000, -0.9162, -0.5436, -0.6446,  0.0000,  0.0000],
+                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
+                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]])
+        >>> # using different paddings for different sides
+        >>> m = nn.ZeroPad2d((1, 1, 2, 0))
+        >>> m(input)
+        tensor([[[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
+                  [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
+                  [ 0.0000, -0.1678, -0.4418,  1.9466,  0.0000],
+                  [ 0.0000,  0.9604, -0.4219, -0.5241,  0.0000],
+                  [ 0.0000, -0.9162, -0.5436, -0.6446,  0.0000]]]])
+    """
+
+    padding: tuple[int, int, int, int]
+
+    def __init__(self, padding: _size_4_t) -> None:
+        super().__init__(padding, 0.0)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"{self.padding}"
+
+
+class ZeroPad3d(ConstantPad3d):
+    r"""Pads the input tensor boundaries with zero.
+
+    For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
+
+    Args:
+        padding (int, tuple): the size of the padding. If is `int`, uses the same
+            padding in all boundaries. If a 6-`tuple`, uses
+            (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
+            :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
+            :math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or
+          :math:`(C, D_{out}, H_{out}, W_{out})`, where
+
+          :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
+
+          :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
+
+          :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
+
+    Examples::
+
+        >>> m = nn.ZeroPad3d(3)
+        >>> input = torch.randn(16, 3, 10, 20, 30)
+        >>> output = m(input)
+        >>> # using different paddings for different sides
+        >>> m = nn.ZeroPad3d((3, 3, 6, 6, 0, 1))
+        >>> output = m(input)
+    """
+
+    padding: tuple[int, int, int, int, int, int]
+
+    def __init__(self, padding: _size_6_t) -> None:
+        super().__init__(padding, 0.0)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"{self.padding}"
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/pixelshuffle.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/pixelshuffle.py
new file mode 100644
index 0000000000000000000000000000000000000000..74c9e0878f0b5ecc48878c63115aafc2128b3afd
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/pixelshuffle.py
@@ -0,0 +1,127 @@
+import torch.nn.functional as F
+from torch import Tensor
+
+from .module import Module
+
+
+__all__ = ["PixelShuffle", "PixelUnshuffle"]
+
+
+class PixelShuffle(Module):
+    r"""Rearrange elements in a tensor according to an upscaling factor.
+
+    Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)`
+    to a tensor of shape :math:`(*, C, H \times r, W \times r)`, where r is an upscale factor.
+
+    This is useful for implementing efficient sub-pixel convolution
+    with a stride of :math:`1/r`.
+
+    See the paper:
+    `Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_
+    by Shi et al. (2016) for more details.
+
+    Args:
+        upscale_factor (int): factor to increase spatial resolution by
+
+    Shape:
+        - Input: :math:`(*, C_{in}, H_{in}, W_{in})`, where * is zero or more batch dimensions
+        - Output: :math:`(*, C_{out}, H_{out}, W_{out})`, where
+
+    .. math::
+        C_{out} = C_{in} \div \text{upscale\_factor}^2
+
+    .. math::
+        H_{out} = H_{in} \times \text{upscale\_factor}
+
+    .. math::
+        W_{out} = W_{in} \times \text{upscale\_factor}
+
+    Examples::
+
+        >>> pixel_shuffle = nn.PixelShuffle(3)
+        >>> input = torch.randn(1, 9, 4, 4)
+        >>> output = pixel_shuffle(input)
+        >>> print(output.size())
+        torch.Size([1, 1, 12, 12])
+
+    .. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network:
+        https://arxiv.org/abs/1609.05158
+    """
+
+    __constants__ = ["upscale_factor"]
+    upscale_factor: int
+
+    def __init__(self, upscale_factor: int) -> None:
+        super().__init__()
+        self.upscale_factor = upscale_factor
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.pixel_shuffle(input, self.upscale_factor)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"upscale_factor={self.upscale_factor}"
+
+
+class PixelUnshuffle(Module):
+    r"""Reverse the PixelShuffle operation.
+
+    Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements
+    in a tensor of shape :math:`(*, C, H \times r, W \times r)` to a tensor of shape
+    :math:`(*, C \times r^2, H, W)`, where r is a downscale factor.
+
+    See the paper:
+    `Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_
+    by Shi et al. (2016) for more details.
+
+    Args:
+        downscale_factor (int): factor to decrease spatial resolution by
+
+    Shape:
+        - Input: :math:`(*, C_{in}, H_{in}, W_{in})`, where * is zero or more batch dimensions
+        - Output: :math:`(*, C_{out}, H_{out}, W_{out})`, where
+
+    .. math::
+        C_{out} = C_{in} \times \text{downscale\_factor}^2
+
+    .. math::
+        H_{out} = H_{in} \div \text{downscale\_factor}
+
+    .. math::
+        W_{out} = W_{in} \div \text{downscale\_factor}
+
+    Examples::
+
+        >>> pixel_unshuffle = nn.PixelUnshuffle(3)
+        >>> input = torch.randn(1, 1, 12, 12)
+        >>> output = pixel_unshuffle(input)
+        >>> print(output.size())
+        torch.Size([1, 9, 4, 4])
+
+    .. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network:
+        https://arxiv.org/abs/1609.05158
+    """
+
+    __constants__ = ["downscale_factor"]
+    downscale_factor: int
+
+    def __init__(self, downscale_factor: int) -> None:
+        super().__init__()
+        self.downscale_factor = downscale_factor
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.pixel_unshuffle(input, self.downscale_factor)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        return f"downscale_factor={self.downscale_factor}"
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/pooling.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/pooling.py
new file mode 100644
index 0000000000000000000000000000000000000000..1dc57c25b168396fa9ceff5b32fd368befa094af
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/pooling.py
@@ -0,0 +1,1550 @@
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.common_types import (
+    _ratio_2_t,
+    _ratio_3_t,
+    _size_1_t,
+    _size_2_opt_t,
+    _size_2_t,
+    _size_3_opt_t,
+    _size_3_t,
+    _size_any_opt_t,
+    _size_any_t,
+)
+
+from .module import Module
+from .utils import _pair, _single, _triple
+
+
+__all__ = [
+    "MaxPool1d",
+    "MaxPool2d",
+    "MaxPool3d",
+    "MaxUnpool1d",
+    "MaxUnpool2d",
+    "MaxUnpool3d",
+    "AvgPool1d",
+    "AvgPool2d",
+    "AvgPool3d",
+    "FractionalMaxPool2d",
+    "FractionalMaxPool3d",
+    "LPPool1d",
+    "LPPool2d",
+    "LPPool3d",
+    "AdaptiveMaxPool1d",
+    "AdaptiveMaxPool2d",
+    "AdaptiveMaxPool3d",
+    "AdaptiveAvgPool1d",
+    "AdaptiveAvgPool2d",
+    "AdaptiveAvgPool3d",
+]
+
+
+class _MaxPoolNd(Module):
+    __constants__ = [
+        "kernel_size",
+        "stride",
+        "padding",
+        "dilation",
+        "return_indices",
+        "ceil_mode",
+    ]
+    return_indices: bool
+    ceil_mode: bool
+
+    def __init__(
+        self,
+        kernel_size: _size_any_t,
+        stride: _size_any_t | None = None,
+        padding: _size_any_t = 0,
+        dilation: _size_any_t = 1,
+        return_indices: bool = False,
+        ceil_mode: bool = False,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride if (stride is not None) else kernel_size
+        self.padding = padding
+        self.dilation = dilation
+        self.return_indices = return_indices
+        self.ceil_mode = ceil_mode
+
+    def extra_repr(self) -> str:
+        return (
+            "kernel_size={kernel_size}, stride={stride}, padding={padding}"
+            ", dilation={dilation}, ceil_mode={ceil_mode}".format(**self.__dict__)
+        )
+
+
+class MaxPool1d(_MaxPoolNd):
+    r"""Applies a 1D max pooling over an input signal composed of several input planes.
+
+    In the simplest case, the output value of the layer with input size :math:`(N, C, L)`
+    and output :math:`(N, C, L_{out})` can be precisely described as:
+
+    .. math::
+        out(N_i, C_j, k) = \max_{m=0, \ldots, \text{kernel\_size} - 1}
+                input(N_i, C_j, stride \times k + m)
+
+    If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides
+    for :attr:`padding` number of points. :attr:`dilation` is the stride between the elements within the
+    sliding window. This `link`_ has a nice visualization of the pooling parameters.
+
+    Note:
+        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding
+        or the input. Sliding windows that would start in the right padded region are ignored.
+
+    Args:
+        kernel_size: The size of the sliding window, must be > 0.
+        stride: The stride of the sliding window, must be > 0. Default value is :attr:`kernel_size`.
+        padding: Implicit negative infinity padding to be added on both sides, must be >= 0 and <= kernel_size / 2.
+        dilation: The stride between elements within a sliding window, must be > 0.
+        return_indices: If ``True``, will return the argmax along with the max values.
+                        Useful for :class:`torch.nn.MaxUnpool1d` later
+        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This
+                   ensures that every element in the input tensor is covered by a sliding window.
+
+    Shape:
+        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.
+        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`,
+
+          where ``ceil_mode = False``
+
+          .. math::
+              L_{out} = \left\lfloor \frac{L_{in} + 2 \times \text{padding} - \text{dilation}
+                   \times (\text{kernel\_size} - 1) - 1}{\text{stride}}\right\rfloor + 1
+
+          where ``ceil_mode = True``
+
+          .. math::
+              L_{out} = \left\lceil \frac{L_{in} + 2 \times \text{padding} - \text{dilation}
+                    \times (\text{kernel\_size} - 1) - 1 + (stride - 1)}{\text{stride}}\right\rceil + 1
+
+        - Ensure that the last pooling starts inside the image, make :math:`L_{out} = L_{out} - 1`
+          when :math:`(L_{out} - 1) * \text{stride} >= L_{in} + \text{padding}`.
+
+    Examples::
+
+        >>> # pool of size=3, stride=2
+        >>> m = nn.MaxPool1d(3, stride=2)
+        >>> input = torch.randn(20, 16, 50)
+        >>> output = m(input)
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+    """
+
+    kernel_size: _size_1_t
+    stride: _size_1_t
+    padding: _size_1_t
+    dilation: _size_1_t
+
+    def forward(self, input: Tensor):
+        """Runs the forward pass."""
+        return F.max_pool1d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.dilation,
+            ceil_mode=self.ceil_mode,
+            return_indices=self.return_indices,
+        )
+
+
+class MaxPool2d(_MaxPoolNd):
+    r"""Applies a 2D max pooling over an input signal composed of several input planes.
+
+    In the simplest case, the output value of the layer with input size :math:`(N, C, H, W)`,
+    output :math:`(N, C, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kH, kW)`
+    can be precisely described as:
+
+    .. math::
+        \begin{aligned}
+            out(N_i, C_j, h, w) ={} & \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \\
+                                    & \text{input}(N_i, C_j, \text{stride[0]} \times h + m,
+                                                   \text{stride[1]} \times w + n)
+        \end{aligned}
+
+    If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides
+    for :attr:`padding` number of points. :attr:`dilation` controls the spacing between the kernel points.
+    It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
+
+    Note:
+        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding
+        or the input. Sliding windows that would start in the right padded region are ignored.
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:
+
+        - a single ``int`` -- in which case the same value is used for the height and width dimension
+        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,
+          and the second `int` for the width dimension
+
+    Args:
+        kernel_size: the size of the window to take a max over
+        stride: the stride of the window. Default value is :attr:`kernel_size`
+        padding: Implicit negative infinity padding to be added on both sides
+        dilation: a parameter that controls the stride of elements in the window
+        return_indices: if ``True``, will return the max indices along with the outputs.
+                        Useful for :class:`torch.nn.MaxUnpool2d` later
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in} + 2 * \text{padding[0]} - \text{dilation[0]}
+                    \times (\text{kernel\_size[0]} - 1) - 1}{\text{stride[0]}} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in} + 2 * \text{padding[1]} - \text{dilation[1]}
+                    \times (\text{kernel\_size[1]} - 1) - 1}{\text{stride[1]}} + 1\right\rfloor
+
+    Examples::
+
+        >>> # pool of square window of size=3, stride=2
+        >>> m = nn.MaxPool2d(3, stride=2)
+        >>> # pool of non-square window
+        >>> m = nn.MaxPool2d((3, 2), stride=(2, 1))
+        >>> input = torch.randn(20, 16, 50, 32)
+        >>> output = m(input)
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+    """
+
+    kernel_size: _size_2_t
+    stride: _size_2_t
+    padding: _size_2_t
+    dilation: _size_2_t
+
+    def forward(self, input: Tensor):
+        """Runs the forward pass."""
+        return F.max_pool2d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.dilation,
+            ceil_mode=self.ceil_mode,
+            return_indices=self.return_indices,
+        )
+
+
+class MaxPool3d(_MaxPoolNd):
+    r"""Applies a 3D max pooling over an input signal composed of several input planes.
+
+    In the simplest case, the output value of the layer with input size :math:`(N, C, D, H, W)`,
+    output :math:`(N, C, D_{out}, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kD, kH, kW)`
+    can be precisely described as:
+
+    .. math::
+        \begin{aligned}
+            \text{out}(N_i, C_j, d, h, w) ={} & \max_{k=0, \ldots, kD-1} \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \\
+                                              & \text{input}(N_i, C_j, \text{stride[0]} \times d + k,
+                                                             \text{stride[1]} \times h + m, \text{stride[2]} \times w + n)
+        \end{aligned}
+
+    If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides
+    for :attr:`padding` number of points. :attr:`dilation` controls the spacing between the kernel points.
+    It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
+
+    Note:
+        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding
+        or the input. Sliding windows that would start in the right padded region are ignored.
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`dilation` can either be:
+
+        - a single ``int`` -- in which case the same value is used for the depth, height and width dimension
+        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,
+          the second `int` for the height dimension and the third `int` for the width dimension
+
+    Args:
+        kernel_size: the size of the window to take a max over
+        stride: the stride of the window. Default value is :attr:`kernel_size`
+        padding: Implicit negative infinity padding to be added on all three sides
+        dilation: a parameter that controls the stride of elements in the window
+        return_indices: if ``True``, will return the max indices along with the outputs.
+                        Useful for :class:`torch.nn.MaxUnpool3d` later
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`, where
+
+          .. math::
+              D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times
+                (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times
+                (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] - \text{dilation}[2] \times
+                (\text{kernel\_size}[2] - 1) - 1}{\text{stride}[2]} + 1\right\rfloor
+
+    Examples::
+
+        >>> # pool of square window of size=3, stride=2
+        >>> m = nn.MaxPool3d(3, stride=2)
+        >>> # pool of non-square window
+        >>> m = nn.MaxPool3d((3, 2, 2), stride=(2, 1, 2))
+        >>> input = torch.randn(20, 16, 50, 44, 31)
+        >>> output = m(input)
+
+    .. _link:
+        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
+    """
+
+    kernel_size: _size_3_t
+    stride: _size_3_t
+    padding: _size_3_t
+    dilation: _size_3_t
+
+    def forward(self, input: Tensor):
+        """Runs the forward pass."""
+        return F.max_pool3d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.dilation,
+            ceil_mode=self.ceil_mode,
+            return_indices=self.return_indices,
+        )
+
+
+class _MaxUnpoolNd(Module):
+    def extra_repr(self) -> str:
+        return f"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}"
+
+
+class MaxUnpool1d(_MaxUnpoolNd):
+    r"""Computes a partial inverse of :class:`MaxPool1d`.
+
+    :class:`MaxPool1d` is not fully invertible, since the non-maximal values are lost.
+
+    :class:`MaxUnpool1d` takes in as input the output of :class:`MaxPool1d`
+    including the indices of the maximal values and computes a partial inverse
+    in which all non-maximal values are set to zero.
+
+    Note:
+        This operation may behave nondeterministically when the input indices has repeat values.
+        See https://github.com/pytorch/pytorch/issues/80827 and :doc:`/notes/randomness` for more information.
+
+    .. note:: :class:`MaxPool1d` can map several input sizes to the same output
+              sizes. Hence, the inversion process can get ambiguous.
+              To accommodate this, you can provide the needed output size
+              as an additional argument :attr:`output_size` in the forward call.
+              See the Inputs and Example below.
+
+    Args:
+        kernel_size (int or tuple): Size of the max pooling window.
+        stride (int or tuple): Stride of the max pooling window.
+            It is set to :attr:`kernel_size` by default.
+        padding (int or tuple): Padding that was added to the input
+
+    Inputs:
+        - `input`: the input Tensor to invert
+        - `indices`: the indices given out by :class:`~torch.nn.MaxPool1d`
+        - `output_size` (optional): the targeted output size
+
+    Shape:
+        - Input: :math:`(N, C, H_{in})` or :math:`(C, H_{in})`.
+        - Output: :math:`(N, C, H_{out})` or :math:`(C, H_{out})`, where
+
+          .. math::
+              H_{out} = (H_{in} - 1) \times \text{stride}[0] - 2 \times \text{padding}[0] + \text{kernel\_size}[0]
+
+          or as given by :attr:`output_size` in the call operator
+
+    Example::
+
+        >>> # xdoctest: +IGNORE_WANT("do other tests modify the global state?")
+        >>> pool = nn.MaxPool1d(2, stride=2, return_indices=True)
+        >>> unpool = nn.MaxUnpool1d(2, stride=2)
+        >>> input = torch.tensor([[[1., 2, 3, 4, 5, 6, 7, 8]]])
+        >>> output, indices = pool(input)
+        >>> unpool(output, indices)
+        tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.]]])
+
+        >>> # Example showcasing the use of output_size
+        >>> input = torch.tensor([[[1., 2, 3, 4, 5, 6, 7, 8, 9]]])
+        >>> output, indices = pool(input)
+        >>> unpool(output, indices, output_size=input.size())
+        tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.,  0.]]])
+
+        >>> unpool(output, indices)
+        tensor([[[ 0.,  2.,  0.,  4.,  0.,  6.,  0., 8.]]])
+    """
+
+    kernel_size: _size_1_t
+    stride: _size_1_t
+    padding: _size_1_t
+
+    def __init__(
+        self,
+        kernel_size: _size_1_t,
+        stride: _size_1_t | None = None,
+        padding: _size_1_t = 0,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = _single(kernel_size)
+        self.stride = _single(stride if (stride is not None) else kernel_size)
+        self.padding = _single(padding)
+
+    def forward(
+        self, input: Tensor, indices: Tensor, output_size: list[int] | None = None
+    ) -> Tensor:
+        """Runs the forward pass."""
+        return F.max_unpool1d(
+            input, indices, self.kernel_size, self.stride, self.padding, output_size
+        )
+
+
+class MaxUnpool2d(_MaxUnpoolNd):
+    r"""Computes a partial inverse of :class:`MaxPool2d`.
+
+    :class:`MaxPool2d` is not fully invertible, since the non-maximal values are lost.
+
+    :class:`MaxUnpool2d` takes in as input the output of :class:`MaxPool2d`
+    including the indices of the maximal values and computes a partial inverse
+    in which all non-maximal values are set to zero.
+
+    Note:
+        This operation may behave nondeterministically when the input indices has repeat values.
+        See https://github.com/pytorch/pytorch/issues/80827 and :doc:`/notes/randomness` for more information.
+
+    .. note:: :class:`MaxPool2d` can map several input sizes to the same output
+              sizes. Hence, the inversion process can get ambiguous.
+              To accommodate this, you can provide the needed output size
+              as an additional argument :attr:`output_size` in the forward call.
+              See the Inputs and Example below.
+
+    Args:
+        kernel_size (int or tuple): Size of the max pooling window.
+        stride (int or tuple): Stride of the max pooling window.
+            It is set to :attr:`kernel_size` by default.
+        padding (int or tuple): Padding that was added to the input
+
+    Inputs:
+        - `input`: the input Tensor to invert
+        - `indices`: the indices given out by :class:`~torch.nn.MaxPool2d`
+        - `output_size` (optional): the targeted output size
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          .. math::
+            H_{out} = (H_{in} - 1) \times \text{stride[0]} - 2 \times \text{padding[0]} + \text{kernel\_size[0]}
+
+          .. math::
+            W_{out} = (W_{in} - 1) \times \text{stride[1]} - 2 \times \text{padding[1]} + \text{kernel\_size[1]}
+
+          or as given by :attr:`output_size` in the call operator
+
+    Example::
+
+        >>> pool = nn.MaxPool2d(2, stride=2, return_indices=True)
+        >>> unpool = nn.MaxUnpool2d(2, stride=2)
+        >>> input = torch.tensor([[[[ 1.,  2.,  3.,  4.],
+                                    [ 5.,  6.,  7.,  8.],
+                                    [ 9., 10., 11., 12.],
+                                    [13., 14., 15., 16.]]]])
+        >>> output, indices = pool(input)
+        >>> unpool(output, indices)
+        tensor([[[[  0.,   0.,   0.,   0.],
+                  [  0.,   6.,   0.,   8.],
+                  [  0.,   0.,   0.,   0.],
+                  [  0.,  14.,   0.,  16.]]]])
+        >>> # Now using output_size to resolve an ambiguous size for the inverse
+        >>> input = torch.tensor([[[[ 1.,  2.,  3.,  4.,  5.],
+                                    [ 6.,  7.,  8.,  9., 10.],
+                                    [11., 12., 13., 14., 15.],
+                                    [16., 17., 18., 19., 20.]]]])
+        >>> output, indices = pool(input)
+        >>> # This call will not work without specifying output_size
+        >>> unpool(output, indices, output_size=input.size())
+        tensor([[[[ 0.,  0.,  0.,  0.,  0.],
+                  [ 0.,  7.,  0.,  9.,  0.],
+                  [ 0.,  0.,  0.,  0.,  0.],
+                  [ 0., 17.,  0., 19.,  0.]]]])
+
+
+    """
+
+    kernel_size: _size_2_t
+    stride: _size_2_t
+    padding: _size_2_t
+
+    def __init__(
+        self,
+        kernel_size: _size_2_t,
+        stride: _size_2_t | None = None,
+        padding: _size_2_t = 0,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride if (stride is not None) else kernel_size)
+        self.padding = _pair(padding)
+
+    def forward(
+        self, input: Tensor, indices: Tensor, output_size: list[int] | None = None
+    ) -> Tensor:
+        """Runs the forward pass."""
+        return F.max_unpool2d(
+            input, indices, self.kernel_size, self.stride, self.padding, output_size
+        )
+
+
+class MaxUnpool3d(_MaxUnpoolNd):
+    r"""Computes a partial inverse of :class:`MaxPool3d`.
+
+    :class:`MaxPool3d` is not fully invertible, since the non-maximal values are lost.
+    :class:`MaxUnpool3d` takes in as input the output of :class:`MaxPool3d`
+    including the indices of the maximal values and computes a partial inverse
+    in which all non-maximal values are set to zero.
+
+    Note:
+        This operation may behave nondeterministically when the input indices has repeat values.
+        See https://github.com/pytorch/pytorch/issues/80827 and :doc:`/notes/randomness` for more information.
+
+    .. note:: :class:`MaxPool3d` can map several input sizes to the same output
+              sizes. Hence, the inversion process can get ambiguous.
+              To accommodate this, you can provide the needed output size
+              as an additional argument :attr:`output_size` in the forward call.
+              See the Inputs section below.
+
+    Args:
+        kernel_size (int or tuple): Size of the max pooling window.
+        stride (int or tuple): Stride of the max pooling window.
+            It is set to :attr:`kernel_size` by default.
+        padding (int or tuple): Padding that was added to the input
+
+    Inputs:
+        - `input`: the input Tensor to invert
+        - `indices`: the indices given out by :class:`~torch.nn.MaxPool3d`
+        - `output_size` (optional): the targeted output size
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`, where
+
+          .. math::
+              D_{out} = (D_{in} - 1) \times \text{stride[0]} - 2 \times \text{padding[0]} + \text{kernel\_size[0]}
+
+          .. math::
+              H_{out} = (H_{in} - 1) \times \text{stride[1]} - 2 \times \text{padding[1]} + \text{kernel\_size[1]}
+
+          .. math::
+              W_{out} = (W_{in} - 1) \times \text{stride[2]} - 2 \times \text{padding[2]} + \text{kernel\_size[2]}
+
+          or as given by :attr:`output_size` in the call operator
+
+    Example::
+
+        >>> # pool of square window of size=3, stride=2
+        >>> pool = nn.MaxPool3d(3, stride=2, return_indices=True)
+        >>> unpool = nn.MaxUnpool3d(3, stride=2)
+        >>> output, indices = pool(torch.randn(20, 16, 51, 33, 15))
+        >>> unpooled_output = unpool(output, indices)
+        >>> unpooled_output.size()
+        torch.Size([20, 16, 51, 33, 15])
+    """
+
+    kernel_size: _size_3_t
+    stride: _size_3_t
+    padding: _size_3_t
+
+    def __init__(
+        self,
+        kernel_size: _size_3_t,
+        stride: _size_3_t | None = None,
+        padding: _size_3_t = 0,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = _triple(kernel_size)
+        self.stride = _triple(stride if (stride is not None) else kernel_size)
+        self.padding = _triple(padding)
+
+    def forward(
+        self, input: Tensor, indices: Tensor, output_size: list[int] | None = None
+    ) -> Tensor:
+        """Runs the forward pass."""
+        return F.max_unpool3d(
+            input, indices, self.kernel_size, self.stride, self.padding, output_size
+        )
+
+
+class _AvgPoolNd(Module):
+    __constants__ = [
+        "kernel_size",
+        "stride",
+        "padding",
+        "ceil_mode",
+        "count_include_pad",
+    ]
+
+    def extra_repr(self) -> str:
+        return f"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}"
+
+
+class AvgPool1d(_AvgPoolNd):
+    r"""Applies a 1D average pooling over an input signal composed of several input planes.
+
+    In the simplest case, the output value of the layer with input size :math:`(N, C, L)`,
+    output :math:`(N, C, L_{out})` and :attr:`kernel_size` :math:`k`
+    can be precisely described as:
+
+    .. math::
+
+        \text{out}(N_i, C_j, l) = \frac{1}{k} \sum_{m=0}^{k-1}
+                               \text{input}(N_i, C_j, \text{stride} \times l + m)
+
+    If :attr:`padding` is non-zero, then the input is implicitly zero-padded on both sides
+    for :attr:`padding` number of points.
+
+    Note:
+        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding
+        or the input. Sliding windows that would start in the right padded region are ignored.
+
+    .. note::
+        pad should be at most half of effective kernel size.
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding` can each be
+    an ``int`` or a one-element tuple.
+
+    Args:
+        kernel_size: the size of the window
+        stride: the stride of the window. Default value is :attr:`kernel_size`
+        padding: implicit zero padding to be added on both sides
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+        count_include_pad: when True, will include the zero-padding in the averaging calculation
+
+    Shape:
+        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.
+        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where
+
+          .. math::
+              L_{out} = \left\lfloor \frac{L_{in} +
+              2 \times \text{padding} - \text{kernel\_size}}{\text{stride}} + 1\right\rfloor
+
+          Per the note above, if ``ceil_mode`` is True and :math:`(L_{out} - 1) \times \text{stride} \geq L_{in}
+          + \text{padding}`, we skip the last window as it would start in the right padded region, resulting in
+          :math:`L_{out}` being reduced by one.
+
+    Examples::
+
+        >>> # pool with window of size=3, stride=2
+        >>> m = nn.AvgPool1d(3, stride=2)
+        >>> m(torch.tensor([[[1., 2, 3, 4, 5, 6, 7]]]))
+        tensor([[[2., 4., 6.]]])
+    """
+
+    kernel_size: _size_1_t
+    stride: _size_1_t
+    padding: _size_1_t
+    ceil_mode: bool
+    count_include_pad: bool
+
+    def __init__(
+        self,
+        kernel_size: _size_1_t,
+        stride: _size_1_t = None,
+        padding: _size_1_t = 0,
+        ceil_mode: bool = False,
+        count_include_pad: bool = True,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = _single(kernel_size)
+        self.stride = _single(stride if stride is not None else kernel_size)
+        self.padding = _single(padding)
+        self.ceil_mode = ceil_mode
+        self.count_include_pad = count_include_pad
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.avg_pool1d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.ceil_mode,
+            self.count_include_pad,
+        )
+
+
+class AvgPool2d(_AvgPoolNd):
+    r"""Applies a 2D average pooling over an input signal composed of several input planes.
+
+    In the simplest case, the output value of the layer with input size :math:`(N, C, H, W)`,
+    output :math:`(N, C, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kH, kW)`
+    can be precisely described as:
+
+    .. math::
+
+        out(N_i, C_j, h, w)  = \frac{1}{kH * kW} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1}
+                               input(N_i, C_j, stride[0] \times h + m, stride[1] \times w + n)
+
+    If :attr:`padding` is non-zero, then the input is implicitly zero-padded on both sides
+    for :attr:`padding` number of points.
+
+    Note:
+        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding
+        or the input. Sliding windows that would start in the right padded region are ignored.
+
+    .. note::
+        pad should be at most half of effective kernel size.
+
+    The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding` can either be:
+
+        - a single ``int`` or a single-element tuple -- in which case the same value is used for the height and width dimension
+        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,
+          and the second `int` for the width dimension
+
+    Args:
+        kernel_size: the size of the window
+        stride: the stride of the window. Default value is :attr:`kernel_size`
+        padding: implicit zero padding to be added on both sides
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+        count_include_pad: when True, will include the zero-padding in the averaging calculation
+        divisor_override: if specified, it will be used as divisor, otherwise size of the pooling region will be used.
+
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in}  + 2 \times \text{padding}[0] -
+                \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in}  + 2 \times \text{padding}[1] -
+                \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor
+
+          Per the note above, if ``ceil_mode`` is True and :math:`(H_{out} - 1)\times \text{stride}[0]\geq H_{in}
+          + \text{padding}[0]`, we skip the last window as it would start in the bottom padded region,
+          resulting in :math:`H_{out}` being reduced by one.
+
+          The same applies for :math:`W_{out}`.
+
+    Examples::
+
+        >>> # pool of square window of size=3, stride=2
+        >>> m = nn.AvgPool2d(3, stride=2)
+        >>> # pool of non-square window
+        >>> m = nn.AvgPool2d((3, 2), stride=(2, 1))
+        >>> input = torch.randn(20, 16, 50, 32)
+        >>> output = m(input)
+    """
+
+    __constants__ = [
+        "kernel_size",
+        "stride",
+        "padding",
+        "ceil_mode",
+        "count_include_pad",
+        "divisor_override",
+    ]
+
+    kernel_size: _size_2_t
+    stride: _size_2_t
+    padding: _size_2_t
+    ceil_mode: bool
+    count_include_pad: bool
+
+    def __init__(
+        self,
+        kernel_size: _size_2_t,
+        stride: _size_2_t | None = None,
+        padding: _size_2_t = 0,
+        ceil_mode: bool = False,
+        count_include_pad: bool = True,
+        divisor_override: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride if (stride is not None) else kernel_size
+        self.padding = padding
+        self.ceil_mode = ceil_mode
+        self.count_include_pad = count_include_pad
+        self.divisor_override = divisor_override
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.avg_pool2d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.ceil_mode,
+            self.count_include_pad,
+            self.divisor_override,
+        )
+
+
+class AvgPool3d(_AvgPoolNd):
+    r"""Applies a 3D average pooling over an input signal composed of several input planes.
+
+    In the simplest case, the output value of the layer with input size :math:`(N, C, D, H, W)`,
+    output :math:`(N, C, D_{out}, H_{out}, W_{out})` and :attr:`kernel_size` :math:`(kD, kH, kW)`
+    can be precisely described as:
+
+    .. math::
+        \begin{aligned}
+            \text{out}(N_i, C_j, d, h, w) ={} & \sum_{k=0}^{kD-1} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} \\
+                                              & \frac{\text{input}(N_i, C_j, \text{stride}[0] \times d + k,
+                                                      \text{stride}[1] \times h + m, \text{stride}[2] \times w + n)}
+                                                     {kD \times kH \times kW}
+        \end{aligned}
+
+    If :attr:`padding` is non-zero, then the input is implicitly zero-padded on all three sides
+    for :attr:`padding` number of points.
+
+    Note:
+        When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding
+        or the input. Sliding windows that would start in the right padded region are ignored.
+
+    .. note::
+        pad should be at most half of effective kernel size.
+
+    The parameters :attr:`kernel_size`, :attr:`stride` can either be:
+
+        - a single ``int`` -- in which case the same value is used for the depth, height and width dimension
+        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,
+          the second `int` for the height dimension and the third `int` for the width dimension
+
+    Args:
+        kernel_size: the size of the window
+        stride: the stride of the window. Default value is :attr:`kernel_size`
+        padding: implicit zero padding to be added on all three sides
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+        count_include_pad: when True, will include the zero-padding in the averaging calculation
+        divisor_override: if specified, it will be used as divisor, otherwise :attr:`kernel_size` will be used
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or
+          :math:`(C, D_{out}, H_{out}, W_{out})`, where
+
+          .. math::
+              D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] -
+                    \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] -
+                    \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] -
+                    \text{kernel\_size}[2]}{\text{stride}[2]} + 1\right\rfloor
+
+          Per the note above, if ``ceil_mode`` is True and :math:`(D_{out} - 1)\times \text{stride}[0]\geq D_{in}
+          + \text{padding}[0]`, we skip the last window as it would start in the padded region,
+          resulting in :math:`D_{out}` being reduced by one.
+
+          The same applies for :math:`W_{out}` and :math:`H_{out}`.
+
+    Examples::
+
+        >>> # pool of square window of size=3, stride=2
+        >>> m = nn.AvgPool3d(3, stride=2)
+        >>> # pool of non-square window
+        >>> m = nn.AvgPool3d((3, 2, 2), stride=(2, 1, 2))
+        >>> input = torch.randn(20, 16, 50, 44, 31)
+        >>> output = m(input)
+    """
+
+    __constants__ = [
+        "kernel_size",
+        "stride",
+        "padding",
+        "ceil_mode",
+        "count_include_pad",
+        "divisor_override",
+    ]
+
+    kernel_size: _size_3_t
+    stride: _size_3_t
+    padding: _size_3_t
+    ceil_mode: bool
+    count_include_pad: bool
+
+    def __init__(
+        self,
+        kernel_size: _size_3_t,
+        stride: _size_3_t | None = None,
+        padding: _size_3_t = 0,
+        ceil_mode: bool = False,
+        count_include_pad: bool = True,
+        divisor_override: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride if (stride is not None) else kernel_size
+        self.padding = padding
+        self.ceil_mode = ceil_mode
+        self.count_include_pad = count_include_pad
+        self.divisor_override = divisor_override
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.avg_pool3d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.ceil_mode,
+            self.count_include_pad,
+            self.divisor_override,
+        )
+
+    def __setstate__(self, d):
+        super().__setstate__(d)
+        self.__dict__.setdefault("padding", 0)
+        self.__dict__.setdefault("ceil_mode", False)
+        self.__dict__.setdefault("count_include_pad", True)
+
+
+class FractionalMaxPool2d(Module):
+    r"""Applies a 2D fractional max pooling over an input signal composed of several input planes.
+
+    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham
+
+    The max-pooling operation is applied in :math:`kH \times kW` regions by a stochastic
+    step size determined by the target output size.
+    The number of output features is equal to the number of input planes.
+
+    .. note:: Exactly one of ``output_size`` or ``output_ratio`` must be defined.
+
+    Args:
+        kernel_size: the size of the window to take a max over.
+                     Can be a single number k (for a square kernel of k x k) or a tuple `(kh, kw)`
+        output_size: the target output size of the image of the form `oH x oW`.
+                     Can be a tuple `(oH, oW)` or a single number oH for a square image `oH x oH`.
+                     Note that we must have :math:`kH + oH - 1 <= H_{in}` and :math:`kW + oW - 1 <= W_{in}`
+        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.
+                      This has to be a number or tuple in the range (0, 1).
+                      Note that we must have :math:`kH + (output\_ratio\_H * H_{in}) - 1 <= H_{in}`
+                      and :math:`kW + (output\_ratio\_W * W_{in}) - 1 <= W_{in}`
+        return_indices: if ``True``, will return the indices along with the outputs.
+                        Useful to pass to :meth:`nn.MaxUnpool2d`. Default: ``False``
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+          :math:`(H_{out}, W_{out})=\text{output\_size}` or
+          :math:`(H_{out}, W_{out})=\text{output\_ratio} \times (H_{in}, W_{in})`.
+
+    Examples:
+        >>> # pool of square window of size=3, and target output size 13x12
+        >>> m = nn.FractionalMaxPool2d(3, output_size=(13, 12))
+        >>> # pool of square window and target output size being half of input image size
+        >>> m = nn.FractionalMaxPool2d(3, output_ratio=(0.5, 0.5))
+        >>> input = torch.randn(20, 16, 50, 32)
+        >>> output = m(input)
+
+    .. _Fractional MaxPooling:
+        https://arxiv.org/abs/1412.6071
+    """
+
+    __constants__ = ["kernel_size", "return_indices", "output_size", "output_ratio"]
+
+    kernel_size: _size_2_t
+    return_indices: bool
+    output_size: _size_2_t
+    output_ratio: _ratio_2_t
+
+    def __init__(
+        self,
+        kernel_size: _size_2_t,
+        output_size: _size_2_t | None = None,
+        output_ratio: _ratio_2_t | None = None,
+        return_indices: bool = False,
+        _random_samples=None,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = _pair(kernel_size)
+        self.return_indices = return_indices
+        self.register_buffer("_random_samples", _random_samples)
+        self.output_size = _pair(output_size) if output_size is not None else None
+        self.output_ratio = _pair(output_ratio) if output_ratio is not None else None
+        if output_size is None and output_ratio is None:
+            raise ValueError(
+                "FractionalMaxPool2d requires specifying either "
+                "an output size, or a pooling ratio"
+            )
+        if output_size is not None and output_ratio is not None:
+            raise ValueError(
+                "only one of output_size and output_ratio may be specified"
+            )
+        if self.output_ratio is not None:
+            if not (0 < self.output_ratio[0] < 1 and 0 < self.output_ratio[1] < 1):
+                raise ValueError(
+                    f"output_ratio must be between 0 and 1 (got {output_ratio})"
+                )
+
+    def forward(self, input: Tensor):
+        return F.fractional_max_pool2d(
+            input,
+            self.kernel_size,
+            self.output_size,
+            self.output_ratio,
+            self.return_indices,
+            _random_samples=self._random_samples,
+        )
+
+
+class FractionalMaxPool3d(Module):
+    r"""Applies a 3D fractional max pooling over an input signal composed of several input planes.
+
+    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham
+
+    The max-pooling operation is applied in :math:`kT \times kH \times kW` regions by a stochastic
+    step size determined by the target output size.
+    The number of output features is equal to the number of input planes.
+
+    .. note:: Exactly one of ``output_size`` or ``output_ratio`` must be defined.
+
+    Args:
+        kernel_size: the size of the window to take a max over.
+                     Can be a single number `k` (for a square kernel of `k x k x k`) or a tuple `(kt x kh x kw)`,
+                     `k` must greater than 0.
+        output_size: the target output size of the image of the form `oT x oH x oW`.
+                     Can be a tuple `(oT, oH, oW)` or a single number oH for a square image `oH x oH x oH`
+        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.
+                      This has to be a number or tuple in the range (0, 1)
+        return_indices: if ``True``, will return the indices along with the outputs.
+                        Useful to pass to :meth:`nn.MaxUnpool3d`. Default: ``False``
+
+    Shape:
+        - Input: :math:`(N, C, T_{in}, H_{in}, W_{in})` or :math:`(C, T_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, T_{out}, H_{out}, W_{out})` or :math:`(C, T_{out}, H_{out}, W_{out})`, where
+          :math:`(T_{out}, H_{out}, W_{out})=\text{output\_size}` or
+          :math:`(T_{out}, H_{out}, W_{out})=\text{output\_ratio} \times (T_{in}, H_{in}, W_{in})`
+
+    Examples:
+        >>> # pool of cubic window of size=3, and target output size 13x12x11
+        >>> m = nn.FractionalMaxPool3d(3, output_size=(13, 12, 11))
+        >>> # pool of cubic window and target output size being half of input size
+        >>> m = nn.FractionalMaxPool3d(3, output_ratio=(0.5, 0.5, 0.5))
+        >>> input = torch.randn(20, 16, 50, 32, 16)
+        >>> output = m(input)
+
+    .. _Fractional MaxPooling:
+        https://arxiv.org/abs/1412.6071
+    """
+
+    __constants__ = ["kernel_size", "return_indices", "output_size", "output_ratio"]
+    kernel_size: _size_3_t
+    return_indices: bool
+    output_size: _size_3_t
+    output_ratio: _ratio_3_t
+
+    def __init__(
+        self,
+        kernel_size: _size_3_t,
+        output_size: _size_3_t | None = None,
+        output_ratio: _ratio_3_t | None = None,
+        return_indices: bool = False,
+        _random_samples=None,
+    ) -> None:
+        super().__init__()
+        if (isinstance(kernel_size, int) and kernel_size <= 0) or (
+            isinstance(kernel_size, (tuple, list))
+            and not all(k > 0 for k in kernel_size)
+        ):
+            raise ValueError(f"kernel_size must greater than 0, but got {kernel_size}")
+        self.kernel_size = _triple(kernel_size)
+        self.return_indices = return_indices
+        self.register_buffer("_random_samples", _random_samples)
+        self.output_size = _triple(output_size) if output_size is not None else None
+        self.output_ratio = _triple(output_ratio) if output_ratio is not None else None
+        if output_size is None and output_ratio is None:
+            raise ValueError(
+                "FractionalMaxPool3d requires specifying either "
+                "an output size, or a pooling ratio"
+            )
+        if output_size is not None and output_ratio is not None:
+            raise ValueError(
+                "only one of output_size and output_ratio may be specified"
+            )
+        if self.output_ratio is not None:
+            if not (
+                0 < self.output_ratio[0] < 1
+                and 0 < self.output_ratio[1] < 1
+                and 0 < self.output_ratio[2] < 1
+            ):
+                raise ValueError(
+                    f"output_ratio must be between 0 and 1 (got {output_ratio})"
+                )
+
+    def forward(self, input: Tensor):
+        return F.fractional_max_pool3d(
+            input,
+            self.kernel_size,
+            self.output_size,
+            self.output_ratio,
+            self.return_indices,
+            _random_samples=self._random_samples,
+        )
+
+
+class _LPPoolNd(Module):
+    __constants__ = ["norm_type", "kernel_size", "stride", "ceil_mode"]
+
+    norm_type: float
+    ceil_mode: bool
+
+    def __init__(
+        self,
+        norm_type: float,
+        kernel_size: _size_any_t,
+        stride: _size_any_t | None = None,
+        ceil_mode: bool = False,
+    ) -> None:
+        super().__init__()
+        self.norm_type = norm_type
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.ceil_mode = ceil_mode
+
+    def extra_repr(self) -> str:
+        return (
+            "norm_type={norm_type}, kernel_size={kernel_size}, stride={stride}, "
+            "ceil_mode={ceil_mode}".format(**self.__dict__)
+        )
+
+
+class LPPool1d(_LPPoolNd):
+    r"""Applies a 1D power-average pooling over an input signal composed of several input planes.
+
+    On each window, the function computed is:
+
+    .. math::
+        f(X) = \sqrt[p]{\sum_{x \in X} x^{p}}
+
+    - At p = :math:`\infty`, one gets Max Pooling
+    - At p = 1, one gets Sum Pooling (which is proportional to Average Pooling)
+
+    .. note:: If the sum to the power of `p` is zero, the gradient of this function is
+              not defined. This implementation will set the gradient to zero in this case.
+
+    Args:
+        kernel_size: a single int, the size of the window
+        stride: a single int, the stride of the window. Default value is :attr:`kernel_size`
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+
+    Note:
+        When :attr:`ceil_mode` is ``True``, sliding windows may go off-bounds if they start within the
+        left padding or the input. Sliding windows that would start in the right padded region are ignored.
+
+    Shape:
+        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.
+        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where
+
+          .. math::
+              L_{out} = \left\lfloor\frac{L_{in} - \text{kernel\_size}}{\text{stride}} + 1\right\rfloor
+
+    Examples::
+        >>> # power-2 pool of window of length 3, with stride 2.
+        >>> m = nn.LPPool1d(2, 3, stride=2)
+        >>> input = torch.randn(20, 16, 50)
+        >>> output = m(input)
+    """
+
+    kernel_size: _size_1_t
+    stride: _size_1_t
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.lp_pool1d(
+            input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode
+        )
+
+
+class LPPool2d(_LPPoolNd):
+    r"""Applies a 2D power-average pooling over an input signal composed of several input planes.
+
+    On each window, the function computed is:
+
+    .. math::
+        f(X) = \sqrt[p]{\sum_{x \in X} x^{p}}
+
+    - At p = :math:`\infty`, one gets Max Pooling
+    - At p = 1, one gets Sum Pooling (which is proportional to average pooling)
+
+    The parameters :attr:`kernel_size`, :attr:`stride` can either be:
+
+        - a single ``int`` -- in which case the same value is used for the height and width dimension
+        - a ``tuple`` of two ints -- in which case, the first `int` is used for the height dimension,
+          and the second `int` for the width dimension
+
+    .. note:: If the sum to the power of `p` is zero, the gradient of this function is
+              not defined. This implementation will set the gradient to zero in this case.
+
+    Args:
+        kernel_size: the size of the window
+        stride: the stride of the window. Default value is :attr:`kernel_size`
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+
+    Note:
+        When :attr:`ceil_mode` is ``True``, sliding windows may go off-bounds if they start within the
+        left padding or the input. Sliding windows that would start in the right padded region are ignored.
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in} - \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in} - \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor
+
+    Examples::
+
+        >>> # power-2 pool of square window of size=3, stride=2
+        >>> m = nn.LPPool2d(2, 3, stride=2)
+        >>> # pool of non-square window of power 1.2
+        >>> m = nn.LPPool2d(1.2, (3, 2), stride=(2, 1))
+        >>> input = torch.randn(20, 16, 50, 32)
+        >>> output = m(input)
+
+    """
+
+    kernel_size: _size_2_t
+    stride: _size_2_t
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.lp_pool2d(
+            input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode
+        )
+
+
+class LPPool3d(_LPPoolNd):
+    r"""Applies a 3D power-average pooling over an input signal composed of several input planes.
+
+    On each window, the function computed is:
+
+    .. math::
+        f(X) = \sqrt[p]{\sum_{x \in X} x^{p}}
+
+    - At p = :math:`\infty`, one gets Max Pooling
+    - At p = 1, one gets Sum Pooling (which is proportional to average pooling)
+
+    The parameters :attr:`kernel_size`, :attr:`stride` can either be:
+
+        - a single ``int`` -- in which case the same value is used for the height, width and depth dimension
+        - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension,
+          the second `int` for the height dimension and the third `int` for the width dimension
+
+    .. note:: If the sum to the power of `p` is zero, the gradient of this function is
+              not defined. This implementation will set the gradient to zero in this case.
+
+    Args:
+        kernel_size: the size of the window
+        stride: the stride of the window. Default value is :attr:`kernel_size`
+        ceil_mode: when True, will use `ceil` instead of `floor` to compute the output shape
+
+    Note:
+        When :attr:`ceil_mode` is ``True``, sliding windows may go off-bounds if they start within the
+        left padding or the input. Sliding windows that would start in the right padded region are ignored.
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or
+          :math:`(C, D_{out}, H_{out}, W_{out})`, where
+
+          .. math::
+              D_{out} = \left\lfloor\frac{D_{in} - \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor
+
+          .. math::
+              H_{out} = \left\lfloor\frac{H_{in} - \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor
+
+          .. math::
+              W_{out} = \left\lfloor\frac{W_{in} - \text{kernel\_size}[2]}{\text{stride}[2]} + 1\right\rfloor
+
+    Examples::
+
+        >>> # power-2 pool of square window of size=3, stride=2
+        >>> m = nn.LPPool3d(2, 3, stride=2)
+        >>> # pool of non-square window of power 1.2
+        >>> m = nn.LPPool3d(1.2, (3, 2, 2), stride=(2, 1, 2))
+        >>> input = torch.randn(20, 16, 50, 44, 31)
+        >>> output = m(input)
+
+    """
+
+    kernel_size: _size_3_t
+    stride: _size_3_t
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.lp_pool3d(
+            input, float(self.norm_type), self.kernel_size, self.stride, self.ceil_mode
+        )
+
+
+class _AdaptiveMaxPoolNd(Module):
+    __constants__ = ["output_size", "return_indices"]
+    return_indices: bool
+
+    def __init__(
+        self, output_size: _size_any_opt_t, return_indices: bool = False
+    ) -> None:
+        super().__init__()
+        self.output_size = output_size
+        self.return_indices = return_indices
+
+    def extra_repr(self) -> str:
+        return f"output_size={self.output_size}"
+
+
+# FIXME (by @ssnl): Improve adaptive pooling docs: specify what the input and
+#   output shapes are, and how the operation computes output.
+
+
+class AdaptiveMaxPool1d(_AdaptiveMaxPoolNd):
+    r"""Applies a 1D adaptive max pooling over an input signal composed of several input planes.
+
+    The output size is :math:`L_{out}`, for any input size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        output_size: the target output size :math:`L_{out}`.
+        return_indices: if ``True``, will return the indices along with the outputs.
+                        Useful to pass to nn.MaxUnpool1d. Default: ``False``
+
+    Shape:
+        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.
+        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where
+          :math:`L_{out}=\text{output\_size}`.
+
+    Examples:
+        >>> # target output size of 5
+        >>> m = nn.AdaptiveMaxPool1d(5)
+        >>> input = torch.randn(1, 64, 8)
+        >>> output = m(input)
+
+    """
+
+    output_size: _size_1_t
+
+    def forward(self, input: Tensor):
+        """Runs the forward pass."""
+        return F.adaptive_max_pool1d(input, self.output_size, self.return_indices)
+
+
+class AdaptiveMaxPool2d(_AdaptiveMaxPoolNd):
+    r"""Applies a 2D adaptive max pooling over an input signal composed of several input planes.
+
+    The output is of size :math:`H_{out} \times W_{out}`, for any input size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        output_size: the target output size of the image of the form :math:`H_{out} \times W_{out}`.
+                     Can be a tuple :math:`(H_{out}, W_{out})` or a single :math:`H_{out}` for a
+                     square image :math:`H_{out} \times H_{out}`. :math:`H_{out}` and :math:`W_{out}`
+                     can be either a ``int``, or ``None`` which means the size will be the same as that
+                     of the input.
+        return_indices: if ``True``, will return the indices along with the outputs.
+                        Useful to pass to nn.MaxUnpool2d. Default: ``False``
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
+          :math:`(H_{out}, W_{out})=\text{output\_size}`.
+
+    Examples:
+        >>> # target output size of 5x7
+        >>> m = nn.AdaptiveMaxPool2d((5, 7))
+        >>> input = torch.randn(1, 64, 8, 9)
+        >>> output = m(input)
+        >>> # target output size of 7x7 (square)
+        >>> m = nn.AdaptiveMaxPool2d(7)
+        >>> input = torch.randn(1, 64, 10, 9)
+        >>> output = m(input)
+        >>> # target output size of 10x7
+        >>> m = nn.AdaptiveMaxPool2d((None, 7))
+        >>> input = torch.randn(1, 64, 10, 9)
+        >>> output = m(input)
+
+    """
+
+    output_size: _size_2_opt_t
+
+    def forward(self, input: Tensor):
+        """Runs the forward pass."""
+        return F.adaptive_max_pool2d(input, self.output_size, self.return_indices)
+
+
+class AdaptiveMaxPool3d(_AdaptiveMaxPoolNd):
+    r"""Applies a 3D adaptive max pooling over an input signal composed of several input planes.
+
+    The output is of size :math:`D_{out} \times H_{out} \times W_{out}`, for any input size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        output_size: the target output size of the image of the form :math:`D_{out} \times H_{out} \times W_{out}`.
+                     Can be a tuple :math:`(D_{out}, H_{out}, W_{out})` or a single
+                     :math:`D_{out}` for a cube :math:`D_{out} \times D_{out} \times D_{out}`.
+                     :math:`D_{out}`, :math:`H_{out}` and :math:`W_{out}` can be either a
+                     ``int``, or ``None`` which means the size will be the same as that of the input.
+
+        return_indices: if ``True``, will return the indices along with the outputs.
+                        Useful to pass to nn.MaxUnpool3d. Default: ``False``
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,
+          where :math:`(D_{out}, H_{out}, W_{out})=\text{output\_size}`.
+
+    Examples:
+        >>> # target output size of 5x7x9
+        >>> m = nn.AdaptiveMaxPool3d((5, 7, 9))
+        >>> input = torch.randn(1, 64, 8, 9, 10)
+        >>> output = m(input)
+        >>> # target output size of 7x7x7 (cube)
+        >>> m = nn.AdaptiveMaxPool3d(7)
+        >>> input = torch.randn(1, 64, 10, 9, 8)
+        >>> output = m(input)
+        >>> # target output size of 7x9x8
+        >>> m = nn.AdaptiveMaxPool3d((7, None, None))
+        >>> input = torch.randn(1, 64, 10, 9, 8)
+        >>> output = m(input)
+
+    """
+
+    output_size: _size_3_opt_t
+
+    def forward(self, input: Tensor):
+        """Runs the forward pass."""
+        return F.adaptive_max_pool3d(input, self.output_size, self.return_indices)
+
+
+class _AdaptiveAvgPoolNd(Module):
+    __constants__ = ["output_size"]
+
+    def __init__(self, output_size: _size_any_opt_t) -> None:
+        super().__init__()
+        self.output_size = output_size
+
+    def extra_repr(self) -> str:
+        return f"output_size={self.output_size}"
+
+
+class AdaptiveAvgPool1d(_AdaptiveAvgPoolNd):
+    r"""Applies a 1D adaptive average pooling over an input signal composed of several input planes.
+
+    The output size is :math:`L_{out}`, for any input size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        output_size: the target output size :math:`L_{out}`.
+
+    Shape:
+        - Input: :math:`(N, C, L_{in})` or :math:`(C, L_{in})`.
+        - Output: :math:`(N, C, L_{out})` or :math:`(C, L_{out})`, where
+          :math:`L_{out}=\text{output\_size}`.
+
+    Examples:
+        >>> # target output size of 5
+        >>> m = nn.AdaptiveAvgPool1d(5)
+        >>> input = torch.randn(1, 64, 8)
+        >>> output = m(input)
+
+    """
+
+    output_size: _size_1_t
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.adaptive_avg_pool1d(input, self.output_size)
+
+
+class AdaptiveAvgPool2d(_AdaptiveAvgPoolNd):
+    r"""Applies a 2D adaptive average pooling over an input signal composed of several input planes.
+
+    The output is of size H x W, for any input size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        output_size: the target output size of the image of the form H x W.
+                     Can be a tuple (H, W) or a single H for a square image H x H.
+                     H and W can be either a ``int``, or ``None`` which means the size will
+                     be the same as that of the input.
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, S_{0}, S_{1})` or :math:`(C, S_{0}, S_{1})`, where
+          :math:`S=\text{output\_size}`.
+
+    Examples:
+        >>> # target output size of 5x7
+        >>> m = nn.AdaptiveAvgPool2d((5, 7))
+        >>> input = torch.randn(1, 64, 8, 9)
+        >>> output = m(input)
+        >>> # target output size of 7x7 (square)
+        >>> m = nn.AdaptiveAvgPool2d(7)
+        >>> input = torch.randn(1, 64, 10, 9)
+        >>> output = m(input)
+        >>> # target output size of 10x7
+        >>> m = nn.AdaptiveAvgPool2d((None, 7))
+        >>> input = torch.randn(1, 64, 10, 9)
+        >>> output = m(input)
+
+    """
+
+    output_size: _size_2_opt_t
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.adaptive_avg_pool2d(input, self.output_size)
+
+
+class AdaptiveAvgPool3d(_AdaptiveAvgPoolNd):
+    r"""Applies a 3D adaptive average pooling over an input signal composed of several input planes.
+
+    The output is of size D x H x W, for any input size.
+    The number of output features is equal to the number of input planes.
+
+    Args:
+        output_size: the target output size of the form D x H x W.
+                     Can be a tuple (D, H, W) or a single number D for a cube D x D x D.
+                     D, H and W can be either a ``int``, or ``None`` which means the size will
+                     be the same as that of the input.
+
+    Shape:
+        - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
+        - Output: :math:`(N, C, S_{0}, S_{1}, S_{2})` or :math:`(C, S_{0}, S_{1}, S_{2})`,
+          where :math:`S=\text{output\_size}`.
+
+    Examples:
+        >>> # target output size of 5x7x9
+        >>> m = nn.AdaptiveAvgPool3d((5, 7, 9))
+        >>> input = torch.randn(1, 64, 8, 9, 10)
+        >>> output = m(input)
+        >>> # target output size of 7x7x7 (cube)
+        >>> m = nn.AdaptiveAvgPool3d(7)
+        >>> input = torch.randn(1, 64, 10, 9, 8)
+        >>> output = m(input)
+        >>> # target output size of 7x9x8
+        >>> m = nn.AdaptiveAvgPool3d((7, None, None))
+        >>> input = torch.randn(1, 64, 10, 9, 8)
+        >>> output = m(input)
+
+    """
+
+    output_size: _size_3_opt_t
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Runs the forward pass."""
+        return F.adaptive_avg_pool3d(input, self.output_size)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/rnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/rnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f6e3b6154725c77501dd9695103a17c00042a7f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/rnn.py
@@ -0,0 +1,1874 @@
+# mypy: allow-untyped-decorators
+# mypy: allow-untyped-defs
+import math
+import numbers
+import warnings
+import weakref
+from typing import overload
+from typing_extensions import deprecated
+
+import torch
+from torch import _VF, Tensor
+from torch.nn import init
+from torch.nn.parameter import Parameter
+from torch.nn.utils.rnn import PackedSequence
+
+from .module import Module
+
+
+__all__ = [
+    "RNNBase",
+    "RNN",
+    "LSTM",
+    "GRU",
+    "RNNCellBase",
+    "RNNCell",
+    "LSTMCell",
+    "GRUCell",
+]
+
+_rnn_impls = {
+    "RNN_TANH": _VF.rnn_tanh,
+    "RNN_RELU": _VF.rnn_relu,
+}
+
+
+def _apply_permutation(tensor: Tensor, permutation: Tensor, dim: int = 1) -> Tensor:
+    return tensor.index_select(dim, permutation)
+
+
+@deprecated(
+    "`apply_permutation` is deprecated, please use `tensor.index_select(dim, permutation)` instead",
+    category=FutureWarning,
+)
+def apply_permutation(tensor: Tensor, permutation: Tensor, dim: int = 1) -> Tensor:
+    return _apply_permutation(tensor, permutation, dim)
+
+
+class RNNBase(Module):
+    r"""Base class for RNN modules (RNN, LSTM, GRU).
+
+    Implements aspects of RNNs shared by the RNN, LSTM, and GRU classes, such as module initialization
+    and utility methods for parameter storage management.
+
+    .. note::
+        The forward method is not implemented by the RNNBase class.
+
+    .. note::
+        LSTM and GRU classes override some methods implemented by RNNBase.
+    """
+
+    __constants__ = [
+        "mode",
+        "input_size",
+        "hidden_size",
+        "num_layers",
+        "bias",
+        "batch_first",
+        "dropout",
+        "bidirectional",
+        "proj_size",
+    ]
+    __jit_unused_properties__ = ["all_weights"]
+
+    mode: str
+    input_size: int
+    hidden_size: int
+    num_layers: int
+    bias: bool
+    batch_first: bool
+    dropout: float
+    bidirectional: bool
+    proj_size: int
+
+    def __init__(
+        self,
+        mode: str,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        bias: bool = True,
+        batch_first: bool = False,
+        dropout: float = 0.0,
+        bidirectional: bool = False,
+        proj_size: int = 0,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.mode = mode
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bias = bias
+        self.batch_first = batch_first
+        self.dropout = float(dropout)
+        self.bidirectional = bidirectional
+        self.proj_size = proj_size
+        self._flat_weight_refs: list[weakref.ReferenceType[Parameter] | None] = []
+        num_directions = 2 if bidirectional else 1
+
+        if (
+            not isinstance(dropout, numbers.Number)
+            or not 0 <= dropout <= 1
+            or isinstance(dropout, bool)
+        ):
+            raise ValueError(
+                "dropout should be a number in range [0, 1] "
+                "representing the probability of an element being "
+                "zeroed"
+            )
+        if dropout > 0 and num_layers == 1:
+            warnings.warn(
+                "dropout option adds dropout after all but last "
+                "recurrent layer, so non-zero dropout expects "
+                f"num_layers greater than 1, but got dropout={dropout} and "
+                f"num_layers={num_layers}",
+                stacklevel=2,
+            )
+
+        if not isinstance(bias, bool):
+            raise TypeError(f"bias should be of type bool, got: {type(bias).__name__}")
+        if not isinstance(batch_first, bool):
+            raise TypeError(
+                f"batch_first should be of type bool, got: {type(batch_first).__name__}"
+            )
+        if not isinstance(input_size, int):
+            raise TypeError(
+                f"input_size should be of type int, got: {type(input_size).__name__}"
+            )
+        if input_size <= 0:
+            raise ValueError("input_size must be greater than zero")
+        if not isinstance(hidden_size, int):
+            raise TypeError(
+                f"hidden_size should be of type int, got: {type(hidden_size).__name__}"
+            )
+        if hidden_size <= 0:
+            raise ValueError("hidden_size must be greater than zero")
+        if num_layers <= 0:
+            raise ValueError("num_layers must be greater than zero")
+        if proj_size < 0:
+            raise ValueError(
+                "proj_size should be a positive integer or zero to disable projections"
+            )
+        if proj_size >= hidden_size:
+            raise ValueError("proj_size has to be smaller than hidden_size")
+
+        if mode == "LSTM":
+            gate_size = 4 * hidden_size
+        elif mode == "GRU":
+            gate_size = 3 * hidden_size
+        elif mode == "RNN_TANH":
+            gate_size = hidden_size
+        elif mode == "RNN_RELU":
+            gate_size = hidden_size
+        else:
+            raise ValueError("Unrecognized RNN mode: " + mode)
+
+        self._flat_weights_names = []
+        self._all_weights = []
+        for layer in range(num_layers):
+            for direction in range(num_directions):
+                real_hidden_size = proj_size if proj_size > 0 else hidden_size
+                layer_input_size = (
+                    input_size if layer == 0 else real_hidden_size * num_directions
+                )
+
+                w_ih = Parameter(
+                    torch.empty((gate_size, layer_input_size), **factory_kwargs)
+                )
+                w_hh = Parameter(
+                    torch.empty((gate_size, real_hidden_size), **factory_kwargs)
+                )
+                b_ih = Parameter(torch.empty(gate_size, **factory_kwargs))
+                # Second bias vector included for CuDNN compatibility. Only one
+                # bias vector is needed in standard definition.
+                b_hh = Parameter(torch.empty(gate_size, **factory_kwargs))
+                layer_params: tuple[Tensor, ...] = ()
+                if self.proj_size == 0:
+                    if bias:
+                        layer_params = (w_ih, w_hh, b_ih, b_hh)
+                    else:
+                        layer_params = (w_ih, w_hh)
+                else:
+                    w_hr = Parameter(
+                        torch.empty((proj_size, hidden_size), **factory_kwargs)
+                    )
+                    if bias:
+                        layer_params = (w_ih, w_hh, b_ih, b_hh, w_hr)
+                    else:
+                        layer_params = (w_ih, w_hh, w_hr)
+
+                suffix = "_reverse" if direction == 1 else ""
+                param_names = ["weight_ih_l{}{}", "weight_hh_l{}{}"]
+                if bias:
+                    param_names += ["bias_ih_l{}{}", "bias_hh_l{}{}"]
+                if self.proj_size > 0:
+                    param_names += ["weight_hr_l{}{}"]
+                param_names = [x.format(layer, suffix) for x in param_names]
+
+                for name, param in zip(param_names, layer_params, strict=True):
+                    setattr(self, name, param)
+                self._flat_weights_names.extend(param_names)
+                self._all_weights.append(param_names)
+
+        self._init_flat_weights()
+
+        self.reset_parameters()
+
+    def _init_flat_weights(self) -> None:
+        self._flat_weights = [
+            getattr(self, wn) if hasattr(self, wn) else None
+            for wn in self._flat_weights_names
+        ]
+        self._flat_weight_refs = [
+            weakref.ref(w) if w is not None else None for w in self._flat_weights
+        ]
+        self.flatten_parameters()
+
+    def __setattr__(self, attr, value) -> None:
+        if hasattr(self, "_flat_weights_names") and attr in self._flat_weights_names:
+            # keep self._flat_weights up to date if you do self.weight = ...
+            idx = self._flat_weights_names.index(attr)
+            self._flat_weights[idx] = value
+        super().__setattr__(attr, value)
+
+    def flatten_parameters(self) -> None:
+        """Reset parameter data pointer so that they can use faster code paths.
+
+        Right now, this works only if the module is on the GPU and cuDNN is enabled.
+        Otherwise, it's a no-op.
+        """
+        # Short-circuits if _flat_weights is only partially instantiated
+        if len(self._flat_weights) != len(self._flat_weights_names):
+            return
+
+        for w in self._flat_weights:
+            if not isinstance(w, Tensor):
+                return
+        # Short-circuits if any tensor in self._flat_weights is not acceptable to cuDNN
+        # or the tensors in _flat_weights are of different dtypes
+
+        first_fw = self._flat_weights[0]  # type: ignore[union-attr]
+        dtype = first_fw.dtype  # type: ignore[union-attr]
+        for fw in self._flat_weights:
+            if (
+                not isinstance(fw, Tensor)
+                or fw.dtype != dtype
+                or not fw.is_cuda
+                or not torch.backends.cudnn.is_acceptable(fw)
+            ):
+                return
+
+        # If any parameters alias, we fall back to the slower, copying code path. This is
+        # a sufficient check, because overlapping parameter buffers that don't completely
+        # alias would break the assumptions of the uniqueness check in
+        # Module.named_parameters().
+        unique_data_ptrs = {
+            p.data_ptr()  # type: ignore[union-attr]
+            for p in self._flat_weights
+        }
+        if len(unique_data_ptrs) != len(self._flat_weights):
+            return
+
+        with torch.cuda.device_of(first_fw):
+            import torch.backends.cudnn.rnn as rnn
+
+            # Note: no_grad() is necessary since _cudnn_rnn_flatten_weight is
+            # an inplace operation on self._flat_weights
+            with torch.no_grad():
+                if torch._use_cudnn_rnn_flatten_weight():
+                    num_weights = 4 if self.bias else 2
+                    if self.proj_size > 0:
+                        num_weights += 1
+                    torch._cudnn_rnn_flatten_weight(
+                        self._flat_weights,  # type: ignore[arg-type]
+                        num_weights,
+                        self.input_size,
+                        rnn.get_cudnn_mode(self.mode),
+                        self.hidden_size,
+                        self.proj_size,
+                        self.num_layers,
+                        self.batch_first,
+                        bool(self.bidirectional),
+                    )
+
+    def _apply(self, fn, recurse=True):
+        self._flat_weight_refs = []
+        ret = super()._apply(fn, recurse)
+
+        # Resets _flat_weights
+        # Note: be v. careful before removing this, as 3rd party device types
+        # likely rely on this behavior to properly .to() modules like LSTM.
+        self._init_flat_weights()
+
+        return ret
+
+    def reset_parameters(self) -> None:
+        stdv = 1.0 / math.sqrt(self.hidden_size) if self.hidden_size > 0 else 0
+        for weight in self.parameters():
+            init.uniform_(weight, -stdv, stdv)
+
+    def check_input(self, input: Tensor, batch_sizes: Tensor | None) -> None:
+        if not torch.jit.is_scripting():
+            if (
+                input.dtype != self._flat_weights[0].dtype  # type: ignore[union-attr]
+                and not torch._C._is_any_autocast_enabled()
+            ):
+                raise ValueError(
+                    f"RNN input dtype ({input.dtype}) does not match weight dtype ({self._flat_weights[0].dtype}). "  # type: ignore[union-attr]
+                    f"Convert input: input.to({self._flat_weights[0].dtype}), or convert model: model.to({input.dtype})"  # type: ignore[union-attr]
+                )
+        expected_input_dim = 2 if batch_sizes is not None else 3
+        if input.dim() != expected_input_dim:
+            raise RuntimeError(
+                f"input must have {expected_input_dim} dimensions, got {input.dim()}"
+            )
+        if self.input_size != input.size(-1):
+            raise RuntimeError(
+                f"input.size(-1) must be equal to input_size. Expected {self.input_size}, got {input.size(-1)}"
+            )
+
+    def get_expected_hidden_size(
+        self, input: Tensor, batch_sizes: Tensor | None
+    ) -> tuple[int, int, int]:
+        if batch_sizes is not None:
+            mini_batch = int(batch_sizes[0])
+        else:
+            mini_batch = input.size(0) if self.batch_first else input.size(1)
+        num_directions = 2 if self.bidirectional else 1
+        if self.proj_size > 0:
+            expected_hidden_size = (
+                self.num_layers * num_directions,
+                mini_batch,
+                self.proj_size,
+            )
+        else:
+            expected_hidden_size = (
+                self.num_layers * num_directions,
+                mini_batch,
+                self.hidden_size,
+            )
+        return expected_hidden_size
+
+    def check_hidden_size(
+        self,
+        hx: Tensor,
+        expected_hidden_size: tuple[int, int, int],
+        msg: str = "Expected hidden size {}, got {}",
+    ) -> None:
+        if hx.size() != expected_hidden_size:
+            raise RuntimeError(msg.format(expected_hidden_size, list(hx.size())))
+
+    def _weights_have_changed(self):
+        # Returns True if the weight tensors have changed since the last forward pass.
+        # This is the case when used with torch.func.functional_call(), for example.
+        weights_changed = False
+        for ref, name in zip(
+            self._flat_weight_refs, self._flat_weights_names, strict=True
+        ):
+            weight = getattr(self, name) if hasattr(self, name) else None
+            if weight is not None and ref is not None and ref() is not weight:
+                weights_changed = True
+                break
+        return weights_changed
+
+    def check_forward_args(
+        self, input: Tensor, hidden: Tensor, batch_sizes: Tensor | None
+    ) -> None:
+        self.check_input(input, batch_sizes)
+        expected_hidden_size = self.get_expected_hidden_size(input, batch_sizes)
+
+        self.check_hidden_size(hidden, expected_hidden_size)
+
+    def permute_hidden(self, hx: Tensor, permutation: Tensor | None):
+        if permutation is None:
+            return hx
+        return _apply_permutation(hx, permutation)
+
+    def extra_repr(self) -> str:
+        s = "{input_size}, {hidden_size}"
+        if self.proj_size != 0:
+            s += ", proj_size={proj_size}"
+        if self.num_layers != 1:
+            s += ", num_layers={num_layers}"
+        if self.bias is not True:
+            s += ", bias={bias}"
+        if self.batch_first is not False:
+            s += ", batch_first={batch_first}"
+        if self.dropout != 0:
+            s += ", dropout={dropout}"
+        if self.bidirectional is not False:
+            s += ", bidirectional={bidirectional}"
+        return s.format(**self.__dict__)
+
+    def _update_flat_weights(self) -> None:
+        if not torch.jit.is_scripting():
+            if self._weights_have_changed():
+                self._init_flat_weights()
+
+    def __getstate__(self):
+        # If weights have been changed, update the _flat_weights in __getstate__ here.
+        self._update_flat_weights()
+        # Don't serialize the weight references.
+        state = self.__dict__.copy()
+        del state["_flat_weight_refs"]
+        return state
+
+    def __setstate__(self, d):
+        super().__setstate__(d)
+        if "all_weights" in d:
+            self._all_weights = d["all_weights"]
+        # In PyTorch 1.8 we added a proj_size member variable to LSTM.
+        # LSTMs that were serialized via torch.save(module) before PyTorch 1.8
+        # don't have it, so to preserve compatibility we set proj_size here.
+        if "proj_size" not in d:
+            self.proj_size = 0
+
+        if not isinstance(self._all_weights[0][0], str):
+            num_layers = self.num_layers
+            num_directions = 2 if self.bidirectional else 1
+            self._flat_weights_names = []
+            self._all_weights = []
+            for layer in range(num_layers):
+                for direction in range(num_directions):
+                    suffix = "_reverse" if direction == 1 else ""
+                    weights = [
+                        "weight_ih_l{}{}",
+                        "weight_hh_l{}{}",
+                        "bias_ih_l{}{}",
+                        "bias_hh_l{}{}",
+                        "weight_hr_l{}{}",
+                    ]
+                    weights = [x.format(layer, suffix) for x in weights]
+                    if self.bias:
+                        if self.proj_size > 0:
+                            self._all_weights += [weights]
+                            self._flat_weights_names.extend(weights)
+                        else:
+                            self._all_weights += [weights[:4]]
+                            self._flat_weights_names.extend(weights[:4])
+                    else:
+                        if self.proj_size > 0:
+                            self._all_weights += [weights[:2]] + [weights[-1:]]
+                            self._flat_weights_names.extend(
+                                weights[:2] + [weights[-1:]]
+                            )
+                        else:
+                            self._all_weights += [weights[:2]]
+                            self._flat_weights_names.extend(weights[:2])
+            self._flat_weights = [
+                getattr(self, wn) if hasattr(self, wn) else None
+                for wn in self._flat_weights_names
+            ]
+
+        self._flat_weight_refs = [
+            weakref.ref(w) if w is not None else None for w in self._flat_weights
+        ]
+
+    @property
+    def all_weights(self) -> list[list[Parameter]]:
+        return [
+            [getattr(self, weight) for weight in weights]
+            for weights in self._all_weights
+        ]
+
+    def _replicate_for_data_parallel(self):
+        replica = super()._replicate_for_data_parallel()
+        # Need to copy these caches, otherwise the replica will share the same
+        # flat weights list.
+        replica._flat_weights = replica._flat_weights[:]
+        replica._flat_weights_names = replica._flat_weights_names[:]
+        return replica
+
+
+class RNN(RNNBase):
+    r"""__init__(input_size,hidden_size,num_layers=1,nonlinearity='tanh',bias=True,batch_first=False,dropout=0.0,bidirectional=False,device=None,dtype=None)
+
+    Apply a multi-layer Elman RNN with :math:`\tanh` or :math:`\text{ReLU}`
+    non-linearity to an input sequence. For each element in the input sequence,
+    each layer computes the following function:
+
+    .. math::
+        h_t = \tanh(x_t W_{ih}^T + b_{ih} + h_{t-1}W_{hh}^T + b_{hh})
+
+    where :math:`h_t` is the hidden state at time `t`, :math:`x_t` is
+    the input at time `t`, and :math:`h_{(t-1)}` is the hidden state of the
+    previous layer at time `t-1` or the initial hidden state at time `0`.
+    If :attr:`nonlinearity` is ``'relu'``, then :math:`\text{ReLU}` is used instead of :math:`\tanh`.
+
+    .. code-block:: python
+
+        # Efficient implementation equivalent to the following with bidirectional=False
+        rnn = nn.RNN(input_size, hidden_size, num_layers)
+        params = dict(rnn.named_parameters())
+        def forward(x, hx=None, batch_first=False):
+            if batch_first:
+                x = x.transpose(0, 1)
+            seq_len, batch_size, _ = x.size()
+            if hx is None:
+                hx = torch.zeros(rnn.num_layers, batch_size, rnn.hidden_size)
+            h_t_minus_1 = hx.clone()
+            h_t = hx.clone()
+            output = []
+            for t in range(seq_len):
+                for layer in range(rnn.num_layers):
+                    input_t = x[t] if layer == 0 else h_t[layer - 1]
+                    h_t[layer] = torch.tanh(
+                        input_t @ params[f"weight_ih_l{layer}"].T
+                        + h_t_minus_1[layer] @ params[f"weight_hh_l{layer}"].T
+                        + params[f"bias_hh_l{layer}"]
+                        + params[f"bias_ih_l{layer}"]
+                    )
+                output.append(h_t[-1].clone())
+                h_t_minus_1 = h_t.clone()
+            output = torch.stack(output)
+            if batch_first:
+                output = output.transpose(0, 1)
+            return output, h_t
+
+    Args:
+        input_size: The number of expected features in the input `x`
+        hidden_size: The number of features in the hidden state `h`
+        num_layers: Number of recurrent layers. E.g., setting ``num_layers=2``
+            would mean stacking two RNNs together to form a `stacked RNN`,
+            with the second RNN taking in outputs of the first RNN and
+            computing the final results. Default: 1
+        nonlinearity: The non-linearity to use. Can be either ``'tanh'`` or ``'relu'``. Default: ``'tanh'``
+        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.
+            Default: ``True``
+        batch_first: If ``True``, then the input and output tensors are provided
+            as `(batch, seq, feature)` instead of `(seq, batch, feature)`.
+            Note that this does not apply to hidden or cell states. See the
+            Inputs/Outputs sections below for details.  Default: ``False``
+        dropout: If non-zero, introduces a `Dropout` layer on the outputs of each
+            RNN layer except the last layer, with dropout probability equal to
+            :attr:`dropout`. Default: 0
+        bidirectional: If ``True``, becomes a bidirectional RNN. Default: ``False``
+
+    Inputs: input, hx
+        * **input**: tensor of shape :math:`(L, H_{in})` for unbatched input,
+          :math:`(L, N, H_{in})` when ``batch_first=False`` or
+          :math:`(N, L, H_{in})` when ``batch_first=True`` containing the features of
+          the input sequence.  The input can also be a packed variable length sequence.
+          See :func:`torch.nn.utils.rnn.pack_padded_sequence` or
+          :func:`torch.nn.utils.rnn.pack_sequence` for details.
+        * **hx**: tensor of shape :math:`(D * \text{num\_layers}, H_{out})` for unbatched input or
+          :math:`(D * \text{num\_layers}, N, H_{out})` containing the initial hidden
+          state for the input sequence batch. Defaults to zeros if not provided.
+
+        where:
+
+        .. math::
+            \begin{aligned}
+                N ={} & \text{batch size} \\
+                L ={} & \text{sequence length} \\
+                D ={} & 2 \text{ if bidirectional=True otherwise } 1 \\
+                H_{in} ={} & \text{input\_size} \\
+                H_{out} ={} & \text{hidden\_size}
+            \end{aligned}
+
+    Outputs: output, h_n
+        * **output**: tensor of shape :math:`(L, D * H_{out})` for unbatched input,
+          :math:`(L, N, D * H_{out})` when ``batch_first=False`` or
+          :math:`(N, L, D * H_{out})` when ``batch_first=True`` containing the output features
+          `(h_t)` from the last layer of the RNN, for each `t`. If a
+          :class:`torch.nn.utils.rnn.PackedSequence` has been given as the input, the output
+          will also be a packed sequence.
+        * **h_n**: tensor of shape :math:`(D * \text{num\_layers}, H_{out})` for unbatched input or
+          :math:`(D * \text{num\_layers}, N, H_{out})` containing the final hidden state
+          for each element in the batch.
+
+    Attributes:
+        weight_ih_l[k]: the learnable input-hidden weights of the k-th layer,
+            of shape `(hidden_size, input_size)` for `k = 0`. Otherwise, the shape is
+            `(hidden_size, num_directions * hidden_size)`
+        weight_hh_l[k]: the learnable hidden-hidden weights of the k-th layer,
+            of shape `(hidden_size, hidden_size)`
+        bias_ih_l[k]: the learnable input-hidden bias of the k-th layer,
+            of shape `(hidden_size)`
+        bias_hh_l[k]: the learnable hidden-hidden bias of the k-th layer,
+            of shape `(hidden_size)`
+
+    .. note::
+        All the weights and biases are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`
+        where :math:`k = \frac{1}{\text{hidden\_size}}`
+
+    .. note::
+        For bidirectional RNNs, forward and backward are directions 0 and 1 respectively.
+        Example of splitting the output layers when ``batch_first=False``:
+        ``output.view(seq_len, batch, num_directions, hidden_size)``.
+
+    .. note::
+        ``batch_first`` argument is ignored for unbatched inputs.
+
+    .. include:: ../cudnn_rnn_determinism.rst
+
+    .. include:: ../cudnn_persistent_rnn.rst
+
+    Examples::
+
+        >>> rnn = nn.RNN(10, 20, 2)
+        >>> input = torch.randn(5, 3, 10)
+        >>> h0 = torch.randn(2, 3, 20)
+        >>> output, hn = rnn(input, h0)
+    """
+
+    @overload
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        nonlinearity: str = "tanh",
+        bias: bool = True,
+        batch_first: bool = False,
+        dropout: float = 0.0,
+        bidirectional: bool = False,
+        device=None,
+        dtype=None,
+    ) -> None: ...
+
+    @overload
+    def __init__(self, *args, **kwargs) -> None: ...
+
+    def __init__(self, *args, **kwargs):
+        if "proj_size" in kwargs:
+            raise ValueError(
+                "proj_size argument is only supported for LSTM, not RNN or GRU"
+            )
+        if len(args) > 3:
+            self.nonlinearity = args[3]
+            args = args[:3] + args[4:]
+        else:
+            self.nonlinearity = kwargs.pop("nonlinearity", "tanh")
+        if self.nonlinearity == "tanh":
+            mode = "RNN_TANH"
+        elif self.nonlinearity == "relu":
+            mode = "RNN_RELU"
+        else:
+            raise ValueError(
+                f"Unknown nonlinearity '{self.nonlinearity}'. Select from 'tanh' or 'relu'."
+            )
+        super().__init__(mode, *args, **kwargs)
+
+    @overload
+    @torch._jit_internal._overload_method  # noqa: F811
+    def forward(
+        self,
+        input: Tensor,
+        hx: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        pass
+
+    @overload
+    @torch._jit_internal._overload_method  # noqa: F811
+    def forward(
+        self,
+        input: PackedSequence,
+        hx: Tensor | None = None,
+    ) -> tuple[PackedSequence, Tensor]:
+        pass
+
+    def forward(self, input, hx=None):  # noqa: F811
+        """
+        Runs the forward pass.
+        """
+        self._update_flat_weights()
+
+        num_directions = 2 if self.bidirectional else 1
+        orig_input = input
+
+        if isinstance(orig_input, PackedSequence):
+            input, batch_sizes, sorted_indices, unsorted_indices = input
+            max_batch_size = batch_sizes[0]
+            # script() is unhappy when max_batch_size is different type in cond branches, so we duplicate
+            if hx is None:
+                hx = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    self.hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+            else:
+                # Each batch of the hidden state should match the input sequence that
+                # the user believes he/she is passing in.
+                hx = self.permute_hidden(hx, sorted_indices)
+        else:
+            batch_sizes = None
+            if input.dim() not in (2, 3):
+                raise ValueError(
+                    f"RNN: Expected input to be 2D or 3D, got {input.dim()}D tensor instead"
+                )
+            is_batched = input.dim() == 3
+            batch_dim = 0 if self.batch_first else 1
+            if not is_batched:
+                input = input.unsqueeze(batch_dim)
+                if hx is not None:
+                    if hx.dim() != 2:
+                        raise RuntimeError(
+                            f"For unbatched 2-D input, hx should also be 2-D but got {hx.dim()}-D tensor"
+                        )
+                    hx = hx.unsqueeze(1)
+            else:
+                if hx is not None and hx.dim() != 3:
+                    raise RuntimeError(
+                        f"For batched 3-D input, hx should also be 3-D but got {hx.dim()}-D tensor"
+                    )
+            max_batch_size = input.size(0) if self.batch_first else input.size(1)
+            sorted_indices = None
+            unsorted_indices = None
+            if hx is None:
+                hx = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    self.hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+            else:
+                # Each batch of the hidden state should match the input sequence that
+                # the user believes he/she is passing in.
+                hx = self.permute_hidden(hx, sorted_indices)
+
+        if hx is None:
+            raise AssertionError("hx must not be None")
+        self.check_forward_args(input, hx, batch_sizes)
+        if self.mode != "RNN_TANH" and self.mode != "RNN_RELU":
+            raise AssertionError(f"mode must be RNN_TANH or RNN_RELU, got {self.mode}")
+        if batch_sizes is None:
+            if self.mode == "RNN_TANH":
+                # pyrefly: ignore [no-matching-overload]
+                result = _VF.rnn_tanh(
+                    input,
+                    hx,
+                    self._flat_weights,  # type: ignore[arg-type]
+                    self.bias,
+                    self.num_layers,
+                    self.dropout,
+                    self.training,
+                    self.bidirectional,
+                    self.batch_first,
+                )
+            else:
+                # pyrefly: ignore [no-matching-overload]
+                result = _VF.rnn_relu(
+                    input,
+                    hx,
+                    self._flat_weights,  # type: ignore[arg-type]
+                    self.bias,
+                    self.num_layers,
+                    self.dropout,
+                    self.training,
+                    self.bidirectional,
+                    self.batch_first,
+                )
+        else:
+            if self.mode == "RNN_TANH":
+                # pyrefly: ignore [no-matching-overload]
+                result = _VF.rnn_tanh(
+                    input,
+                    batch_sizes,
+                    hx,
+                    self._flat_weights,  # type: ignore[arg-type]
+                    self.bias,
+                    self.num_layers,
+                    self.dropout,
+                    self.training,
+                    self.bidirectional,
+                )
+            else:
+                # pyrefly: ignore [no-matching-overload]
+                result = _VF.rnn_relu(
+                    input,
+                    batch_sizes,
+                    hx,
+                    self._flat_weights,  # type: ignore[arg-type]
+                    self.bias,
+                    self.num_layers,
+                    self.dropout,
+                    self.training,
+                    self.bidirectional,
+                )
+
+        output = result[0]
+        hidden = result[1]
+
+        if isinstance(orig_input, PackedSequence):
+            output_packed = PackedSequence(
+                output,
+                batch_sizes,
+                sorted_indices,
+                unsorted_indices,
+            )
+            return output_packed, self.permute_hidden(hidden, unsorted_indices)
+
+        if not is_batched:  # type: ignore[possibly-undefined]
+            output = output.squeeze(batch_dim)  # type: ignore[possibly-undefined]
+            hidden = hidden.squeeze(1)
+
+        return output, self.permute_hidden(hidden, unsorted_indices)
+
+
+# XXX: LSTM and GRU implementation is different from RNNBase, this is because:
+# 1. we want to support nn.LSTM and nn.GRU in TorchScript and TorchScript in
+#    its current state could not support the python Union Type or Any Type
+# 2. TorchScript static typing does not allow a Function or Callable type in
+#    Dict values, so we have to separately call _VF instead of using _rnn_impls
+# 3. This is temporary only and in the transition state that we want to make it
+#    on time for the release
+#
+# More discussion details in https://github.com/pytorch/pytorch/pull/23266
+#
+# TODO: remove the overriding implementations for LSTM and GRU when TorchScript
+# support expressing these two modules generally.
+
+
+class LSTM(RNNBase):
+    r"""__init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,proj_size=0,device=None,dtype=None)
+
+    Apply a multi-layer long short-term memory (LSTM) RNN to an input sequence.
+    For each element in the input sequence, each layer computes the following
+    function:
+
+    .. math::
+        \begin{array}{ll} \\
+            i_t = \sigma(W_{ii} x_t + b_{ii} + W_{hi} h_{t-1} + b_{hi}) \\
+            f_t = \sigma(W_{if} x_t + b_{if} + W_{hf} h_{t-1} + b_{hf}) \\
+            g_t = \tanh(W_{ig} x_t + b_{ig} + W_{hg} h_{t-1} + b_{hg}) \\
+            o_t = \sigma(W_{io} x_t + b_{io} + W_{ho} h_{t-1} + b_{ho}) \\
+            c_t = f_t \odot c_{t-1} + i_t \odot g_t \\
+            h_t = o_t \odot \tanh(c_t) \\
+        \end{array}
+
+    where :math:`h_t` is the hidden state at time `t`, :math:`c_t` is the cell
+    state at time `t`, :math:`x_t` is the input at time `t`, :math:`h_{t-1}`
+    is the hidden state of the layer at time `t-1` or the initial hidden
+    state at time `0`, and :math:`i_t`, :math:`f_t`, :math:`g_t`,
+    :math:`o_t` are the input, forget, cell, and output gates, respectively.
+    :math:`\sigma` is the sigmoid function, and :math:`\odot` is the Hadamard product.
+
+    In a multilayer LSTM, the input :math:`x^{(l)}_t` of the :math:`l` -th layer
+    (:math:`l \ge 2`) is the hidden state :math:`h^{(l-1)}_t` of the previous layer multiplied by
+    dropout :math:`\delta^{(l-1)}_t` where each :math:`\delta^{(l-1)}_t` is a Bernoulli random
+    variable which is :math:`0` with probability :attr:`dropout`.
+
+    If ``proj_size > 0`` is specified, LSTM with projections will be used. This changes
+    the LSTM cell in the following way. First, the dimension of :math:`h_t` will be changed from
+    ``hidden_size`` to ``proj_size`` (dimensions of :math:`W_{hi}` will be changed accordingly).
+    Second, the output hidden state of each layer will be multiplied by a learnable projection
+    matrix: :math:`h_t = W_{hr}h_t`. Note that as a consequence of this, the output
+    of LSTM network will be of different shape as well. See Inputs/Outputs sections below for exact
+    dimensions of all variables. You can find more details in https://arxiv.org/abs/1402.1128.
+
+    Args:
+        input_size: The number of expected features in the input `x`
+        hidden_size: The number of features in the hidden state `h`
+        num_layers: Number of recurrent layers. E.g., setting ``num_layers=2``
+            would mean stacking two LSTMs together to form a `stacked LSTM`,
+            with the second LSTM taking in outputs of the first LSTM and
+            computing the final results. Default: 1
+        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.
+            Default: ``True``
+        batch_first: If ``True``, then the input and output tensors are provided
+            as `(batch, seq, feature)` instead of `(seq, batch, feature)`.
+            Note that this does not apply to hidden or cell states. See the
+            Inputs/Outputs sections below for details.  Default: ``False``
+        dropout: If non-zero, introduces a `Dropout` layer on the outputs of each
+            LSTM layer except the last layer, with dropout probability equal to
+            :attr:`dropout`. Default: 0
+        bidirectional: If ``True``, becomes a bidirectional LSTM. Default: ``False``
+        proj_size: If ``> 0``, will use LSTM with projections of corresponding size. Default: 0
+
+    Inputs: input, (h_0, c_0)
+        * **input**: tensor of shape :math:`(L, H_{in})` for unbatched input,
+          :math:`(L, N, H_{in})` when ``batch_first=False`` or
+          :math:`(N, L, H_{in})` when ``batch_first=True`` containing the features of
+          the input sequence.  The input can also be a packed variable length sequence.
+          See :func:`torch.nn.utils.rnn.pack_padded_sequence` or
+          :func:`torch.nn.utils.rnn.pack_sequence` for details.
+        * **h_0**: tensor of shape :math:`(D * \text{num\_layers}, H_{out})` for unbatched input or
+          :math:`(D * \text{num\_layers}, N, H_{out})` containing the
+          initial hidden state for each element in the input sequence.
+          Defaults to zeros if (h_0, c_0) is not provided.
+        * **c_0**: tensor of shape :math:`(D * \text{num\_layers}, H_{cell})` for unbatched input or
+          :math:`(D * \text{num\_layers}, N, H_{cell})` containing the
+          initial cell state for each element in the input sequence.
+          Defaults to zeros if (h_0, c_0) is not provided.
+
+        where:
+
+        .. math::
+            \begin{aligned}
+                N ={} & \text{batch size} \\
+                L ={} & \text{sequence length} \\
+                D ={} & 2 \text{ if bidirectional=True otherwise } 1 \\
+                H_{in} ={} & \text{input\_size} \\
+                H_{cell} ={} & \text{hidden\_size} \\
+                H_{out} ={} & \text{proj\_size if } \text{proj\_size}>0 \text{ otherwise hidden\_size} \\
+            \end{aligned}
+
+    Outputs: output, (h_n, c_n)
+        * **output**: tensor of shape :math:`(L, D * H_{out})` for unbatched input,
+          :math:`(L, N, D * H_{out})` when ``batch_first=False`` or
+          :math:`(N, L, D * H_{out})` when ``batch_first=True`` containing the output features
+          `(h_t)` from the last layer of the LSTM, for each `t`. If a
+          :class:`torch.nn.utils.rnn.PackedSequence` has been given as the input, the output
+          will also be a packed sequence. When ``bidirectional=True``, `output` will contain
+          a concatenation of the forward and reverse hidden states at each time step in the sequence.
+        * **h_n**: tensor of shape :math:`(D * \text{num\_layers}, H_{out})` for unbatched input or
+          :math:`(D * \text{num\_layers}, N, H_{out})` containing the
+          final hidden state for each element in the sequence. When ``bidirectional=True``,
+          `h_n` will contain a concatenation of the final forward and reverse hidden states, respectively.
+        * **c_n**: tensor of shape :math:`(D * \text{num\_layers}, H_{cell})` for unbatched input or
+          :math:`(D * \text{num\_layers}, N, H_{cell})` containing the
+          final cell state for each element in the sequence. When ``bidirectional=True``,
+          `c_n` will contain a concatenation of the final forward and reverse cell states, respectively.
+
+    Attributes:
+        weight_ih_l[k] : the learnable input-hidden weights of the :math:`\text{k}^{th}` layer
+            `(W_ii|W_if|W_ig|W_io)`, of shape `(4*hidden_size, input_size)` for `k = 0`.
+            Otherwise, the shape is `(4*hidden_size, num_directions * hidden_size)`. If
+            ``proj_size > 0`` was specified, the shape will be
+            `(4*hidden_size, num_directions * proj_size)` for `k > 0`
+        weight_hh_l[k] : the learnable hidden-hidden weights of the :math:`\text{k}^{th}` layer
+            `(W_hi|W_hf|W_hg|W_ho)`, of shape `(4*hidden_size, hidden_size)`. If ``proj_size > 0``
+            was specified, the shape will be `(4*hidden_size, proj_size)`.
+        bias_ih_l[k] : the learnable input-hidden bias of the :math:`\text{k}^{th}` layer
+            `(b_ii|b_if|b_ig|b_io)`, of shape `(4*hidden_size)`
+        bias_hh_l[k] : the learnable hidden-hidden bias of the :math:`\text{k}^{th}` layer
+            `(b_hi|b_hf|b_hg|b_ho)`, of shape `(4*hidden_size)`
+        weight_hr_l[k] : the learnable projection weights of the :math:`\text{k}^{th}` layer
+            of shape `(proj_size, hidden_size)`. Only present when ``proj_size > 0`` was
+            specified.
+        weight_ih_l[k]_reverse: Analogous to `weight_ih_l[k]` for the reverse direction.
+            Only present when ``bidirectional=True``.
+        weight_hh_l[k]_reverse:  Analogous to `weight_hh_l[k]` for the reverse direction.
+            Only present when ``bidirectional=True``.
+        bias_ih_l[k]_reverse:  Analogous to `bias_ih_l[k]` for the reverse direction.
+            Only present when ``bidirectional=True``.
+        bias_hh_l[k]_reverse:  Analogous to `bias_hh_l[k]` for the reverse direction.
+            Only present when ``bidirectional=True``.
+        weight_hr_l[k]_reverse:  Analogous to `weight_hr_l[k]` for the reverse direction.
+            Only present when ``bidirectional=True`` and ``proj_size > 0`` was specified.
+
+    .. note::
+        All the weights and biases are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`
+        where :math:`k = \frac{1}{\text{hidden\_size}}`
+
+    .. note::
+        For bidirectional LSTMs, forward and backward are directions 0 and 1 respectively.
+        Example of splitting the output layers when ``batch_first=False``:
+        ``output.view(seq_len, batch, num_directions, hidden_size)``.
+
+    .. note::
+        For bidirectional LSTMs, `h_n` is not equivalent to the last element of `output`; the
+        former contains the final forward and reverse hidden states, while the latter contains the
+        final forward hidden state and the initial reverse hidden state.
+
+    .. note::
+        ``batch_first`` argument is ignored for unbatched inputs.
+
+    .. note::
+        ``proj_size`` should be smaller than ``hidden_size``.
+
+    .. include:: ../cudnn_rnn_determinism.rst
+
+    .. include:: ../cudnn_persistent_rnn.rst
+
+    Examples::
+
+        >>> rnn = nn.LSTM(10, 20, 2)
+        >>> input = torch.randn(5, 3, 10)
+        >>> h0 = torch.randn(2, 3, 20)
+        >>> c0 = torch.randn(2, 3, 20)
+        >>> output, (hn, cn) = rnn(input, (h0, c0))
+    """
+
+    @overload
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        bias: bool = True,
+        batch_first: bool = False,
+        dropout: float = 0.0,
+        bidirectional: bool = False,
+        proj_size: int = 0,
+        device=None,
+        dtype=None,
+    ) -> None: ...
+
+    @overload
+    def __init__(self, *args, **kwargs) -> None: ...
+
+    def __init__(self, *args, **kwargs):
+        super().__init__("LSTM", *args, **kwargs)
+
+    def get_expected_cell_size(
+        self, input: Tensor, batch_sizes: Tensor | None
+    ) -> tuple[int, int, int]:
+        if batch_sizes is not None:
+            mini_batch = int(batch_sizes[0])
+        else:
+            mini_batch = input.size(0) if self.batch_first else input.size(1)
+        num_directions = 2 if self.bidirectional else 1
+        expected_hidden_size = (
+            self.num_layers * num_directions,
+            mini_batch,
+            self.hidden_size,
+        )
+        return expected_hidden_size
+
+    # In the future, we should prevent mypy from applying contravariance rules here.
+    # See torch/nn/modules/module.py::_forward_unimplemented
+    # pyrefly: ignore [bad-override]
+    def check_forward_args(
+        self,
+        input: Tensor,
+        hidden: tuple[Tensor, Tensor],  # type: ignore[override]
+        batch_sizes: Tensor | None,
+    ) -> None:
+        self.check_input(input, batch_sizes)
+        self.check_hidden_size(
+            hidden[0],
+            self.get_expected_hidden_size(input, batch_sizes),
+            "Expected hidden[0] size {}, got {}",
+        )
+        self.check_hidden_size(
+            hidden[1],
+            self.get_expected_cell_size(input, batch_sizes),
+            "Expected hidden[1] size {}, got {}",
+        )
+
+    # Same as above, see torch/nn/modules/module.py::_forward_unimplemented
+    def permute_hidden(  # type: ignore[override]
+        self,
+        hx: tuple[Tensor, Tensor],
+        permutation: Tensor | None,
+    ) -> tuple[Tensor, Tensor]:
+        if permutation is None:
+            return hx
+        return _apply_permutation(hx[0], permutation), _apply_permutation(
+            hx[1], permutation
+        )
+
+    # Same as above, see torch/nn/modules/module.py::_forward_unimplemented
+    @overload  # type: ignore[override]
+    @torch._jit_internal._overload_method  # noqa: F811
+    def forward(
+        self,
+        input: Tensor,
+        hx: tuple[Tensor, Tensor] | None = None,
+    ) -> tuple[Tensor, tuple[Tensor, Tensor]]:  # noqa: F811
+        pass
+
+    # Same as above, see torch/nn/modules/module.py::_forward_unimplemented
+    @overload
+    @torch._jit_internal._overload_method  # noqa: F811
+    def forward(
+        self,
+        input: PackedSequence,
+        hx: tuple[Tensor, Tensor] | None = None,
+    ) -> tuple[PackedSequence, tuple[Tensor, Tensor]]:  # noqa: F811
+        pass
+
+    def forward(self, input, hx=None):  # noqa: F811
+        self._update_flat_weights()
+
+        orig_input = input
+        # xxx: isinstance check needs to be in conditional for TorchScript to compile
+        batch_sizes = None
+        num_directions = 2 if self.bidirectional else 1
+        real_hidden_size = self.proj_size if self.proj_size > 0 else self.hidden_size
+        if isinstance(orig_input, PackedSequence):
+            input, batch_sizes, sorted_indices, unsorted_indices = input
+            max_batch_size = batch_sizes[0]
+            if hx is None:
+                h_zeros = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    real_hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+                c_zeros = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    self.hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+                hx = (h_zeros, c_zeros)
+            else:
+                # Each batch of the hidden state should match the input sequence that
+                # the user believes he/she is passing in.
+                hx = self.permute_hidden(hx, sorted_indices)
+        else:
+            if input.dim() not in (2, 3):
+                raise ValueError(
+                    f"LSTM: Expected input to be 2D or 3D, got {input.dim()}D instead"
+                )
+            is_batched = input.dim() == 3
+            batch_dim = 0 if self.batch_first else 1
+            if not is_batched:
+                input = input.unsqueeze(batch_dim)
+            max_batch_size = input.size(0) if self.batch_first else input.size(1)
+            sorted_indices = None
+            unsorted_indices = None
+            if hx is None:
+                h_zeros = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    real_hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+                c_zeros = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    self.hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+                hx = (h_zeros, c_zeros)
+                self.check_forward_args(input, hx, batch_sizes)
+            else:
+                if is_batched:
+                    if hx[0].dim() != 3 or hx[1].dim() != 3:
+                        msg = (
+                            "For batched 3-D input, hx and cx should "
+                            f"also be 3-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors"
+                        )
+                        raise RuntimeError(msg)
+                else:
+                    if hx[0].dim() != 2 or hx[1].dim() != 2:
+                        msg = (
+                            "For unbatched 2-D input, hx and cx should "
+                            f"also be 2-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors"
+                        )
+                        raise RuntimeError(msg)
+                    hx = (hx[0].unsqueeze(1), hx[1].unsqueeze(1))
+                # Each batch of the hidden state should match the input sequence that
+                # the user believes he/she is passing in.
+                self.check_forward_args(input, hx, batch_sizes)
+                hx = self.permute_hidden(hx, sorted_indices)
+
+        if batch_sizes is None:
+            # pyrefly: ignore [no-matching-overload]
+            result = _VF.lstm(
+                input,
+                hx,
+                self._flat_weights,  # type: ignore[arg-type]
+                self.bias,
+                self.num_layers,
+                self.dropout,
+                self.training,
+                self.bidirectional,
+                self.batch_first,
+            )
+        else:
+            # pyrefly: ignore [no-matching-overload]
+            result = _VF.lstm(
+                input,
+                batch_sizes,
+                hx,
+                self._flat_weights,  # type: ignore[arg-type]
+                self.bias,
+                self.num_layers,
+                self.dropout,
+                self.training,
+                self.bidirectional,
+            )
+        output = result[0]
+        hidden = result[1:]
+        # xxx: isinstance check needs to be in conditional for TorchScript to compile
+        if isinstance(orig_input, PackedSequence):
+            output_packed = PackedSequence(
+                output,
+                batch_sizes,
+                sorted_indices,
+                unsorted_indices,
+            )
+            return output_packed, self.permute_hidden(hidden, unsorted_indices)
+        else:
+            if not is_batched:  # type: ignore[possibly-undefined]
+                output = output.squeeze(batch_dim)  # type: ignore[possibly-undefined]
+                hidden = (hidden[0].squeeze(1), hidden[1].squeeze(1))
+            return output, self.permute_hidden(hidden, unsorted_indices)
+
+
+class GRU(RNNBase):
+    r"""__init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,device=None,dtype=None)
+
+    Apply a multi-layer gated recurrent unit (GRU) RNN to an input sequence.
+    For each element in the input sequence, each layer computes the following
+    function:
+
+    .. math::
+        \begin{array}{ll}
+            r_t = \sigma(W_{ir} x_t + b_{ir} + W_{hr} h_{(t-1)} + b_{hr}) \\
+            z_t = \sigma(W_{iz} x_t + b_{iz} + W_{hz} h_{(t-1)} + b_{hz}) \\
+            n_t = \tanh(W_{in} x_t + b_{in} + r_t \odot (W_{hn} h_{(t-1)}+ b_{hn})) \\
+            h_t = (1 - z_t) \odot n_t + z_t \odot h_{(t-1)}
+        \end{array}
+
+    where :math:`h_t` is the hidden state at time `t`, :math:`x_t` is the input
+    at time `t`, :math:`h_{(t-1)}` is the hidden state of the layer
+    at time `t-1` or the initial hidden state at time `0`, and :math:`r_t`,
+    :math:`z_t`, :math:`n_t` are the reset, update, and new gates, respectively.
+    :math:`\sigma` is the sigmoid function, and :math:`\odot` is the Hadamard product.
+
+    In a multilayer GRU, the input :math:`x^{(l)}_t` of the :math:`l` -th layer
+    (:math:`l \ge 2`) is the hidden state :math:`h^{(l-1)}_t` of the previous layer multiplied by
+    dropout :math:`\delta^{(l-1)}_t` where each :math:`\delta^{(l-1)}_t` is a Bernoulli random
+    variable which is :math:`0` with probability :attr:`dropout`.
+
+    Args:
+        input_size: The number of expected features in the input `x`
+        hidden_size: The number of features in the hidden state `h`
+        num_layers: Number of recurrent layers. E.g., setting ``num_layers=2``
+            would mean stacking two GRUs together to form a `stacked GRU`,
+            with the second GRU taking in outputs of the first GRU and
+            computing the final results. Default: 1
+        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.
+            Default: ``True``
+        batch_first: If ``True``, then the input and output tensors are provided
+            as `(batch, seq, feature)` instead of `(seq, batch, feature)`.
+            Note that this does not apply to hidden or cell states. See the
+            Inputs/Outputs sections below for details.  Default: ``False``
+        dropout: If non-zero, introduces a `Dropout` layer on the outputs of each
+            GRU layer except the last layer, with dropout probability equal to
+            :attr:`dropout`. Default: 0
+        bidirectional: If ``True``, becomes a bidirectional GRU. Default: ``False``
+
+    Inputs: input, h_0
+        * **input**: tensor of shape :math:`(L, H_{in})` for unbatched input,
+          :math:`(L, N, H_{in})` when ``batch_first=False`` or
+          :math:`(N, L, H_{in})` when ``batch_first=True`` containing the features of
+          the input sequence.  The input can also be a packed variable length sequence.
+          See :func:`torch.nn.utils.rnn.pack_padded_sequence` or
+          :func:`torch.nn.utils.rnn.pack_sequence` for details.
+        * **h_0**: tensor of shape :math:`(D * \text{num\_layers}, H_{out})` or
+          :math:`(D * \text{num\_layers}, N, H_{out})`
+          containing the initial hidden state for the input sequence. Defaults to zeros if not provided.
+
+        where:
+
+        .. math::
+            \begin{aligned}
+                N ={} & \text{batch size} \\
+                L ={} & \text{sequence length} \\
+                D ={} & 2 \text{ if bidirectional=True otherwise } 1 \\
+                H_{in} ={} & \text{input\_size} \\
+                H_{out} ={} & \text{hidden\_size}
+            \end{aligned}
+
+    Outputs: output, h_n
+        * **output**: tensor of shape :math:`(L, D * H_{out})` for unbatched input,
+          :math:`(L, N, D * H_{out})` when ``batch_first=False`` or
+          :math:`(N, L, D * H_{out})` when ``batch_first=True`` containing the output features
+          `(h_t)` from the last layer of the GRU, for each `t`. If a
+          :class:`torch.nn.utils.rnn.PackedSequence` has been given as the input, the output
+          will also be a packed sequence.
+        * **h_n**: tensor of shape :math:`(D * \text{num\_layers}, H_{out})` or
+          :math:`(D * \text{num\_layers}, N, H_{out})` containing the final hidden state
+          for the input sequence.
+
+    Attributes:
+        weight_ih_l[k] : the learnable input-hidden weights of the :math:`\text{k}^{th}` layer
+            (W_ir|W_iz|W_in), of shape `(3*hidden_size, input_size)` for `k = 0`.
+            Otherwise, the shape is `(3*hidden_size, num_directions * hidden_size)`
+        weight_hh_l[k] : the learnable hidden-hidden weights of the :math:`\text{k}^{th}` layer
+            (W_hr|W_hz|W_hn), of shape `(3*hidden_size, hidden_size)`
+        bias_ih_l[k] : the learnable input-hidden bias of the :math:`\text{k}^{th}` layer
+            (b_ir|b_iz|b_in), of shape `(3*hidden_size)`
+        bias_hh_l[k] : the learnable hidden-hidden bias of the :math:`\text{k}^{th}` layer
+            (b_hr|b_hz|b_hn), of shape `(3*hidden_size)`
+
+    .. note::
+        All the weights and biases are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`
+        where :math:`k = \frac{1}{\text{hidden\_size}}`
+
+    .. note::
+        For bidirectional GRUs, forward and backward are directions 0 and 1 respectively.
+        Example of splitting the output layers when ``batch_first=False``:
+        ``output.view(seq_len, batch, num_directions, hidden_size)``.
+
+    .. note::
+        ``batch_first`` argument is ignored for unbatched inputs.
+
+    .. note::
+        The calculation of new gate :math:`n_t` subtly differs from the original paper and other frameworks.
+        In the original implementation, the Hadamard product :math:`(\odot)` between :math:`r_t` and the
+        previous hidden state :math:`h_{(t-1)}` is done before the multiplication with the weight matrix
+        `W` and addition of bias:
+
+        .. math::
+            \begin{aligned}
+                n_t = \tanh(W_{in} x_t + b_{in} + W_{hn} ( r_t \odot h_{(t-1)} ) + b_{hn})
+            \end{aligned}
+
+        This is in contrast to PyTorch implementation, which is done after :math:`W_{hn} h_{(t-1)}`
+
+        .. math::
+            \begin{aligned}
+                n_t = \tanh(W_{in} x_t + b_{in} + r_t \odot (W_{hn} h_{(t-1)}+ b_{hn}))
+            \end{aligned}
+
+        This implementation differs on purpose for efficiency.
+
+    .. include:: ../cudnn_persistent_rnn.rst
+
+    Examples::
+
+        >>> rnn = nn.GRU(10, 20, 2)
+        >>> input = torch.randn(5, 3, 10)
+        >>> h0 = torch.randn(2, 3, 20)
+        >>> output, hn = rnn(input, h0)
+    """
+
+    @overload
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        bias: bool = True,
+        batch_first: bool = False,
+        dropout: float = 0.0,
+        bidirectional: bool = False,
+        device=None,
+        dtype=None,
+    ) -> None: ...
+
+    @overload
+    def __init__(self, *args, **kwargs) -> None: ...
+
+    def __init__(self, *args, **kwargs):
+        if "proj_size" in kwargs:
+            raise ValueError(
+                "proj_size argument is only supported for LSTM, not RNN or GRU"
+            )
+        super().__init__("GRU", *args, **kwargs)
+
+    @overload  # type: ignore[override]
+    @torch._jit_internal._overload_method  # noqa: F811
+    def forward(
+        self,
+        input: Tensor,
+        hx: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor]:  # noqa: F811
+        pass
+
+    @overload
+    @torch._jit_internal._overload_method  # noqa: F811
+    def forward(
+        self,
+        input: PackedSequence,
+        hx: Tensor | None = None,
+    ) -> tuple[PackedSequence, Tensor]:  # noqa: F811
+        pass
+
+    def forward(self, input, hx=None):  # noqa: F811
+        self._update_flat_weights()
+
+        orig_input = input
+        # xxx: isinstance check needs to be in conditional for TorchScript to compile
+        if isinstance(orig_input, PackedSequence):
+            input, batch_sizes, sorted_indices, unsorted_indices = input
+            max_batch_size = batch_sizes[0]
+            if hx is None:
+                num_directions = 2 if self.bidirectional else 1
+                hx = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    self.hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+            else:
+                # Each batch of the hidden state should match the input sequence that
+                # the user believes he/she is passing in.
+                hx = self.permute_hidden(hx, sorted_indices)
+        else:
+            batch_sizes = None
+            if input.dim() not in (2, 3):
+                raise ValueError(
+                    f"GRU: Expected input to be 2D or 3D, got {input.dim()}D instead"
+                )
+            is_batched = input.dim() == 3
+            batch_dim = 0 if self.batch_first else 1
+            if not is_batched:
+                input = input.unsqueeze(batch_dim)
+                if hx is not None:
+                    if hx.dim() != 2:
+                        raise RuntimeError(
+                            f"For unbatched 2-D input, hx should also be 2-D but got {hx.dim()}-D tensor"
+                        )
+                    hx = hx.unsqueeze(1)
+            else:
+                if hx is not None and hx.dim() != 3:
+                    raise RuntimeError(
+                        f"For batched 3-D input, hx should also be 3-D but got {hx.dim()}-D tensor"
+                    )
+            max_batch_size = input.size(0) if self.batch_first else input.size(1)
+            sorted_indices = None
+            unsorted_indices = None
+            if hx is None:
+                num_directions = 2 if self.bidirectional else 1
+                hx = torch.zeros(
+                    self.num_layers * num_directions,
+                    max_batch_size,
+                    self.hidden_size,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
+            else:
+                # Each batch of the hidden state should match the input sequence that
+                # the user believes he/she is passing in.
+                hx = self.permute_hidden(hx, sorted_indices)
+
+        self.check_forward_args(input, hx, batch_sizes)
+        if batch_sizes is None:
+            # pyrefly: ignore [no-matching-overload]
+            result = _VF.gru(
+                input,
+                hx,
+                self._flat_weights,  # type: ignore[arg-type]
+                self.bias,
+                self.num_layers,
+                self.dropout,
+                self.training,
+                self.bidirectional,
+                self.batch_first,
+            )
+        else:
+            # pyrefly: ignore [no-matching-overload]
+            result = _VF.gru(
+                input,
+                batch_sizes,
+                hx,
+                self._flat_weights,  # type: ignore[arg-type]
+                self.bias,
+                self.num_layers,
+                self.dropout,
+                self.training,
+                self.bidirectional,
+            )
+        output = result[0]
+        hidden = result[1]
+
+        # xxx: isinstance check needs to be in conditional for TorchScript to compile
+        if isinstance(orig_input, PackedSequence):
+            output_packed = PackedSequence(
+                output,
+                batch_sizes,
+                sorted_indices,
+                unsorted_indices,
+            )
+            return output_packed, self.permute_hidden(hidden, unsorted_indices)
+        else:
+            if not is_batched:  # type: ignore[possibly-undefined]
+                output = output.squeeze(batch_dim)  # type: ignore[possibly-undefined]
+                hidden = hidden.squeeze(1)
+
+            return output, self.permute_hidden(hidden, unsorted_indices)
+
+
+class RNNCellBase(Module):
+    __constants__ = ["input_size", "hidden_size", "bias"]
+
+    input_size: int
+    hidden_size: int
+    bias: bool
+    weight_ih: Tensor
+    weight_hh: Tensor
+    # WARNING: bias_ih and bias_hh purposely not defined here.
+    # See https://github.com/pytorch/pytorch/issues/39670
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool,
+        num_chunks: int,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.bias = bias
+        self.weight_ih = Parameter(
+            torch.empty((num_chunks * hidden_size, input_size), **factory_kwargs)
+        )
+        self.weight_hh = Parameter(
+            torch.empty((num_chunks * hidden_size, hidden_size), **factory_kwargs)
+        )
+        if bias:
+            self.bias_ih = Parameter(
+                torch.empty(num_chunks * hidden_size, **factory_kwargs)
+            )
+            self.bias_hh = Parameter(
+                torch.empty(num_chunks * hidden_size, **factory_kwargs)
+            )
+        else:
+            self.register_parameter("bias_ih", None)
+            self.register_parameter("bias_hh", None)
+
+        self.reset_parameters()
+
+    def extra_repr(self) -> str:
+        s = "{input_size}, {hidden_size}"
+        if "bias" in self.__dict__ and self.bias is not True:
+            s += ", bias={bias}"
+        if "nonlinearity" in self.__dict__ and self.nonlinearity != "tanh":
+            s += ", nonlinearity={nonlinearity}"
+        return s.format(**self.__dict__)
+
+    def reset_parameters(self) -> None:
+        stdv = 1.0 / math.sqrt(self.hidden_size) if self.hidden_size > 0 else 0
+        for weight in self.parameters():
+            init.uniform_(weight, -stdv, stdv)
+
+
+class RNNCell(RNNCellBase):
+    r"""An Elman RNN cell with tanh or ReLU non-linearity.
+
+    .. math::
+
+        h' = \tanh(W_{ih} x + b_{ih}  +  W_{hh} h + b_{hh})
+
+    If :attr:`nonlinearity` is `'relu'`, then ReLU is used in place of tanh.
+
+    Args:
+        input_size: The number of expected features in the input `x`
+        hidden_size: The number of features in the hidden state `h`
+        bias: If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.
+            Default: ``True``
+        nonlinearity: The non-linearity to use. Can be either ``'tanh'`` or ``'relu'``. Default: ``'tanh'``
+
+    Inputs: input, hidden
+        - **input**: tensor containing input features
+        - **hidden**: tensor containing the initial hidden state
+          Defaults to zero if not provided.
+
+    Outputs: h'
+        - **h'** of shape `(batch, hidden_size)`: tensor containing the next hidden state
+          for each element in the batch
+
+    Shape:
+        - input: :math:`(N, H_{in})` or :math:`(H_{in})` tensor containing input features where
+          :math:`H_{in}` = `input_size`.
+        - hidden: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the initial hidden
+          state where :math:`H_{out}` = `hidden_size`. Defaults to zero if not provided.
+        - output: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the next hidden state.
+
+    Attributes:
+        weight_ih: the learnable input-hidden weights, of shape
+            `(hidden_size, input_size)`
+        weight_hh: the learnable hidden-hidden weights, of shape
+            `(hidden_size, hidden_size)`
+        bias_ih: the learnable input-hidden bias, of shape `(hidden_size)`
+        bias_hh: the learnable hidden-hidden bias, of shape `(hidden_size)`
+
+    .. note::
+        All the weights and biases are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`
+        where :math:`k = \frac{1}{\text{hidden\_size}}`
+
+    Examples::
+
+        >>> rnn = nn.RNNCell(10, 20)
+        >>> input = torch.randn(6, 3, 10)
+        >>> hx = torch.randn(3, 20)
+        >>> output = []
+        >>> for i in range(6):
+        ...     hx = rnn(input[i], hx)
+        ...     output.append(hx)
+    """
+
+    __constants__ = ["input_size", "hidden_size", "bias", "nonlinearity"]
+    nonlinearity: str
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+        nonlinearity: str = "tanh",
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(input_size, hidden_size, bias, num_chunks=1, **factory_kwargs)
+        self.nonlinearity = nonlinearity
+
+    def forward(self, input: Tensor, hx: Tensor | None = None) -> Tensor:
+        if input.dim() not in (1, 2):
+            raise ValueError(
+                f"RNNCell: Expected input to be 1D or 2D, got {input.dim()}D instead"
+            )
+        if hx is not None and hx.dim() not in (1, 2):
+            raise ValueError(
+                f"RNNCell: Expected hidden to be 1D or 2D, got {hx.dim()}D instead"
+            )
+        is_batched = input.dim() == 2
+        if not is_batched:
+            input = input.unsqueeze(0)
+
+        if hx is None:
+            hx = torch.zeros(
+                input.size(0), self.hidden_size, dtype=input.dtype, device=input.device
+            )
+        else:
+            hx = hx.unsqueeze(0) if not is_batched else hx
+
+        if self.nonlinearity == "tanh":
+            ret = _VF.rnn_tanh_cell(
+                input,
+                hx,
+                self.weight_ih,
+                self.weight_hh,
+                self.bias_ih,
+                self.bias_hh,
+            )
+        elif self.nonlinearity == "relu":
+            ret = _VF.rnn_relu_cell(
+                input,
+                hx,
+                self.weight_ih,
+                self.weight_hh,
+                self.bias_ih,
+                self.bias_hh,
+            )
+        else:
+            ret = input  # TODO: remove when jit supports exception flow
+            raise RuntimeError(f"Unknown nonlinearity: {self.nonlinearity}")
+
+        if not is_batched:
+            ret = ret.squeeze(0)
+
+        return ret
+
+
+class LSTMCell(RNNCellBase):
+    r"""A long short-term memory (LSTM) cell.
+
+    .. math::
+
+        \begin{array}{ll}
+        i = \sigma(W_{ii} x + b_{ii} + W_{hi} h + b_{hi}) \\
+        f = \sigma(W_{if} x + b_{if} + W_{hf} h + b_{hf}) \\
+        g = \tanh(W_{ig} x + b_{ig} + W_{hg} h + b_{hg}) \\
+        o = \sigma(W_{io} x + b_{io} + W_{ho} h + b_{ho}) \\
+        c' = f \odot c + i \odot g \\
+        h' = o \odot \tanh(c') \\
+        \end{array}
+
+    where :math:`\sigma` is the sigmoid function, and :math:`\odot` is the Hadamard product.
+
+    Args:
+        input_size: The number of expected features in the input `x`
+        hidden_size: The number of features in the hidden state `h`
+        bias: If ``False``, then the layer does not use bias weights `b_ih` and
+            `b_hh`. Default: ``True``
+
+    Inputs: input, (h_0, c_0)
+        - **input** of shape `(batch, input_size)` or `(input_size)`: tensor containing input features
+        - **h_0** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the initial hidden state
+        - **c_0** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the initial cell state
+
+          If `(h_0, c_0)` is not provided, both **h_0** and **c_0** default to zero.
+
+    Outputs: (h_1, c_1)
+        - **h_1** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the next hidden state
+        - **c_1** of shape `(batch, hidden_size)` or `(hidden_size)`: tensor containing the next cell state
+
+    Attributes:
+        weight_ih: the learnable input-hidden weights, of shape
+            `(4*hidden_size, input_size)`
+        weight_hh: the learnable hidden-hidden weights, of shape
+            `(4*hidden_size, hidden_size)`
+        bias_ih: the learnable input-hidden bias, of shape `(4*hidden_size)`
+        bias_hh: the learnable hidden-hidden bias, of shape `(4*hidden_size)`
+
+    .. note::
+        All the weights and biases are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`
+        where :math:`k = \frac{1}{\text{hidden\_size}}`
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    Examples::
+
+        >>> rnn = nn.LSTMCell(10, 20)  # (input_size, hidden_size)
+        >>> input = torch.randn(2, 3, 10)  # (time_steps, batch, input_size)
+        >>> hx = torch.randn(3, 20)  # (batch, hidden_size)
+        >>> cx = torch.randn(3, 20)
+        >>> output = []
+        >>> for i in range(input.size()[0]):
+        ...     hx, cx = rnn(input[i], (hx, cx))
+        ...     output.append(hx)
+        >>> output = torch.stack(output, dim=0)
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(input_size, hidden_size, bias, num_chunks=4, **factory_kwargs)
+
+    def forward(
+        self, input: Tensor, hx: tuple[Tensor, Tensor] | None = None
+    ) -> tuple[Tensor, Tensor]:
+        if input.dim() not in (1, 2):
+            raise ValueError(
+                f"LSTMCell: Expected input to be 1D or 2D, got {input.dim()}D instead"
+            )
+        if hx is not None:
+            for idx, value in enumerate(hx):
+                if value.dim() not in (1, 2):
+                    raise ValueError(
+                        f"LSTMCell: Expected hx[{idx}] to be 1D or 2D, got {value.dim()}D instead"
+                    )
+        is_batched = input.dim() == 2
+        if not is_batched:
+            input = input.unsqueeze(0)
+
+        if hx is None:
+            zeros = torch.zeros(
+                input.size(0), self.hidden_size, dtype=input.dtype, device=input.device
+            )
+            hx = (zeros, zeros)
+        else:
+            hx = (hx[0].unsqueeze(0), hx[1].unsqueeze(0)) if not is_batched else hx
+
+        ret = _VF.lstm_cell(
+            input,
+            hx,
+            self.weight_ih,
+            self.weight_hh,
+            self.bias_ih,
+            self.bias_hh,
+        )
+
+        if not is_batched:
+            ret = (ret[0].squeeze(0), ret[1].squeeze(0))
+        return ret
+
+
+class GRUCell(RNNCellBase):
+    r"""A gated recurrent unit (GRU) cell.
+
+    .. math::
+
+        \begin{array}{ll}
+        r = \sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\
+        z = \sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\
+        n = \tanh(W_{in} x + b_{in} + r \odot (W_{hn} h + b_{hn})) \\
+        h' = (1 - z) \odot n + z \odot h
+        \end{array}
+
+    where :math:`\sigma` is the sigmoid function, and :math:`\odot` is the Hadamard product.
+
+    Args:
+        input_size: The number of expected features in the input `x`
+        hidden_size: The number of features in the hidden state `h`
+        bias: If ``False``, then the layer does not use bias weights `b_ih` and
+            `b_hh`. Default: ``True``
+
+    Inputs: input, hidden
+        - **input** : tensor containing input features
+        - **hidden** : tensor containing the initial hidden
+          state for each element in the batch.
+          Defaults to zero if not provided.
+
+    Outputs: h'
+        - **h'** : tensor containing the next hidden state
+          for each element in the batch
+
+    Shape:
+        - input: :math:`(N, H_{in})` or :math:`(H_{in})` tensor containing input features where
+          :math:`H_{in}` = `input_size`.
+        - hidden: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the initial hidden
+          state where :math:`H_{out}` = `hidden_size`. Defaults to zero if not provided.
+        - output: :math:`(N, H_{out})` or :math:`(H_{out})` tensor containing the next hidden state.
+
+    Attributes:
+        weight_ih: the learnable input-hidden weights, of shape
+            `(3*hidden_size, input_size)`
+        weight_hh: the learnable hidden-hidden weights, of shape
+            `(3*hidden_size, hidden_size)`
+        bias_ih: the learnable input-hidden bias, of shape `(3*hidden_size)`
+        bias_hh: the learnable hidden-hidden bias, of shape `(3*hidden_size)`
+
+    .. note::
+        All the weights and biases are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`
+        where :math:`k = \frac{1}{\text{hidden\_size}}`
+
+    On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
+
+    Examples::
+
+        >>> rnn = nn.GRUCell(10, 20)
+        >>> input = torch.randn(6, 3, 10)
+        >>> hx = torch.randn(3, 20)
+        >>> output = []
+        >>> for i in range(6):
+        ...     hx = rnn(input[i], hx)
+        ...     output.append(hx)
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(input_size, hidden_size, bias, num_chunks=3, **factory_kwargs)
+
+    def forward(self, input: Tensor, hx: Tensor | None = None) -> Tensor:
+        if input.dim() not in (1, 2):
+            raise ValueError(
+                f"GRUCell: Expected input to be 1D or 2D, got {input.dim()}D instead"
+            )
+        if hx is not None and hx.dim() not in (1, 2):
+            raise ValueError(
+                f"GRUCell: Expected hidden to be 1D or 2D, got {hx.dim()}D instead"
+            )
+        is_batched = input.dim() == 2
+        if not is_batched:
+            input = input.unsqueeze(0)
+
+        if hx is None:
+            hx = torch.zeros(
+                input.size(0), self.hidden_size, dtype=input.dtype, device=input.device
+            )
+        else:
+            hx = hx.unsqueeze(0) if not is_batched else hx
+
+        ret = _VF.gru_cell(
+            input,
+            hx,
+            self.weight_ih,
+            self.weight_hh,
+            self.bias_ih,
+            self.bias_hh,
+        )
+
+        if not is_batched:
+            ret = ret.squeeze(0)
+
+        return ret
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/sparse.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/sparse.py
new file mode 100644
index 0000000000000000000000000000000000000000..486e2fc88631c4c90a084d6275e044d69f91caba
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/sparse.py
@@ -0,0 +1,543 @@
+# mypy: allow-untyped-defs
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F, init
+from torch.nn.parameter import Parameter
+
+from .module import Module
+
+
+__all__ = ["Embedding", "EmbeddingBag"]
+
+
+class Embedding(Module):
+    r"""A simple lookup table that stores embeddings of a fixed dictionary and size.
+
+    This module is often used to store word embeddings and retrieve them using indices.
+    The input to the module is a list of indices, and the output is the corresponding
+    word embeddings.
+
+    Args:
+        num_embeddings (int): size of the dictionary of embeddings
+        embedding_dim (int): the size of each embedding vector
+        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;
+                                     therefore, the embedding vector at :attr:`padding_idx` is not updated during training,
+                                     i.e. it remains as a fixed "pad". For a newly constructed Embedding,
+                                     the embedding vector at :attr:`padding_idx` will default to all zeros,
+                                     but can be updated to another value to be used as the padding vector.
+        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`
+                                    is renormalized to have norm :attr:`max_norm`.
+        norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.
+        scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse of frequency of
+                                                the words in the mini-batch. Default ``False``.
+        sparse (bool, optional): If ``True``, gradient w.r.t. :attr:`weight` matrix will be a sparse tensor.
+                                 See Notes for more details regarding sparse gradients.
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape (num_embeddings, embedding_dim)
+                         initialized from :math:`\mathcal{N}(0, 1)`
+
+    Shape:
+        - Input: :math:`(*)`, IntTensor or LongTensor of arbitrary shape containing the indices to extract
+        - Output: :math:`(*, H)`, where `*` is the input shape and :math:`H=\text{embedding\_dim}`
+
+    .. note::
+        Keep in mind that only a limited number of optimizers support
+        sparse gradients: currently it's :class:`optim.SGD` (`CUDA` and `CPU`),
+        :class:`optim.SparseAdam` (`CUDA` and `CPU`) and :class:`optim.Adagrad` (`CPU`)
+
+    .. note::
+        When :attr:`max_norm` is not ``None``, :class:`Embedding`'s forward method will modify the
+        :attr:`weight` tensor in-place. Since tensors needed for gradient computations cannot be
+        modified in-place, performing a differentiable operation on ``Embedding.weight`` before
+        calling :class:`Embedding`'s forward method requires cloning ``Embedding.weight`` when
+        :attr:`max_norm` is not ``None``. For example::
+
+            n, d, m = 3, 5, 7
+            embedding = nn.Embedding(n, d, max_norm=1.0)
+            W = torch.randn((m, d), requires_grad=True)
+            idx = torch.tensor([1, 2])
+            a = (
+                embedding.weight.clone() @ W.t()
+            )  # weight must be cloned for this to be differentiable
+            b = embedding(idx) @ W.t()  # modifies weight in-place
+            out = a.unsqueeze(0) + b.unsqueeze(1)
+            loss = out.sigmoid().prod()
+            loss.backward()
+
+    Examples::
+
+        >>> # an Embedding module containing 10 tensors of size 3
+        >>> embedding = nn.Embedding(10, 3)
+        >>> # a batch of 2 samples of 4 indices each
+        >>> input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> embedding(input)
+        tensor([[[-0.0251, -1.6902,  0.7172],
+                 [-0.6431,  0.0748,  0.6969],
+                 [ 1.4970,  1.3448, -0.9685],
+                 [-0.3677, -2.7265, -0.1685]],
+
+                [[ 1.4970,  1.3448, -0.9685],
+                 [ 0.4362, -0.4004,  0.9400],
+                 [-0.6431,  0.0748,  0.6969],
+                 [ 0.9124, -2.3616,  1.1151]]])
+
+
+        >>> # example with padding_idx
+        >>> embedding = nn.Embedding(10, 3, padding_idx=0)
+        >>> input = torch.LongTensor([[0, 2, 0, 5]])
+        >>> embedding(input)
+        tensor([[[ 0.0000,  0.0000,  0.0000],
+                 [ 0.1535, -2.0309,  0.9315],
+                 [ 0.0000,  0.0000,  0.0000],
+                 [-0.1655,  0.9897,  0.0635]]])
+
+        >>> # example of changing `pad` vector
+        >>> padding_idx = 0
+        >>> embedding = nn.Embedding(3, 3, padding_idx=padding_idx)
+        >>> embedding.weight
+        Parameter containing:
+        tensor([[ 0.0000,  0.0000,  0.0000],
+                [-0.7895, -0.7089, -0.0364],
+                [ 0.6778,  0.5803,  0.2678]], requires_grad=True)
+        >>> with torch.no_grad():
+        ...     embedding.weight[padding_idx] = torch.ones(3)
+        >>> embedding.weight
+        Parameter containing:
+        tensor([[ 1.0000,  1.0000,  1.0000],
+                [-0.7895, -0.7089, -0.0364],
+                [ 0.6778,  0.5803,  0.2678]], requires_grad=True)
+    """
+
+    __constants__ = [
+        "num_embeddings",
+        "embedding_dim",
+        "padding_idx",
+        "max_norm",
+        "norm_type",
+        "scale_grad_by_freq",
+        "sparse",
+    ]
+
+    num_embeddings: int
+    embedding_dim: int
+    padding_idx: int | None
+    max_norm: float | None
+    norm_type: float
+    scale_grad_by_freq: bool
+    weight: Tensor
+    freeze: bool
+    sparse: bool
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        padding_idx: int | None = None,
+        max_norm: float | None = None,
+        norm_type: float = 2.0,
+        scale_grad_by_freq: bool = False,
+        sparse: bool = False,
+        _weight: Tensor | None = None,
+        _freeze: bool = False,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        if padding_idx is not None:
+            if padding_idx > 0:
+                if padding_idx >= self.num_embeddings:
+                    raise AssertionError("Padding_idx must be within num_embeddings")
+            elif padding_idx < 0:
+                if padding_idx < -self.num_embeddings:
+                    raise AssertionError("Padding_idx must be within num_embeddings")
+                padding_idx = self.num_embeddings + padding_idx
+        self.padding_idx = padding_idx
+        self.max_norm = max_norm
+        self.norm_type = norm_type
+        self.scale_grad_by_freq = scale_grad_by_freq
+        if _weight is None:
+            self.weight = Parameter(
+                torch.empty((num_embeddings, embedding_dim), **factory_kwargs),
+                requires_grad=not _freeze,
+            )
+            self.reset_parameters()
+        else:
+            if list(_weight.shape) != [num_embeddings, embedding_dim]:
+                raise AssertionError(
+                    "Shape of weight does not match num_embeddings and embedding_dim"
+                )
+            self.weight = Parameter(_weight, requires_grad=not _freeze)
+
+        self.sparse = sparse
+
+    def reset_parameters(self) -> None:
+        init.normal_(self.weight)
+        self._fill_padding_idx_with_zero()
+
+    def _fill_padding_idx_with_zero(self) -> None:
+        if self.padding_idx is not None:
+            with torch.no_grad():
+                self.weight[self.padding_idx].fill_(0)
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.embedding(
+            input,
+            self.weight,
+            self.padding_idx,
+            self.max_norm,
+            self.norm_type,
+            self.scale_grad_by_freq,
+            self.sparse,
+        )
+
+    def extra_repr(self) -> str:
+        s = "{num_embeddings}, {embedding_dim}"
+        if self.padding_idx is not None:
+            s += ", padding_idx={padding_idx}"
+        if self.max_norm is not None:
+            s += ", max_norm={max_norm}"
+        if self.norm_type != 2:
+            s += ", norm_type={norm_type}"
+        if self.scale_grad_by_freq is not False:
+            s += ", scale_grad_by_freq={scale_grad_by_freq}"
+        if self.sparse is not False:
+            s += ", sparse=True"
+        return s.format(**self.__dict__)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        embeddings,
+        freeze=True,
+        padding_idx=None,
+        max_norm=None,
+        norm_type=2.0,
+        scale_grad_by_freq=False,
+        sparse=False,
+    ):
+        r"""Create Embedding instance from given 2-dimensional FloatTensor.
+
+        Args:
+            embeddings (Tensor): FloatTensor containing weights for the Embedding.
+                First dimension is being passed to Embedding as ``num_embeddings``, second as ``embedding_dim``.
+            freeze (bool, optional): If ``True``, the tensor does not get updated in the learning process.
+                Equivalent to ``embedding.weight.requires_grad = False``. Default: ``True``
+            padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;
+                                         therefore, the embedding vector at :attr:`padding_idx` is not updated during training,
+                                         i.e. it remains as a fixed "pad".
+            max_norm (float, optional): See module initialization documentation.
+            norm_type (float, optional): See module initialization documentation. Default ``2``.
+            scale_grad_by_freq (bool, optional): See module initialization documentation. Default ``False``.
+            sparse (bool, optional): See module initialization documentation.
+
+        Examples::
+
+            >>> # FloatTensor containing pretrained weights
+            >>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
+            >>> embedding = nn.Embedding.from_pretrained(weight)
+            >>> # Get embeddings for index 1
+            >>> input = torch.LongTensor([1])
+            >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+            >>> embedding(input)
+            tensor([[ 4.0000,  5.1000,  6.3000]])
+        """
+        if embeddings.dim() != 2:
+            raise AssertionError("Embeddings parameter is expected to be 2-dimensional")
+        rows, cols = embeddings.shape
+        embedding = cls(
+            num_embeddings=rows,
+            embedding_dim=cols,
+            _weight=embeddings,
+            _freeze=freeze,
+            padding_idx=padding_idx,
+            max_norm=max_norm,
+            norm_type=norm_type,
+            scale_grad_by_freq=scale_grad_by_freq,
+            sparse=sparse,
+        )
+        return embedding
+
+
+class EmbeddingBag(Module):
+    r"""Compute sums or means of 'bags' of embeddings, without instantiating the intermediate embeddings.
+
+    For bags of constant length, no :attr:`per_sample_weights`, no indices equal to :attr:`padding_idx`,
+    and with 2D inputs, this class
+
+        * with ``mode="sum"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.sum(dim=1)``,
+        * with ``mode="mean"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.mean(dim=1)``,
+        * with ``mode="max"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.max(dim=1)``.
+
+    However, :class:`~torch.nn.EmbeddingBag` is much more time and memory efficient than using a chain of these
+    operations.
+
+    EmbeddingBag also supports per-sample weights as an argument to the forward
+    pass. This scales the output of the Embedding before performing a weighted
+    reduction as specified by ``mode``. If :attr:`per_sample_weights` is passed, the
+    only supported ``mode`` is ``"sum"``, which computes a weighted sum according to
+    :attr:`per_sample_weights`.
+
+    Args:
+        num_embeddings (int): size of the dictionary of embeddings
+        embedding_dim (int): the size of each embedding vector
+        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`
+                                    is renormalized to have norm :attr:`max_norm`.
+        norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.
+        scale_grad_by_freq (bool, optional): if given, this will scale gradients by the inverse of frequency of
+                                                the words in the mini-batch. Default ``False``.
+                                                Note: this option is not supported when ``mode="max"``.
+        mode (str, optional): ``"sum"``, ``"mean"`` or ``"max"``. Specifies the way to reduce the bag.
+                                 ``"sum"`` computes the weighted sum, taking :attr:`per_sample_weights`
+                                 into consideration. ``"mean"`` computes the average of the values
+                                 in the bag, ``"max"`` computes the max value over each bag.
+                                 Default: ``"mean"``
+        sparse (bool, optional): if ``True``, gradient w.r.t. :attr:`weight` matrix will be a sparse tensor. See
+                                 Notes for more details regarding sparse gradients. Note: this option is not
+                                 supported when ``mode="max"``.
+        include_last_offset (bool, optional): if ``True``, the size of offsets is equal to the number of bags + 1.
+                                              The last element is the size of the input, or the ending index position
+                                              of the last bag (sequence). This matches the CSR format. Ignored when
+                                              input is 2D. Default ``False``.
+        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the
+                                     gradient; therefore, the embedding vector at :attr:`padding_idx` is not updated
+                                     during training, i.e. it remains as a fixed "pad". For a newly constructed
+                                     EmbeddingBag, the embedding vector at :attr:`padding_idx` will default to all
+                                     zeros, but can be updated to another value to be used as the padding vector.
+                                     Note that the embedding vector at :attr:`padding_idx` is excluded from the
+                                     reduction.
+
+    Attributes:
+        weight (Tensor): the learnable weights of the module of shape `(num_embeddings, embedding_dim)`
+                         initialized from :math:`\mathcal{N}(0, 1)`.
+
+    Examples::
+
+        >>> # an EmbeddingBag module containing 10 tensors of size 3
+        >>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum')
+        >>> # a batch of 2 samples of 4 indices each
+        >>> input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)
+        >>> offsets = torch.tensor([0, 4], dtype=torch.long)
+        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+        >>> embedding_sum(input, offsets)
+        tensor([[-0.8861, -5.4350, -0.0523],
+                [ 1.1306, -2.5798, -1.0044]])
+
+        >>> # Example with padding_idx
+        >>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum', padding_idx=2)
+        >>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9], dtype=torch.long)
+        >>> offsets = torch.tensor([0, 4], dtype=torch.long)
+        >>> embedding_sum(input, offsets)
+        tensor([[ 0.0000,  0.0000,  0.0000],
+                [-0.7082,  3.2145, -2.6251]])
+
+        >>> # An EmbeddingBag can be loaded from an Embedding like so
+        >>> embedding = nn.Embedding(10, 3, padding_idx=2)
+        >>> embedding_sum = nn.EmbeddingBag.from_pretrained(
+                embedding.weight,
+                padding_idx=embedding.padding_idx,
+                mode='sum')
+    """
+
+    __constants__ = [
+        "num_embeddings",
+        "embedding_dim",
+        "max_norm",
+        "norm_type",
+        "scale_grad_by_freq",
+        "mode",
+        "sparse",
+        "include_last_offset",
+        "padding_idx",
+    ]
+
+    num_embeddings: int
+    embedding_dim: int
+    max_norm: float | None
+    norm_type: float
+    scale_grad_by_freq: bool
+    weight: Tensor
+    mode: str
+    sparse: bool
+    include_last_offset: bool
+    padding_idx: int | None
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        max_norm: float | None = None,
+        norm_type: float = 2.0,
+        scale_grad_by_freq: bool = False,
+        mode: str = "mean",
+        sparse: bool = False,
+        _weight: Tensor | None = None,
+        include_last_offset: bool = False,
+        padding_idx: int | None = None,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.max_norm = max_norm
+        self.norm_type = norm_type
+        self.scale_grad_by_freq = scale_grad_by_freq
+        if padding_idx is not None:
+            if padding_idx > 0:
+                if padding_idx >= self.num_embeddings:
+                    raise AssertionError("padding_idx must be within num_embeddings")
+            elif padding_idx < 0:
+                if padding_idx < -self.num_embeddings:
+                    raise AssertionError("padding_idx must be within num_embeddings")
+                padding_idx = self.num_embeddings + padding_idx
+        self.padding_idx = padding_idx
+        if _weight is None:
+            self.weight = Parameter(
+                torch.empty((num_embeddings, embedding_dim), **factory_kwargs)
+            )
+            self.reset_parameters()
+        else:
+            if list(_weight.shape) != [num_embeddings, embedding_dim]:
+                raise AssertionError(
+                    "Shape of weight does not match num_embeddings and embedding_dim"
+                )
+            self.weight = Parameter(_weight)
+        self.mode = mode
+        self.sparse = sparse
+        self.include_last_offset = include_last_offset
+
+    def reset_parameters(self) -> None:
+        init.normal_(self.weight)
+        self._fill_padding_idx_with_zero()
+
+    def _fill_padding_idx_with_zero(self) -> None:
+        if self.padding_idx is not None:
+            with torch.no_grad():
+                self.weight[self.padding_idx].fill_(0)
+
+    def forward(
+        self,
+        input: Tensor,
+        offsets: Tensor | None = None,
+        per_sample_weights: Tensor | None = None,
+    ) -> Tensor:
+        """Forward pass of EmbeddingBag.
+
+        Args:
+            input (Tensor): Tensor containing bags of indices into the embedding matrix.
+            offsets (Tensor, optional): Only used when :attr:`input` is 1D. :attr:`offsets` determines
+                the starting index position of each bag (sequence) in :attr:`input`.
+            per_sample_weights (Tensor, optional): a tensor of float / double weights, or None
+                to indicate all weights should be taken to be ``1``. If specified, :attr:`per_sample_weights`
+                must have exactly the same shape as input and is treated as having the same
+                :attr:`offsets`, if those are not ``None``. Only supported for ``mode='sum'``.
+
+        Returns:
+            Tensor output shape of `(B, embedding_dim)`.
+
+        .. note::
+
+            A few notes about ``input`` and ``offsets``:
+
+            - :attr:`input` and :attr:`offsets` have to be of the same type, either int or long
+
+            - If :attr:`input` is 2D of shape `(B, N)`, it will be treated as ``B`` bags (sequences)
+              each of fixed length ``N``, and this will return ``B`` values aggregated in a way
+              depending on the :attr:`mode`. :attr:`offsets` is ignored and required to be ``None`` in this case.
+
+            - If :attr:`input` is 1D of shape `(N)`, it will be treated as a concatenation of
+              multiple bags (sequences).  :attr:`offsets` is required to be a 1D tensor containing the
+              starting index positions of each bag in :attr:`input`. Therefore, for :attr:`offsets` of shape `(B)`,
+              :attr:`input` will be viewed as having ``B`` bags. Empty bags (i.e., having 0-length) will have
+              returned vectors filled by zeros.
+        """
+        return F.embedding_bag(
+            input,
+            self.weight,
+            offsets,
+            self.max_norm,
+            self.norm_type,
+            self.scale_grad_by_freq,
+            self.mode,
+            self.sparse,
+            per_sample_weights,
+            self.include_last_offset,
+            self.padding_idx,
+        )
+
+    def extra_repr(self) -> str:
+        s = "{num_embeddings}, {embedding_dim}"
+        if self.max_norm is not None:
+            s += ", max_norm={max_norm}"
+        if self.norm_type != 2:
+            s += ", norm_type={norm_type}"
+        if self.scale_grad_by_freq is not False:
+            s += ", scale_grad_by_freq={scale_grad_by_freq}"
+        s += ", mode={mode}"
+        if self.padding_idx is not None:
+            s += ", padding_idx={padding_idx}"
+        return s.format(**{k: repr(v) for k, v in self.__dict__.items()})
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        embeddings: Tensor,
+        freeze: bool = True,
+        max_norm: float | None = None,
+        norm_type: float = 2.0,
+        scale_grad_by_freq: bool = False,
+        mode: str = "mean",
+        sparse: bool = False,
+        include_last_offset: bool = False,
+        padding_idx: int | None = None,
+    ) -> "EmbeddingBag":
+        r"""Create EmbeddingBag instance from given 2-dimensional FloatTensor.
+
+        Args:
+            embeddings (Tensor): FloatTensor containing weights for the EmbeddingBag.
+                First dimension is being passed to EmbeddingBag as 'num_embeddings', second as 'embedding_dim'.
+            freeze (bool, optional): If ``True``, the tensor does not get updated in the learning process.
+                Equivalent to ``embeddingbag.weight.requires_grad = False``. Default: ``True``
+            max_norm (float, optional): See module initialization documentation. Default: ``None``
+            norm_type (float, optional): See module initialization documentation. Default ``2``.
+            scale_grad_by_freq (bool, optional): See module initialization documentation. Default ``False``.
+            mode (str, optional): See module initialization documentation. Default: ``"mean"``
+            sparse (bool, optional): See module initialization documentation. Default: ``False``.
+            include_last_offset (bool, optional): See module initialization documentation. Default: ``False``.
+            padding_idx (int, optional): See module initialization documentation. Default: ``None``.
+
+        Examples::
+
+            >>> # FloatTensor containing pretrained weights
+            >>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
+            >>> embeddingbag = nn.EmbeddingBag.from_pretrained(weight)
+            >>> # Get embeddings for index 1
+            >>> input = torch.LongTensor([[1, 0]])
+            >>> # xdoctest: +IGNORE_WANT("non-deterministic")
+            >>> embeddingbag(input)
+            tensor([[ 2.5000,  3.7000,  4.6500]])
+        """
+        if embeddings.dim() != 2:
+            raise AssertionError("Embeddings parameter is expected to be 2-dimensional")
+        rows, cols = embeddings.shape
+        embeddingbag = cls(
+            num_embeddings=rows,
+            embedding_dim=cols,
+            _weight=embeddings,
+            max_norm=max_norm,
+            norm_type=norm_type,
+            scale_grad_by_freq=scale_grad_by_freq,
+            mode=mode,
+            sparse=sparse,
+            include_last_offset=include_last_offset,
+            padding_idx=padding_idx,
+        )
+        embeddingbag.weight.requires_grad = not freeze
+        return embeddingbag
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/transformer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..7f47013c14585db727d69b337b00dbc89a22cbdc
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/transformer.py
@@ -0,0 +1,1249 @@
+# mypy: allow-untyped-defs
+import copy
+import warnings
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.init import xavier_uniform_
+
+from .activation import MultiheadAttention
+from .container import ModuleList
+from .dropout import Dropout
+from .linear import Linear
+from .module import Module
+from .normalization import LayerNorm
+
+
+__all__ = [
+    "Transformer",
+    "TransformerEncoder",
+    "TransformerDecoder",
+    "TransformerEncoderLayer",
+    "TransformerDecoderLayer",
+]
+
+
+def _generate_square_subsequent_mask(
+    sz: int,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
+) -> Tensor:
+    r"""Generate a square causal mask for the sequence.
+
+    The masked positions are filled with float('-inf'). Unmasked positions are filled with float(0.0).
+    """
+    return torch.triu(
+        torch.full((sz, sz), float("-inf"), dtype=dtype, device=device),
+        diagonal=1,
+    )
+
+
+def _get_seq_len(src: Tensor, batch_first: bool) -> int | None:
+    if src.is_nested:
+        return None
+    else:
+        src_size = src.size()
+        if len(src_size) == 2:
+            # unbatched: S, E
+            return src_size[0]
+        else:
+            # batched: B, S, E if batch_first else S, B, E
+            seq_len_pos = 1 if batch_first else 0
+            return src_size[seq_len_pos]
+
+
+class Transformer(Module):
+    r"""A basic transformer layer.
+
+
+    This Transformer layer implements the original Transformer architecture described
+    in the `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_ paper. The
+    intent of this layer is as a reference implementation for foundational understanding
+    and thus it contains only limited features relative to newer Transformer architectures.
+    Given the fast pace of innovation in transformer-like architectures, we recommend
+    exploring this `tutorial <https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html>`_
+    to build an efficient transformer layer from building blocks in core or using higher
+    level libraries from the `PyTorch Ecosystem <https://landscape.pytorch.org/>`_.
+
+    Args:
+        d_model: the number of expected features in the encoder/decoder inputs (default=512).
+        nhead: the number of heads in the multiheadattention models (default=8).
+        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
+        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of encoder/decoder intermediate layer, can be a string
+            ("relu" or "gelu") or a unary callable. Default: relu
+        custom_encoder: custom encoder (default=None).
+        custom_decoder: custom decoder (default=None).
+        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+        norm_first: if ``True``, encoder and decoder layers will perform LayerNorms before
+            other attention and feedforward operations, otherwise after. Default: ``False`` (after).
+        bias: If set to ``False``, ``Linear`` and ``LayerNorm`` layers will not learn an additive
+            bias. Default: ``True``.
+
+    Examples:
+        >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
+        >>> src = torch.rand((10, 32, 512))
+        >>> tgt = torch.rand((20, 32, 512))
+        >>> out = transformer_model(src, tgt)
+
+    Note: A full example to apply nn.Transformer module for the word language model is available in
+    https://github.com/pytorch/examples/tree/master/word_language_model
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        nhead: int = 8,
+        num_encoder_layers: int = 6,
+        num_decoder_layers: int = 6,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: str | Callable[[Tensor], Tensor] = F.relu,
+        custom_encoder: Any | None = None,
+        custom_decoder: Any | None = None,
+        layer_norm_eps: float = 1e-5,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        torch._C._log_api_usage_once(f"torch.nn.modules.{self.__class__.__name__}")
+
+        if custom_encoder is not None:
+            self.encoder = custom_encoder
+        else:
+            encoder_layer = TransformerEncoderLayer(
+                d_model,
+                nhead,
+                dim_feedforward,
+                dropout,
+                activation,
+                layer_norm_eps,
+                batch_first,
+                norm_first,
+                bias,
+                **factory_kwargs,
+            )
+            encoder_norm = LayerNorm(
+                d_model,
+                eps=layer_norm_eps,
+                bias=bias,
+                **factory_kwargs,
+            )
+            self.encoder = TransformerEncoder(
+                encoder_layer, num_encoder_layers, encoder_norm
+            )
+
+        if custom_decoder is not None:
+            self.decoder = custom_decoder
+        else:
+            decoder_layer = TransformerDecoderLayer(
+                d_model,
+                nhead,
+                dim_feedforward,
+                dropout,
+                activation,
+                layer_norm_eps,
+                batch_first,
+                norm_first,
+                bias,
+                **factory_kwargs,
+            )
+            decoder_norm = LayerNorm(
+                d_model,
+                eps=layer_norm_eps,
+                bias=bias,
+                **factory_kwargs,
+            )
+            self.decoder = TransformerDecoder(
+                decoder_layer, num_decoder_layers, decoder_norm
+            )
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+        self.batch_first = batch_first
+
+    def forward(
+        self,
+        src: Tensor,
+        tgt: Tensor,
+        src_mask: Tensor | None = None,
+        tgt_mask: Tensor | None = None,
+        memory_mask: Tensor | None = None,
+        src_key_padding_mask: Tensor | None = None,
+        tgt_key_padding_mask: Tensor | None = None,
+        memory_key_padding_mask: Tensor | None = None,
+        src_is_causal: bool | None = None,
+        tgt_is_causal: bool | None = None,
+        memory_is_causal: bool = False,
+    ) -> Tensor:
+        r"""Take in and process masked source/target sequences.
+
+        .. note::
+
+            If a boolean tensor is provided for any of the [src/tgt/memory]_mask arguments, positions with a ``True`` value are
+            not allowed to participate in the attention,
+            which is the opposite of the definition for :attr:`attn_mask`
+            in :func:`torch.nn.functional.scaled_dot_product_attention`.
+
+        Args:
+            src: the sequence to the encoder (required).
+            tgt: the sequence to the decoder (required).
+            src_mask: the additive mask for the src sequence (optional).
+            tgt_mask: the additive mask for the tgt sequence (optional).
+            memory_mask: the additive mask for the encoder output (optional).
+            src_key_padding_mask: the Tensor mask for src keys per batch (optional).
+            tgt_key_padding_mask: the Tensor mask for tgt keys per batch (optional).
+            memory_key_padding_mask: the Tensor mask for memory keys per batch (optional).
+            src_is_causal: If specified, applies a causal mask as ``src_mask``.
+                Default: ``None``; try to detect a causal mask.
+                Warning:
+                ``src_is_causal`` provides a hint that ``src_mask`` is
+                the causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+            tgt_is_causal: If specified, applies a causal mask as ``tgt_mask``.
+                Default: ``None``; try to detect a causal mask.
+                Warning:
+                ``tgt_is_causal`` provides a hint that ``tgt_mask`` is
+                the causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+            memory_is_causal: If specified, applies a causal mask as
+                ``memory_mask``.
+                Default: ``False``.
+                Warning:
+                ``memory_is_causal`` provides a hint that
+                ``memory_mask`` is the causal mask. Providing incorrect
+                hints can result in incorrect execution, including
+                forward and backward compatibility.
+
+        Shape:
+            - src: :math:`(S, E)` for unbatched input, :math:`(S, N, E)` if `batch_first=False` or
+              `(N, S, E)` if `batch_first=True`.
+            - tgt: :math:`(T, E)` for unbatched input, :math:`(T, N, E)` if `batch_first=False` or
+              `(N, T, E)` if `batch_first=True`.
+            - src_mask: :math:`(S, S)` or :math:`(N\cdot\text{num\_heads}, S, S)`.
+            - tgt_mask: :math:`(T, T)` or :math:`(N\cdot\text{num\_heads}, T, T)`.
+            - memory_mask: :math:`(T, S)`.
+            - src_key_padding_mask: :math:`(S)` for unbatched input otherwise :math:`(N, S)`.
+            - tgt_key_padding_mask: :math:`(T)` for unbatched input otherwise :math:`(N, T)`.
+            - memory_key_padding_mask: :math:`(S)` for unbatched input otherwise :math:`(N, S)`.
+
+            Note: [src/tgt/memory]_mask ensures that position :math:`i` is allowed to attend the unmasked
+            positions. If a BoolTensor is provided, positions with ``True``
+            are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+            is provided, it will be added to the attention weight.
+            [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
+            the attention. If a BoolTensor is provided, the positions with the
+            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+
+            - output: :math:`(T, E)` for unbatched input, :math:`(T, N, E)` if `batch_first=False` or
+              `(N, T, E)` if `batch_first=True`.
+
+            Note: Due to the multi-head attention architecture in the transformer model,
+            the output sequence length of a transformer is same as the input sequence
+            (i.e. target) length of the decoder.
+
+            where :math:`S` is the source sequence length, :math:`T` is the target sequence length, :math:`N` is the
+            batch size, :math:`E` is the feature number
+
+        Examples:
+            >>> # xdoctest: +SKIP
+            >>> output = transformer_model(
+            ...     src, tgt, src_mask=src_mask, tgt_mask=tgt_mask
+            ... )
+        """
+        is_batched = src.dim() == 3
+        if not self.batch_first and src.size(1) != tgt.size(1) and is_batched:
+            raise RuntimeError("the batch number of src and tgt must be equal")
+        elif self.batch_first and src.size(0) != tgt.size(0) and is_batched:
+            raise RuntimeError("the batch number of src and tgt must be equal")
+
+        if src.size(-1) != self.d_model or tgt.size(-1) != self.d_model:
+            raise RuntimeError(
+                "the feature number of src and tgt must be equal to d_model"
+            )
+
+        memory = self.encoder(
+            src,
+            mask=src_mask,
+            src_key_padding_mask=src_key_padding_mask,
+            is_causal=src_is_causal,
+        )
+        output = self.decoder(
+            tgt,
+            memory,
+            tgt_mask=tgt_mask,
+            memory_mask=memory_mask,
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=memory_key_padding_mask,
+            tgt_is_causal=tgt_is_causal,
+            memory_is_causal=memory_is_causal,
+        )
+        return output
+
+    @staticmethod
+    def generate_square_subsequent_mask(
+        sz: int,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> Tensor:
+        r"""Generate a square causal mask for the sequence.
+
+        The masked positions are filled with float('-inf'). Unmasked positions are filled with float(0.0).
+        """
+        return _generate_square_subsequent_mask(sz, dtype=dtype, device=device)
+
+    def _reset_parameters(self) -> None:
+        r"""Initiate parameters in the transformer model."""
+        for p in self.parameters():
+            if p.dim() > 1:
+                xavier_uniform_(p)
+
+
+class TransformerEncoder(Module):
+    r"""TransformerEncoder is a stack of N encoder layers.
+
+    This TransformerEncoder layer implements the original architecture described
+    in the `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_ paper. The
+    intent of this layer is as a reference implementation for foundational understanding
+    and thus it contains only limited features relative to newer Transformer architectures.
+    Given the fast pace of innovation in transformer-like architectures, we recommend
+    exploring this `tutorial <https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html>`_
+    to build efficient layers from building blocks in core or using higher
+    level libraries from the `PyTorch Ecosystem <https://landscape.pytorch.org/>`_.
+
+    .. warning::
+        All layers in the TransformerEncoder are initialized with the same parameters.
+        It is recommended to manually initialize the layers after creating the TransformerEncoder instance.
+
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+        enable_nested_tensor: if True, input will automatically convert to nested tensor
+            (and convert back on output). This will improve the overall performance of
+            TransformerEncoder when padding rate is high. Default: ``True`` (enabled).
+
+    Examples:
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+
+    __constants__ = ["norm"]
+
+    def __init__(
+        self,
+        encoder_layer: "TransformerEncoderLayer",
+        num_layers: int,
+        norm: Module | None = None,
+        enable_nested_tensor: bool = True,
+        mask_check: bool = True,
+    ) -> None:
+        super().__init__()
+        torch._C._log_api_usage_once(f"torch.nn.modules.{self.__class__.__name__}")
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        # this attribute saves the value providedat object construction
+        self.enable_nested_tensor = enable_nested_tensor
+        # this attribute controls whether nested tensors are used
+        self.use_nested_tensor = enable_nested_tensor
+        self.mask_check = mask_check
+
+        enc_layer = "encoder_layer"
+        why_not_sparsity_fast_path = ""
+        if not isinstance(encoder_layer, torch.nn.TransformerEncoderLayer):
+            why_not_sparsity_fast_path = f"{enc_layer} was not TransformerEncoderLayer"
+        elif encoder_layer.norm_first:
+            why_not_sparsity_fast_path = f"{enc_layer}.norm_first was True"
+        elif not encoder_layer.self_attn.batch_first:
+            why_not_sparsity_fast_path = (
+                f"{enc_layer}.self_attn.batch_first was not True"
+                + "(use batch_first for better inference performance)"
+            )
+        elif not encoder_layer.self_attn._qkv_same_embed_dim:
+            why_not_sparsity_fast_path = (
+                f"{enc_layer}.self_attn._qkv_same_embed_dim was not True"
+            )
+        elif encoder_layer.self_attn.in_proj_bias is None:
+            why_not_sparsity_fast_path = f"{enc_layer}.self_attn was passed bias=False"
+        elif not encoder_layer.activation_relu_or_gelu:
+            why_not_sparsity_fast_path = (
+                f"{enc_layer}.activation_relu_or_gelu was not True"
+            )
+        elif encoder_layer.norm1.eps != encoder_layer.norm2.eps:
+            why_not_sparsity_fast_path = (
+                f"{enc_layer}.norm1.eps was not equal to {enc_layer}.norm2.eps"
+            )
+        elif encoder_layer.self_attn.num_heads % 2 == 1:
+            why_not_sparsity_fast_path = f"{enc_layer}.self_attn.num_heads is odd"
+
+        if enable_nested_tensor and why_not_sparsity_fast_path:
+            warnings.warn(
+                f"enable_nested_tensor is True, but self.use_nested_tensor is False because {why_not_sparsity_fast_path}",
+                stacklevel=2,
+            )
+            self.use_nested_tensor = False
+
+    def forward(
+        self,
+        src: Tensor,
+        mask: Tensor | None = None,
+        src_key_padding_mask: Tensor | None = None,
+        is_causal: bool | None = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+            is_causal: If specified, applies a causal mask as ``mask``.
+                Default: ``None``; try to detect a causal mask.
+                Warning:
+                ``is_causal`` provides a hint that ``mask`` is the
+                causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+
+        Shape:
+            see the docs in :class:`~torch.nn.Transformer`.
+        """
+        src_key_padding_mask = F._canonical_mask(
+            mask=src_key_padding_mask,
+            mask_name="src_key_padding_mask",
+            other_type=F._none_or_dtype(mask),
+            other_name="mask",
+            target_type=src.dtype,
+        )
+
+        mask = F._canonical_mask(
+            mask=mask,
+            mask_name="mask",
+            other_type=None,
+            other_name="",
+            target_type=src.dtype,
+            check_other=False,
+        )
+
+        output = src
+        convert_to_nested = False
+        first_layer = self.layers[0]
+        src_key_padding_mask_for_layers = src_key_padding_mask
+        why_not_sparsity_fast_path = ""
+        str_first_layer = "self.layers[0]"
+        batch_first = first_layer.self_attn.batch_first
+        is_fastpath_enabled = torch.backends.mha.get_fastpath_enabled()
+        do_mask_check = getattr(self, "mask_check", True)
+
+        if not is_fastpath_enabled:
+            why_not_sparsity_fast_path = (
+                "torch.backends.mha.get_fastpath_enabled() was not True"
+            )
+        elif not hasattr(self, "use_nested_tensor"):
+            why_not_sparsity_fast_path = "use_nested_tensor attribute not present"
+        elif not self.use_nested_tensor:
+            why_not_sparsity_fast_path = (
+                "self.use_nested_tensor (set in init) was not True"
+            )
+        elif first_layer.training:
+            why_not_sparsity_fast_path = f"{str_first_layer} was in training mode"
+        elif src.dim() != 3:
+            why_not_sparsity_fast_path = (
+                f"input not batched; expected src.dim() of 3 but got {src.dim()}"
+            )
+        elif src_key_padding_mask is None:
+            why_not_sparsity_fast_path = "src_key_padding_mask was None"
+        # This check avoids a call to torch._nested_tensor_from_mask_left_aligned() that
+        # breaks in torch.compile.
+        elif do_mask_check and torch.compiler.is_compiling():
+            why_not_sparsity_fast_path = (
+                "mask_check enabled with torch.compile or torch.export"
+            )
+        elif do_mask_check and not torch._nested_tensor_from_mask_left_aligned(
+            src, src_key_padding_mask.logical_not()
+        ):
+            why_not_sparsity_fast_path = "mask_check enabled, and src and src_key_padding_mask was not left aligned"
+        elif output.is_nested:
+            why_not_sparsity_fast_path = "NestedTensor input is not supported"
+        elif mask is not None:
+            why_not_sparsity_fast_path = (
+                "src_key_padding_mask and mask were both supplied"
+            )
+        elif torch.is_autocast_enabled():
+            why_not_sparsity_fast_path = "autocast is enabled"
+
+        if not why_not_sparsity_fast_path:
+            tensor_args = (
+                src,
+                first_layer.self_attn.in_proj_weight,
+                first_layer.self_attn.in_proj_bias,
+                first_layer.self_attn.out_proj.weight,
+                first_layer.self_attn.out_proj.bias,
+                first_layer.norm1.weight,
+                first_layer.norm1.bias,
+                first_layer.norm2.weight,
+                first_layer.norm2.bias,
+                first_layer.linear1.weight,
+                first_layer.linear1.bias,
+                first_layer.linear2.weight,
+                first_layer.linear2.bias,
+            )
+            _supported_device_type = [
+                "cpu",
+                "cuda",
+                "xpu",
+                torch.utils.backend_registration._privateuse1_backend_name,
+            ]
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_sparsity_fast_path = "some Tensor argument has_torch_function"
+            elif src.device.type not in _supported_device_type:
+                why_not_sparsity_fast_path = (
+                    f"src device is neither one of {_supported_device_type}"
+                )
+            elif torch.is_grad_enabled() and any(x.requires_grad for x in tensor_args):
+                why_not_sparsity_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+
+            if (not why_not_sparsity_fast_path) and (src_key_padding_mask is not None):
+                convert_to_nested = True
+                output = torch._nested_tensor_from_mask(
+                    output, src_key_padding_mask.logical_not(), mask_check=False
+                )
+                src_key_padding_mask_for_layers = None
+
+        seq_len = _get_seq_len(src, batch_first)
+        is_causal = _detect_is_causal_mask(mask, is_causal, seq_len)
+
+        for mod in self.layers:
+            output = mod(
+                output,
+                src_mask=mask,
+                is_causal=is_causal,
+                src_key_padding_mask=src_key_padding_mask_for_layers,
+            )
+
+        if convert_to_nested:
+            output = output.to_padded_tensor(0.0, src.size())
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerDecoder(Module):
+    r"""TransformerDecoder is a stack of N decoder layers.
+
+    This TransformerDecoder layer implements the original architecture described
+    in the `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_ paper. The
+    intent of this layer is as a reference implementation for foundational understanding
+    and thus it contains only limited features relative to newer Transformer architectures.
+    Given the fast pace of innovation in transformer-like architectures, we recommend
+    exploring this `tutorial <https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html>`_
+    to build efficient layers from building blocks in core or using higher
+    level libraries from the `PyTorch Ecosystem <https://landscape.pytorch.org/>`_.
+
+    .. warning::
+        All layers in the TransformerDecoder are initialized with the same parameters.
+        It is recommended to manually initialize the layers after creating the TransformerDecoder instance.
+
+    Args:
+        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
+        num_layers: the number of sub-decoder-layers in the decoder (required).
+        norm: the layer normalization component (optional).
+
+    Examples:
+        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
+        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
+        >>> memory = torch.rand(10, 32, 512)
+        >>> tgt = torch.rand(20, 32, 512)
+        >>> out = transformer_decoder(tgt, memory)
+    """
+
+    __constants__ = ["norm"]
+
+    def __init__(
+        self,
+        decoder_layer: "TransformerDecoderLayer",
+        num_layers: int,
+        norm: Module | None = None,
+    ) -> None:
+        super().__init__()
+        torch._C._log_api_usage_once(f"torch.nn.modules.{self.__class__.__name__}")
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        tgt_mask: Tensor | None = None,
+        memory_mask: Tensor | None = None,
+        tgt_key_padding_mask: Tensor | None = None,
+        memory_key_padding_mask: Tensor | None = None,
+        tgt_is_causal: bool | None = None,
+        memory_is_causal: bool = False,
+    ) -> Tensor:
+        r"""Pass the inputs (and mask) through the decoder layer in turn.
+
+        Args:
+            tgt: the sequence to the decoder (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+            tgt_is_causal: If specified, applies a causal mask as ``tgt mask``.
+                Default: ``None``; try to detect a causal mask.
+                Warning:
+                ``tgt_is_causal`` provides a hint that ``tgt_mask`` is
+                the causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+            memory_is_causal: If specified, applies a causal mask as
+                ``memory mask``.
+                Default: ``False``.
+                Warning:
+                ``memory_is_causal`` provides a hint that
+                ``memory_mask`` is the causal mask. Providing incorrect
+                hints can result in incorrect execution, including
+                forward and backward compatibility.
+
+        Shape:
+            see the docs in :class:`~torch.nn.Transformer`.
+        """
+        output = tgt
+
+        seq_len = _get_seq_len(tgt, self.layers[0].self_attn.batch_first)
+        tgt_is_causal = _detect_is_causal_mask(tgt_mask, tgt_is_causal, seq_len)
+
+        for mod in self.layers:
+            output = mod(
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+                tgt_is_causal=tgt_is_causal,
+                memory_is_causal=memory_is_causal,
+            )
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerEncoderLayer(Module):
+    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
+
+    This TransformerEncoderLayer implements the original architecture described
+    in the `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_ paper. The
+    intent of this layer is as a reference implementation for foundational understanding
+    and thus it contains only limited features relative to newer Transformer architectures.
+    Given the fast pace of innovation in transformer-like architectures, we recommend
+    exploring this `tutorial <https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html>`_
+    to build efficient layers from building blocks in core or using higher
+    level libraries from the `PyTorch Ecosystem <https://landscape.pytorch.org/>`_.
+
+    TransformerEncoderLayer can handle either traditional torch.tensor inputs,
+    or Nested Tensor inputs.  Derived classes are expected to similarly accept
+    both input formats.  (Not all combinations of inputs are currently
+    supported by TransformerEncoderLayer while Nested Tensor is in prototype
+    state.)
+
+    If you are implementing a custom layer, you may derive it either from
+    the Module or TransformerEncoderLayer class.  If your custom layer
+    supports both torch.Tensors and Nested Tensors inputs, make its
+    implementation a derived class of TransformerEncoderLayer. If your custom
+    Layer supports only torch.Tensor inputs, derive its implementation from
+    Module.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of the intermediate layer, can be a string
+            ("relu" or "gelu") or a unary callable. Default: relu
+        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+        norm_first: if ``True``, layer norm is done prior to attention and feedforward
+            operations, respectively. Otherwise it's done after. Default: ``False`` (after).
+        bias: If set to ``False``, ``Linear`` and ``LayerNorm`` layers will not learn an additive
+            bias. Default: ``True``.
+
+    Examples:
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = encoder_layer(src)
+
+    Alternatively, when ``batch_first`` is ``True``:
+        >>> encoder_layer = nn.TransformerEncoderLayer(
+        ...     d_model=512, nhead=8, batch_first=True
+        ... )
+        >>> src = torch.rand(32, 10, 512)
+        >>> out = encoder_layer(src)
+
+    Fast path:
+        forward() will use a special optimized implementation described in
+        `FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness`_ if all of the following
+        conditions are met:
+
+        - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor
+          argument ``requires_grad``
+        - training is disabled (using ``.eval()``)
+        - batch_first is ``True`` and the input is batched (i.e., ``src.dim() == 3``)
+        - activation is one of: ``"relu"``, ``"gelu"``, ``torch.functional.relu``, or ``torch.functional.gelu``
+        - at most one of ``src_mask`` and ``src_key_padding_mask`` is passed
+        - if src is a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_, neither ``src_mask``
+          nor ``src_key_padding_mask`` is passed
+        - the two ``LayerNorm`` instances have a consistent ``eps`` value (this will naturally be the case
+          unless the caller has manually modified one without modifying the other)
+
+        If the optimized implementation is in use, a
+        `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be
+        passed for ``src`` to represent padding more efficiently than using a padding
+        mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ will be
+        returned, and an additional speedup proportional to the fraction of the input that
+        is padding can be expected.
+
+        .. _`FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness`:
+         https://arxiv.org/abs/2205.14135
+
+    """
+
+    __constants__ = ["norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: str | Callable[[Tensor], Tensor] = F.relu,
+        layer_norm_eps: float = 1e-5,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            bias=bias,
+            batch_first=batch_first,
+            **factory_kwargs,
+        )
+        # Implementation of Feedforward model
+        self.linear1 = Linear(d_model, dim_feedforward, bias=bias, **factory_kwargs)
+        self.dropout = Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs)
+
+        self.norm_first = norm_first
+        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
+        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            activation = _get_activation_fn(activation)
+
+        # We can't test self.activation in forward() in TorchScript,
+        # so stash some information about it instead.
+        if activation is F.relu or isinstance(activation, torch.nn.ReLU):
+            self.activation_relu_or_gelu = 1
+        elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
+            self.activation_relu_or_gelu = 2
+        else:
+            self.activation_relu_or_gelu = 0
+        self.activation = activation
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        if not hasattr(self, "activation"):
+            self.activation = F.relu
+
+    def forward(
+        self,
+        src: Tensor,
+        src_mask: Tensor | None = None,
+        src_key_padding_mask: Tensor | None = None,
+        is_causal: bool = False,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+            is_causal: If specified, applies a causal mask as ``src mask``.
+                Default: ``False``.
+                Warning:
+                ``is_causal`` provides a hint that ``src_mask`` is the
+                causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+
+        Shape:
+            see the docs in :class:`~torch.nn.Transformer`.
+        """
+        src_key_padding_mask = F._canonical_mask(
+            mask=src_key_padding_mask,
+            mask_name="src_key_padding_mask",
+            other_type=F._none_or_dtype(src_mask),
+            other_name="src_mask",
+            target_type=src.dtype,
+        )
+
+        src_mask = F._canonical_mask(
+            mask=src_mask,
+            mask_name="src_mask",
+            other_type=None,
+            other_name="",
+            target_type=src.dtype,
+            check_other=False,
+        )
+
+        is_fastpath_enabled = torch.backends.mha.get_fastpath_enabled()
+
+        why_not_sparsity_fast_path = ""
+        if not is_fastpath_enabled:
+            why_not_sparsity_fast_path = (
+                "torch.backends.mha.get_fastpath_enabled() was not True"
+            )
+        elif src.dim() != 3:
+            why_not_sparsity_fast_path = (
+                f"input not batched; expected src.dim() of 3 but got {src.dim()}"
+            )
+        elif self.training:
+            why_not_sparsity_fast_path = "training is enabled"
+        elif not self.self_attn.batch_first:
+            why_not_sparsity_fast_path = "self_attn.batch_first was not True"
+        elif self.self_attn.in_proj_bias is None:
+            why_not_sparsity_fast_path = "self_attn was passed bias=False"
+        elif not self.self_attn._qkv_same_embed_dim:
+            why_not_sparsity_fast_path = "self_attn._qkv_same_embed_dim was not True"
+        elif not self.activation_relu_or_gelu:
+            why_not_sparsity_fast_path = "activation_relu_or_gelu was not True"
+        elif self.norm1.eps != self.norm2.eps:
+            why_not_sparsity_fast_path = "norm1.eps is not equal to norm2.eps"
+        elif src.is_nested and (
+            src_key_padding_mask is not None or src_mask is not None
+        ):
+            why_not_sparsity_fast_path = "neither src_key_padding_mask nor src_mask are not supported with NestedTensor input"
+        elif self.self_attn.num_heads % 2 == 1:
+            why_not_sparsity_fast_path = "num_head is odd"
+        elif torch.is_autocast_enabled():
+            why_not_sparsity_fast_path = "autocast is enabled"
+        elif any(
+            len(getattr(m, "_forward_hooks", {}))
+            + len(getattr(m, "_forward_pre_hooks", {}))
+            for m in self.modules()
+        ):
+            why_not_sparsity_fast_path = "forward pre-/hooks are attached to the module"
+        if not why_not_sparsity_fast_path:
+            tensor_args = (
+                src,
+                self.self_attn.in_proj_weight,
+                self.self_attn.in_proj_bias,
+                self.self_attn.out_proj.weight,
+                self.self_attn.out_proj.bias,
+                self.norm1.weight,
+                self.norm1.bias,
+                self.norm2.weight,
+                self.norm2.bias,
+                self.linear1.weight,
+                self.linear1.bias,
+                self.linear2.weight,
+                self.linear2.bias,
+            )
+
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            _supported_device_type = [
+                "cpu",
+                "cuda",
+                "xpu",
+                torch.utils.backend_registration._privateuse1_backend_name,
+            ]
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_sparsity_fast_path = "some Tensor argument has_torch_function"
+            elif not all(
+                (x.device.type in _supported_device_type) for x in tensor_args
+            ):
+                why_not_sparsity_fast_path = (
+                    "some Tensor argument's device is neither one of "
+                    f"{_supported_device_type}"
+                )
+            elif torch.is_grad_enabled() and any(x.requires_grad for x in tensor_args):
+                why_not_sparsity_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+
+            if not why_not_sparsity_fast_path:
+                merged_mask, mask_type = self.self_attn.merge_masks(
+                    src_mask, src_key_padding_mask, src
+                )
+                return torch._transformer_encoder_layer_fwd(
+                    src,
+                    self.self_attn.embed_dim,
+                    self.self_attn.num_heads,
+                    self.self_attn.in_proj_weight,
+                    self.self_attn.in_proj_bias,
+                    self.self_attn.out_proj.weight,
+                    self.self_attn.out_proj.bias,
+                    self.activation_relu_or_gelu == 2,
+                    self.norm_first,
+                    self.norm1.eps,
+                    self.norm1.weight,
+                    self.norm1.bias,
+                    self.norm2.weight,
+                    self.norm2.bias,
+                    self.linear1.weight,
+                    self.linear1.bias,
+                    self.linear2.weight,
+                    self.linear2.bias,
+                    merged_mask,
+                    mask_type,
+                )
+
+        # see Fig. 1 of https://arxiv.org/pdf/2002.04745v1.pdf
+        x = src
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x), src_mask, src_key_padding_mask, is_causal=is_causal
+            )
+            x = x + self._ff_block(self.norm2(x))
+        else:
+            x = self.norm1(
+                x
+                + self._sa_block(x, src_mask, src_key_padding_mask, is_causal=is_causal)
+            )
+            x = self.norm2(x + self._ff_block(x))
+
+        return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Tensor | None,
+        key_padding_mask: Tensor | None,
+        is_causal: bool = False,
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+            is_causal=is_causal,
+        )[0]
+        return self.dropout1(x)
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout2(x)
+
+
+class TransformerDecoderLayer(Module):
+    r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
+
+    This TransformerDecoderLayer implements the original architecture described
+    in the `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_ paper. The
+    intent of this layer is as a reference implementation for foundational understanding
+    and thus it contains only limited features relative to newer Transformer architectures.
+    Given the fast pace of innovation in transformer-like architectures, we recommend
+    exploring this `tutorial <https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html>`_
+    to build efficient layers from building blocks in core or using higher
+    level libraries from the `PyTorch Ecosystem <https://landscape.pytorch.org/>`_.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of the intermediate layer, can be a string
+            ("relu" or "gelu") or a unary callable. Default: relu
+        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+        norm_first: if ``True``, layer norm is done prior to self attention, multihead
+            attention and feedforward operations, respectively. Otherwise it's done after.
+            Default: ``False`` (after).
+        bias: If set to ``False``, ``Linear`` and ``LayerNorm`` layers will not learn an additive
+            bias. Default: ``True``.
+
+    Examples:
+        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
+        >>> memory = torch.rand(10, 32, 512)
+        >>> tgt = torch.rand(20, 32, 512)
+        >>> out = decoder_layer(tgt, memory)
+
+    Alternatively, when ``batch_first`` is ``True``:
+        >>> decoder_layer = nn.TransformerDecoderLayer(
+        ...     d_model=512, nhead=8, batch_first=True
+        ... )
+        >>> memory = torch.rand(32, 10, 512)
+        >>> tgt = torch.rand(32, 20, 512)
+        >>> out = decoder_layer(tgt, memory)
+    """
+
+    __constants__ = ["norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: str | Callable[[Tensor], Tensor] = F.relu,
+        layer_norm_eps: float = 1e-5,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            bias=bias,
+            **factory_kwargs,
+        )
+        self.multihead_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            bias=bias,
+            **factory_kwargs,
+        )
+        # Implementation of Feedforward model
+        self.linear1 = Linear(d_model, dim_feedforward, bias=bias, **factory_kwargs)
+        self.dropout = Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs)
+
+        self.norm_first = norm_first
+        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
+        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
+        self.norm3 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+        self.dropout3 = Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            self.activation = _get_activation_fn(activation)
+        else:
+            self.activation = activation
+
+    def __setstate__(self, state):
+        if "activation" not in state:
+            state["activation"] = F.relu
+        super().__setstate__(state)
+
+    def forward(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        tgt_mask: Tensor | None = None,
+        memory_mask: Tensor | None = None,
+        tgt_key_padding_mask: Tensor | None = None,
+        memory_key_padding_mask: Tensor | None = None,
+        tgt_is_causal: bool = False,
+        memory_is_causal: bool = False,
+    ) -> Tensor:
+        r"""Pass the inputs (and mask) through the decoder layer.
+
+        Args:
+            tgt: the sequence to the decoder layer (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+            tgt_is_causal: If specified, applies a causal mask as ``tgt mask``.
+                Default: ``False``.
+                Warning:
+                ``tgt_is_causal`` provides a hint that ``tgt_mask`` is
+                the causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+            memory_is_causal: If specified, applies a causal mask as
+                ``memory mask``.
+                Default: ``False``.
+                Warning:
+                ``memory_is_causal`` provides a hint that
+                ``memory_mask`` is the causal mask. Providing incorrect
+                hints can result in incorrect execution, including
+                forward and backward compatibility.
+
+        Shape:
+            see the docs in :class:`~torch.nn.Transformer`.
+        """
+        # see Fig. 1 of https://arxiv.org/pdf/2002.04745v1.pdf
+
+        x = tgt
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x), tgt_mask, tgt_key_padding_mask, tgt_is_causal
+            )
+            x = x + self._mha_block(
+                self.norm2(x),
+                memory,
+                memory_mask,
+                memory_key_padding_mask,
+                memory_is_causal,
+            )
+            x = x + self._ff_block(self.norm3(x))
+        else:
+            x = self.norm1(
+                x + self._sa_block(x, tgt_mask, tgt_key_padding_mask, tgt_is_causal)
+            )
+            x = self.norm2(
+                x
+                + self._mha_block(
+                    x, memory, memory_mask, memory_key_padding_mask, memory_is_causal
+                )
+            )
+            x = self.norm3(x + self._ff_block(x))
+
+        return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Tensor | None,
+        key_padding_mask: Tensor | None,
+        is_causal: bool = False,
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            is_causal=is_causal,
+            need_weights=False,
+        )[0]
+        return self.dropout1(x)
+
+    # multihead attention block
+    def _mha_block(
+        self,
+        x: Tensor,
+        mem: Tensor,
+        attn_mask: Tensor | None,
+        key_padding_mask: Tensor | None,
+        is_causal: bool = False,
+    ) -> Tensor:
+        x = self.multihead_attn(
+            x,
+            mem,
+            mem,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            is_causal=is_causal,
+            need_weights=False,
+        )[0]
+        return self.dropout2(x)
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout3(x)
+
+
+def _get_clones(module, N):
+    # FIXME: copy.deepcopy() is not defined on nn.module
+    return ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}")
+
+
+def _detect_is_causal_mask(
+    mask: Tensor | None,
+    is_causal: bool | None = None,
+    size: int | None = None,
+) -> bool:
+    """Return whether the given attention mask is causal.
+
+    Warning:
+    If ``is_causal`` is not ``None``, its value will be returned as is.  If a
+    user supplies an incorrect ``is_causal`` hint,
+
+    ``is_causal=False`` when the mask is in fact a causal attention.mask
+       may lead to reduced performance relative to what would be achievable
+       with ``is_causal=True``;
+    ``is_causal=True`` when the mask is in fact not a causal attention.mask
+       may lead to incorrect and unpredictable execution - in some scenarios,
+       a causal mask may be applied based on the hint, in other execution
+       scenarios the specified mask may be used.  The choice may not appear
+       to be deterministic, in that a number of factors like alignment,
+       hardware SKU, etc influence the decision whether to use a mask or
+       rely on the hint.
+    ``size`` if not None, check whether the mask is a causal mask of the provided size
+       Otherwise, checks for any causal mask.
+    """
+    # Prevent type refinement
+    make_causal = is_causal is True
+
+    if is_causal is None and mask is not None:
+        sz = size if size is not None else mask.size(-2)
+        causal_comparison = _generate_square_subsequent_mask(
+            sz, device=mask.device, dtype=mask.dtype
+        )
+
+        # Do not use `torch.equal` so we handle batched masks by
+        # broadcasting the comparison.
+        if mask.size() == causal_comparison.size():
+            make_causal = bool((mask == causal_comparison).all())
+        else:
+            make_causal = False
+
+    return make_causal
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/upsampling.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/upsampling.py
new file mode 100644
index 0000000000000000000000000000000000000000..29e58bc6a9f3779924584e2934874a1333b3e501
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/upsampling.py
@@ -0,0 +1,298 @@
+# mypy: allow-untyped-defs
+
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.common_types import _ratio_2_t, _ratio_any_t, _size_2_t, _size_any_t
+
+from .module import Module
+
+
+__all__ = ["Upsample", "UpsamplingNearest2d", "UpsamplingBilinear2d"]
+
+
+class Upsample(Module):
+    r"""Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data.
+
+    The input data is assumed to be of the form
+    `minibatch x channels x [optional depth] x [optional height] x width`.
+    Hence, for spatial inputs, we expect a 4D Tensor and for volumetric inputs, we expect a 5D Tensor.
+
+    The algorithms available for upsampling are nearest neighbor and linear,
+    bilinear, bicubic and trilinear for 3D, 4D and 5D input Tensor,
+    respectively.
+
+    One can either give a :attr:`scale_factor` or the target output :attr:`size` to
+    calculate the output size. (You cannot give both, as it is ambiguous)
+
+    Args:
+        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int], optional):
+            output spatial sizes
+        scale_factor (float or Tuple[float] or Tuple[float, float] or Tuple[float, float, float], optional):
+            multiplier for spatial size. Has to match input size if it is a tuple.
+        mode (str, optional): the upsampling algorithm: one of ``'nearest'``,
+            ``'linear'``, ``'bilinear'``, ``'bicubic'`` and ``'trilinear'``.
+            Default: ``'nearest'``
+        align_corners (bool, optional): if ``True``, the corner pixels of the input
+            and output tensors are aligned, and thus preserving the values at
+            those pixels. This only has effect when :attr:`mode` is
+            ``'linear'``, ``'bilinear'``, ``'bicubic'``, or ``'trilinear'``.
+            Default: ``False``
+        recompute_scale_factor (bool, optional): recompute the scale_factor for use in the
+            interpolation calculation. If `recompute_scale_factor` is ``True``, then
+            `scale_factor` must be passed in and `scale_factor` is used to compute the
+            output `size`. The computed output `size` will be used to infer new scales for
+            the interpolation. Note that when `scale_factor` is floating-point, it may differ
+            from the recomputed `scale_factor` due to rounding and precision issues.
+            If `recompute_scale_factor` is ``False``, then `size` or `scale_factor` will
+            be used directly for interpolation.
+
+    Shape:
+        - Input: :math:`(N, C, W_{in})`, :math:`(N, C, H_{in}, W_{in})` or :math:`(N, C, D_{in}, H_{in}, W_{in})`
+        - Output: :math:`(N, C, W_{out})`, :math:`(N, C, H_{out}, W_{out})`
+          or :math:`(N, C, D_{out}, H_{out}, W_{out})`, where
+
+    .. math::
+        D_{out} = \left\lfloor D_{in} \times \text{scale\_factor} \right\rfloor
+
+    .. math::
+        H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor
+
+    .. math::
+        W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor
+
+    .. warning::
+        With ``align_corners = True``, the linearly interpolating modes
+        (`linear`, `bilinear`, `bicubic`, and `trilinear`) don't proportionally
+        align the output and input pixels, and thus the output values can depend
+        on the input size. This was the default behavior for these modes up to
+        version 0.3.1. Since then, the default behavior is
+        ``align_corners = False``. See below for concrete examples on how this
+        affects the outputs.
+
+    .. note::
+        If you want downsampling/general resizing, you should use :func:`~nn.functional.interpolate`.
+
+    Examples::
+
+        >>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
+        >>> input
+        tensor([[[[1., 2.],
+                  [3., 4.]]]])
+
+        >>> m = nn.Upsample(scale_factor=2, mode='nearest')
+        >>> m(input)
+        tensor([[[[1., 1., 2., 2.],
+                  [1., 1., 2., 2.],
+                  [3., 3., 4., 4.],
+                  [3., 3., 4., 4.]]]])
+
+        >>> # xdoctest: +IGNORE_WANT("other tests seem to modify printing styles")
+        >>> m = nn.Upsample(scale_factor=2, mode='bilinear')  # align_corners=False
+        >>> m(input)
+        tensor([[[[1.0000, 1.2500, 1.7500, 2.0000],
+                  [1.5000, 1.7500, 2.2500, 2.5000],
+                  [2.5000, 2.7500, 3.2500, 3.5000],
+                  [3.0000, 3.2500, 3.7500, 4.0000]]]])
+
+        >>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        >>> m(input)
+        tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],
+                  [1.6667, 2.0000, 2.3333, 2.6667],
+                  [2.3333, 2.6667, 3.0000, 3.3333],
+                  [3.0000, 3.3333, 3.6667, 4.0000]]]])
+
+        >>> # Try scaling the same data in a larger tensor
+        >>> input_3x3 = torch.zeros(3, 3).view(1, 1, 3, 3)
+        >>> input_3x3[:, :, :2, :2].copy_(input)
+        tensor([[[[1., 2.],
+                  [3., 4.]]]])
+        >>> input_3x3
+        tensor([[[[1., 2., 0.],
+                  [3., 4., 0.],
+                  [0., 0., 0.]]]])
+
+        >>> # xdoctest: +IGNORE_WANT("seems to fail when other tests are run in the same session")
+        >>> m = nn.Upsample(scale_factor=2, mode='bilinear')  # align_corners=False
+        >>> # Notice that values in top left corner are the same with the small input (except at boundary)
+        >>> m(input_3x3)
+        tensor([[[[1.0000, 1.2500, 1.7500, 1.5000, 0.5000, 0.0000],
+                  [1.5000, 1.7500, 2.2500, 1.8750, 0.6250, 0.0000],
+                  [2.5000, 2.7500, 3.2500, 2.6250, 0.8750, 0.0000],
+                  [2.2500, 2.4375, 2.8125, 2.2500, 0.7500, 0.0000],
+                  [0.7500, 0.8125, 0.9375, 0.7500, 0.2500, 0.0000],
+                  [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])
+
+        >>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        >>> # Notice that values in top left corner are now changed
+        >>> m(input_3x3)
+        tensor([[[[1.0000, 1.4000, 1.8000, 1.6000, 0.8000, 0.0000],
+                  [1.8000, 2.2000, 2.6000, 2.2400, 1.1200, 0.0000],
+                  [2.6000, 3.0000, 3.4000, 2.8800, 1.4400, 0.0000],
+                  [2.4000, 2.7200, 3.0400, 2.5600, 1.2800, 0.0000],
+                  [1.2000, 1.3600, 1.5200, 1.2800, 0.6400, 0.0000],
+                  [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])
+    """
+
+    __constants__ = [
+        "size",
+        "scale_factor",
+        "mode",
+        "align_corners",
+        "name",
+        "recompute_scale_factor",
+    ]
+    name: str
+    size: _size_any_t | None
+    scale_factor: _ratio_any_t | None
+    mode: str
+    align_corners: bool | None
+    recompute_scale_factor: bool | None
+
+    def __init__(
+        self,
+        size: _size_any_t | None = None,
+        scale_factor: _ratio_any_t | None = None,
+        mode: str = "nearest",
+        align_corners: bool | None = None,
+        recompute_scale_factor: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.name = type(self).__name__
+        self.size = size
+        if isinstance(scale_factor, tuple):
+            self.scale_factor = tuple(float(factor) for factor in scale_factor)
+        else:
+            self.scale_factor = float(scale_factor) if scale_factor else None
+        self.mode = mode
+        self.align_corners = align_corners
+        self.recompute_scale_factor = recompute_scale_factor
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Runs the forward pass.
+        """
+        return F.interpolate(
+            input,
+            self.size,
+            self.scale_factor,
+            self.mode,
+            self.align_corners,
+            recompute_scale_factor=self.recompute_scale_factor,
+        )
+
+    def __setstate__(self, state):
+        if "recompute_scale_factor" not in state:
+            state["recompute_scale_factor"] = True
+
+        super().__setstate__(state)
+
+    def extra_repr(self) -> str:
+        """
+        Return the extra representation of the module.
+        """
+        if self.scale_factor is not None:
+            info = "scale_factor=" + repr(self.scale_factor)
+        else:
+            info = "size=" + repr(self.size)
+        info += ", mode=" + repr(self.mode)
+        return info
+
+
+class UpsamplingNearest2d(Upsample):
+    r"""Applies a 2D nearest neighbor upsampling to an input signal composed of several input channels.
+
+    To specify the scale, it takes either the :attr:`size` or the :attr:`scale_factor`
+    as it's constructor argument.
+
+    When :attr:`size` is given, it is the output size of the image `(h, w)`.
+
+    Args:
+        size (int or Tuple[int, int], optional): output spatial sizes
+        scale_factor (float or Tuple[float, float], optional): multiplier for
+            spatial size.
+
+    .. warning::
+        This class is deprecated in favor of :func:`~nn.functional.interpolate`.
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})`
+        - Output: :math:`(N, C, H_{out}, W_{out})` where
+
+    .. math::
+          H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor
+
+    .. math::
+          W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor
+
+    Examples::
+
+        >>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
+        >>> input
+        tensor([[[[1., 2.],
+                  [3., 4.]]]])
+
+        >>> m = nn.UpsamplingNearest2d(scale_factor=2)
+        >>> m(input)
+        tensor([[[[1., 1., 2., 2.],
+                  [1., 1., 2., 2.],
+                  [3., 3., 4., 4.],
+                  [3., 3., 4., 4.]]]])
+    """
+
+    def __init__(
+        self,
+        size: _size_2_t | None = None,
+        scale_factor: _ratio_2_t | None = None,
+    ) -> None:
+        super().__init__(size, scale_factor, mode="nearest")
+
+
+class UpsamplingBilinear2d(Upsample):
+    r"""Applies a 2D bilinear upsampling to an input signal composed of several input channels.
+
+    To specify the scale, it takes either the :attr:`size` or the :attr:`scale_factor`
+    as it's constructor argument.
+
+    When :attr:`size` is given, it is the output size of the image `(h, w)`.
+
+    Args:
+        size (int or Tuple[int, int], optional): output spatial sizes
+        scale_factor (float or Tuple[float, float], optional): multiplier for
+            spatial size.
+
+    .. warning::
+        This class is deprecated in favor of :func:`~nn.functional.interpolate`. It is
+        equivalent to ``nn.functional.interpolate(..., mode='bilinear', align_corners=True)``.
+
+    Shape:
+        - Input: :math:`(N, C, H_{in}, W_{in})`
+        - Output: :math:`(N, C, H_{out}, W_{out})` where
+
+    .. math::
+        H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor
+
+    .. math::
+        W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor
+
+    Examples::
+
+        >>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
+        >>> input
+        tensor([[[[1., 2.],
+                  [3., 4.]]]])
+
+        >>> # xdoctest: +IGNORE_WANT("do other tests modify the global state?")
+        >>> m = nn.UpsamplingBilinear2d(scale_factor=2)
+        >>> m(input)
+        tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],
+                  [1.6667, 2.0000, 2.3333, 2.6667],
+                  [2.3333, 2.6667, 3.0000, 3.3333],
+                  [3.0000, 3.3333, 3.6667, 4.0000]]]])
+    """
+
+    def __init__(
+        self,
+        size: _size_2_t | None = None,
+        scale_factor: _ratio_2_t | None = None,
+    ) -> None:
+        super().__init__(size, scale_factor, mode="bilinear", align_corners=True)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..63a1f4a9bb84c05caffcc22ee1bf9fc540df39e5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/modules/utils.py
@@ -0,0 +1,82 @@
+# mypy: allow-untyped-defs
+import collections
+from itertools import repeat
+from typing import Any
+
+
+__all__ = ["consume_prefix_in_state_dict_if_present"]
+
+
+def _ntuple(n, name="parse"):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return tuple(x)
+        return tuple(repeat(x, n))
+
+    parse.__name__ = name
+    return parse
+
+
+_single = _ntuple(1, "_single")
+_pair = _ntuple(2, "_pair")
+_triple = _ntuple(3, "_triple")
+_quadruple = _ntuple(4, "_quadruple")
+
+
+def _reverse_repeat_tuple(t, n):
+    r"""Reverse the order of `t` and repeat each element for `n` times.
+
+    This can be used to translate padding arg used by Conv and Pooling modules
+    to the ones used by `F.pad`.
+    """
+    return tuple(x for x in reversed(t) for _ in range(n))
+
+
+def _list_with_default(out_size: list[int], defaults: list[int]) -> list[int]:
+    import torch
+
+    if isinstance(out_size, (int, torch.SymInt)):
+        return out_size
+    if len(defaults) <= len(out_size):
+        raise ValueError(f"Input dimension should be at least {len(out_size) + 1}")
+    return [
+        v if v is not None else d
+        for v, d in zip(out_size, defaults[-len(out_size) :], strict=False)
+    ]
+
+
+def consume_prefix_in_state_dict_if_present(
+    state_dict: dict[str, Any],
+    prefix: str,
+) -> None:
+    r"""Strip the prefix in state_dict in place, if any.
+
+    .. note::
+        Given a `state_dict` from a DP/DDP model, a local model can load it by applying
+        `consume_prefix_in_state_dict_if_present(state_dict, "module.")` before calling
+        :meth:`torch.nn.Module.load_state_dict`.
+
+    Args:
+        state_dict (OrderedDict): a state-dict to be loaded to the model.
+        prefix (str): prefix.
+    """
+    keys = list(state_dict.keys())
+    for key in keys:
+        if key.startswith(prefix):
+            newkey = key[len(prefix) :]
+            state_dict[newkey] = state_dict.pop(key)
+
+    # also strip the prefix in metadata if any.
+    if hasattr(state_dict, "_metadata"):
+        keys = list(state_dict._metadata.keys())
+        for key in keys:
+            # for the metadata dict, the key can be:
+            # '': for the DDP module, which we want to remove.
+            # 'module': for the actual model.
+            # 'module.xx.xx': for the rest.
+            if len(key) == 0:
+                continue
+            # handling both, 'module' case and  'module.' cases
+            if key == prefix.replace(".", "") or key.startswith(prefix):
+                newkey = key[len(prefix) :]
+                state_dict._metadata[newkey] = state_dict._metadata.pop(key)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ad8648d10aadc8dec59ea7ebc54aa77cd60ee4f5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/__init__.py
@@ -0,0 +1,27 @@
+from typing_extensions import deprecated
+
+from torch.nn.parallel.data_parallel import data_parallel, DataParallel
+from torch.nn.parallel.distributed import DistributedDataParallel
+from torch.nn.parallel.parallel_apply import parallel_apply
+from torch.nn.parallel.replicate import replicate
+from torch.nn.parallel.scatter_gather import gather, scatter
+
+
+__all__ = [
+    "replicate",
+    "scatter",
+    "parallel_apply",
+    "gather",
+    "data_parallel",
+    "DataParallel",
+    "DistributedDataParallel",
+]
+
+
+@deprecated(
+    "`torch.nn.parallel.DistributedDataParallelCPU` is deprecated, "
+    "please use `torch.nn.parallel.DistributedDataParallel` instead.",
+    category=FutureWarning,
+)
+class DistributedDataParallelCPU(DistributedDataParallel):
+    pass
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/_functions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5e82285e9747491d427f9db8b10b6d987deaf53
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/_functions.py
@@ -0,0 +1,176 @@
+import warnings
+from itertools import chain
+
+import torch
+from torch._utils import _get_device_index
+from torch.autograd import Function
+from torch.nn.parallel import comm
+
+
+class Broadcast(Function):
+    @staticmethod
+    def forward(ctx, target_gpus, *inputs):
+        if not all(i.device.type != "cpu" for i in inputs):
+            raise AssertionError("Broadcast function not implemented for CPU tensors")
+        target_gpus = [_get_device_index(x, True) for x in target_gpus]
+        ctx.target_gpus = target_gpus
+        if len(inputs) == 0:
+            return ()
+        ctx.num_inputs = len(inputs)
+        ctx.input_device = inputs[0].get_device()
+
+        ctx.complex_mask = [inp.is_complex() for inp in inputs]
+
+        outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)
+
+        for device_outputs in outputs:
+            for i, is_complex in enumerate(ctx.complex_mask):
+                if is_complex:
+                    device_outputs[i] = torch.view_as_complex(device_outputs[i])
+
+        non_differentiables = []
+        for idx, input_requires_grad in enumerate(ctx.needs_input_grad[1:]):
+            if not input_requires_grad:
+                non_differentiables.extend(output[idx] for output in outputs)
+        ctx.mark_non_differentiable(*non_differentiables)
+        return tuple(chain.from_iterable(outputs))
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        grads = ReduceAddCoalesced.apply(
+            ctx.input_device, ctx.num_inputs, *grad_outputs
+        )
+
+        return (None,) + grads
+
+
+class ReduceAddCoalesced(Function):
+    @staticmethod
+    def forward(ctx, destination, num_inputs, *grads):
+        ctx.target_gpus = [
+            grads[i].get_device() for i in range(0, len(grads), num_inputs)
+        ]
+
+        complex_mask = [grads[i].is_complex() for i in range(num_inputs)]
+        ctx.complex_mask = complex_mask
+
+        grads_converted = tuple(
+            torch.view_as_real(g) if g.is_complex() else g for g in grads
+        )
+
+        grads_ = [
+            grads_converted[i : i + num_inputs]
+            for i in range(0, len(grads_converted), num_inputs)
+        ]
+        results = comm.reduce_add_coalesced(grads_, destination)
+
+        results = tuple(
+            torch.view_as_complex(r) if is_complex else r
+            for r, is_complex in zip(results, complex_mask)
+        )
+
+        return results
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        return (
+            None,
+            None,
+        ) + Broadcast.apply(ctx.target_gpus, *grad_outputs)
+
+
+class Gather(Function):
+    @staticmethod
+    def forward(ctx, target_device, dim, *inputs):
+        if not all(i.device.type != "cpu" for i in inputs):
+            raise AssertionError("Gather function not implemented for CPU tensors")
+        if target_device == "cpu":
+            ctx.target_device = "cpu"
+        else:
+            target_device = _get_device_index(target_device, True)
+            ctx.target_device = target_device
+        ctx.dim = dim
+        ctx.input_gpus = tuple(i.get_device() for i in inputs)
+        if all(t.dim() == 0 for t in inputs) and dim == 0:
+            inputs = tuple(t.view(1) for t in inputs)
+            warnings.warn(
+                "Was asked to gather along dimension 0, but all "
+                "input tensors were scalars; will instead unsqueeze "
+                "and return a vector.",
+                stacklevel=2,
+            )
+            ctx.unsqueezed_scalar = True
+        else:
+            ctx.unsqueezed_scalar = False
+        ctx.input_sizes = tuple(i.size(ctx.dim) for i in inputs)
+
+        is_complex = len(inputs) > 0 and inputs[0].is_complex()
+
+        output = comm.gather(inputs, ctx.dim, ctx.target_device)
+
+        if is_complex:
+            output = torch.view_as_complex(output)
+
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        scattered_grads = Scatter.apply(
+            ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output
+        )
+        if ctx.unsqueezed_scalar:
+            scattered_grads = tuple(g[0] for g in scattered_grads)
+        return (None, None) + scattered_grads
+
+
+class Scatter(Function):
+    @staticmethod
+    def forward(ctx, target_gpus, chunk_sizes, dim, input):
+        target_gpus = [_get_device_index(x, True) for x in target_gpus]
+        ctx.dim = dim
+        ctx.input_device = input.get_device() if input.device.type != "cpu" else -1
+        streams = None
+        if torch.accelerator.is_available() and ctx.input_device == -1:
+            # Perform CPU to GPU copies in a background stream
+            streams = [_get_stream(torch.device(device)) for device in target_gpus]
+
+        is_complex = input.is_complex()
+
+        outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams)
+
+        if is_complex:
+            outputs = tuple(torch.view_as_complex(o) for o in outputs)
+
+        # Synchronize with the copy stream
+        if streams is not None:
+            for i, output in enumerate(outputs):
+                with torch.accelerator.device_index(target_gpus[i]):
+                    main_stream = torch.accelerator.current_stream()
+                    main_stream.wait_stream(streams[i])
+                    output.record_stream(main_stream)
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *grad_output):
+        return None, None, None, Gather.apply(ctx.input_device, ctx.dim, *grad_output)
+
+
+# background streams used for copying
+_streams: list[torch.Stream | None] | None = None
+
+
+def _get_stream(device: torch.device):
+    """Get a background stream for copying between CPU and target device."""
+    global _streams
+    if device.type == "cpu" or not torch.accelerator.is_available():
+        return None
+    if torch.accelerator.current_accelerator().type != device.type:
+        raise AssertionError(
+            f"Expected current accelerator type {torch.accelerator.current_accelerator().type} "
+            f"to match device type {device.type}"
+        )
+    if _streams is None:
+        _streams = [None] * torch.accelerator.device_count()
+    if _streams[device.index] is None:
+        _streams[device.index] = torch.Stream(device.index)
+    return _streams[device.index]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/comm.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/comm.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d79196bc2082d767351ee0fe3f075ab4997bc6d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/comm.py
@@ -0,0 +1,264 @@
+# mypy: allow-untyped-defs
+import warnings
+
+import torch
+from torch._utils import (
+    _flatten_dense_tensors,
+    _get_device_index,
+    _handle_complex,
+    _reorder_tensors_as,
+    _take_tensors,
+    _unflatten_dense_tensors,
+)
+from torch.cuda import nccl
+
+
+def broadcast(tensor, devices=None, *, out=None):
+    r"""Broadcasts a tensor to specified GPU devices.
+
+    Args:
+        tensor (Tensor): tensor to broadcast. Can be on CPU or GPU.
+        devices (Iterable[torch.device, str or int], optional): an iterable of
+          GPU devices, among which to broadcast.
+        out (Sequence[Tensor], optional, keyword-only): the GPU tensors to
+          store output results.
+
+    .. note::
+        Exactly one of :attr:`devices` and :attr:`out` must be specified.
+
+    Returns:
+        - If :attr:`devices` is specified,
+            a tuple containing copies of :attr:`tensor`, placed on
+            :attr:`devices`.
+        - If :attr:`out` is specified,
+            a tuple containing :attr:`out` tensors, each containing a copy of
+            :attr:`tensor`.
+    """
+    tensor = _handle_complex(tensor)
+    if not ((devices is None) ^ (out is None)):
+        raise RuntimeError(
+            f"Exactly one of 'devices' and 'out' must be specified, but got devices={devices} and out={out}"
+        )
+    if devices is not None:
+        devices = [_get_device_index(d) for d in devices]
+        return torch._C._broadcast(tensor, devices)
+    else:
+        # pyrefly: ignore [bad-argument-type]
+        return torch._C._broadcast_out(tensor, out)
+
+
+def broadcast_coalesced(tensors, devices, buffer_size=10485760):
+    """Broadcast a sequence of tensors to the specified GPUs.
+
+    Small tensors are first coalesced into a buffer to reduce the number of synchronizations.
+
+    Args:
+        tensors (sequence): tensors to broadcast. Must be on the same device,
+          either CPU or GPU.
+        devices (Iterable[torch.device, str or int]): an iterable of GPU
+          devices, among which to broadcast.
+        buffer_size (int): maximum size of the buffer used for coalescing
+
+    Returns:
+        A tuple containing copies of :attr:`tensor`, placed on :attr:`devices`.
+    """
+    devices = [_get_device_index(d) for d in devices]
+    tensors = [_handle_complex(t) for t in tensors]
+    return torch._C._broadcast_coalesced(tensors, devices, buffer_size)
+
+
+def reduce_add(inputs, destination=None):
+    """Sum tensors from multiple GPUs.
+
+    All inputs should have matching shapes, dtype, and layout. The output tensor
+    will be of the same shape, dtype, and layout.
+
+    Args:
+        inputs (Iterable[Tensor]): an iterable of tensors to add.
+        destination (int, optional): a device on which the output will be
+            placed (default: current device).
+
+    Returns:
+        A tensor containing an elementwise sum of all inputs, placed on the
+        :attr:`destination` device.
+    """
+    destination = _get_device_index(destination, optional=True)
+    input_size = inputs[0].size()
+    root_index = None  # index of input tensor that already is on the correct device
+    for i, inp in enumerate(inputs):
+        if inp.device.type == "cpu":
+            raise AssertionError(
+                f"reduce_add expects all inputs to be on GPUs, but input {i} is on CPU"
+            )
+        if inp.get_device() == destination:
+            root_index = i
+        if inp.size() != input_size:
+            got = "x".join(str(x) for x in inp.size())
+            expected = "x".join(str(x) for x in input_size)
+            raise ValueError(
+                f"input {i} has invalid size: got {got}, but expected {expected}"
+            )
+    if root_index is None:
+        raise RuntimeError(
+            "reduce_add expects destination to be on the same GPU with one of the tensors"
+        )
+
+    if len(inputs) == 1:
+        return inputs[0]
+
+    if nccl.is_available(inputs):
+        result = torch.empty_like(inputs[root_index])
+        nccl.reduce(inputs, output=result, root=root_index)
+    else:
+        destination_device = torch.device(inputs[root_index].device.type, destination)
+        nonroot = [t for i, t in enumerate(inputs) if i != root_index]
+        # make a new tensor w/o clone
+        result = inputs[root_index] + nonroot[0].to(
+            device=destination_device, non_blocking=True
+        )
+        for other in nonroot[1:]:
+            result.add_(other.to(device=destination_device, non_blocking=True))
+    return result
+
+
+def reduce_add_coalesced(inputs, destination=None, buffer_size=10485760):
+    """Sum tensors from multiple GPUs.
+
+    Small tensors are first coalesced into a buffer to reduce the number
+    of synchronizations.
+
+    Args:
+        inputs (Iterable[Iterable[Tensor]]): iterable of iterables that
+            contain tensors from a single device.
+        destination (int, optional): a device on which the output will be
+            placed (default: current device).
+        buffer_size (int): maximum size of the buffer used for coalescing
+
+    Returns:
+        A tuple of tensors containing an elementwise sum of each group of
+        inputs, placed on the ``destination`` device.
+    """
+    # TODO: When `len(inputs) == 1` and all inputs are on `destination`, just
+    #       return `inputs`.
+    dense_tensors: list[list] = [[] for _ in inputs]  # shape (num_gpus, num_tensors)
+    output = []
+    ref_order = []
+    # process sparse ones first since they may have different sizes on different gpus
+    for tensor_at_gpus in zip(*inputs, strict=True):
+        if all(t.is_sparse for t in tensor_at_gpus):
+            result = reduce_add(tensor_at_gpus, destination)  # this will be sparse too
+            output.append(result)
+            ref_order.append(tensor_at_gpus[0])
+        else:
+            for coll, t in zip(dense_tensors, tensor_at_gpus, strict=True):
+                coll.append(t.to_dense() if t.is_sparse else t)
+            ref_order.append(dense_tensors[0][-1])
+    itrs = [_take_tensors(tensors, buffer_size) for tensors in dense_tensors]
+    # now the dense ones, which have consistent sizes
+    for chunks in zip(*itrs, strict=True):
+        flat_tensors = [
+            _flatten_dense_tensors(chunk) for chunk in chunks
+        ]  # (num_gpus,)
+        flat_result = reduce_add(flat_tensors, destination)
+        for t in _unflatten_dense_tensors(flat_result, chunks[0]):
+            # The unflattened tensors do not share storage, and we don't expose
+            # base flat tensor anyways, so give them different version counters.
+            # See NOTE [ Version Counter in comm.*_coalesced ]
+            output.append(t.data)
+    return tuple(_reorder_tensors_as(output, ref_order))
+
+
+def scatter(tensor, devices=None, chunk_sizes=None, dim=0, streams=None, *, out=None):
+    """Scatters tensor across multiple GPUs.
+
+    Args:
+        tensor (Tensor): tensor to scatter. Can be on CPU or GPU.
+        devices (Iterable[torch.device, str or int], optional): an iterable of
+          GPU devices, among which to scatter.
+        chunk_sizes (Iterable[int], optional): sizes of chunks to be placed on
+          each device. It should match :attr:`devices` in length and sums to
+          ``tensor.size(dim)``. If not specified, :attr:`tensor` will be divided
+          into equal chunks.
+        dim (int, optional): A dimension along which to chunk :attr:`tensor`.
+          Default: ``0``.
+        streams (Iterable[torch.cuda.Stream], optional): an iterable of Streams, among
+          which to execute the scatter. If not specified, the default stream will
+          be utilized.
+        out (Sequence[Tensor], optional, keyword-only): the GPU tensors to
+          store output results. Sizes of these tensors must match that of
+          :attr:`tensor`, except for :attr:`dim`, where the total size must
+          sum to ``tensor.size(dim)``.
+
+    .. note::
+        Exactly one of :attr:`devices` and :attr:`out` must be specified. When
+        :attr:`out` is specified, :attr:`chunk_sizes` must not be specified and
+        will be inferred from sizes of :attr:`out`.
+
+    Returns:
+        - If :attr:`devices` is specified,
+            a tuple containing chunks of :attr:`tensor`, placed on
+            :attr:`devices`.
+        - If :attr:`out` is specified,
+            a tuple containing :attr:`out` tensors, each containing a chunk of
+            :attr:`tensor`.
+    """
+    tensor = _handle_complex(tensor)
+    if out is None:
+        # pyrefly: ignore [not-iterable]
+        devices = [_get_device_index(d) for d in devices]
+        return tuple(torch._C._scatter(tensor, devices, chunk_sizes, dim, streams))
+    else:
+        if devices is not None:
+            raise RuntimeError(
+                f"'devices' must not be specified when 'out' is specified, but got devices={devices}"
+            )
+        if chunk_sizes is not None:
+            raise RuntimeError(
+                f"'chunk_sizes' must not be specified when 'out' is specified, but got chunk_sizes={chunk_sizes}"
+            )
+        return tuple(torch._C._scatter_out(tensor, out, dim, streams))
+
+
+def gather(tensors, dim=0, destination=None, *, out=None):
+    r"""Gathers tensors from multiple GPU devices.
+
+    Args:
+        tensors (Iterable[Tensor]): an iterable of tensors to gather.
+          Tensor sizes in all dimensions other than :attr:`dim` have to match.
+        dim (int, optional): a dimension along which the tensors will be
+          concatenated. Default: ``0``.
+        destination (torch.device, str, or int, optional): the output device.
+          Can be CPU or CUDA. Default: the current CUDA device.
+        out (Tensor, optional, keyword-only): the tensor to store gather result.
+          Its sizes must match those of :attr:`tensors`, except for :attr:`dim`,
+          where the size must equal ``sum(tensor.size(dim) for tensor in tensors)``.
+          Can be on CPU or CUDA.
+
+    .. note::
+        :attr:`destination` must not be specified when :attr:`out` is specified.
+
+    Returns:
+        - If :attr:`destination` is specified,
+            a tensor located on :attr:`destination` device, that is a result of
+            concatenating :attr:`tensors` along :attr:`dim`.
+        - If :attr:`out` is specified,
+            the :attr:`out` tensor, now containing results of concatenating
+            :attr:`tensors` along :attr:`dim`.
+    """
+    tensors = [_handle_complex(t) for t in tensors]
+    if out is None:
+        if destination == -1:
+            warnings.warn(
+                "Using -1 to represent CPU tensor is deprecated. Please use a "
+                'device object or string instead, e.g., "cpu".',
+                FutureWarning,
+                stacklevel=2,
+            )
+        destination = _get_device_index(destination, allow_cpu=True, optional=True)
+        return torch._C._gather(tensors, dim, destination)
+    else:
+        if destination is not None:
+            raise RuntimeError(
+                f"'destination' must not be specified when 'out' is specified, but got destination={destination}"
+            )
+        return torch._C._gather_out(tensors, out, dim)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py
new file mode 100644
index 0000000000000000000000000000000000000000..16537624663b1e350c26fe3276f503e42c26cafa
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py
@@ -0,0 +1,289 @@
+# mypy: allow-untyped-defs
+import operator
+import warnings
+from collections.abc import Sequence
+from itertools import chain
+from typing import Any, Generic, TypeVar
+
+import torch
+from torch._utils import (
+    _get_all_device_indices,
+    _get_available_device_type,
+    _get_device_index,
+    _get_devices_properties,
+)
+from torch.nn.modules import Module
+from torch.nn.parallel.parallel_apply import parallel_apply
+from torch.nn.parallel.replicate import replicate
+from torch.nn.parallel.scatter_gather import gather, scatter_kwargs
+
+
+__all__ = ["DataParallel", "data_parallel"]
+
+
+def _check_balance(device_ids: Sequence[int | torch.device]) -> None:
+    imbalance_warn = """
+    There is an imbalance between your GPUs. You may want to exclude GPU {} which
+    has less than 75% of the memory or cores of GPU {}. You can do so by setting
+    the device_ids argument to DataParallel, or by setting the CUDA_VISIBLE_DEVICES
+    environment variable."""
+    device_ids = [_get_device_index(x, True) for x in device_ids]
+    dev_props = _get_devices_properties(device_ids)
+
+    def warn_imbalance(get_prop) -> bool:
+        values = [get_prop(props) for props in dev_props]
+        min_pos, min_val = min(enumerate(values), key=operator.itemgetter(1))
+        max_pos, max_val = max(enumerate(values), key=operator.itemgetter(1))
+        if min_val / max_val < 0.75:
+            warnings.warn(
+                imbalance_warn.format(device_ids[min_pos], device_ids[max_pos]),
+                stacklevel=2,
+            )
+            return True
+        return False
+
+    if warn_imbalance(lambda props: props.total_memory):
+        return
+    if warn_imbalance(lambda props: props.multi_processor_count):
+        return
+
+
+T = TypeVar("T", bound=Module)
+
+
+class DataParallel(Module, Generic[T]):
+    r"""Implements data parallelism at the module level.
+
+    This container parallelizes the application of the given :attr:`module` by
+    splitting the input across the specified devices by chunking in the batch
+    dimension (other objects will be copied once per device). In the forward
+    pass, the module is replicated on each device, and each replica handles a
+    portion of the input. During the backwards pass, gradients from each replica
+    are summed into the original module.
+
+    The batch size should be larger than the number of GPUs used.
+
+    .. warning::
+        It is recommended to use :class:`~torch.nn.parallel.DistributedDataParallel`,
+        instead of this class, to do multi-GPU training, even if there is only a single
+        node. See: :ref:`cuda-nn-ddp-instead` and :ref:`ddp`.
+
+    Arbitrary positional and keyword inputs are allowed to be passed into
+    DataParallel but some types are specially handled. tensors will be
+    **scattered** on dim specified (default 0). tuple, list and dict types will
+    be shallow copied. The other types will be shared among different threads
+    and can be corrupted if written to in the model's forward pass.
+
+    The parallelized :attr:`module` must have its parameters and buffers on
+    ``device_ids[0]`` before running this :class:`~torch.nn.DataParallel`
+    module.
+
+    .. warning::
+        In each forward, :attr:`module` is **replicated** on each device, so any
+        updates to the running module in ``forward`` will be lost. For example,
+        if :attr:`module` has a counter attribute that is incremented in each
+        ``forward``, it will always stay at the initial value because the update
+        is done on the replicas which are destroyed after ``forward``. However,
+        :class:`~torch.nn.DataParallel` guarantees that the replica on
+        ``device[0]`` will have its parameters and buffers sharing storage with
+        the base parallelized :attr:`module`. So **in-place** updates to the
+        parameters or buffers on ``device[0]`` will be recorded. E.g.,
+        :class:`~torch.nn.BatchNorm2d` and :func:`~torch.nn.utils.spectral_norm`
+        rely on this behavior to update the buffers.
+
+    .. warning::
+        Forward and backward hooks defined on :attr:`module` and its submodules
+        will be invoked ``len(device_ids)`` times, each with inputs located on
+        a particular device. Particularly, the hooks are only guaranteed to be
+        executed in correct order with respect to operations on corresponding
+        devices. For example, it is not guaranteed that hooks set via
+        :meth:`~torch.nn.Module.register_forward_pre_hook` be executed before
+        `all` ``len(device_ids)`` :meth:`~torch.nn.Module.forward` calls, but
+        that each such hook be executed before the corresponding
+        :meth:`~torch.nn.Module.forward` call of that device.
+
+    .. warning::
+        When :attr:`module` returns a scalar (i.e., 0-dimensional tensor) in
+        :func:`forward`, this wrapper will return a vector of length equal to
+        number of devices used in data parallelism, containing the result from
+        each device.
+
+    .. note::
+        There is a subtlety in using the
+        ``pack sequence -> recurrent network -> unpack sequence`` pattern in a
+        :class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`.
+        See :ref:`pack-rnn-unpack-with-data-parallelism` section in FAQ for
+        details.
+
+
+    Args:
+        module (Module): module to be parallelized
+        device_ids (list of int or torch.device): CUDA devices (default: all devices)
+        output_device (int or torch.device): device location of output (default: device_ids[0])
+
+    Attributes:
+        module (Module): the module to be parallelized
+
+    Example::
+
+        >>> # xdoctest: +SKIP
+        >>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
+        >>> output = net(input_var)  # input_var can be on any device, including CPU
+    """
+
+    # TODO: update notes/cuda.rst when this class handles 8+ GPUs well
+
+    def __init__(
+        self,
+        module: T,
+        device_ids: Sequence[int | torch.device] | None = None,
+        output_device: int | torch.device | None = None,
+        dim: int = 0,
+    ) -> None:
+        super().__init__()
+        torch._C._log_api_usage_once("torch.nn.parallel.DataParallel")
+        device_type = _get_available_device_type()
+        if device_type is None or device_type == "mps":
+            self.module = module
+            self.device_ids = []
+            return
+
+        if device_ids is None:
+            device_ids = _get_all_device_indices()
+
+        if device_ids is None:
+            raise RuntimeError("no available devices were found")
+
+        if output_device is None:
+            output_device = device_ids[0]
+
+        self.dim = dim
+        self.module = module
+        self.device_ids = [_get_device_index(x, True) for x in device_ids]
+        self.output_device = _get_device_index(output_device, True)
+        self.src_device_obj = torch.device(device_type, self.device_ids[0])
+
+        if device_type == "cuda":
+            _check_balance(self.device_ids)
+
+        if len(self.device_ids) == 1:
+            self.module.to(self.src_device_obj)
+
+    def forward(self, *inputs: Any, **kwargs: Any) -> Any:
+        with torch.autograd.profiler.record_function("DataParallel.forward"):
+            if not self.device_ids:
+                return self.module(*inputs, **kwargs)
+
+            # pyrefly: ignore [bad-argument-type]
+            for t in chain(self.module.parameters(), self.module.buffers()):
+                if t.device != self.src_device_obj:
+                    raise RuntimeError(
+                        "module must have its parameters and buffers "
+                        f"on device {self.src_device_obj} (device_ids[0]) but found one of "
+                        f"them on device: {t.device}"
+                    )
+
+            inputs, module_kwargs = self.scatter(inputs, kwargs, self.device_ids)
+            # for forward function without any inputs, empty list and dict will be created
+            # so the module can be executed on one device which is the first one in device_ids
+            if not inputs and not module_kwargs:
+                inputs = ((),)
+                module_kwargs = ({},)
+
+            if len(self.device_ids) == 1:
+                return self.module(*inputs[0], **module_kwargs[0])
+            replicas = self.replicate(self.module, self.device_ids[: len(inputs)])
+            outputs = self.parallel_apply(replicas, inputs, module_kwargs)
+            return self.gather(outputs, self.output_device)
+
+    def replicate(self, module: T, device_ids: Sequence[int | torch.device]) -> list[T]:
+        return replicate(module, device_ids, not torch.is_grad_enabled())
+
+    def scatter(
+        self,
+        inputs: tuple[Any, ...],
+        kwargs: dict[str, Any] | None,
+        device_ids: Sequence[int | torch.device],
+    ) -> Any:
+        return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
+
+    def parallel_apply(
+        self, replicas: Sequence[T], inputs: Sequence[Any], kwargs: Any
+    ) -> list[Any]:
+        return parallel_apply(
+            replicas, inputs, kwargs, self.device_ids[: len(replicas)]
+        )
+
+    def gather(self, outputs: Any, output_device: int | torch.device) -> Any:
+        return gather(outputs, output_device, dim=self.dim)
+
+
+def data_parallel(
+    module: Module,
+    inputs: Any,
+    device_ids: Sequence[int | torch.device] | None = None,
+    output_device: int | torch.device | None = None,
+    dim: int = 0,
+    module_kwargs: Any | None = None,
+) -> torch.Tensor:
+    r"""Evaluate module(input) in parallel across the GPUs given in device_ids.
+
+    This is the functional version of the DataParallel module.
+
+    Args:
+        module (Module): the module to evaluate in parallel
+        inputs (Tensor): inputs to the module
+        device_ids (list of int or torch.device): GPU ids on which to replicate module
+        output_device (list of int or torch.device): GPU location of the output  Use -1 to indicate the CPU.
+            (default: device_ids[0])
+    Returns:
+        a Tensor containing the result of module(input) located on
+        output_device
+    """
+    if not isinstance(inputs, tuple):
+        inputs = (inputs,) if inputs is not None else ()
+
+    device_type = _get_available_device_type()
+
+    if device_type is None:
+        raise RuntimeError("device type could not be determined")
+
+    if device_ids is None:
+        device_ids = _get_all_device_indices()
+
+    if device_ids is None:
+        raise RuntimeError("no available devices were found")
+
+    if output_device is None:
+        output_device = device_ids[0]
+
+    device_ids = [_get_device_index(x, True) for x in device_ids]
+    output_device = _get_device_index(output_device, True)
+    # pyrefly: ignore [bad-argument-type, no-matching-overload]
+    src_device_obj = torch.device(device_type, device_ids[0])
+
+    # pyrefly: ignore [bad-argument-type]
+    for t in chain(module.parameters(), module.buffers()):
+        if t.device != src_device_obj:
+            raise RuntimeError(
+                "module must have its parameters and buffers "
+                f"on device {src_device_obj} (device_ids[0]) but found one of "
+                f"them on device: {t.device}"
+            )
+
+    inputs, module_kwargs = scatter_kwargs(inputs, module_kwargs, device_ids, dim)
+    # for module without any inputs, empty list and dict will be created
+    # so the module can be executed on one device which is the first one in device_ids
+    if not inputs and not module_kwargs:
+        inputs = ((),)
+        module_kwargs = ({},)
+
+    if module_kwargs is None:
+        raise AssertionError("module_kwargs should not be None after scatter_kwargs")
+
+    if len(device_ids) == 1:
+        return module(*inputs[0], **module_kwargs[0])
+    used_device_ids = device_ids[: len(inputs)]
+    replicas = replicate(module, used_device_ids)
+    outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
+    return gather(outputs, output_device, dim)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/distributed.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/distributed.py
new file mode 100644
index 0000000000000000000000000000000000000000..50f230b995fbacd5c93a237e85da8840e102d8db
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/distributed.py
@@ -0,0 +1,2608 @@
+# mypy: allow-untyped-defs
+import copy
+import functools
+import inspect
+import itertools
+import logging
+import os
+import sys
+import warnings
+import weakref
+from collections import defaultdict, deque
+from collections.abc import Callable
+from contextlib import contextmanager
+from dataclasses import dataclass, fields, is_dataclass
+from enum import auto, Enum
+from typing import Any, Optional, TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+from torch._utils import _get_device_index
+from torch.autograd import Function, Variable
+from torch.distributed.algorithms.join import Join, Joinable, JoinHook
+from torch.nn.modules import Module
+from torch.nn.parallel.scatter_gather import gather, scatter_kwargs
+from torch.utils._pytree import tree_flatten, tree_unflatten
+
+
+RPC_AVAILABLE = False
+
+# Default bucket size in MiB for gradient reduction
+_DEFAULT_BUCKET_CAP_MB = 25
+# Conversion factor from MiB to bytes
+_MB_TO_BYTES = 1024 * 1024
+
+
+@dataclass(frozen=True)
+class _BucketCapacityConfig:
+    """Configuration for DDP gradient reduction bucket capacities.
+
+    This immutable dataclass encapsulates bucket size settings for
+    DistributedDataParallel. Use the `create()` factory method to construct.
+
+    The bucket capacity determines how parameters are grouped for AllReduce:
+    - Smaller buckets: more frequent, smaller AllReduce operations
+    - Larger buckets: less frequent, larger AllReduce operations
+
+    Attributes:
+        bucket_bytes_cap: The primary bucket size limit in bytes. This is passed
+            to the C++ Reducer as the default bucket size.
+
+        per_bucket_bytes_caps: Per-bucket size limits in bytes when user specifies
+            bucket_cap_mb_list. Empty tuple means use uniform bucket_bytes_cap.
+
+        first_bucket_bytes_cap: Size limit for the first bucket in bytes. Can be
+            smaller than bucket_bytes_cap to reduce latency (first bucket contains
+            parameters whose gradients complete last, so a smaller first bucket
+            allows AllReduce to start sooner).
+    """
+
+    bucket_bytes_cap: int
+    per_bucket_bytes_caps: tuple[int, ...]
+    first_bucket_bytes_cap: int
+
+    @classmethod
+    def create(
+        cls,
+        bucket_cap_mb: int | None,
+        bucket_cap_mb_list: list[int] | None,
+        use_python_reducer: bool,
+    ) -> "_BucketCapacityConfig":
+        """Factory method to create a BucketCapacityConfig from user inputs.
+
+        Args:
+            bucket_cap_mb: Single bucket size limit in MiB. If None, uses default
+                of 25 MiB. Ignored if bucket_cap_mb_list is provided.
+            bucket_cap_mb_list: Per-bucket size limits in MiB. When provided, each
+                bucket gets its own size limit.
+            use_python_reducer: Whether Python reducer is being used (for validation).
+
+        Returns:
+            A new BucketCapacityConfig instance.
+
+        Raises:
+            AssertionError: If bucket_cap_mb_list is provided with Python reducer.
+        """
+        is_using_default = bucket_cap_mb is None
+        effective_bucket_cap_mb = (
+            bucket_cap_mb if bucket_cap_mb is not None else _DEFAULT_BUCKET_CAP_MB
+        )
+
+        # Process per-bucket size list if provided
+        per_bucket_bytes_caps: tuple[int, ...] = ()
+        if bucket_cap_mb_list:
+            if use_python_reducer:
+                raise AssertionError(
+                    "when using bucket_cap_mb_list, python reducer is not supported"
+                )
+            per_bucket_bytes_caps = tuple(
+                int(cap_mb * _MB_TO_BYTES) for cap_mb in bucket_cap_mb_list
+            )
+            effective_bucket_cap_mb = max(bucket_cap_mb_list)
+            is_using_default = False
+
+        bucket_bytes_cap = int(effective_bucket_cap_mb * _MB_TO_BYTES)
+
+        # First bucket optimization: use smaller size when using defaults
+        # to reduce latency for early-computed gradients
+        first_bucket_bytes_cap = (
+            dist._DEFAULT_FIRST_BUCKET_BYTES if is_using_default else bucket_bytes_cap
+        )
+
+        return cls(
+            bucket_bytes_cap=bucket_bytes_cap,
+            per_bucket_bytes_caps=per_bucket_bytes_caps,
+            first_bucket_bytes_cap=first_bucket_bytes_cap,
+        )
+
+    @property
+    def bucket_cap_mb(self) -> int:
+        """Return bucket size in MiB (for backward compatibility)."""
+        return self.bucket_bytes_cap // _MB_TO_BYTES
+
+    @property
+    def has_custom_per_bucket_caps(self) -> bool:
+        """Return True if user specified per-bucket size limits."""
+        return len(self.per_bucket_bytes_caps) > 0
+
+    def compute_bucket_size_limits(
+        self,
+        static_graph: bool,
+        find_unused_parameters: bool,
+    ) -> tuple[list[int], list[int]]:
+        """Compute bucket size limits for initial bucketing and rebuilding.
+
+        Args:
+            static_graph: Whether the computation graph is static.
+            find_unused_parameters: Whether to find unused parameters.
+
+        Returns:
+            A tuple of (bucket_size_limits, bucket_size_limits_for_rebuilding):
+            - bucket_size_limits: Used for initial bucket assignment
+            - bucket_size_limits_for_rebuilding: Passed to C++ Reducer (empty list
+              means use C++ Reducer's default logic)
+        """
+        # Case 1: User provided explicit per-bucket size limits
+        if self.has_custom_per_bucket_caps:
+            limits = list(self.per_bucket_bytes_caps)
+            return (limits, limits)
+
+        # Case 2: Compute default bucket size limits
+        # When static_graph or not finding unused params, disable initial bucketing
+        # by using maxsize (put all params in one bucket initially)
+        if static_graph or not find_unused_parameters:
+            bucket_size_limits = [sys.maxsize]
+        elif self.first_bucket_bytes_cap < self.bucket_bytes_cap:
+            # Use smaller first bucket for latency optimization
+            bucket_size_limits = [self.first_bucket_bytes_cap, self.bucket_bytes_cap]
+        else:
+            bucket_size_limits = [self.bucket_bytes_cap]
+
+        # Empty list lets C++ Reducer use its default logic
+        return (bucket_size_limits, [])
+
+
+if dist.is_available():
+    from torch.distributed.distributed_c10d import (
+        _get_default_group,
+        _rank_not_in_group,
+        ReduceOp,
+    )
+    from torch.distributed.utils import (
+        _alloc_storage,
+        _cast_forward_inputs,
+        _free_storage,
+        _sync_module_states,
+        _to_kwargs,
+        _verify_param_shape_across_processes,
+    )
+if dist.rpc.is_available():
+    RPC_AVAILABLE = True
+    from torch.distributed.rpc import RRef
+
+if TYPE_CHECKING:
+    from torch.utils.hooks import RemovableHandle
+
+
+__all__ = ["DistributedDataParallel"]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _MixedPrecision:
+    """
+    This configures DDP-native mixed precision training.
+
+    Attributes:
+        param_dtype (torch.dtype): This specifies the dtype for model
+            parameters, inputs (when ``cast_forward_inputs`` is set to
+            ``True``), and therefore the dtype for computation.
+            However, outside the forward and backward passes, parameters are in
+            full precision. Model checkpointing always happens in full
+            precision.
+        reduce_dtype (torch.dtype): This specifies the dtype for gradient
+            reduction, which is permitted to differ from ``param_dtype``.
+        buffer_dtype (torch.dtype): This specifies the dtype for buffers.
+
+    .. note:: This API is experimental and subject to change.
+
+    .. note:: Only floating point tensors are cast to their specified dtypes.
+
+    .. note:: ``state_dict`` checkpoints parameters and buffers in full
+        precision.
+
+    .. note:: Each low precision dtype must be specified explicitly. For
+        example, ``_MixedPrecision(reduce_dtype=torch.float16)`` only specifies
+        the reduction dtype to be low precision, and DDP will not cast
+        parameters or buffers.
+
+    .. note:: If a ``reduce_dtype`` is not specified, then gradient reduction
+        happens in ``param_dtype`` if specified or the original parameter dtype
+        otherwise. For example, ``_MixedPrecision(param_dtype=torch.float16)``
+        would result in communication occurring in fp16.
+    """
+
+    param_dtype: torch.dtype | None = None
+    reduce_dtype: torch.dtype | None = None
+    buffer_dtype: torch.dtype | None = None
+    # TODO (rohan-varma): keep_low_precision_grads: bool = False
+    # TODO (rohan-varma): APIs to allow users to run batchnorm and layernorm
+    # in full precision. For DDP, this can be implemented by not performing the
+    # parameter cast for BN and LN units.
+
+
+def _cast_buffers(mixed_precision_config, root_module):
+    """Casts buffers to the given ``buffer_dtype``."""
+    for buf in root_module.buffers():
+        if hasattr(buf, "_ddp_ignored") and buf._ddp_ignored:
+            continue
+
+        buf.data = buf.to(dtype=mixed_precision_config.buffer_dtype)
+
+
+def _setup_mixed_precision_params(mixed_precision_config, root_module):
+    """Create and free storage for the mixed precision parameters."""
+    for param in root_module.parameters():
+        # Do not setup mixed precision for DDP ignored parameters.
+        if hasattr(param, "_ddp_ignored") and param._ddp_ignored:
+            continue
+
+        if not hasattr(param, "_mp_param"):
+            param._mp_param = torch.zeros_like(
+                param,
+                device=param.device,
+                dtype=mixed_precision_config.param_dtype,
+                requires_grad=param.requires_grad,
+            )
+            _free_storage(param._mp_param)
+            # _fp_param will point to the full precision param so it can be switched
+            # back to at the end of forward / backward.
+            param._fp_param = param.data
+
+
+def _tree_flatten_with_rref(output):
+    output_is_rref = RPC_AVAILABLE and isinstance(output, RRef)
+    if output_is_rref:
+        output_tensor_list, treespec = tree_flatten(output.local_value())
+    else:
+        output_tensor_list, treespec = tree_flatten(output)
+    # Need to return flattened tensors, spec to re-pack them, as well
+    # as if the return type was actually an RRef to reconstruct.
+    return output_tensor_list, treespec, output_is_rref
+
+
+def _tree_unflatten_with_rref(output, treespec, output_is_rref):
+    output = tree_unflatten(output, treespec)
+    if output_is_rref:
+        output = RRef(output)
+    return output
+
+
+def _find_tensors(obj):
+    r"""Recursively find all tensors contained in the specified object."""
+    if RPC_AVAILABLE and isinstance(obj, RRef):
+        # If the current node is the owner of the RRef, unwrap it and try to
+        # find Tensors.
+        # TODO: Expand to remote RRefs.
+        if obj.is_owner():
+            return _find_tensors(obj.local_value())
+    if isinstance(obj, torch.Tensor):
+        return [obj]
+    if isinstance(obj, (list, tuple)):
+        return itertools.chain.from_iterable(map(_find_tensors, obj))
+    if isinstance(obj, dict):
+        return itertools.chain.from_iterable(map(_find_tensors, obj.values()))
+    if is_dataclass(obj):
+        return itertools.chain.from_iterable(
+            map(_find_tensors, (getattr(obj, f.name) for f in fields(obj)))
+        )
+
+    return []
+
+
+def _dump_DDP_relevant_env_vars():
+    relevant_env_vars = [
+        "RANK",
+        "LOCAL_RANK",
+        "WORLD_SIZE",
+        "MASTER_PORT",
+        "MASTER_ADDR",
+        "CUDA_VISIBLE_DEVICES",
+        "GLOO_SOCKET_IFNAME",
+        "GLOO_DEVICE_TRANSPORT",
+        "NCCL_SOCKET_IFNAME",
+        "TORCH_NCCL_BLOCKING_WAIT",
+        "NCCL_DEBUG",
+        "NCCL_DEBUG_SUBSYS",
+        "NCCL_IB_DISABLE",
+        # More NCCL env vars:
+        "NCCL_P2P_DISABLE",
+        "NCCL_P2P_LEVEL",
+        "NCCL_SHM_DISABLE",
+        "NCCL_SOCKET_NTHREADS",
+        "NCCL_NSOCKS_PERTHREAD",
+        "NCCL_BUFFSIZE",
+        "NCCL_NTHREADS",
+        "NCCL_RINGS",
+        "NCCL_MAX_NCHANNELS",
+        "NCCL_MIN_NCHANNELS",
+        "NCCL_CHECKS_DISABLE",
+        "NCCL_CHECK_POINTERS",
+        "NCCL_LAUNCH_MODE",
+        "NCCL_IB_HCA",
+        "NCCL_IB_TIMEOUT",
+        "NCCL_IB_RETRY_CNT",
+        "NCCL_IB_GID_INDEX",
+        "NCCL_IB_SL",
+        "NCCL_IB_TC",
+        "NCCL_IB_AR_THRESHOLD",
+        "NCCL_IB_CUDA_SUPPORT",
+        "NCCL_NET_GDR_LEVEL",
+        "NCCL_NET_GDR_READ",
+        "NCCL_SINGLE_RING_THRESHOLD",
+        "NCCL_LL_THRESHOLD",
+        "NCCL_TREE_THRESHOLD",
+        "NCCL_ALGO",
+        "NCCL_PROTO",
+        "NCCL_IGNORE_CPU_AFFINITY",
+        "NCCL_DEBUG_FILE",
+        "NCCL_COLLNET_ENABLE",
+        "NCCL_TOPO_FILE",
+        "NCCL_TOPO_DUMP_FILE",
+        "TORCH_NCCL_ASYNC_ERROR_HANDLING",
+    ]
+    formatted_output = ""
+    for var in relevant_env_vars:
+        value = os.environ.get(var, "N/A")
+        formatted_output += f"env:{var}={value}\n"
+    print(formatted_output)
+
+
+class _BufferCommHookLocation(Enum):
+    PRE_FORWARD = auto()
+    POST_FORWARD = auto()
+
+
+@dataclass
+class _BufferCommHook:
+    buffer_comm_hook: Callable
+    buffer_comm_hook_state: Any
+    buffer_comm_hook_location: _BufferCommHookLocation
+
+
+# Add a DDPSink to run various functions when backwards starts, such as
+# queueing call back of out-most backward/graph task,
+# this helps call back is fired after all gradients' calculation
+# is completed.
+class _DDPSink(Function):
+    @staticmethod
+    # pyrefly: ignore [bad-override]
+    def forward(ctx, ddp_weakref, *inputs):
+        # set_materialize_grads(False) will ensure that None gradients stay as
+        # None and are not filled with zeros.
+        ctx.set_materialize_grads(False)
+        ctx.ddp_weakref = ddp_weakref
+        ret = inputs
+        if ddp_weakref()._ddp_sink_clone:
+            ret = tuple(
+                inp.clone() if isinstance(inp, torch.Tensor) else inp for inp in inputs
+            )
+        return ret
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        # Enqueue delay allreduce for static graph training on the first
+        # iteration.
+        ddp_weakref = ctx.ddp_weakref()
+        reducer = ddp_weakref.reducer
+        static_graph = ddp_weakref.static_graph
+        delay_ar_enqueued = (
+            static_graph and ddp_weakref._static_graph_delay_allreduce_enqueued
+        )
+        if static_graph and not delay_ar_enqueued:
+            Variable._execution_engine.queue_callback(  # type: ignore[call-arg,misc]
+                reducer._delay_all_reduce
+            )
+            ddp_weakref._static_graph_delay_allreduce_enqueued = True
+
+        return (None, *grad_outputs)
+
+
+class _DDPJoinHook(JoinHook):
+    def __init__(self, ddp, divide_by_initial_world_size):
+        """Set config variables for internal usage."""
+        if not isinstance(ddp, DistributedDataParallel):
+            raise AssertionError(
+                "DDP join hook requires passing in a DistributedDataParallel "
+                f"instance as the state, got {type(ddp).__name__}"
+            )
+        if ddp.logger is None:
+            raise AssertionError("ddp.logger must not be None")
+        ddp.logger._set_uneven_input_join()
+        self.ddp = ddp
+        self.ddp._divide_by_initial_world_size = divide_by_initial_world_size
+        super().__init__()
+
+    def main_hook(self):
+        """Shadow the DDP collective communication operations in the forward and backward passes."""
+        ddp = self.ddp
+        # Buckets are rebuilt only once during a training period
+        ddp.reducer._rebuild_buckets()
+
+        # Schedule a broadcast if we are syncing module buffers in the
+        # forward pass
+        # TODO: make DDP uneven inputs context manager support buffer
+        # comm hook (https://github.com/pytorch/pytorch/issues/65436)
+        ddp._check_and_sync_module_buffers()
+
+        # Check if need to sync in the backward pass
+        should_sync_backwards = ddp._check_global_requires_backward_grad_sync(
+            is_joined_rank=True
+        )
+        # Forward parameter sync is disabled in the next iteration if we
+        # are skipping gradient sync this iteration, so set
+        # `require_forward_param_sync` accordingly
+        ddp.require_forward_param_sync = should_sync_backwards
+        if not should_sync_backwards:
+            return
+
+        # Schedule one allreduce per gradient bucket to match the backward
+        # pass allreduce
+        ddp._match_all_reduce_for_bwd_pass()
+
+        # Check if we need to allreduce locally unused parameters
+        if ddp.find_unused_parameters:
+            ddp._match_unused_params_allreduce()
+
+        # Rebuilt parameters are pushed only once during a training period
+        ddp.reducer._push_all_rebuilt_params()
+
+    def post_hook(self, is_last_joiner: bool):
+        """Sync the final model to ensure that the model is the same across all processes."""
+        self.ddp._sync_final_model(is_last_joiner)
+
+
+class DistributedDataParallel(Module, Joinable):
+    r"""Implement distributed data parallelism based on ``torch.distributed`` at module level.
+
+    This container provides data parallelism by synchronizing gradients
+    across each model replica. The devices to synchronize across are
+    specified by the input ``process_group``, which is the entire world
+    by default. Note that ``DistributedDataParallel`` does not chunk or
+    otherwise shard the input across participating GPUs; the user is
+    responsible for defining how to do so, for example through the use
+    of a :class:`DistributedSampler`.
+
+    See also: :ref:`distributed-basics` and :ref:`cuda-nn-ddp-instead`.
+    The same constraints on input as in :class:`torch.nn.DataParallel` apply.
+
+    Creation of this class requires that ``torch.distributed`` to be already
+    initialized, by calling :func:`torch.distributed.init_process_group`.
+
+    ``DistributedDataParallel`` is proven to be significantly faster than
+    :class:`torch.nn.DataParallel` for single-node multi-GPU data
+    parallel training.
+
+    To use ``DistributedDataParallel`` on a host with N GPUs, you should spawn
+    up ``N`` processes, ensuring that each process exclusively works on a single
+    GPU from 0 to N-1. This can be done by either setting
+    ``CUDA_VISIBLE_DEVICES`` for every process or by calling the following API for GPUs,
+
+        >>> # xdoctest: +SKIP("undefined variables")
+        >>> torch.cuda.set_device(i)
+
+    or calling the unified API for :ref:`accelerator<accelerators>`,
+
+        >>> # xdoctest: +SKIP("undefined variables")
+        >>> torch.accelerator.set_device_index(i)
+
+    where i is from 0 to N-1. In each process, you should refer the following
+    to construct this module:
+
+        >>> # xdoctest: +SKIP("undefined variables")
+        >>> if torch.accelerator.is_available():
+        >>>     device_type = torch.accelerator.current_accelerator().type
+        >>>     vendor_backend = torch.distributed.get_default_backend_for_device(device_type)
+        >>>
+        >>> torch.distributed.init_process_group(
+        >>>     backend=vendor_backend, world_size=N, init_method='...'
+        >>> )
+        >>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)
+
+    Or you can use the latest API for initialization:
+
+        >>> torch.distributed.init_process_group(device_id=i)
+
+    In order to spawn up multiple processes per node, you can use either
+    ``torch.distributed.launch`` or ``torch.multiprocessing.spawn``.
+
+    .. note::
+        Please refer to `PyTorch Distributed Overview <https://pytorch.org/tutorials/beginner/dist_overview.html>`__
+        for a brief introduction to all features related to distributed training.
+
+    .. note::
+        ``DistributedDataParallel`` can be used in conjunction with
+        :class:`torch.distributed.optim.ZeroRedundancyOptimizer` to reduce
+        per-rank optimizer states memory footprint. Please refer to
+        `ZeroRedundancyOptimizer recipe <https://pytorch.org/tutorials/recipes/zero_redundancy_optimizer.html>`__
+        for more details.
+
+    .. note:: ``nccl`` backend is currently the fastest and highly recommended
+        backend when using GPUs. This applies to both single-node and
+        multi-node distributed training.
+
+    .. note:: This module also supports mixed-precision distributed training.
+        This means that your model can have different types of parameters such
+        as mixed types of ``fp16`` and ``fp32``, the gradient reduction on these
+        mixed types of parameters will just work fine.
+
+    .. note:: If you use ``torch.save`` on one process to checkpoint the module,
+        and ``torch.load`` on some other processes to recover it, make sure that
+        ``map_location`` is configured properly for every process. Without
+        ``map_location``, ``torch.load`` would recover the module to devices
+        where the module was saved from.
+
+    .. note:: When a model is trained on ``M`` nodes with ``batch=N``, the
+        gradient will be ``M`` times smaller when compared to the same model
+        trained on a single node with ``batch=M*N`` if the loss is summed (NOT
+        averaged as usual) across instances in a batch (because the gradients
+        between different nodes are averaged). You should take this into
+        consideration when you want to obtain a mathematically equivalent
+        training process compared to the local training counterpart. But in most
+        cases, you can just treat a DistributedDataParallel wrapped model, a
+        DataParallel wrapped model and an ordinary model on a single GPU as the
+        same (E.g. using the same learning rate for equivalent batch size).
+
+    .. note::
+        Parameters are never broadcast between processes. The module performs
+        an all-reduce step on gradients and assumes that they will be modified
+        by the optimizer in all processes in the same way. Buffers
+        (e.g. BatchNorm stats) are broadcast from the module in process of rank
+        0, to all other replicas in the system in every iteration.
+
+    .. note::
+        If you are using DistributedDataParallel in conjunction with the
+        :ref:`distributed-rpc-framework`, you should always use
+        :meth:`torch.distributed.autograd.backward` to compute gradients and
+        :class:`torch.distributed.optim.DistributedOptimizer` for optimizing
+        parameters.
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined variables")
+            >>> import torch.distributed.autograd as dist_autograd
+            >>> from torch.nn.parallel import DistributedDataParallel as DDP
+            >>> import torch
+            >>> from torch import optim
+            >>> from torch.distributed.optim import DistributedOptimizer
+            >>> import torch.distributed.rpc as rpc
+            >>> from torch.distributed.rpc import RRef
+            >>>
+            >>> t1 = torch.rand((3, 3), requires_grad=True)
+            >>> t2 = torch.rand((3, 3), requires_grad=True)
+            >>> rref = rpc.remote("worker1", torch.add, args=(t1, t2))
+            >>> ddp_model = DDP(my_model)
+            >>>
+            >>> # Setup optimizer
+            >>> optimizer_params = [rref]
+            >>> for param in ddp_model.parameters():
+            >>>     optimizer_params.append(RRef(param))
+            >>>
+            >>> dist_optim = DistributedOptimizer(
+            >>>     optim.SGD,
+            >>>     optimizer_params,
+            >>>     lr=0.05,
+            >>> )
+            >>>
+            >>> with dist_autograd.context() as context_id:
+            >>>     pred = ddp_model(rref.to_here())
+            >>>     loss = loss_func(pred, target)
+            >>>     dist_autograd.backward(context_id, [loss])
+            >>>     dist_optim.step(context_id)
+
+    .. note::
+        DistributedDataParallel currently offers limited support for gradient
+        checkpointing with :meth:`torch.utils.checkpoint`.
+        If the checkpoint is done with use_reentrant=False (recommended), DDP
+        will work as expected without any limitations.
+        If, however, the checkpoint is done with use_reentrant=True (the default),
+        DDP will work as expected when there are no unused parameters in the model
+        and each layer is checkpointed at most once (make sure you are not passing
+        `find_unused_parameters=True` to DDP). We currently do not support the
+        case where a layer is checkpointed multiple times, or when there unused
+        parameters in the checkpointed model.
+
+    .. note::
+        To let a non-DDP model load a state dict from a DDP model,
+        :meth:`~torch.nn.modules.utils.consume_prefix_in_state_dict_if_present`
+        needs to be applied to strip the prefix "module." in the DDP state dict before loading.
+
+    .. warning::
+        Constructor, forward method, and differentiation of the output (or a
+        function of the output of this module) are distributed synchronization
+        points. Take that into account in case different processes might be
+        executing different code.
+
+    .. warning::
+        This module assumes all parameters are registered in the model by the
+        time it is created. No parameters should be added nor removed later.
+        Same applies to buffers.
+
+    .. warning::
+        This module assumes all parameters are registered in the model of each
+        distributed processes are in the same order. The module itself will
+        conduct gradient ``allreduce`` following the reverse order of the
+        registered parameters of the model. In other words, it is users'
+        responsibility to ensure that each distributed process has the exact
+        same model and thus the exact same parameter registration order.
+
+    .. warning::
+        This module allows parameters with non-rowmajor-contiguous strides.
+        For example, your model may contain some parameters whose
+        :class:`torch.memory_format` is ``torch.contiguous_format``
+        and others whose format is ``torch.channels_last``.  However,
+        corresponding parameters in different processes must have the
+        same strides.
+
+    .. warning::
+        This module doesn't work with :func:`torch.autograd.grad` (i.e. it will
+        only work if gradients are to be accumulated in ``.grad`` attributes of
+        parameters).
+
+    .. warning::
+        If you plan on using this module with a ``nccl`` backend or a ``gloo``
+        backend (that uses Infiniband), together with a DataLoader that uses
+        multiple workers, please change the multiprocessing start method to
+        ``forkserver`` or ``spawn``. Unfortunately
+        Gloo (that uses Infiniband) and NCCL2 are not fork safe, and you will
+        likely experience deadlocks if you don't change this setting.
+
+    .. warning::
+        You should never try to change your model's parameters after wrapping
+        up your model with ``DistributedDataParallel``. Because, when
+        wrapping up your model with ``DistributedDataParallel``, the constructor
+        of ``DistributedDataParallel`` will register the additional gradient
+        reduction functions on all the parameters of the model itself at the
+        time of construction. If you change the model's parameters afterwards,
+        gradient reduction functions no longer match the correct set of
+        parameters.
+
+    .. warning::
+        Using ``DistributedDataParallel`` in conjunction with the
+        :ref:`distributed-rpc-framework` is experimental and subject to change.
+
+    Args:
+        module (Module): module to be parallelized
+        device_ids (list of int or torch.device): CUDA devices.
+                   1) For single-device modules, ``device_ids`` can
+                   contain exactly one device id, which represents the only
+                   CUDA device where the input module corresponding to this process resides.
+                   Alternatively, ``device_ids`` can also be ``None``.
+                   2) For multi-device modules and CPU modules,
+                   ``device_ids`` must be ``None``.
+
+                   When ``device_ids`` is ``None`` for both cases,
+                   both the input data for the forward pass and the actual module
+                   must be placed on the correct device.
+                   (default: ``None``)
+        output_device (int or torch.device): Device location of output for
+                      single-device CUDA modules. For multi-device modules and
+                      CPU modules, it must be ``None``, and the module itself
+                      dictates the output location. (default: ``device_ids[0]``
+                      for single-device modules)
+        broadcast_buffers (bool): Flag that enables syncing (broadcasting)
+                          buffers of the module at beginning of the ``forward``
+                          function. (default: ``True``)
+        init_sync (bool): Whether to sync during initialization to verify param
+                          shapes and broadcast parameters and buffers.
+                          WARNING: if this is set to False the user is required
+                          to ensure themselves that the weights are the same on
+                          all ranks.
+                          (default: ``True``)
+        process_group: The process group to be used for distributed data
+                       all-reduction. If ``None``, the default process group, which
+                       is created by :func:`torch.distributed.init_process_group`,
+                       will be used. (default: ``None``)
+        bucket_cap_mb: ``DistributedDataParallel`` will bucket parameters into
+                       multiple buckets so that gradient reduction of each
+                       bucket can potentially overlap with backward computation.
+                       :attr:`bucket_cap_mb` controls the bucket size in
+                       MebiBytes (MiB). If ``None``, a default size of 25 MiB
+                       will be used. (default: ``None``)
+        find_unused_parameters (bool): Traverse the autograd graph from all
+                               tensors contained in the return value of the
+                               wrapped module's ``forward`` function. Parameters
+                               that don't receive gradients as part of this
+                               graph are preemptively marked as being ready to
+                               be reduced. In addition, parameters that may have
+                               been used in the wrapped module's ``forward``
+                               function but were not part of loss computation and
+                               thus would also not receive gradients are
+                               preemptively marked as ready to be reduced.
+                               (default: ``False``)
+        check_reduction: This argument is deprecated.
+        gradient_as_bucket_view (bool): When set to ``True``, gradients will be views
+                      pointing to different offsets of ``allreduce`` communication
+                      buckets. This can reduce peak memory usage, where the
+                      saved memory size will be equal to the total gradients
+                      size. Moreover, it avoids the overhead of copying between
+                      gradients and ``allreduce`` communication buckets. When
+                      gradients are views, ``detach_()`` cannot be called on the
+                      gradients. If hitting such errors, please fix it by
+                      referring to the :meth:`~torch.optim.Optimizer.zero_grad`
+                      function in ``torch/optim/optimizer.py`` as a solution.
+                      Note that gradients will be views after first iteration, so
+                      the peak memory saving should be checked after first iteration.
+        static_graph (bool): When set to ``True``, DDP knows the trained graph is
+                     static. Static graph means 1) The set of used and unused
+                     parameters will not change during the whole training loop; in
+                     this case, it does not matter whether users set
+                     ``find_unused_parameters = True`` or not. 2) How the graph is trained
+                     will not change during the whole training loop (meaning there is
+                     no control flow depending on iterations).
+                     When static_graph is set to be ``True``, DDP will support cases that
+                     can not be supported in the past:
+                     1) Reentrant backwards.
+                     2) Activation checkpointing multiple times.
+                     3) Activation checkpointing when model has unused parameters.
+                     4) There are model parameters that are outside of forward function.
+                     5) Potentially improve performance when there are unused parameters,
+                     as DDP will not search graph in each iteration to detect unused
+                     parameters when static_graph is set to be ``True``.
+                     To check whether you can set static_graph to be ``True``, one way is to
+                     check ddp logging data at the end of your previous model training,
+                     if ``ddp_logging_data.get("can_set_static_graph") == True``, mostly you
+                     can set ``static_graph = True`` as well.
+
+                     Example::
+                         >>> # xdoctest: +SKIP("undefined variables")
+                         >>> model_DDP = torch.nn.parallel.DistributedDataParallel(model)
+                         >>> # Training loop
+                         >>> ...
+                         >>> ddp_logging_data = model_DDP._get_ddp_logging_data()
+                         >>> static_graph = ddp_logging_data.get("can_set_static_graph")
+        delay_all_reduce_named_params (list of tuple of str and torch.nn.Parameter): a list
+                    of named parameters whose all reduce will be delayed when the gradient of
+                    the parameter specified in ``param_to_hook_all_reduce`` is ready. Other
+                    arguments of DDP do not apply to named params specified in this argument
+                    as these named params will be ignored by DDP reducer.
+        param_to_hook_all_reduce (torch.nn.Parameter): a parameter to hook delayed all reduce
+                    of parameters specified in ``delay_all_reduce_named_params``.
+        skip_all_reduce_unused_params: When set to True, DDP will skip reducing unused parameters.
+                    This requires that unused parameters remain the same across all ranks throughout
+                    the entire training process. If this condition is not met, it may cause
+                    desynchronization and result in training hang.
+        batched_grad_copy (bool): When set to ``True``, individual per-parameter
+                    gradient-to-bucket copy and division operations are deferred
+                    and flushed as a single ``_foreach_copy_`` plus one flat
+                    ``div_`` when a bucket becomes ready. This reduces per-parameter
+                    kernel launches down to 2 kernels per bucket, which can improve
+                    throughput for models with many small parameters. The
+                    optimization is most effective with
+                    ``optimizer.zero_grad(set_to_none=True)`` (the default), where
+                    ``gradient_as_bucket_view`` alone cannot avoid copies because
+                    the bucket view alias is destroyed every iteration.
+                    (default: ``False``)
+
+
+    Attributes:
+        module (Module): the module to be parallelized.
+
+    Example::
+
+        >>> # xdoctest: +SKIP("undefined variables")
+        >>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
+        >>> net = torch.nn.parallel.DistributedDataParallel(model)
+    """
+
+    # used to track whether the given thread is inside ddp forward for torchdynamo purposes
+    _active_ddp_module: Optional["DistributedDataParallel"] = None
+
+    def __init__(
+        self,
+        module,
+        device_ids=None,
+        output_device=None,
+        dim=0,
+        broadcast_buffers=True,
+        init_sync=True,
+        process_group=None,
+        bucket_cap_mb=None,
+        find_unused_parameters=False,
+        check_reduction=False,
+        gradient_as_bucket_view=False,
+        static_graph=False,
+        delay_all_reduce_named_params=None,
+        param_to_hook_all_reduce=None,
+        mixed_precision: _MixedPrecision | None = None,
+        device_mesh=None,
+        skip_all_reduce_unused_params=False,
+        bucket_cap_mb_list: list[int] | None = None,
+        batched_grad_copy=False,
+    ):
+        super().__init__()
+        Joinable.__init__(self)
+        self._use_python_reducer = (
+            torch._dynamo.utils.get_optimize_ddp_mode() == "python_reducer"
+        )
+        self.logger: dist.Logger | None = None
+        if bool(delay_all_reduce_named_params is not None) != bool(
+            param_to_hook_all_reduce is not None
+        ):
+            self._log_and_throw(
+                ValueError,
+                "delay_all_reduce_named_params and param_to_hook_all_reduce "
+                "need to be set at the same time.",
+            )
+
+        if process_group and device_mesh is not None:
+            raise RuntimeError(
+                "Cannot specify both process_group and device_mesh arguments."
+            )
+        elif process_group is None and device_mesh is None:
+            self.process_group = _get_default_group()
+        elif device_mesh is None:
+            # pyrefly: ignore [bad-assignment]
+            self.process_group = process_group
+        else:
+            if device_mesh.ndim != 1:
+                raise RuntimeError(
+                    f"Only 1D device mesh is supported, but got {device_mesh}."
+                )
+            self.device_mesh = device_mesh
+            self.process_group = device_mesh.get_group(mesh_dim=0)
+
+            root_mesh = device_mesh._get_root_mesh()
+            # if a root mesh is not the same as device_mesh,
+            # meaning the device_mesh is sliced out from the root mesh.
+            if root_mesh != device_mesh:
+                # TODO: This is a temporary work around to enable DDP + TP.
+                # We should do the logic in DDP so that the 2D implementation is
+                # sound and the state_dict works out of the box.
+                # This has to be done before check UninitializedParameter.
+                from torch.distributed.tensor.parallel.ddp import (
+                    _pre_dp_module_transform,
+                )
+
+                _pre_dp_module_transform(module)
+
+        self._delay_all_reduce_params = []
+        if hasattr(module, "_ddp_params_and_buffers_to_ignore"):
+            self.parameters_to_ignore = set(module._ddp_params_and_buffers_to_ignore)
+        else:
+            self.parameters_to_ignore = set()
+        if delay_all_reduce_named_params is not None:
+            for name, param in delay_all_reduce_named_params:
+                self.parameters_to_ignore.add(name)
+                self._delay_all_reduce_params.append(param)
+
+        self._module_parameters = [
+            p
+            for n, p in module.named_parameters()
+            if n not in self.parameters_to_ignore
+        ]
+        if not any(p.requires_grad for p in self._module_parameters):
+            if len(self._delay_all_reduce_params):
+                logger.info("Delay the AllReduce of all parameters.")
+            else:
+                self._log_and_throw(
+                    RuntimeError,
+                    "DistributedDataParallel is not needed when a module "
+                    "doesn't have any parameter that requires a gradient.",
+                )
+
+        if device_ids is not None and len(device_ids) > 1:
+            self._log_and_throw(
+                ValueError,
+                "device_ids can only be None or contain a single element.",
+            )
+
+        self.is_multi_device_module = (
+            len({p.device for p in self._module_parameters}) > 1
+        )
+        distinct_device_types = {
+            p.device.type for p in self._module_parameters if p.device is not None
+        }
+        if len(distinct_device_types) != 1:
+            self._log_and_throw(
+                ValueError,
+                "DistributedDataParallel's input module must be on "
+                f"the same type of devices, but input module parameters locate in {distinct_device_types}.",
+            )
+
+        self.device_type = next(iter(distinct_device_types))
+
+        if (
+            device_ids is None
+            or len(device_ids) == 0  # For backward compatibility.
+            or self.device_type == "cpu"
+            or self.is_multi_device_module
+        ):
+            if device_ids or output_device:
+                devices = {p.device for p in self._module_parameters}
+                self._log_and_throw(
+                    ValueError,
+                    "DistributedDataParallel device_ids and output_device arguments "
+                    "only work with single-device/multiple-device GPU modules or CPU modules, "
+                    f"but got device_ids {device_ids}, output_device {output_device}, "
+                    f"and module parameters {devices}.",
+                )
+
+            self.device_ids = None
+            self.output_device = None
+        else:
+            self.device_ids = [_get_device_index(x, True) for x in device_ids]
+
+            if output_device is None:
+                output_device = device_ids[0]
+
+            self.output_device = _get_device_index(output_device, True)
+
+        self.static_graph = False
+        self.dim = dim
+        self.module = module
+        self.device = next(iter(self._module_parameters)).device
+        self.broadcast_buffers = broadcast_buffers
+        self.find_unused_parameters = find_unused_parameters
+        self.require_backward_grad_sync = True
+        self.require_forward_param_sync = True
+        self.gradient_as_bucket_view = gradient_as_bucket_view
+        self.batched_grad_copy = batched_grad_copy
+        self.mixed_precision = mixed_precision
+        if self.mixed_precision is not None:
+            logger.warning("Received mixed precision config %s", self.mixed_precision)
+
+        if check_reduction:
+            # This argument is no longer used since the reducer
+            # will ensure reduction completes even if some parameters
+            # do not receive gradients.
+            warnings.warn(
+                "The `check_reduction` argument in `DistributedDataParallel` "
+                "module is deprecated. Please avoid using it.",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        # Check that a module does not have Uninitialized parameters
+        for param in self._module_parameters:
+            if isinstance(param, torch.nn.parameter.UninitializedParameter):
+                self._log_and_throw(
+                    RuntimeError,
+                    "Modules with uninitialized parameters can't be used with `DistributedDataParallel`. "
+                    "Run a dummy forward pass to correctly initialize the modules",
+                )
+        # used for intra-node param sync and inter-node sync as well
+        self.broadcast_bucket_size = 250 * 1024 * 1024
+
+        # Initialize bucket capacity configuration using the factory method
+        self._bucket_config = _BucketCapacityConfig.create(
+            bucket_cap_mb=bucket_cap_mb,
+            bucket_cap_mb_list=bucket_cap_mb_list,
+            use_python_reducer=self._use_python_reducer,
+        )
+        # Expose config values as instance attributes for backward compatibility
+        # TODO: Remove in the future
+        self.bucket_bytes_cap = self._bucket_config.bucket_bytes_cap
+        self.bucket_bytes_cap_list = list(self._bucket_config.per_bucket_bytes_caps)
+        self.bucket_bytes_cap_default = (
+            self._bucket_config.first_bucket_bytes_cap
+            < self._bucket_config.bucket_bytes_cap
+        )
+        bucket_cap_mb = self._bucket_config.bucket_cap_mb
+
+        # Whether to perform input tensor CPU to GPU copies on a side-stream
+        self.use_side_stream_for_tensor_copies = (
+            os.environ.get("PYTORCH_DDP_USE_SIDE_STREAM", "1") == "1"
+        )
+
+        # Initialize gradient buffers and register all reduce hook
+        self._delay_grad_buffer: torch.Tensor | None = None
+        self._delay_grad_views: list[torch.Tensor] = []
+        self._delay_all_reduce_all_params = False
+        if len(self._delay_all_reduce_params) != 0:
+            self._register_delay_all_reduce_hook(
+                bucket_cap_mb=bucket_cap_mb,
+                param_to_hook_all_reduce=param_to_hook_all_reduce,
+                device_ids=device_ids,
+            )
+            if self._delay_all_reduce_all_params:
+                return
+
+        self.skip_all_reduce_unused_params = skip_all_reduce_unused_params
+
+        # Build parameters for reducer.
+        parameters, expect_sparse_gradient = self._build_params_for_reducer()
+
+        # All collectives during initialization are gated by this flag.
+        if init_sync:
+            # Verify model equivalence.
+            _verify_param_shape_across_processes(self.process_group, parameters)
+            # Sync params and buffers. Ensures all DDP models start off at the same value.
+            _sync_module_states(
+                module=self.module,
+                process_group=self.process_group,
+                broadcast_bucket_size=self.broadcast_bucket_size,
+                src=0,
+                params_and_buffers_to_ignore=self.parameters_to_ignore,
+                broadcast_buffers=self.broadcast_buffers,
+            )
+
+        # In debug mode, build a mapping of parameter index -> parameter.
+        param_to_name_mapping = self._build_debug_param_to_name_mapping(parameters)
+
+        # Builds reducer.
+        self._ddp_init_helper(
+            parameters,
+            expect_sparse_gradient,
+            param_to_name_mapping,
+            static_graph,
+        )
+        self._comm_hooks: list[tuple[Callable, object]] = []
+
+        if self.mixed_precision is not None:
+            _setup_mixed_precision_params(self.mixed_precision, self.module)
+            _cast_buffers(self.mixed_precision, self.module)
+            # Stream used for async low precision copies.
+            self._mp_stream = torch.Stream()
+            self._submodule_to_event = defaultdict(deque)  # type: ignore[var-annotated]
+            # Add forward pre-hook to root module to kick off copies to lower
+            # precision.
+            self.module.register_forward_pre_hook(
+                self._root_copy_hook, prepend=False, with_kwargs=True
+            )
+            # Add forward pre hook to all submodules to wait for copy events
+            # before running computation.
+            for module in self.module.modules():
+                module.register_forward_pre_hook(
+                    self._module_wait_for_copy_hook,
+                    prepend=False,
+                    with_kwargs=True,
+                )
+            # Set up callbacks in backward to upcast and use full precision
+            # params. TODO (rohan-varma): Make this compose with general
+            # comm hooks and apply_optimizer_in_backward. Importing inline to
+            # avoid circular import issue.
+            from torch.distributed.algorithms.ddp_comm_hooks.mixed_precision_hooks import (
+                _AllreduceUpcastHookState,
+                _reducer_allreduce_and_upcast_hook,
+            )
+
+            upcast_hook_state = _AllreduceUpcastHookState(
+                ddp_weakref=weakref.ref(self),
+                upcast_stream=torch.Stream(),
+            )
+            self.register_comm_hook(
+                upcast_hook_state,
+                _reducer_allreduce_and_upcast_hook,
+            )
+            # Inform reducer of reduced precision param dtype for correctness
+            # of type checks between gradient and bucket.
+            self.reducer._set_mixed_precision_param_dtype(  # type: ignore[attr-defined]
+                self.mixed_precision.param_dtype
+            )
+
+        self._has_rebuilt_buckets = False
+
+        if static_graph:
+            self._set_static_graph()
+
+        self._lazy_init_ran = False
+
+        # Register the AccumulateGrad post hooks if optimize_ddp is
+        # True. The hooks will be deregistered if compiled_autograd is not
+        # enabled.
+        self._accum_grad_hooks: list[RemovableHandle] = []
+        if self._use_python_reducer:
+            # pyrefly: ignore [bad-assignment]
+            torch._inductor.config._fuse_ddp_communication = True
+            # pyrefly: ignore [bad-assignment]
+            torch._inductor.config._fuse_ddp_bucket_size = bucket_cap_mb
+            # Directly adding this to the trace rule will disturb the users
+            # who are using DDPOptimizer.
+            torch._dynamo.trace_rules.LEGACY_MOD_INLINELIST.add(
+                "torch.nn.parallel.distributed"
+            )
+            torch._dynamo.trace_rules.get_legacy_mod_inlinelist.cache_clear()
+            # NOTE: we should init these lazily
+            self._register_accum_grad_hook()
+
+        # Whether or not DDPSink performs a clone.
+        self._ddp_sink_clone = True
+
+    def _register_accum_grad_hook(self):
+        import torch.distributed._functional_collectives as fcol
+
+        def compiled_accum_grad_hook(
+            param,
+            *,
+            param_index: int,
+        ):
+            if not self.require_backward_grad_sync:
+                return
+
+            if param.grad is None:
+                return
+
+            if self._comm_hooks:
+                for hook, state in self._comm_hooks:
+                    hook(state, (param.grad, param))
+            else:
+                gradient = param.grad / self.process_group.size()
+                gradient = fcol.all_reduce(gradient, "sum", self.process_group)
+                param.grad.copy_(gradient)
+
+        for index, param in enumerate(self._module_parameters):
+            if not param.requires_grad:
+                continue
+            self._accum_grad_hooks.append(
+                param.register_post_accumulate_grad_hook(
+                    functools.partial(
+                        compiled_accum_grad_hook,
+                        param_index=index,
+                    )
+                )
+            )
+
+    def _delayed_all_reduce_hook(self, grad):
+        world_size = dist.get_world_size(self.process_group)
+
+        self._delay_grad_buffer.div_(world_size)  # type: ignore[union-attr]
+        _ = dist.all_reduce(
+            self._delay_grad_buffer, group=self.process_group, async_op=True
+        )
+        return grad
+
+    def _register_delay_all_reduce_hook(
+        self,
+        bucket_cap_mb,
+        param_to_hook_all_reduce,
+        device_ids,
+    ):
+        # 1. Create gradient buffer
+        device = torch.device("cpu") if device_ids is None else device_ids[0]
+        self._delay_grad_buffer = torch.zeros(
+            sum(p.numel() for p in self._delay_all_reduce_params),
+            device=device,
+        )
+
+        # 2. Broadcast the parameters
+        detached_params = [p.detach() for p in self._delay_all_reduce_params]
+        dist._broadcast_coalesced(self.process_group, detached_params, bucket_cap_mb, 0)
+
+        # 3. Hook all reduce to the specified parameter
+        param_to_hook_all_reduce.register_hook(self._delayed_all_reduce_hook)
+
+        # 4. Build tensor views for gradients
+        offset = 0
+        for param in self._delay_all_reduce_params:
+            grad_view = self._delay_grad_buffer[offset : (offset + param.numel())].view(
+                param.shape
+            )
+            self._delay_grad_views.append(grad_view)
+            offset = offset + param.numel()
+
+        # 5. Check whether the all reduce of all params requiring grad is delayed.
+        for module_name, module in self.module.named_modules():
+            for param_name, param in module.named_parameters(recurse=False):
+                if param.requires_grad:
+                    full_name = f"{module_name}.{param_name}"
+                    if full_name not in self.parameters_to_ignore:
+                        # There is at least a param whose all reduce will not be delayed.
+                        # In this case, we should not set self._delay_all_reduce_all_params
+                        # to True.
+                        return
+        self._delay_all_reduce_all_params = True
+
+    def _setup_in_backward_optimizers(self):
+        # Check if user has used apply_optim_in_backward to overlap optimizer
+        # step + DDP backward. Current constraints:
+        # 1. Only allreduce is supported at the moment, no custom communication.
+        # 2. For DDP-managed parameters that have their optimizer run in
+        # backward, their gradients are set to ``None``. If your use case
+        # requires DDP parameters grad not to be set to ``None`` after their
+        # in-backward optimizer runs, please ping
+        # https://github.com/pytorch/pytorch/issues/90052.
+        # NOTE: we use self._module_parameters instead of .parameters() since
+        # the former excludes ignored (non-DDP managed) parameters.
+        if any(hasattr(p, "_in_backward_optimizers") for p in self._module_parameters):
+            torch._C._log_api_usage_once("ddp.optimizer_in_backward")
+            # Remove hooks that apply_optim_in_backward had registered because
+            # DDP customizes how optimizer is overlapped with backward due to
+            # the allreduce.
+            param_to_handle_map = (
+                dist.optim.apply_optimizer_in_backward.param_to_optim_hook_handle_map
+            )
+            for p in self._module_parameters:
+                for handle in param_to_handle_map.get(p, []):
+                    handle.remove()
+
+            # Need a weakref to DDP instance to run all_reduce (from reducer)
+            # and get managed DDP parameters.
+            ddp_weakref = weakref.ref(self)
+            # Note: importing in function, otherwise this will cause a circular
+            # import.
+            from torch.distributed.algorithms.ddp_comm_hooks.optimizer_overlap_hooks import (
+                _apply_optim_in_backward_hook,
+            )
+
+            self.register_comm_hook(
+                ddp_weakref,
+                _apply_optim_in_backward_hook(
+                    gradient_is_bucket_view=self.gradient_as_bucket_view
+                ),
+            )
+
+            self.reducer._set_optimizer_in_backward()  # type: ignore[attr-defined]
+
+    def _fire_reducer_autograd_hook(self, idx, *unused):
+        """
+        Fire the reducer's autograd hook to allreduce params in a Reducer bucket.
+
+        Note that this is only used during mixed precision training as the
+        Reducer's hooks installed during construction time would not be called
+        as we're working in the low precision parameter setting.
+        """
+        self.reducer._autograd_hook(idx)  # type: ignore[attr-defined]
+
+    def _root_copy_hook(self, *args: Any, **kwargs: Any) -> None:
+        """
+        For DDP mixed precision, put low precision copies on separate stream and create events to wait for them.
+
+        When training with DDP mixed precision, this root pre-forward hook kicks
+        off low precision copies on a separate stream and creates respective
+        events to wait for them.
+        """
+        # Clear out previous iteration submodule to event. This is because we
+        # may have populated some events for modules that didn't end up being
+        # used.
+        self._submodule_to_event = defaultdict(deque)  # type: ignore[var-annotated]
+        with self._mp_stream:
+            for submodule in self.module.modules():
+                for param in submodule.parameters(recurse=False):
+                    # Do not cast DDP ignored parameters.
+                    if hasattr(param, "_ddp_ignored") and param._ddp_ignored:
+                        continue
+                    _alloc_storage(param._mp_param, param.size())
+                    # copy() implicitly casts to low precision
+                    with torch.no_grad():
+                        param._mp_param.copy_(param.data)
+                        # TODO: when zero_grad(set_to_none=False) or in grad
+                        # accumulation case, accumulated grads can be in fp32
+                        # which can cause errors when running DDP backwards due
+                        # to mismatched incoming and accumulated gradient types.
+                        # So we manually cast the accumulated grad down for now,
+                        # in the future we may shift to FSDP style gradient
+                        # accumulation management where the accumulated gradient
+                        # is saved and .grad field is set to None, bypassing
+                        # this issue.
+                        if param.grad is not None:
+                            param.grad.data = param.grad.to(
+                                self.mixed_precision.param_dtype  # type: ignore[union-attr]
+                            )
+                    param.data = param._mp_param
+                copy_event = torch.Event()
+                copy_event.record()
+                self._submodule_to_event[submodule].append(copy_event)
+
+    def _module_wait_for_copy_hook(
+        self,
+        module,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Before carrying out computation, wait on the appropriate event to ensure low precision copies have finished."""
+        try:
+            event = self._submodule_to_event[module].popleft()
+        except IndexError:
+            # copy event has already been waited on
+            return
+
+        event.wait(stream=torch.accelerator.current_stream())
+        for p in module.parameters(recurse=False):
+            # Don't register hooks if param does not require grad
+            if not p.requires_grad or (hasattr(p, "_ddp_ignored") and p._ddp_ignored):
+                continue
+            # We need to register autograd hook here instead of DDP's ctor
+            # since we're working with the low precision param. Register them
+            # via obtaining the gradient accumulator.
+            tmp = p.expand_as(p)
+            grad_acc = tmp.grad_fn.next_functions[0][0]
+
+            hook = grad_acc.register_hook(
+                functools.partial(self._fire_reducer_autograd_hook, p._idx)
+            )
+            p._ddp_mp_hook_state = (grad_acc, hook)
+
+    def _log_and_throw(self, err_type, err_msg):
+        if self.logger is not None:
+            self.logger.set_error_and_log(f"{str(err_type)}: {err_msg}")
+        raise err_type(err_msg)
+
+    def _ddp_init_helper(
+        self,
+        parameters,
+        expect_sparse_gradient,
+        param_to_name_mapping,
+        static_graph,
+    ):
+        """
+        DDP init helper function to manage parameters, grad hooks, logging, and SyncBatchNorm.
+
+        Initialization helper function that does the following:
+        (1) bucketing the parameters for reductions
+        (2) resetting the bucketing states
+        (3) registering the grad hooks
+        (4) Logging construction-time DDP logging data
+        (5) passing a handle of DDP to SyncBatchNorm Layer
+        """
+        # Notice, the parameters order is not in the order in which they are used,
+        # especially in models with control flow.
+        #
+        # Alongside parameters are not presented in the real execution order,
+        # if a certain model happens to also
+        #   1) have other collectives comm ops in its backward graph.
+        #   2) have unused parameter in subset ranks of the whole world.
+        # bucketing could insert ALL-REDUCE comm op too early on the rank with unused parameter,
+        # matching up with other collectives comm ops on other ranks unexpectedly.
+        #
+        # In order to handle this corner case, when the parameters are not in the real execution order,
+        # we don't do bucketing, thus only one ALL-REDUCE is inserted after all the gradients
+        # of the whole graph are computed.
+        #
+        # Notice, here we only disable bucketing for the first iteration.
+        # After the first iteration, it's OK to rebuild buckets,
+        # because "bucket rebuild" bucketizes parameters based on its real execution order in backward graph.
+
+        # Compute bucket size limits for initial bucketing and rebuilding
+        # See BucketCapacityConfig.compute_bucket_size_limits for detailed explanation
+        (
+            bucket_size_limits,
+            bucket_size_limits_for_rebuilding,
+        ) = self._bucket_config.compute_bucket_size_limits(
+            static_graph=static_graph,
+            find_unused_parameters=self.find_unused_parameters,
+        )
+
+        (
+            bucket_indices,
+            per_bucket_size_limits,
+        ) = dist._compute_bucket_assignment_by_size(
+            parameters,
+            bucket_size_limits,
+            expect_sparse_gradient,
+        )
+
+        # Remember index for parameters if we are in mixed precision, as we
+        # need to pass in index to Reducer's autograd hook via python.
+        if self.mixed_precision is not None:
+            for i, p in enumerate(parameters):
+                p._idx = i
+
+        # Note: reverse list of buckets because we want to approximate the
+        # order in which their gradients are produced, and assume they
+        # are used in the forward pass in the order they are defined.
+        self.reducer = dist.Reducer(
+            parameters,
+            list(reversed(bucket_indices)),
+            list(reversed(per_bucket_size_limits)),
+            self.process_group,
+            expect_sparse_gradient,
+            # The bucket size limit is specified in the constructor.
+            # Additionally, we allow for a single small bucket for parameters
+            # that are defined first, such that their gradients don't spill into
+            # a much larger bucket, adding unnecessary latency after gradient
+            # computation finishes. Experiments showed 1MB is a reasonable value.
+            self.bucket_bytes_cap,
+            self.find_unused_parameters,
+            self.gradient_as_bucket_view,
+            param_to_name_mapping,
+            # First bucket can be smaller for reduced latency. See
+            # BucketCapacityConfig.first_bucket_bytes_cap for details.
+            self._bucket_config.first_bucket_bytes_cap,
+            self.skip_all_reduce_unused_params,
+            self._use_python_reducer,
+            bucket_size_limits_for_rebuilding,
+            self.batched_grad_copy,
+        )
+
+        self.logger = dist.Logger(self.reducer)
+        # Set as a weak reference to avoid reference cycle between
+        # logger and reducer.
+        self.reducer.set_logger(self.logger)
+
+        has_sync_bn = False
+        for submodule in self.module.modules():
+            if isinstance(submodule, torch.nn.SyncBatchNorm):
+                has_sync_bn = True
+                break
+
+        # Set logging data that can be got during construction time.
+        self.logger.set_construction_data_and_log(
+            self.module.__class__.__name__,
+            [] if self.device_ids is None else self.device_ids,
+            -1 if self.output_device is None else self.output_device,
+            self.broadcast_buffers,
+            has_sync_bn,
+            static_graph,
+        )
+
+        # passing a handle to torch.nn.SyncBatchNorm layer
+        self._passing_sync_batchnorm_handle(self.module)
+
+    def __getstate__(self):
+        self._check_default_group()
+        attrs = copy.copy(self.__dict__)
+        del attrs["process_group"]
+        del attrs["reducer"]
+        del attrs["logger"]
+        return attrs
+
+    def __setstate__(self, state):
+        # If serializable, then the process group should be the default one
+        self.process_group = _get_default_group()
+        super().__setstate__(state)
+        self.__dict__.setdefault("require_forward_param_sync", True)
+        self.__dict__.setdefault("require_backward_grad_sync", True)
+        parameters, expect_sparse_gradient = self._build_params_for_reducer()
+        # In debug mode, build a mapping of parameter index -> parameter.
+        param_to_name_mapping = self._build_debug_param_to_name_mapping(parameters)
+        # Builds reducer.
+        self._ddp_init_helper(
+            parameters,
+            expect_sparse_gradient,
+            param_to_name_mapping,
+            self.static_graph,
+        )
+        if self.static_graph:
+            self.reducer._set_static_graph()
+            if self.logger is None:
+                raise AssertionError(
+                    "self.logger must not be None when static_graph is True"
+                )
+            self.logger._set_static_graph()
+
+    def _build_params_for_reducer(self):
+        # Build tuple of (module, parameter) for all parameters that require grads.
+        modules_and_parameters = [
+            (module, parameter)
+            for module_name, module in self.module.named_modules()
+            for parameter in [
+                param
+                # Note that we access module.named_parameters instead of
+                # parameters(module). parameters(module) is only needed in the
+                # single-process multi device case, where it accesses replicated
+                # parameters through _former_parameters.
+                for param_name, param in module.named_parameters(recurse=False)
+                if param.requires_grad
+                and f"{module_name}.{param_name}" not in self.parameters_to_ignore
+            ]
+        ]
+
+        # Deduplicate any parameters that might be shared across child modules.
+        memo = set()
+        modules_and_parameters = [
+            # "p not in memo" is the deduplication check.
+            # "not memo.add(p)" is always True, and it's only there to cause "add(p)" if needed.
+            (m, p)
+            for m, p in modules_and_parameters
+            if p not in memo and not memo.add(p)  # type: ignore[func-returns-value]
+        ]
+
+        # Build list of parameters.
+        parameters = [parameter for _, parameter in modules_and_parameters]
+
+        # Checks if a module will produce a sparse gradient.
+        def produces_sparse_gradient(module):
+            if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)):
+                return module.sparse
+            return False
+
+        # Build list of booleans indicating whether or not to expect sparse
+        # gradients for the corresponding parameters.
+        expect_sparse_gradient = [
+            produces_sparse_gradient(module) for module, _ in modules_and_parameters
+        ]
+
+        self._assign_modules_buffers()
+
+        return parameters, expect_sparse_gradient
+
+    def _assign_modules_buffers(self):
+        """
+        Assign self.module.named_buffers to self.modules_buffers.
+
+        Assigns module buffers to self.modules_buffers which are then used to
+        broadcast across ranks when broadcast_buffers=True. Note that this
+        must be called every time buffers need to be synced because buffers can
+        be reassigned by user module,
+        see https://github.com/pytorch/pytorch/issues/63916.
+        """
+        # Collect buffers for modules, filtering out buffers that should be ignored.
+        named_module_buffers = [
+            (buffer, buffer_name)
+            for buffer_name, buffer in self.module.named_buffers()
+            if buffer_name not in self.parameters_to_ignore
+        ]
+        self.modules_buffers = [
+            buffer for (buffer, buffer_name) in named_module_buffers
+        ]
+        # Dict[str, tensor] representing module buffers not ignored by DDP.
+        self.named_module_buffers = {
+            buffer_name: buffer for (buffer, buffer_name) in named_module_buffers
+        }
+
+    def _build_debug_param_to_name_mapping(self, parameters):
+        param_to_param_index = {parameters[i]: i for i in range(len(parameters))}
+        param_set = set(parameters)
+        param_index_to_param_fqn = {}
+        for module_name, module in self.module.named_modules():
+            for param_name, param in module.named_parameters(recurse=False):
+                fqn = f"{module_name}.{param_name}"
+                # Bypass ignored parameters since those are not reduced by DDP
+                # to begin with.
+                if fqn not in self.parameters_to_ignore and param.requires_grad:
+                    if param not in param_set:
+                        self._log_and_throw(
+                            ValueError,
+                            f"Param with name {fqn} found in module parameters, but not DDP parameters."
+                            " This indicates a bug in DDP, please report an issue to PyTorch.",
+                        )
+                    param_index = param_to_param_index[param]
+                    param_index_to_param_fqn[param_index] = fqn
+
+        # Ensure we covered all parameters
+        if len(param_set) != len(param_index_to_param_fqn):
+            self._log_and_throw(
+                ValueError,
+                (
+                    "Expected param to name mapping to cover all parameters, but"
+                    f" got conflicting lengths: {len(param_set)} vs "
+                    f"{len(param_index_to_param_fqn)}. This indicates a bug in DDP"
+                    ", please report an issue to PyTorch."
+                ),
+            )
+
+        return param_index_to_param_fqn
+
+    def _get_parameters(self, m, recurse=True):
+        """Return a generator of module parameters."""
+
+        def model_parameters(m):
+            ps = (
+                m._former_parameters.values()
+                if hasattr(m, "_former_parameters")
+                else m.parameters(recurse=False)
+            )
+            yield from ps
+
+        for mod in m.modules() if recurse else [m]:
+            yield from model_parameters(mod)
+
+    def _check_default_group(self):
+        pickle_not_supported = False
+        try:
+            if self.process_group != _get_default_group():
+                pickle_not_supported = True
+        except RuntimeError:
+            pickle_not_supported = True
+
+        if pickle_not_supported:
+            self._log_and_throw(
+                RuntimeError,
+                "DDP Pickling/Unpickling are only supported "
+                "when using DDP with the default process "
+                "group. That is, when you have called "
+                "init_process_group and have not passed "
+                "process_group argument to DDP constructor",
+            )
+
+    @contextmanager
+    def no_sync(self):
+        r"""
+        Context manager to disable gradient synchronizations across DDP processes.
+
+        Within this context, gradients will be accumulated on module
+        variables, which will later be synchronized in the first
+        forward-backward pass exiting the context.
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined variables")
+            >>> ddp = torch.nn.parallel.DistributedDataParallel(model, pg)
+            >>> with ddp.no_sync():
+            >>>     for input in inputs:
+            >>>         ddp(input).backward()  # no synchronization, accumulate grads
+            >>> ddp(another_input).backward()  # synchronize grads
+
+        .. warning::
+            The forward pass should be included inside the context manager, or
+            else gradients will still be synchronized.
+        """
+        old_require_backward_grad_sync = self.require_backward_grad_sync
+        self.require_backward_grad_sync = False
+        try:
+            yield
+        finally:
+            self.require_backward_grad_sync = old_require_backward_grad_sync
+
+    @classmethod
+    def _get_active_ddp_module(cls):
+        """`TorchDynamo` requires DDP's status and module for cooperative optimization."""
+        return cls._active_ddp_module
+
+    # note, this ctxmgr function is marked 'skip' in torchdynamo, so dynamo only kicks in
+    # for the 'module_to_run' underneath
+    # see torch._dynamo/eval_frame.py TorchPatcher.patch for more details
+    @contextmanager
+    @torch._disable_dynamo(recursive=False)
+    def _inside_ddp_forward(self):
+        # Save and restore the previous _active_ddp_module to handle nested
+        # DDP correctly (e.g., TorchRec wraps embeddings in an inner DDP inside
+        # an outer DDP).  Without this, the inner DDP's exit would clear the
+        # flag to None, causing DDPOptimizer to miss compiled regions that run
+        # after the inner forward.
+        old = DistributedDataParallel._active_ddp_module
+        DistributedDataParallel._active_ddp_module = self
+        try:
+            yield
+        finally:
+            DistributedDataParallel._active_ddp_module = old
+
+    def _run_ddp_forward(self, *inputs, **kwargs):
+        if self._use_python_reducer:
+            return self.module(*inputs, **kwargs)  # type: ignore[index]
+        else:
+            with self._inside_ddp_forward():
+                return self.module(*inputs, **kwargs)  # type: ignore[index]
+
+    def _clear_grad_buffer(self):
+        # Making param.grad points to the grad buffers before backward is based on the
+        # assumption that the grad accumulation is done in place in autograd engine,
+        # for some edge cases, if the grad accumulation in autograd engine is not in
+        # place, then the param.grad and grad buffers are detached.
+        if self._delay_grad_buffer is not None:
+            # We batch zero_grad for all params by resetting the whole grad
+            # buffer when the grad of all params is set to None.
+            all_param_grad_none = all(
+                param.grad is None for param in self._delay_all_reduce_params
+            )
+
+            for index, param in enumerate(self._delay_all_reduce_params):
+                if param.grad is None:
+                    param.grad = self._delay_grad_views[index]
+                    if not all_param_grad_none:
+                        param.grad.zero_()
+
+            if all_param_grad_none:
+                self._delay_grad_buffer.zero_()
+
+    def _lazy_init(self):
+        # Initialization for DDP that occurs after construction, but lazily
+        # before the first forward pass.
+        self._setup_in_backward_optimizers()
+        self._lazy_init_ran = True
+
+    def _pre_forward(self, *inputs, **kwargs):
+        if self._use_python_reducer:
+            return inputs, kwargs
+
+        if not self._lazy_init_ran and not torch.compiler.is_compiling():
+            self._lazy_init()
+
+        if self._delay_all_reduce_all_params:
+            return inputs, kwargs
+
+        if torch.is_grad_enabled() and self.require_backward_grad_sync:
+            if self.logger is None:
+                raise AssertionError("self.logger must not be None")
+            self.logger.set_runtime_stats_and_log()
+            self.reducer.prepare_for_forward()
+
+        # Notify the join context that this process has not joined, if
+        # needed
+        work = Join.notify_join_context(self)
+        if work:
+            self.reducer._set_forward_pass_work_handle(
+                work,
+                self._divide_by_initial_world_size,  # type: ignore[arg-type]
+            )
+
+        # Calling _rebuild_buckets before forward computation,
+        # It may allocate new buckets before deallocating old buckets
+        # inside _rebuild_buckets. To save peak memory usage,
+        # call _rebuild_buckets before the peak memory usage increases
+        # during forward computation.
+        # This should be called only once during whole training period.
+        if torch.is_grad_enabled() and self.reducer._rebuild_buckets():
+            logger.info("Reducer buckets have been rebuilt in this iteration.")
+            self._has_rebuilt_buckets = True
+
+        # sync params according to location (before/after forward) user
+        # specified as part of hook, if hook was specified.
+        if self._check_sync_bufs_pre_fwd():
+            self._sync_buffers()
+
+        if self._join_config.enable:
+            # Notify joined ranks whether they should sync in backwards pass or not.
+            self._check_global_requires_backward_grad_sync(is_joined_rank=False)
+
+        if self.device_ids:
+            moved_inputs, moved_kwargs = _to_kwargs(
+                inputs,
+                kwargs,
+                torch.device(self.device_type, self.device_ids[0]),
+                self.use_side_stream_for_tensor_copies,
+            )
+            args, kwargs = moved_inputs[0], moved_kwargs[0]
+            # Cast inputs to reduced precision if needed.
+            if self.mixed_precision is not None:
+                args, kwargs = _cast_forward_inputs(
+                    self.mixed_precision.param_dtype,
+                    *args,
+                    **kwargs,
+                )
+            return args, kwargs
+        else:
+            # Cast inputs to reduced precision if needed.
+            # TODO (rohan-varma) test this codepath.
+            if self.mixed_precision is not None:
+                inputs, kwargs = _cast_forward_inputs(
+                    self.mixed_precision.param_dtype,
+                    *inputs,
+                    **kwargs,
+                )
+            return inputs, kwargs
+
+    def _post_forward(self, output):
+        if self._use_python_reducer:
+            return output
+
+        if self._delay_all_reduce_all_params:
+            self._clear_grad_buffer()
+            return output
+
+        # sync params according to location (before/after forward) user
+        # specified as part of hook, if hook was specified.
+        if self._check_sync_bufs_post_fwd():
+            self._sync_buffers()
+
+        if torch.is_grad_enabled() and self.require_backward_grad_sync:
+            self.require_forward_param_sync = True
+            # We'll return the output object verbatim since it is a freeform
+            # object. We need to find any tensors in this object, though,
+            # because we need to figure out which parameters were used during
+            # this forward pass, to ensure we short circuit reduction for any
+            # unused parameters. Only if `find_unused_parameters` is set.
+            if self.find_unused_parameters and not self.static_graph:
+                # Do not need to populate this for static graph.
+                self.reducer.prepare_for_backward(list(_find_tensors(output)))
+            else:
+                self.reducer.prepare_for_backward([])
+        else:
+            self.require_forward_param_sync = False
+
+        # TODO: DDPSink is currently enabled for unused parameter detection and
+        # static graph training for first iteration.
+        if (self.find_unused_parameters and not self.static_graph) or (
+            self.static_graph and not self._static_graph_delay_allreduce_enqueued
+        ):
+            (
+                output_tensor_list,
+                treespec,
+                output_is_rref,
+            ) = _tree_flatten_with_rref(output)
+            output_placeholders: list[torch.Tensor | None] = [
+                None for _ in range(len(output_tensor_list))
+            ]
+            # Do not touch tensors that have no grad_fn, which can cause issues
+            # such as https://github.com/pytorch/pytorch/issues/60733
+            for i, output in enumerate(output_tensor_list):
+                if torch.is_tensor(output) and output.grad_fn is None:
+                    output_placeholders[i] = output
+
+            # When find_unused_parameters=True, makes tensors which require grad
+            # run through the DDPSink backward pass. When not all outputs are
+            # used in loss, this makes those corresponding tensors receive
+            # undefined gradient which the reducer then handles to ensure
+            # param.grad field is not touched and we don't error out.
+            passthrough_tensor_list = _DDPSink.apply(
+                weakref.ref(self),
+                *output_tensor_list,
+            )
+            for i in range(len(output_placeholders)):
+                if output_placeholders[i] is None:
+                    output_placeholders[i] = passthrough_tensor_list[i]
+
+            # Reconstruct output data structure.
+            output = _tree_unflatten_with_rref(
+                output_placeholders, treespec, output_is_rref
+            )
+
+        # At the end of the forward pass, reset the grad buffer and grad views
+        self._clear_grad_buffer()
+        return output
+
+    def forward(self, *inputs, **kwargs):
+        with torch.autograd.profiler.record_function("DistributedDataParallel.forward"):
+            inputs, kwargs = self._pre_forward(*inputs, **kwargs)
+            output = (
+                self.module.forward(*inputs, **kwargs)
+                if self._delay_all_reduce_all_params
+                else self._run_ddp_forward(*inputs, **kwargs)
+            )
+            return self._post_forward(output)
+
+    def scatter(self, inputs, kwargs, device_ids):
+        return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
+
+    def to_kwargs(self, inputs, kwargs, device_id):
+        # Kept for BC
+        return _to_kwargs(
+            inputs,
+            kwargs,
+            torch.device(self.device_type, device_id),
+            self.use_side_stream_for_tensor_copies,
+        )
+
+    def gather(self, outputs, output_device):
+        return gather(outputs, output_device, dim=self.dim)
+
+    def train(self, mode=True):
+        super().train(mode)
+        return self
+
+    # When running in join mode, schedules an allreduce to notify joined ranks
+    # of whether backwards pass synchronization will run this iteration or not.
+    def _check_global_requires_backward_grad_sync(self, is_joined_rank):
+        if not is_joined_rank and self.require_backward_grad_sync:
+            requires_sync_tensor = torch.ones(1, device=self.device)
+        else:
+            requires_sync_tensor = torch.zeros(1, device=self.device)
+
+        work = dist.all_reduce(
+            requires_sync_tensor, group=self.process_group, async_op=True
+        )
+
+        # (kwen2501) This if condition is a plain translation of previous
+        # behavior, i.e. in the `is_joined_rank=False` case, `work.wait()`
+        # is not called and it doesn't care about the result. I am guessing
+        # that it just wants to fire a matching all-reduce and does not want
+        # the main stream to wait.
+        if is_joined_rank:
+            work.wait()
+            should_sync_backwards = requires_sync_tensor.item() != 0
+            return should_sync_backwards
+        else:
+            return None  # Return value is not/should not be used.
+
+    # When running in join mode, checks and performs sync of module buffers if
+    # the models have buffers that should be synchronized in the forward pass.
+    def _check_and_sync_module_buffers(self):
+        if self._check_sync_bufs_pre_fwd():
+            authoritative_rank = self._find_common_rank(self._distributed_rank, False)
+            self._sync_module_buffers(authoritative_rank)
+
+    # When running in join model, agrees upon a common rank and broadcast model
+    # parameters to all other ranks.
+    def _sync_final_model(self, is_last_joiner):
+        # Agree upon the process that will be the authoritative model copy.
+        # The current rank is a candidate for being the authoritative copy if
+        # is_last_joiner=True. We break ties via picking the larger rank.
+        self._authoritative_rank = self._find_common_rank(
+            self._distributed_rank, is_last_joiner
+        )
+        _sync_module_states(
+            module=self.module,
+            process_group=self.process_group,
+            broadcast_bucket_size=self.broadcast_bucket_size,
+            src=self._authoritative_rank,
+            params_and_buffers_to_ignore=self.parameters_to_ignore,
+            broadcast_buffers=self.broadcast_buffers,
+        )
+
+    # Schedule comm ops to match those scheduled in the reducer's backward
+    # pass.
+    def _match_all_reduce_for_bwd_pass(self):
+        comm_work = []
+        # Schedule comm in the same order as Reducer schedules them, i.e.
+        # the order of the buckets. Retrieving the bucket order from the reducer
+        # ensures that we keep the same order in join mode, such as when bucket
+        # order is rebuilt dynamically.
+
+        # Returns grad_buckets in order, but real tensors are substituted with
+        # zero tensors of the same shape.
+        grad_buckets = self.reducer._get_zeros_like_grad_buckets()
+        for grad_bucket in grad_buckets:
+            # Joined processes contribute zero gradient. In the case that
+            # divide_by_initial_world_size=True, we divide grads by the static
+            # world size, if not, the dividing factor is reduced by the number
+            # of joined processes.
+            work = self.reducer._run_comm_hook(grad_bucket)
+            comm_work.append(work)
+        for work in comm_work:
+            work.wait()
+
+    # Allreduces the used parameter mapping across ranks.
+    def _match_unused_params_allreduce(self):
+        locally_used_param_map = self.reducer._get_local_used_map()
+        self.process_group.allreduce(locally_used_param_map)
+
+    def join(
+        self,
+        divide_by_initial_world_size: bool = True,
+        enable: bool = True,
+        throw_on_early_termination: bool = False,
+    ):
+        r"""
+        Context manager for training with uneven inputs across processes in DDP.
+
+        This context manager will keep track of already-joined DDP processes,
+        and "shadow" the forward and backward passes by inserting collective
+        communication operations to match with the ones created by non-joined
+        DDP processes. This will ensure each collective call has a corresponding
+        call by already-joined DDP processes, preventing hangs or errors that
+        would otherwise happen when training with uneven inputs across
+        processes. Alternatively, if the flag ``throw_on_early_termination`` is
+        specified to be ``True``, all trainers will throw an error once one rank
+        runs out of inputs, allowing these errors to be caught and handled
+        according to application logic.
+
+        Once all DDP processes have joined, the context manager will broadcast
+        the model corresponding to the last joined process to all processes to
+        ensure the model is the same across all processes
+        (which is guaranteed by DDP).
+
+        To use this to enable training with uneven inputs across processes,
+        simply wrap this context manager around your training loop. No further
+        modifications to the model or data loading is required.
+
+        .. warning::
+            If the model or training loop this context manager is wrapped around
+            has additional distributed collective operations, such as
+            ``SyncBatchNorm`` in the model's forward pass, then the flag
+            ``throw_on_early_termination`` must be enabled. This is because this
+            context manager is not aware of non-DDP collective communication.
+            This flag will cause all ranks to throw when any one rank
+            exhausts inputs, allowing these errors to be caught and recovered
+            from across all ranks.
+
+        Args:
+            divide_by_initial_world_size (bool): If ``True``, will divide
+                gradients by the initial ``world_size`` DDP training was launched
+                with. If ``False``, will compute the effective world size
+                (number of ranks that have not depleted their inputs yet) and
+                divide gradients by that during allreduce. Set
+                ``divide_by_initial_world_size=True`` to ensure every input
+                sample including the uneven inputs have equal weight in terms of
+                how much they contribute to the global gradient. This is
+                achieved by always dividing the gradient by the initial
+                ``world_size`` even when we encounter uneven inputs. If you set
+                this to ``False``, we divide the gradient by the remaining
+                number of nodes. This ensures parity with training on a smaller
+                ``world_size`` although it also means the uneven inputs would
+                contribute more towards the global gradient. Typically, you
+                would want to set this to ``True`` for cases where the last few
+                inputs of your training job are uneven. In extreme cases, where
+                there is a large discrepancy in the number of inputs, setting
+                this to ``False`` might provide better results.
+            enable (bool): Whether to enable uneven input detection or not. Pass
+                in ``enable=False`` to disable in cases where you know that
+                inputs are even across participating processes. Default is
+                ``True``.
+            throw_on_early_termination (bool): Whether to throw an error
+                or continue training when at least one rank has exhausted
+                inputs. If ``True``, will throw upon the first rank reaching end
+                of data. If ``False``, will continue training with a smaller
+                effective world size until all ranks are joined. Note that if
+                this flag is specified, then the flag
+                ``divide_by_initial_world_size`` would be ignored. Default
+                is ``False``.
+
+
+        Example::
+
+            >>> # xdoctest: +SKIP("Distributed")
+            >>> import torch
+            >>> import torch.distributed as dist
+            >>> import os
+            >>> import torch.multiprocessing as mp
+            >>> import torch.nn as nn
+            >>> # On each spawned worker
+            >>> def worker(rank):
+            >>>     dist.init_process_group("nccl", rank=rank, world_size=2)
+            >>>     torch.cuda.set_device(rank)
+            >>>     model = nn.Linear(1, 1, bias=False).to(rank)
+            >>>     model = torch.nn.parallel.DistributedDataParallel(
+            >>>         model, device_ids=[rank], output_device=rank
+            >>>     )
+            >>>     # Rank 1 gets one more input than rank 0.
+            >>>     inputs = [torch.tensor([1]).float() for _ in range(10 + rank)]
+            >>>     with model.join():
+            >>>         for _ in range(5):
+            >>>             for inp in inputs:
+            >>>                 loss = model(inp).sum()
+            >>>                 loss.backward()
+            >>>     # Without the join() API, the below synchronization will hang
+            >>>     # blocking for rank 1's allreduce to complete.
+            >>>     torch.cuda.synchronize(device=rank)
+        """
+        return Join(
+            [self],
+            enable,
+            throw_on_early_termination,
+            divide_by_initial_world_size=divide_by_initial_world_size,
+        )
+
+    def join_hook(
+        self,
+        **kwargs,
+    ):
+        r"""
+        DDP join hook enables training on uneven inputs by mirroring communications in forward and backward passes.
+
+        Arguments:
+            kwargs (dict): a :class:`dict` containing any keyword arguments
+                to modify the behavior of the join hook at run time; all
+                :class:`Joinable` instances sharing the same join context
+                manager are forwarded the same value for ``kwargs``.
+
+        The hook supports the following keyword arguments:
+            divide_by_initial_world_size (bool, optional):
+                If ``True``, then gradients are divided by the initial world
+                size that DDP was launched with.
+                If ``False``, then gradients are divided by the effective world
+                size (i.e. the number of non-joined processes), meaning that
+                the uneven inputs contribute more toward the global gradient.
+                Typically, this should be set to ``True`` if the degree of
+                unevenness is small but can be set to ``False`` in extreme
+                cases for possibly better results.
+                Default is ``True``.
+        """
+        divide_by_initial_world_size = kwargs.get("divide_by_initial_world_size", True)
+        return _DDPJoinHook(
+            self, divide_by_initial_world_size=divide_by_initial_world_size
+        )
+
+    @property
+    def join_device(self):
+        return self.device
+
+    @property
+    def join_process_group(self):
+        return self.process_group
+
+    def _register_buffer_comm_hook(
+        self,
+        state,
+        hook: Callable,
+        comm_hook_location=_BufferCommHookLocation.POST_FORWARD,
+    ):
+        r"""
+        Allow custom registration of hooks that define how buffer are synchronized across ranks.
+
+        The hook takes in an optional state and is passed in a Dict[str, Tensor]
+        corresponding to buffer names and the buffers, and can run arbitrary reductions
+        on buffers as opposed to DDP's default broadcast from rank 0. This is useful for
+        example if a counter needs to be summed or averaged across ranks every iteration.
+
+        Args:
+            state (Any): Optional state that is passed to the hook.
+            hook (Callable): Callable with the following signature:
+                         ``hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]``
+            comm_hook_location (_BufferCommHookLocation): Enum value indicating
+                            where to run the hook.
+                            _BufferCommHookLocation.PRE_FORWARD means that the
+                            hook will run _before_ the forward pass, and
+                            _BufferCommHookLocation.POST_FORWARD means that the
+                            hook will run _after_ the forward pass.
+
+            NOTE: To maximize performance, users can return a
+                list[torch.futures.Future] from their hook, and DDP will
+                install and await these hooks appropriately at the end of
+                the backward pass. This will ensure all buffers are
+                synchronized by the end of the backward pass. If this
+                setting is used, it is recommended to pass
+                comm_hook_location=_BufferCommHookLocation.POST_FORWARD,
+                which will trigger the hook after the forward pass.
+                If _BufferCommHookLocation.PRE_FORWARD is used, users must
+                ensure appropriate synchronization when manipulating GPU
+                buffers in the forward pass.
+        """
+        if not callable(hook):
+            raise AssertionError(f"hook must be callable, got {type(hook).__name__}")
+        self.buffer_hook = _BufferCommHook(
+            buffer_comm_hook=hook,
+            buffer_comm_hook_state=state,
+            buffer_comm_hook_location=comm_hook_location,
+        )
+
+    def register_comm_hook(self, state: object, hook: Callable) -> None:
+        r"""
+        Register communication hook for user-defined DDP aggregation of gradients across multiple workers.
+
+        This hook would be very useful for researchers to try out new ideas. For
+        example, this hook can be used to implement several algorithms like GossipGrad
+        and gradient compression which involve different communication strategies for
+        parameter syncs while running Distributed DataParallel training.
+
+        Args:
+            state (object): Passed to the hook to maintain any state information during the training process.
+                            Examples include error feedback in gradient compression,
+                            peers to communicate with next in GossipGrad, etc.
+
+                            It is locally stored by each worker
+                            and shared by all the gradient tensors on the worker.
+            hook (Callable): Callable with the following signature:
+                             ``hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]``:
+
+                             This function is called once the bucket is ready. The
+                             hook can perform whatever processing is needed and return
+                             a Future indicating completion of any async work (ex: allreduce).
+                             If the hook doesn't perform any communication, it still
+                             must return a completed Future. The Future should hold the
+                             new value of grad bucket's tensors. Once a bucket is ready,
+                             c10d reducer would call this hook and use the tensors returned
+                             by the Future and copy grads to individual parameters.
+                             Note that the future's return type must be a single tensor.
+
+                             We also provide an API called ``get_future`` to retrieve a
+                             Future associated with the completion of ``c10d.ProcessGroup.Work``.
+                             ``get_future`` is currently supported for NCCL and also supported for most
+                             operations on GLOO and MPI, except for peer to peer operations (send/recv).
+
+        .. warning ::
+            Grad bucket's tensors will not be predivided by world_size. User is responsible
+            to divide by the world_size in case of operations like allreduce.
+
+        .. warning ::
+            DDP communication hook can only be registered once and should be registered
+            before calling backward.
+
+        .. warning ::
+            The Future object that hook returns should contain a single tensor
+            that has the same shape with the tensors inside grad bucket.
+
+        .. warning ::
+            ``get_future`` API supports NCCL, and partially GLOO and MPI backends (no support
+            for peer-to-peer operations like send/recv) and will return a ``torch.futures.Future``.
+
+        Example::
+            Below is an example of a noop hook that returns the same tensor.
+
+            >>> # xdoctest: +SKIP('undefined name')
+            >>> def noop(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
+            >>>     fut = torch.futures.Future()
+            >>>     fut.set_result(bucket.buffer())
+            >>>     return fut
+            >>> ddp.register_comm_hook(state=None, hook=noop)
+
+        Example::
+            Below is an example of a Parallel SGD algorithm where gradients are encoded before
+            allreduce, and then decoded after allreduce.
+
+            >>> # xdoctest: +SKIP('undefined name')
+            >>> def encode_and_decode(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
+            >>>     encoded_tensor = encode(bucket.buffer())  # encode gradients
+            >>>     fut = torch.distributed.all_reduce(encoded_tensor).get_future()
+            >>>     # Define the then callback to decode.
+            >>>     def decode(fut):
+            >>>         decoded_tensor = decode(fut.value()[0])  # decode gradients
+            >>>         return decoded_tensor
+            >>>     return fut.then(decode)
+            >>> ddp.register_comm_hook(state=None, hook=encode_and_decode)
+        """
+        self._check_comm_hook(hook)
+        if self.logger is None:
+            raise AssertionError("self.logger must not be None")
+        self.logger._set_comm_hook_name(hook.__qualname__)
+        self._comm_hooks.append((hook, state))
+        dist._register_comm_hook(self.reducer, state, hook)
+
+    def _register_builtin_comm_hook(self, comm_hook_type):
+        r"""
+        Register a built-in communication hook that specifies how DDP aggregates gradients across multiple workers.
+
+        The built-in hooks aim to provide efficient C++ implementations for certain hooks,
+        which might not be as efficient if implemented in Python using a Python communication hook.
+
+        Args:
+            comm_hook_type (dist.BuiltinCommHookType): type of communication hook, such as ALLREDUCE, FP16_COMPRESS, etc.
+
+        .. warning ::
+            DDP communication hook can only be registered once and should be registered
+            before calling backward.
+
+        Example::
+            Below is an example of a FP16 compression where gradients are
+            compressed into 16-bit floating-point numbers before allreduce, and
+            then decompressed after allreduce.
+
+            >>> # xdoctest: +SKIP('undefined name')
+            >>> ddp._register_builtin_comm_hook(dist.BuiltinCommHookType.FP16_COMPRESS)
+
+        """
+        if self.logger is None:
+            raise AssertionError("self.logger must not be None")
+        self.logger._set_comm_hook_name(str(comm_hook_type))
+        dist._register_builtin_comm_hook(self.reducer, comm_hook_type)
+
+    def _register_fused_optim(self, optim: type, *args, optim_params=None, **kwargs):
+        r"""
+        Register an optimizer in DDP to optimize parameter immediately after its gradient reduction.
+
+        Registers an optimizer with DDP such that the optimization for a
+        parameter will run immediately when that parameter's gradient is
+        finished with reduction, instead of waiting for all parameters'
+        gradients to finish reduction. This can result in a training speedup
+        depending on your workload since the optimizer can run while gradient
+        reduction for other parameters are still ongoing. In addition, this has
+        the potential to reduce peak memory consumption during training, as it
+        only needs to load the per-parameter optimizer states of a single
+        parameter at a time, instead of loading all per-parameter optimizer
+        states at once.
+
+        Args:
+            optim (Type): a ``torch.optim.Optimizer`` class to be registered
+            as a fused optimizer.
+            *args (Sequence[Any]): Arguments to forward to `optim`.
+            optim_params (Optional[Iterable[torch.Tensor]]): Set of parameters
+            to optimize, similar to `params` argument of traditional `torch.optim`
+            Optimizers. If this is omitted, all DDP model parameters will be
+            optimized.
+            **kwargs: (Dict[str, Any]): Keyword arguments to forward to `optim`.
+
+        .. warning ::
+            _register_fused_optim should only be called once on a DDP instance,
+            and registering multiple fused optimizers for the same DDP model
+            is not currently supported. Please ping
+            https://github.com/pytorch/pytorch/issues/71595 if this is necessary
+            for your use case.
+
+        .. warning ::
+            _register_fused_optim and register_comm_hook currently do not
+            compose together, meaning that custom DDP communication hooks are
+            not supported with overlapped optimizers. Please ping
+            https://github.com/pytorch/pytorch/issues/71595 if this is necessary
+            for your use case.
+
+        .. warning ::
+            Gradient accumulation and DDP `no_sync` are currently not supported
+            with overlapped optimizer. Please ping
+            https://github.com/pytorch/pytorch/issues/71595 if this is necessary
+            for your use case.
+
+        Example::
+
+            >>> # xdoctest: +SKIP("No rendezvous handler")
+            >>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
+            >>> net = torch.nn.parallel.DistributedDataParallel(model, pg)
+            >>> lr = 1e-2
+            >>> betas = (0.9, 0.99)
+            >>> eps = 1e-6
+            >>> net._register_fused_optim(torch.optim.Adam, lr, betas=betas, eps=eps)
+            >>> # Example with subset of parameters
+            >>> params_to_opt = [list(net.parameters())[0]]
+            >>> net._register_fused_optim(
+            ...   torch.optim.Adam, lr, optim_params=params_to_opt,  betas=betas, eps=eps
+            ... )
+        """
+        # Note: importing in function, otherwise this will cause a circular
+        # import as optimizer_overlap module needs to import DistributedDataParallel.
+        from torch.distributed.algorithms._optimizer_overlap import _as_overlapped_optim
+
+        overlapped_optim = _as_overlapped_optim(optim, optim_params, *args, **kwargs)
+        try:
+            overlapped_optim.register_ddp(self)
+        except NotImplementedError as e:
+            raise RuntimeError(
+                f"{optim} does not support overlapped DDP. Please file an issue to PyTorch or the respective owner of {optim}."
+            ) from e
+
+    def _distributed_broadcast_coalesced(
+        self, tensors, buffer_size, authoritative_rank=0
+    ):
+        dist._broadcast_coalesced(
+            self.process_group, tensors, buffer_size, authoritative_rank
+        )
+
+    def _check_sync_bufs_post_fwd(self):
+        return (
+            self.will_sync_module_buffers()
+            and hasattr(self, "buffer_hook")
+            and self.buffer_hook.buffer_comm_hook_location
+            == _BufferCommHookLocation.POST_FORWARD
+        )
+
+    def _check_sync_bufs_pre_fwd(self):
+        return self.will_sync_module_buffers() and (
+            not hasattr(self, "buffer_hook")
+            or self.buffer_hook.buffer_comm_hook_location
+            == _BufferCommHookLocation.PRE_FORWARD
+        )
+
+    def will_sync_module_buffers(self):
+        return (
+            self.require_forward_param_sync
+            and self.broadcast_buffers
+            and len(self.modules_buffers) > 0
+        )
+
+    def _find_common_rank(self, input_rank, rank_cond):
+        # -1 indicates that this rank is not under consideration to be the
+        # common_rank
+        rank_to_use = torch.tensor(
+            [input_rank if rank_cond else -1],
+            device=self.device,
+        )
+        dist.all_reduce(rank_to_use, op=ReduceOp.MAX, group=self.process_group)
+        if rank_to_use.item() == -1:
+            self._log_and_throw(
+                ValueError,
+                "BUG! Expected rank_cond to be true for at least one process."
+                " This indicates a bug in PyTorch, please report an issue.",
+            )
+        return rank_to_use.item()
+
+    def _sync_buffers(self):
+        with torch.no_grad():
+            # module buffer sync
+            # Synchronize buffers across processes.
+            # If we are running DDP with the join manager, we have to agree
+            # upon a rank to sync module buffers from, since rank 0 may
+            # already have been joined and have stale module buffers.
+            if self._join_config.enable:
+                authoritative_rank = self._find_common_rank(
+                    self._distributed_rank, True
+                )
+            else:
+                # The process with rank 0 is considered the authoritative copy.
+                authoritative_rank = 0
+            # Update self.modules_buffers in case any buffers were
+            # reassigned.
+            self._assign_modules_buffers()
+            with torch.autograd._unsafe_preserve_version_counter(
+                tuple(self.modules_buffers)
+            ):
+                self._sync_module_buffers(authoritative_rank)
+
+    def _sync_module_buffers(self, authoritative_rank):
+        if not hasattr(self, "buffer_hook"):
+            self._default_broadcast_coalesced(authoritative_rank=authoritative_rank)
+        else:
+            hook = self.buffer_hook.buffer_comm_hook
+            state = self.buffer_hook.buffer_comm_hook_state
+            futs = hook(state, self.named_module_buffers)
+            if futs is not None:
+                self.reducer._install_post_backward_futures(futs)
+
+    def _default_broadcast_coalesced(
+        self, bufs=None, bucket_size=None, authoritative_rank=0
+    ):
+        """
+        Broadcasts buffers from rank 0 to rest of workers.
+
+        If bufs, bucket_size are None, default values self.modules_buffers
+        and self.broadcast_bucket_size are used instead.
+        """
+        if bufs is None:
+            bufs = self.modules_buffers
+        if bucket_size is None:
+            bucket_size = self.broadcast_bucket_size
+
+        self._distributed_broadcast_coalesced(bufs, bucket_size, authoritative_rank)
+
+    def _passing_sync_batchnorm_handle(self, module):
+        for layer in module.modules():
+            if isinstance(layer, torch.nn.modules.SyncBatchNorm):
+                if self.device_type == "cpu":
+                    self._log_and_throw(
+                        ValueError,
+                        "SyncBatchNorm layers only work with GPU modules",
+                    )
+
+    def _check_comm_hook(self, hook):
+        if not callable(hook):
+            self._log_and_throw(TypeError, "Communication hook must be callable.")
+
+        sig = inspect.signature(hook)
+        if (
+            sig.parameters["bucket"].annotation != inspect._empty
+            and sig.parameters["bucket"].annotation != dist.GradBucket
+        ):
+            self._log_and_throw(
+                ValueError,
+                "Communication hook: bucket annotation should be dist.GradBucket.",
+            )
+
+        if (
+            sig.return_annotation != inspect._empty
+            and sig.return_annotation != torch.futures.Future[torch.Tensor]
+        ):
+            self._log_and_throw(
+                ValueError,
+                "Communication hook: return annotation should be torch.futures.Future[torch.Tensor].",
+            )
+
+        if hook.__name__ in ["bf16_compress_hook", "bf16_compress_wrapper_hook"]:
+            cuda_supported = (
+                torch.version.cuda is not None
+            ) or torch.version.hip is not None
+            nccl_supported = (
+                dist.is_available()
+                and dist.is_nccl_available()
+                and torch.cuda.nccl.version() >= (2, 10)
+            )
+            xpu_xccl_supported = (
+                dist.is_available()
+                and dist.is_xccl_available()
+                and torch.xpu.is_available()
+            )
+
+            if not ((cuda_supported and nccl_supported) or xpu_xccl_supported):
+                self._log_and_throw(
+                    TypeError,
+                    "BF16 all reduce communication hook required CUDA 11+ and NCCL 2.10+ or XPU and XCCL",
+                )
+
+    @property
+    def _distributed_rank(self):
+        return dist.get_rank(self.process_group)
+
+    @staticmethod
+    def _get_data_parallel_params(module, named_params=False):
+        """Return a generator of parameters managed by a given DDP unit."""
+        for param in (
+            module.parameters() if not named_params else module.named_parameters()
+        ):
+            if not hasattr(param, "_ddp_ignored"):
+                yield param
+
+    @staticmethod
+    def _set_params_and_buffers_to_ignore_for_model(
+        module, params_and_buffers_to_ignore
+    ):
+        """
+        Set parameters and buffers to be ignored by DDP.
+
+        Expected format for parameters is the fully qualified name: {module_name}.{param_name}, and
+        similarly, {module_name}.{buffer_name} for buffers. For example:
+        params_to_ignore = []
+        # NB: model here is vanilla PyTorch module, not yet wrapped with DDP.
+        for module_name, module in model.named_modules():
+            for param_name, param in module.named_parameters(recurse=False):
+                if should_ignore(param):
+                    # Create expected format
+                    fqn = f"{module_name}.{param_name}"
+                    params_to_ignore.append(fqn)
+        torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model(
+            model,
+            params_to_ignore
+        )
+        """
+        # This is a workaround to set parameters and buffers DDP should ignore
+        # during synchronization. It will be removed when the API is finalized
+        # as part of addressing https://github.com/pytorch/pytorch/issues/43690.
+        module._ddp_params_and_buffers_to_ignore = params_and_buffers_to_ignore
+        for name, param in module.named_parameters():
+            if name in params_and_buffers_to_ignore:
+                param._ddp_ignored = True
+        for name, buffer in module.named_buffers():
+            if name in params_and_buffers_to_ignore:
+                buffer._ddp_ignored = True
+
+    def _get_ddp_logging_data(self):
+        r"""
+        Return a dictionary of logging data for debugging and analysis.
+
+        This interface can be called after DistributedDataParallel() is
+        constructed. It returns a dictionary of logging data. It could help
+        for debugging and analysis. The logging data includes DistributedDataParallel
+        constructor input parameters, some internal states of DistributedDataParallel
+        and performance metrics. Simply print the dictionary and see what
+        these metrics are.
+        This is a prototype interface and subject to change in the future.
+        """
+        if self.logger is None:
+            raise AssertionError("self.logger must not be None")
+        ddp_logging_data = self.logger._get_ddp_logging_data()
+        return {**ddp_logging_data.strs_map, **ddp_logging_data.ints_map}
+
+    def _set_ddp_runtime_logging_sample_rate(self, sample_rate):
+        r"""
+        Set sample_rate of collecting runtime stats.
+
+        This interface allows users to set sample_rate of collecting
+        runtime stats. The runtime stats will be recorded for the
+        first 10 iterations, after 10 iterations runtime stats will be
+        recorded once every "sample_rate" training iterations. In
+        default, runtime stats are recorded for the first 10 iterations,
+        after 10 iterations runtime stats are recorded once every
+        "kDDPRuntimeLoggingSampleRate=100" training iterations.
+        This is a prototype interface and subject to change in the future.
+        """
+        if sample_rate < 1:
+            self._log_and_throw(
+                ValueError,
+                "DDP runtime logging sample rate should be equal or greater than 1",
+            )
+        self.reducer._set_ddp_runtime_logging_sample_rate(sample_rate)
+
+    def _set_static_graph(self):
+        """
+        Set static graph for DDP.
+
+        It is recommended to set static graph in the DDP constructor, which will
+        call this private API internally.
+        """
+        # If self.static_graph has been set, no need to set it again
+        if self.static_graph:
+            warnings.warn(
+                "You've set static_graph to be True, no need to set it again.",
+                stacklevel=2,
+            )
+            return
+        self.static_graph = True
+        self._static_graph_delay_allreduce_enqueued = False
+        self.reducer._set_static_graph()
+        if self.logger is None:
+            raise AssertionError("self.logger must not be None")
+        self.logger._set_static_graph()
+        if self.find_unused_parameters:
+            warnings.warn(
+                "You passed find_unused_parameters=true to DistributedDataParallel, "
+                "`_set_static_graph` will detect unused parameters automatically, so "
+                "you do not need to set find_unused_parameters=true, just be sure these "
+                "unused parameters will not change during training loop while calling "
+                "`_set_static_graph`.",
+                stacklevel=2,
+            )
+
+    def _remove_autograd_hooks(self):
+        """Remove autograd hooks registered by the reducer on the model parameters."""
+        self.reducer._remove_autograd_hooks()
+
+    def _check_reducer_finalized(self):
+        """
+        Check if the reducer has processed all buckets and finalized the backward appropriately.
+
+        It is useful to call this method after calling .backward() in your training loop
+        in order to avoid subsequent hard to debug errors down the road due to the
+        reducer not finalizing backward.
+        """
+        self.reducer._check_reducer_finalized()
+
+    def _set_sparse_metadata(self, global_unique_ids):
+        self.reducer._set_sparse_metadata(global_unique_ids)
+
+    def _update_process_group(self, new_process_group):
+        """
+        Dynamically updates the process group for DDP so that we can shrink/expand DDP
+        world size without having to reinitialize DDP.
+
+        NOTE: If you are using custom communications hooks via, register_comm_hook,
+        you need to update the process groups for those hooks separately.
+        """
+        # Force a rebuild of buckets for a new process group. This ensures all ranks
+        # are synchronized in terms of when they will rebuild buckets and also
+        # re-evaluates previous assumptions of buckets given the world size might have
+        # changed.
+        self._has_rebuilt_buckets = False
+        self.reducer._reset_state()
+
+        if not _rank_not_in_group(new_process_group):
+            self.process_group = new_process_group
+            self.reducer._update_process_group(new_process_group)
+
+    def _set_ddp_sink_clone(self, val: bool):
+        """
+        Sets whether or not DDPSink should clone the output tensors or not.
+        The default is True since if the loss is modified in place we run
+        into the view is modified in-place error.
+
+        Although, cloning the tensors can add significant memory and
+        performance hit if the number and size of tensors are large. As
+        a result, this can be set to False if you are not modifying the
+        loss in place.
+        """
+        self._ddp_sink_clone = val
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae8ff24c104465dcb3fb884c56c7e775986cdf73
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py
@@ -0,0 +1,146 @@
+import threading
+from collections.abc import Sequence
+from typing import Any, cast
+
+import torch
+from torch._utils import ExceptionWrapper
+from torch.cuda._utils import _get_device_index
+from torch.nn.modules import Module
+
+
+__all__ = ["get_a_var", "parallel_apply"]
+
+
+def get_a_var(
+    obj: torch.Tensor | list[Any] | tuple[Any, ...] | dict[Any, Any],
+) -> torch.Tensor | None:
+    if isinstance(obj, torch.Tensor):
+        return obj
+
+    if isinstance(obj, (list, tuple)):
+        for result in map(get_a_var, obj):
+            if isinstance(result, torch.Tensor):
+                return result
+    if isinstance(obj, dict):
+        for result in map(get_a_var, obj.items()):
+            if isinstance(result, torch.Tensor):
+                return result
+    return None
+
+
+def parallel_apply(
+    modules: Sequence[Module],
+    inputs: Sequence[Any],
+    kwargs_tup: Sequence[dict[str, Any]] | None = None,
+    devices: Sequence[int | torch.device | None] | None = None,
+) -> list[Any]:
+    r"""Apply each `module` in :attr:`modules` in parallel on each of :attr:`devices`.
+
+    Args:
+        modules (Module): modules to be parallelized
+        inputs (tensor): inputs to the modules
+        devices (list of int or torch.device): CUDA devices
+
+    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
+    :attr:`devices` (if given) should all have same length. Moreover, each
+    element of :attr:`inputs` can either be a single object as the only argument
+    to a module, or a collection of positional arguments.
+    """
+    if len(modules) != len(inputs):
+        raise AssertionError(
+            f"The number of modules {len(modules)} is not equal to "
+            f"the number of inputs {len(inputs)}"
+        )
+    if kwargs_tup is not None:
+        if len(modules) != len(kwargs_tup):
+            raise AssertionError(
+                f"The number of modules {len(modules)} is not equal to "
+                f"the number of kwargs_tup {len(kwargs_tup)}"
+            )
+    else:
+        kwargs_tup = (cast(dict[str, Any], {}),) * len(modules)
+    if devices is not None:
+        if len(modules) != len(devices):
+            raise AssertionError(
+                f"The number of modules {len(modules)} is not equal to "
+                f"the number of devices {len(devices)}"
+            )
+    else:
+        devices = [None] * len(modules)
+    devices = [_get_device_index(x, True) for x in devices]
+    streams = [torch.accelerator.current_stream(x) for x in devices]
+    if not torch.accelerator.is_available():
+        raise AssertionError("No available accelerator found.")
+    device_type = torch.accelerator.current_accelerator().type  # type: ignore[union-attr]
+    lock = threading.Lock()
+    results = {}
+    grad_enabled, autocast_enabled = (
+        torch.is_grad_enabled(),
+        torch.is_autocast_enabled(),
+    )
+
+    def _worker(
+        i: int,
+        module: Module,
+        input: Any,
+        kwargs: dict[str, Any],
+        device: int | torch.device | None = None,
+        stream: torch.Stream | None = None,
+    ) -> None:
+        torch.set_grad_enabled(grad_enabled)
+        if device is None:
+            t = get_a_var(input)
+            if t is None:
+                with lock:
+                    results[i] = ExceptionWrapper(
+                        where=f"in replica {i}, no device was provided and no tensor input was found; "
+                        "device cannot be resolved"
+                    )
+                return
+            device = t.get_device()
+        if isinstance(device, torch.device):
+            device = device.index
+        if stream is None:
+            stream = torch.accelerator.current_stream(device)
+        try:
+            with (
+                torch.accelerator.device_index(device),
+                stream,
+                torch.amp.autocast(device_type, enabled=autocast_enabled),
+            ):
+                # this also avoids accidental slicing of `input` if it is a Tensor
+                if not isinstance(input, (list, tuple)):
+                    input = (input,)
+                output = module(*input, **kwargs)
+            with lock:
+                results[i] = output
+        except Exception:
+            with lock:
+                results[i] = ExceptionWrapper(
+                    where=f"in replica {i} on device {device}"
+                )
+
+    if len(modules) > 1:
+        threads = [
+            threading.Thread(
+                target=_worker, args=(i, module, input, kwargs, device, stream)
+            )
+            for i, (module, input, kwargs, device, stream) in enumerate(
+                zip(modules, inputs, kwargs_tup, devices, streams, strict=True)
+            )
+        ]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+    else:
+        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0], streams[0])
+
+    outputs = []
+    for i in range(len(inputs)):
+        output = results[i]
+        if isinstance(output, ExceptionWrapper):
+            output.reraise()
+        outputs.append(output)
+    return outputs
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/replicate.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/replicate.py
new file mode 100644
index 0000000000000000000000000000000000000000..269af40ade3890c5eaaa6ed7396dd5210687596e
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/replicate.py
@@ -0,0 +1,214 @@
+from collections import OrderedDict
+from collections.abc import Iterator, Sequence
+from typing import cast, TYPE_CHECKING, TypeVar
+from typing_extensions import TypeIs
+
+import torch
+from torch._utils import _get_device_index
+from torch.nn.modules import Module
+from torch.nn.parallel import comm
+
+
+if TYPE_CHECKING:
+    from torch._C import ScriptMethod
+    from torch.jit import ScriptModule
+    from torch.jit._state import EnabledProxy
+
+
+__all__ = ["replicate"]
+
+
+def _is_script_module(module: Module) -> TypeIs["ScriptModule"]:
+    import torch.jit
+
+    return isinstance(module, torch.jit.ScriptModule)
+
+
+def _is_script_method(module: object) -> TypeIs["ScriptMethod"]:
+    import torch.jit
+
+    return isinstance(module, torch._C.ScriptMethod)
+
+
+def _init_script_module() -> "ScriptModule":
+    import torch.jit
+
+    return torch.jit.ScriptModule()
+
+
+def _is_jit_enabled() -> "EnabledProxy":
+    import torch.jit._state
+
+    return torch.jit._state._enabled
+
+
+# Check if we can safely replicate the module.
+# there are two types of module:
+# 1. python modules
+# 2. ScriptModule
+#
+# currently a module cannot be replicated properly if the descendants of
+# any ScriptModule contains python module (type 1 above)
+def _replicatable_module(module: Module, memo: set[Module] | None = None) -> bool:
+    # module.modules() contains module itself as the first element
+    def descendant_modules(module: Module) -> Iterator[Module]:
+        gen = module.modules()
+        next(gen)
+        return gen
+
+    if not _is_jit_enabled():
+        return True
+    if memo is None:
+        memo = set()
+
+    # memoize visited modules
+    memo.add(module)
+    if _is_script_module(module):
+        memo.update(descendant_modules(module))
+        return all(
+            _is_script_module(descendant) for descendant in descendant_modules(module)
+        )
+
+    for child in module.children():
+        # since any unreplicatable module will cause the check to return
+        # False early, visited modules here can be safely ignored.
+        if child in memo:
+            continue
+        if not _replicatable_module(child, memo):
+            return False
+
+    return True
+
+
+def _broadcast_coalesced_reshape(
+    tensors: Sequence[torch.Tensor],
+    devices: Sequence[int | torch.device],
+    detach: bool = False,
+) -> list[list[torch.Tensor]]:
+    from torch.nn.parallel._functions import Broadcast
+
+    if len(tensors) == 0:
+        return []
+
+    if detach:
+        complex_mask = [
+            not isinstance(t, torch.nn.UninitializedParameter) and t.is_complex()
+            for t in tensors
+        ]
+
+        outputs = comm.broadcast_coalesced(tensors, devices)
+
+        for device_outputs in outputs:
+            for i, is_complex in enumerate(complex_mask):
+                if is_complex:
+                    device_outputs[i] = torch.view_as_complex(device_outputs[i])
+
+        return outputs
+    else:
+        tensor_copies = Broadcast.apply(devices, *tensors)
+        return [
+            list(tensor_copies[i : i + len(tensors)])
+            for i in range(0, len(tensor_copies), len(tensors))
+        ]
+
+
+T = TypeVar("T", bound=Module)
+
+
+def replicate(
+    network: T,
+    devices: Sequence[int | torch.device],
+    detach: bool = False,
+) -> list[T]:
+    if not _replicatable_module(network):
+        raise RuntimeError(
+            "Cannot replicate network where python modules are children of ScriptModule"
+        )
+
+    if not devices:
+        return []
+
+    devices = [_get_device_index(x, True) for x in devices]
+    num_replicas = len(devices)
+
+    params = list(network.parameters())
+    param_indices = {param: idx for idx, param in enumerate(params)}
+    param_copies = _broadcast_coalesced_reshape(params, devices, detach)
+
+    buffers = list(network.buffers())
+    buffers_rg: list[torch.Tensor] = []
+    buffers_not_rg: list[torch.Tensor] = []
+    for buf in buffers:
+        if buf.requires_grad and not detach:
+            buffers_rg.append(buf)
+        else:
+            buffers_not_rg.append(buf)
+
+    buffer_indices_rg = {buf: idx for idx, buf in enumerate(buffers_rg)}
+    buffer_indices_not_rg = {buf: idx for idx, buf in enumerate(buffers_not_rg)}
+
+    buffer_copies_rg = _broadcast_coalesced_reshape(buffers_rg, devices, detach=detach)
+    buffer_copies_not_rg = _broadcast_coalesced_reshape(
+        buffers_not_rg, devices, detach=True
+    )
+
+    modules = list(network.modules())
+    module_copies: list[list[Module]] = [[] for _ in devices]
+    module_indices: dict[Module, int] = {}
+
+    for i, module in enumerate(modules):
+        module_indices[module] = i
+        for j in range(num_replicas):
+            replica = module._replicate_for_data_parallel()
+            # This is a temporary fix for DDP. DDP needs to access the
+            # replicated model parameters. It used to do so through
+            # `mode.parameters()`. The fix added in #33907 for DP stops the
+            # `parameters()` API from exposing the replicated parameters.
+            # Hence, we add a `_former_parameters` dict here to support DDP.
+            replica._former_parameters = OrderedDict()
+
+            module_copies[j].append(replica)
+
+    for i, module in enumerate(modules):
+        for key, child in module._modules.items():
+            if child is None:
+                for j in range(num_replicas):
+                    replica = module_copies[j][i]
+                    replica._modules[key] = None
+            else:
+                module_idx = module_indices[child]
+                for j in range(num_replicas):
+                    replica = module_copies[j][i]
+                    setattr(replica, key, module_copies[j][module_idx])
+        for key, param in module._parameters.items():
+            if param is None:
+                for j in range(num_replicas):
+                    replica = module_copies[j][i]
+                    replica._parameters[key] = None
+            else:
+                param_idx = param_indices[param]
+                for j in range(num_replicas):
+                    replica = module_copies[j][i]
+                    param_copy = param_copies[j][param_idx]
+                    # parameters in replicas are no longer leaves,
+                    # so setattr them as non-parameter attributes
+                    setattr(replica, key, param_copy)
+                    # expose the parameter for DDP
+                    replica._former_parameters[key] = param_copy  # type: ignore[operator, index]
+        for key, buf in module._buffers.items():  # type: ignore[assignment]
+            if buf is None:
+                for j in range(num_replicas):
+                    replica = module_copies[j][i]
+                    replica._buffers[key] = None
+            else:
+                if buf.requires_grad and not detach:
+                    buffer_copies = buffer_copies_rg
+                    buffer_idx = buffer_indices_rg[buf]
+                else:
+                    buffer_copies = buffer_copies_not_rg
+                    buffer_idx = buffer_indices_not_rg[buf]
+                for j in range(num_replicas):
+                    replica = module_copies[j][i]
+                    setattr(replica, key, buffer_copies[j][buffer_idx])
+
+    return [cast(T, module_copies[j][0]) for j in range(num_replicas)]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/scatter_gather.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/scatter_gather.py
new file mode 100644
index 0000000000000000000000000000000000000000..bcd975a4fe2f7673b618b60f326b9aae6d22f446
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parallel/scatter_gather.py
@@ -0,0 +1,150 @@
+# mypy: allow-untyped-defs
+from collections.abc import Sequence
+from typing import Any, overload, TypeVar
+from typing_extensions import deprecated
+
+import torch
+from torch.nn.parallel._functions import Gather, Scatter
+
+
+__all__ = ["scatter", "scatter_kwargs", "gather"]
+
+
+@deprecated(
+    "`is_namedtuple` is deprecated, please use the python checks instead",
+    category=FutureWarning,
+)
+def is_namedtuple(obj: Any) -> bool:
+    # Check if type was created from collections.namedtuple or a typing.NamedTuple.
+    return _is_namedtuple(obj)
+
+
+def _is_namedtuple(obj: Any) -> bool:
+    # Check if type was created from collections.namedtuple or a typing.NamedTuple.
+    return (
+        isinstance(obj, tuple) and hasattr(obj, "_asdict") and hasattr(obj, "_fields")
+    )
+
+
+T = TypeVar("T", dict, list, tuple)
+
+
+# For some reason, 'scatter' returns a tuple when given a single Tensor input but a list otherwise.
+@overload
+def scatter(
+    inputs: torch.Tensor,
+    target_gpus: Sequence[int | torch.device],
+    dim: int = ...,
+) -> tuple[torch.Tensor, ...]: ...
+
+
+@overload
+def scatter(
+    inputs: T,
+    target_gpus: Sequence[int | torch.device],
+    dim: int = ...,
+) -> list[T]: ...
+
+
+def scatter(inputs, target_gpus, dim=0):
+    r"""Slice tensors into approximately equal chunks and distributes them across given GPUs.
+
+    Duplicates references to objects that are not tensors.
+    """
+
+    def scatter_map(obj):
+        if isinstance(obj, torch.Tensor):
+            return Scatter.apply(target_gpus, None, dim, obj)
+        if _is_namedtuple(obj):
+            return [
+                type(obj)(*args)
+                # pyrefly: ignore [bad-argument-type, no-matching-overload]
+                for args in zip(*map(scatter_map, obj), strict=False)
+            ]
+        if isinstance(obj, tuple) and len(obj) > 0:
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
+            return list(zip(*map(scatter_map, obj), strict=False))
+        if isinstance(obj, list) and len(obj) > 0:
+            # pyrefly: ignore [bad-argument-type, no-matching-overload]
+            return [list(i) for i in zip(*map(scatter_map, obj), strict=False)]
+        if isinstance(obj, dict) and len(obj) > 0:
+            return [
+                type(obj)(i)
+                # pyrefly: ignore [bad-argument-type, no-matching-overload]
+                for i in zip(*map(scatter_map, obj.items()), strict=False)
+            ]
+        return [obj for _ in target_gpus]
+
+    # After scatter_map is called, a scatter_map cell will exist. This cell
+    # has a reference to the actual function scatter_map, which has references
+    # to a closure that has a reference to the scatter_map cell (because the
+    # fn is recursive). To avoid this reference cycle, we set the function to
+    # None, clearing the cell
+    try:
+        res = scatter_map(inputs)
+    finally:
+        scatter_map = None  # type: ignore[assignment]
+    return res
+
+
+def scatter_kwargs(
+    inputs: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+    target_gpus: Sequence[int | torch.device],
+    dim: int = 0,
+) -> tuple[tuple[Any, ...], tuple[dict[str, Any], ...]]:
+    r"""Scatter with support for kwargs dictionary."""
+    scattered_inputs = scatter(inputs, target_gpus, dim) if inputs else []
+    scattered_kwargs = scatter(kwargs, target_gpus, dim) if kwargs else []
+    if len(scattered_inputs) < len(scattered_kwargs):
+        scattered_inputs.extend(
+            () for _ in range(len(scattered_kwargs) - len(scattered_inputs))
+        )
+    elif len(scattered_kwargs) < len(inputs):
+        scattered_kwargs.extend(
+            {} for _ in range(len(scattered_inputs) - len(scattered_kwargs))
+        )
+    return tuple(scattered_inputs), tuple(scattered_kwargs)
+
+
+def gather(outputs: Any, target_device: int | torch.device, dim: int = 0) -> Any:
+    r"""Gather tensors from different GPUs on a specified device.
+
+    This function is useful for gathering the results of a distributed computation.
+    It takes a sequence of objects, one for each GPU, and returns a single object
+    on the specified device.
+
+    Args:
+        outputs (Any): A sequence of objects (potentially tensors) to gather.
+        target_device (Union[int, torch.device]): The device to gather the tensors to.
+            Use 'cpu' for CPU to avoid a deprecation warning.
+        dim (int, optional): The dimension along which to gather. Default: 0.
+
+    Returns:
+        Any: A gathered object (potentially tensor) on the specified device.
+    """
+
+    def gather_map(outputs):
+        out = outputs[0]
+        if isinstance(out, torch.Tensor):
+            return Gather.apply(target_device, dim, *outputs)
+        if out is None:
+            return None
+        if isinstance(out, dict):
+            if not all(len(out) == len(d) for d in outputs):
+                raise ValueError("All dicts must have the same number of keys")
+            # pyrefly: ignore [not-callable]
+            return type(out)((k, gather_map([d[k] for d in outputs])) for k in out)
+        if _is_namedtuple(out):
+            # pyrefly: ignore [bad-argument-type]
+            return type(out)._make(map(gather_map, zip(*outputs, strict=True)))
+        # pyrefly: ignore [bad-argument-type]
+        return type(out)(map(gather_map, zip(*outputs, strict=True)))
+
+    # Recursive function calls like this create reference cycles.
+    # Setting the function to None clears the refcycle.
+    try:
+        res = gather_map(outputs)
+    finally:
+        gather_map = None  # type: ignore[assignment]
+    return res
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parameter.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parameter.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5af0581d61fe47acdb08e75afc5a51ef06f42f2
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parameter.py
@@ -0,0 +1,310 @@
+from collections import OrderedDict
+from typing import Any
+
+import torch
+from torch._C import _disabled_torch_function_impl
+
+
+__all__ = [
+    "Parameter",
+    "UninitializedParameter",
+    "is_lazy",
+    "Buffer",
+    "UninitializedBuffer",
+    "UninitializedTensorMixin",
+]
+
+
+# Metaclass to combine _TensorMeta and the instance check override for Parameter.
+class _ParameterMeta(torch._C._TensorMeta):
+    # Make `isinstance(t, Parameter)` return True for custom tensor instances that have the _is_param flag.
+    def __instancecheck__(self, instance) -> bool:
+        if self is Parameter:
+            if isinstance(instance, torch.Tensor) and getattr(
+                instance, "_is_param", False
+            ):
+                return True
+        return super().__instancecheck__(instance)
+
+
+class Parameter(torch.Tensor, metaclass=_ParameterMeta):
+    r"""A kind of Tensor that is to be considered a module parameter.
+
+    Parameters are :class:`~torch.Tensor` subclasses, that have a
+    very special property when used with :class:`Module` s - when they're
+    assigned as Module attributes they are automatically added to the list of
+    its parameters, and will appear e.g. in :meth:`~Module.parameters` iterator.
+    Assigning a Tensor doesn't have such effect. This is because one might
+    want to cache some temporary state, like last hidden state of the RNN, in
+    the model. If there was no such class as :class:`Parameter`, these
+    temporaries would get registered too.
+
+    Args:
+        data (Tensor): parameter tensor.
+        requires_grad (bool, optional): if the parameter requires gradient. Note that
+            the torch.no_grad() context does NOT affect the default behavior of
+            Parameter creation--the Parameter will still have `requires_grad=True` in
+            :class:`~no_grad` mode. See :ref:`locally-disable-grad-doc` for more
+            details. Default: `True`
+    """
+
+    def __new__(cls, data=None, requires_grad=True):
+        if data is None:
+            data = torch.empty(0)
+        if type(data) is torch.Tensor or type(data) is Parameter:
+            # For ease of BC maintenance, keep this path for standard Tensor.
+            # Eventually (tm), we should change the behavior for standard Tensor to match.
+            return torch.Tensor._make_subclass(cls, data, requires_grad)
+
+        # Path for custom tensors: set a flag on the instance to indicate parameter-ness.
+        t = data.detach().requires_grad_(requires_grad)
+        if type(t) is not type(data):
+            raise RuntimeError(
+                f"Creating a Parameter from an instance of type {type(data).__name__} "
+                "requires that detach() returns an instance of the same type, but return "
+                f"type {type(t).__name__} was found instead. To use the type as a "
+                "Parameter, please correct the detach() semantics defined by "
+                "its __torch_dispatch__() implementation."
+            )
+        t._is_param = True
+        return t
+
+    # Note: the 3 methods below only apply to standard Tensor. Parameters of custom tensor types
+    # are still considered that custom tensor type and these methods will not be called for them.
+    def __deepcopy__(self, memo):
+        if id(self) in memo:
+            return memo[id(self)]
+        else:
+            result = type(self)(
+                self.data.clone(memory_format=torch.preserve_format), self.requires_grad
+            )
+            memo[id(self)] = result
+            return result
+
+    # pyrefly: ignore [bad-override]
+    def __repr__(self) -> str:
+        return "Parameter containing:\n" + super().__repr__()
+
+    def __reduce_ex__(self, proto):
+        state = torch._utils._get_obj_state(self)
+
+        # See Note [Don't serialize hooks]
+        hooks = OrderedDict()
+        if not state:
+            return (
+                torch._utils._rebuild_parameter,
+                (self.data, self.requires_grad, hooks),
+            )
+
+        return (
+            torch._utils._rebuild_parameter_with_state,
+            (self.data, self.requires_grad, hooks, state),
+        )
+
+    # pyrefly: ignore [bad-override]
+    __torch_function__ = _disabled_torch_function_impl
+
+
+class UninitializedTensorMixin:
+    _allowed_methods = [
+        torch.Tensor.__hash__,
+        torch.Tensor.size,
+        torch.Tensor.copy_,
+        torch.Tensor.is_complex,
+        torch.Tensor.is_floating_point,
+        torch.Tensor.half,
+        torch.Tensor.float,
+        torch.Tensor.double,
+        torch.Tensor.char,
+        torch.Tensor.short,
+        torch.Tensor.int,
+        torch.Tensor.long,
+        torch.Tensor.cuda,
+        torch.Tensor.cpu,
+        torch.Tensor.to,
+        torch.Tensor.get_device,
+        torch._has_compatible_shallow_copy_type,
+    ]
+
+    def materialize(self, shape, device=None, dtype=None) -> None:
+        r"""Create a Parameter or Tensor with the same properties of the uninitialized one.
+
+        Given a shape, it materializes a parameter in the same device
+        and with the same `dtype` as the current one or the specified ones in the
+        arguments.
+
+        Args:
+            shape : (tuple): the shape for the materialized tensor.
+            device (:class:`torch.device`): the desired device of the parameters
+                and buffers in this module. Optional.
+            dtype (:class:`torch.dtype`): the desired floating point type of
+                the floating point parameters and buffers in this module. Optional.
+        """
+        if device is None:
+            device = self.data.device
+        if dtype is None:
+            dtype = self.data.dtype
+        self.data = torch.empty(shape, device=device, dtype=dtype)
+        # pyrefly: ignore [missing-attribute]
+        self.__class__ = self.cls_to_become
+
+    @property
+    def shape(self):
+        raise RuntimeError(
+            "Can't access the shape of an uninitialized parameter or buffer. "
+            "This error usually happens in `load_state_dict` when trying to load "
+            "an uninitialized parameter into an initialized one. "
+            "Call `forward` to initialize the parameters before accessing their attributes."
+        )
+
+    def share_memory_(self):
+        raise RuntimeError(
+            "Can't share memory on an uninitialized parameter or buffer. "
+            "Call `forward` to initialize the parameters before calling "
+            "`module.share_memory()`."
+        )
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}>"
+
+    def __reduce_ex__(self, proto):
+        # See Note [Don't serialize hooks]
+        # pyrefly: ignore [missing-attribute]
+        return (self.__class__, (self.requires_grad,))
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        # method-wrapper is to detect access to Tensor properties that are
+        # wrapped in descriptors
+        if func in cls._allowed_methods or func.__class__.__name__ == "method-wrapper":
+            if kwargs is None:
+                kwargs = {}
+            # pyrefly: ignore [missing-attribute]
+            return super().__torch_function__(func, types, args, kwargs)
+        raise ValueError(
+            f"Attempted to use an uninitialized parameter in {func}. "
+            "This error happens when you are using a `LazyModule` or "
+            f"explicitly manipulating `torch.nn.parameter.{cls.__name__}` "
+            "objects. When using LazyModules Call `forward` with a dummy batch "
+            "to initialize the parameters before calling torch functions"
+        )
+
+
+def is_lazy(param: Any) -> bool:
+    """
+    Returns whether ``param`` is an ``UninitializedParameter`` or ``UninitializedBuffer``.
+
+    Args:
+        param (Any): the input to check.
+    """
+    return isinstance(param, UninitializedTensorMixin)
+
+
+# pyrefly: ignore [inconsistent-inheritance]
+class UninitializedParameter(UninitializedTensorMixin, Parameter):
+    r"""A parameter that is not initialized.
+
+    Uninitialized Parameters are a special case of :class:`torch.nn.Parameter`
+    where the shape of the data is still unknown.
+
+    Unlike a :class:`torch.nn.Parameter`, uninitialized parameters
+    hold no data and attempting to access some properties, like their shape,
+    will throw a runtime error. The only operations that can be performed on a uninitialized
+    parameter are changing its datatype, moving it to a different device and
+    converting it to a regular :class:`torch.nn.Parameter`.
+
+    The default device or dtype to use when the parameter is materialized can be set
+    during construction using e.g. ``device='cuda'``.
+    """
+
+    cls_to_become = Parameter
+
+    def __new__(cls, requires_grad=True, device=None, dtype=None) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        data = torch.empty(0, **factory_kwargs)
+        # pyrefly: ignore [bad-return]
+        return torch.Tensor._make_subclass(cls, data, requires_grad)
+
+    def __deepcopy__(self, memo):
+        if id(self) in memo:
+            return memo[id(self)]
+        else:
+            result = type(self)(self.requires_grad, self.data.device, self.data.dtype)
+            memo[id(self)] = result
+            return result
+
+
+# Metaclass to combine _TensorMeta and the instance check override for Buffer.
+class _BufferMeta(torch._C._TensorMeta):
+    # Make `isinstance(t, Buffer)` return True for custom tensor instances that have the _is_buffer flag.
+    def __instancecheck__(self, instance) -> bool:
+        if self is Buffer:
+            if isinstance(instance, torch.Tensor) and getattr(
+                instance, "_is_buffer", False
+            ):
+                return True
+        return super().__instancecheck__(instance)
+
+
+class Buffer(torch.Tensor, metaclass=_BufferMeta):
+    r"""A kind of Tensor that should not be considered a model
+    parameter. For example, BatchNorm's ``running_mean`` is not a parameter, but is part of the module's state.
+
+    Buffers are :class:`~torch.Tensor` subclasses, that have a
+    very special property when used with :class:`Module` s -- when they're
+    assigned as Module attributes they are automatically added to the list of
+    its buffers, and will appear e.g. in :meth:`~torch.nn.Module.buffers` iterator.
+    Assigning a Tensor doesn't have such effect. One can still assign a Tensor as explicitly by using
+    the :meth:`~torch.nn.Module.register_buffer` function.
+
+    Args:
+        data (Tensor): buffer tensor.
+        persistent (bool, optional): whether the buffer is part of the module's
+            :attr:`state_dict`. Default: ``True``
+    """
+
+    def __new__(cls, data=None, *, persistent=True):
+        if data is None:
+            data = torch.empty(0)
+
+        t = data.detach().requires_grad_(data.requires_grad)
+        # pyrefly: ignore [missing-attribute]
+        t.persistent = persistent
+        # pyrefly: ignore [missing-attribute]
+        t._is_buffer = True
+        return t
+
+    # pyrefly: ignore [bad-override]
+    __torch_function__ = _disabled_torch_function_impl
+
+
+class UninitializedBuffer(UninitializedTensorMixin, torch.Tensor):
+    r"""A buffer that is not initialized.
+
+    Uninitialized Buffer is a a special case of :class:`torch.Tensor`
+    where the shape of the data is still unknown.
+
+    Unlike a :class:`torch.Tensor`, uninitialized parameters
+    hold no data and attempting to access some properties, like their shape,
+    will throw a runtime error. The only operations that can be performed on a uninitialized
+    parameter are changing its datatype, moving it to a different device and
+    converting it to a regular :class:`torch.Tensor`.
+
+    The default device or dtype to use when the buffer is materialized can be set
+    during construction using e.g. ``device='cuda'``.
+    """
+
+    cls_to_become = torch.Tensor
+
+    def __new__(
+        cls, requires_grad=False, device=None, dtype=None, persistent=True
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        data = torch.empty(0, **factory_kwargs)
+        ret = torch.Tensor._make_subclass(cls, data, requires_grad)
+        # pyrefly: ignore [missing-attribute]
+        ret.persistent = persistent
+        # pyrefly: ignore [missing-attribute]
+        ret._is_buffer = True
+        # pyrefly: ignore [bad-return]
+        return ret
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parameter.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parameter.pyi
new file mode 100644
index 0000000000000000000000000000000000000000..3d1cddb7e8b8bf7dd4aced5f645a8670cf266b3c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/parameter.pyi
@@ -0,0 +1,43 @@
+from typing_extensions import TypeIs
+
+from torch import device, dtype, Tensor
+
+class Parameter(Tensor):
+    def __init__(self, data: Tensor = ..., requires_grad: bool = ...) -> None: ...
+
+def is_lazy(
+    param: Tensor,
+) -> TypeIs[UninitializedParameter | UninitializedBuffer]: ...
+
+class UninitializedParameter(Tensor):
+    def __init__(self, data: Tensor = ..., requires_grad: bool = ...) -> None: ...
+    def materialize(
+        self,
+        shape: tuple[int, ...],
+        device: device | None = None,
+        dtype: dtype | None = None,
+    ) -> None: ...
+
+class Buffer(Tensor):
+    persistent: bool
+    def __init__(
+        self,
+        data: Tensor = ...,
+        requires_grad: bool = ...,
+        persistent: bool = ...,
+    ) -> None: ...
+
+class UninitializedBuffer(Tensor):
+    persistent: bool
+    def __init__(
+        self,
+        data: Tensor = ...,
+        requires_grad: bool = ...,
+        persistent: bool = ...,
+    ) -> None: ...
+    def materialize(
+        self,
+        shape: tuple[int, ...],
+        device: device | None = None,
+        dtype: dtype | None = None,
+    ) -> None: ...
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..766b09382aa78e65aba915e4e6faf7979c500d1b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/__init__.py
@@ -0,0 +1,19 @@
+# flake8: noqa: F401
+r"""QAT Dynamic Modules.
+
+This package is in the process of being deprecated.
+Please, use `torch.ao.nn.qat.dynamic` instead.
+"""
+
+from torch.nn.qat import dynamic, modules  # noqa: F403
+from torch.nn.qat.modules import *  # noqa: F403
+
+
+__all__ = [
+    "Linear",
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "Embedding",
+    "EmbeddingBag",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..56838a1cfcae74c171479a717ff4e80ae7c53d71
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/__init__.py
@@ -0,0 +1,8 @@
+# flake8: noqa: F401
+r"""QAT Dynamic Modules.
+
+This package is in the process of being deprecated.
+Please, use `torch.ao.nn.qat.dynamic` instead.
+"""
+
+from torch.nn.qat.dynamic.modules import *  # noqa: F403
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3790f21256074a9a158d3ad15fa2ee0044fd0e6
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/modules/__init__.py
@@ -0,0 +1,4 @@
+from torch.nn.qat.dynamic.modules.linear import Linear
+
+
+__all__ = ["Linear"]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/modules/linear.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/modules/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a5c80ea213c6287d2134e5344ddb19e538d0d5c
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/dynamic/modules/linear.py
@@ -0,0 +1,11 @@
+# flake8: noqa: F401
+r"""QAT Modules.
+
+This file is in the process of migration to `torch/ao/nn/qat/dynamic`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/qat/dynamic/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.qat.dynamic.modules.linear import Linear
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7f55fbdf789a7c845896f5b3bb0570ef74b5890
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/__init__.py
@@ -0,0 +1,21 @@
+# flake8: noqa: F401
+r"""QAT Modules.
+
+This package is in the process of being deprecated.
+Please, use `torch.ao.nn.qat.modules` instead.
+"""
+
+from torch.ao.nn.qat.modules.conv import Conv1d, Conv2d, Conv3d
+from torch.ao.nn.qat.modules.embedding_ops import Embedding, EmbeddingBag
+from torch.ao.nn.qat.modules.linear import Linear
+from torch.nn.qat.modules import conv, embedding_ops, linear
+
+
+__all__ = [
+    "Linear",
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "Embedding",
+    "EmbeddingBag",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/conv.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/conv.py
new file mode 100644
index 0000000000000000000000000000000000000000..31aa101d8f60a5a5a1c61dae98d55db6f187a1c5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/conv.py
@@ -0,0 +1,11 @@
+# flake8: noqa: F401
+r"""QAT Modules.
+
+This file is in the process of migration to `torch/ao/nn/qat`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/qat/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.qat.modules.conv import Conv1d, Conv2d, Conv3d
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/embedding_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/embedding_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a0964739f9e6ac96db8d3c8a7d156d371cc0725
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/embedding_ops.py
@@ -0,0 +1,14 @@
+# flake8: noqa: F401
+r"""QAT Modules.
+
+This file is in the process of migration to `torch/ao/nn/qat`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/qat/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.qat.modules.embedding_ops import Embedding, EmbeddingBag
+
+
+__all__ = ["Embedding", "EmbeddingBag"]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/linear.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e822eba7e0617f7d950f6b269398e729524a28f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/qat/modules/linear.py
@@ -0,0 +1,11 @@
+# flake8: noqa: F401
+r"""QAT Modules.
+
+This file is in the process of migration to `torch/ao/nn/qat`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/qat/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.qat.modules.linear import Linear
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7628c5c15992efa600ea5520aed955ba42c6146
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/__init__.py
@@ -0,0 +1 @@
+from torch.nn.quantizable.modules import *  # noqa: F403
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f0e0ca1743566174b7abfd3663dfa90b744ba56f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/__init__.py
@@ -0,0 +1,9 @@
+from torch.ao.nn.quantizable.modules.activation import MultiheadAttention
+from torch.ao.nn.quantizable.modules.rnn import LSTM, LSTMCell
+
+
+__all__ = [
+    "LSTM",
+    "LSTMCell",
+    "MultiheadAttention",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/activation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/activation.py
new file mode 100644
index 0000000000000000000000000000000000000000..28f3eee958115d05d161af8acf1b2308c02c3248
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/activation.py
@@ -0,0 +1,11 @@
+# flake8: noqa: F401
+r"""Quantizable Modules.
+
+This file is in the process of migration to `torch/ao/nn/quantizable`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantizable/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantizable.modules.activation import MultiheadAttention
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/rnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/rnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e355efcdd63d595580bdac22d46882b91b7d118
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantizable/modules/rnn.py
@@ -0,0 +1,11 @@
+# flake8: noqa: F401
+r"""Quantizable Modules.
+
+This file is in the process of migration to `torch/ao/nn/quantizable`, and
+is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantizable/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantizable.modules.rnn import LSTM, LSTMCell
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e2bbbc13202db1cbddaad4b05241a62190adc46
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/__init__.py
@@ -0,0 +1,39 @@
+from torch.nn.quantized import dynamic, functional, modules  # noqa: F403
+from torch.nn.quantized.modules import *  # noqa: F403
+from torch.nn.quantized.modules import MaxPool2d
+
+
+__all__ = [
+    "BatchNorm2d",
+    "BatchNorm3d",
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "ConvTranspose1d",
+    "ConvTranspose2d",
+    "ConvTranspose3d",
+    "DeQuantize",
+    "Dropout",
+    "ELU",
+    "Embedding",
+    "EmbeddingBag",
+    "GroupNorm",
+    "Hardswish",
+    "InstanceNorm1d",
+    "InstanceNorm2d",
+    "InstanceNorm3d",
+    "LayerNorm",
+    "LeakyReLU",
+    "Linear",
+    "LSTM",
+    "MultiheadAttention",
+    "PReLU",
+    "Quantize",
+    "ReLU6",
+    "Sigmoid",
+    "Softmax",
+    # Wrapper modules
+    "FloatFunctional",
+    "FXFloatFunctional",
+    "QFunctional",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..61faa90bd95cc7e255be2df82c617b5bab46b044
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/__init__.py
@@ -0,0 +1 @@
+from torch.nn.quantized._reference.modules import *  # noqa: F403
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9caa8e58f19393bf3c19421dbaa528b1f86b996
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/__init__.py
@@ -0,0 +1,39 @@
+# flake8: noqa: F401
+r"""Quantized Reference Modules.
+
+This module is in the process of migration to
+`torch/ao/nn/quantized/reference`, and is kept here for
+compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/reference`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.reference.modules.conv import (
+    Conv1d,
+    Conv2d,
+    Conv3d,
+    ConvTranspose1d,
+    ConvTranspose2d,
+    ConvTranspose3d,
+)
+from torch.ao.nn.quantized.reference.modules.linear import Linear
+from torch.ao.nn.quantized.reference.modules.rnn import GRUCell, LSTM, LSTMCell, RNNCell
+from torch.ao.nn.quantized.reference.modules.sparse import Embedding, EmbeddingBag
+
+
+__all__ = [
+    "Linear",
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "ConvTranspose1d",
+    "ConvTranspose2d",
+    "ConvTranspose3d",
+    "RNNCell",
+    "LSTMCell",
+    "GRUCell",
+    "LSTM",
+    "Embedding",
+    "EmbeddingBag",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/conv.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/conv.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac806d0fa60d657f32603e7d761687f3c4a64215
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/conv.py
@@ -0,0 +1,21 @@
+# flake8: noqa: F401
+r"""Quantized Reference Modules.
+
+This module is in the process of migration to
+`torch/ao/nn/quantized/reference`, and is kept here for
+compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/reference`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.reference.modules.conv import (
+    _ConvNd,
+    _ConvTransposeNd,
+    Conv1d,
+    Conv2d,
+    Conv3d,
+    ConvTranspose1d,
+    ConvTranspose2d,
+    ConvTranspose3d,
+)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/linear.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..6be6d5a140bb58f76b0e6061eb4ccb37d385757f
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/linear.py
@@ -0,0 +1,12 @@
+# flake8: noqa: F401
+r"""Quantized Reference Modules.
+
+This module is in the process of migration to
+`torch/ao/nn/quantized/reference`, and is kept here for
+compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/reference`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.reference.modules.linear import Linear
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/rnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/rnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..0573b3309b64b92ef3bb59200117b4ef5a62680b
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/rnn.py
@@ -0,0 +1,19 @@
+# flake8: noqa: F401
+r"""Quantized Reference Modules.
+
+This module is in the process of migration to
+`torch/ao/nn/quantized/reference`, and is kept here for
+compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/reference`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.reference.modules.rnn import (
+    GRUCell,
+    LSTM,
+    LSTMCell,
+    RNNBase,
+    RNNCell,
+    RNNCellBase,
+)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/sparse.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/sparse.py
new file mode 100644
index 0000000000000000000000000000000000000000..85bf997d478ae1d0b4631541433e01f8f2943633
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/sparse.py
@@ -0,0 +1,12 @@
+# flake8: noqa: F401
+r"""Quantized Reference Modules.
+
+This module is in the process of migration to
+`torch/ao/nn/quantized/reference`, and is kept here for
+compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/reference`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.reference.modules.sparse import Embedding, EmbeddingBag
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6a0f30f2fb92c2717cb9f794d69e1e905bff6ed
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/_reference/modules/utils.py
@@ -0,0 +1,18 @@
+# flake8: noqa: F401
+r"""Quantized Reference Modules.
+
+This module is in the process of migration to
+`torch/ao/nn/quantized/reference`, and is kept here for
+compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/reference`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.reference.modules.utils import (
+    _get_weight_qparam_keys,
+    _quantize_and_dequantize_weight,
+    _quantize_weight,
+    _save_weight_qparams,
+    ReferenceQuantizedModule,
+)
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b08cd1bc7149c5506db3a952fff488eb06749f5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/__init__.py
@@ -0,0 +1 @@
+from torch.ao.nn.quantized.dynamic import *  # noqa: F403
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ae09e82c3bb85f754d554af6eb7ac36f29e56ce
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/__init__.py
@@ -0,0 +1,43 @@
+# flake8: noqa: F401
+r"""Quantized Dynamic Modules.
+
+This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
+and is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/dynamic`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.dynamic.modules import conv, linear, rnn
+from torch.ao.nn.quantized.dynamic.modules.conv import (
+    Conv1d,
+    Conv2d,
+    Conv3d,
+    ConvTranspose1d,
+    ConvTranspose2d,
+    ConvTranspose3d,
+)
+from torch.ao.nn.quantized.dynamic.modules.linear import Linear
+from torch.ao.nn.quantized.dynamic.modules.rnn import (
+    GRU,
+    GRUCell,
+    LSTM,
+    LSTMCell,
+    RNNCell,
+)
+
+
+__all__ = [
+    "Linear",
+    "LSTM",
+    "GRU",
+    "LSTMCell",
+    "RNNCell",
+    "GRUCell",
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "ConvTranspose1d",
+    "ConvTranspose2d",
+    "ConvTranspose3d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/conv.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/conv.py
new file mode 100644
index 0000000000000000000000000000000000000000..b5b81a68a88917cddd0fe4b6987bc8ff6eddef01
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/conv.py
@@ -0,0 +1,28 @@
+# flake8: noqa: F401
+r"""Quantized Dynamic Modules.
+
+This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
+and is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/dynamic/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.dynamic.modules.conv import (
+    Conv1d,
+    Conv2d,
+    Conv3d,
+    ConvTranspose1d,
+    ConvTranspose2d,
+    ConvTranspose3d,
+)
+
+
+__all__ = [
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "ConvTranspose1d",
+    "ConvTranspose2d",
+    "ConvTranspose3d",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/linear.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..b23fae2c06aa8c829b0aeab7d29a4c84ad5c35b7
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/linear.py
@@ -0,0 +1,11 @@
+# flake8: noqa: F401
+r"""Quantized Dynamic Modules.
+
+This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
+and is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/dynamic/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.dynamic.modules.linear import Linear
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/rnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/rnn.py
new file mode 100644
index 0000000000000000000000000000000000000000..d5ca396a2d44030f61c8fdf64c8fe4aed3e0cf5d
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/dynamic/modules/rnn.py
@@ -0,0 +1,34 @@
+# flake8: noqa: F401
+r"""Quantized Dynamic Modules.
+
+This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
+and is kept here for compatibility while the migration process is ongoing.
+If you are adding a new entry/functionality, please, add it to the
+appropriate file under the `torch/ao/nn/quantized/dynamic/modules`,
+while adding an import statement here.
+"""
+
+from torch.ao.nn.quantized.dynamic.modules.rnn import (
+    GRU,
+    GRUCell,
+    LSTM,
+    LSTMCell,
+    pack_weight_bias,
+    PackedParameter,
+    RNNBase,
+    RNNCell,
+    RNNCellBase,
+)
+
+
+__all__ = [
+    "pack_weight_bias",
+    "PackedParameter",
+    "RNNBase",
+    "LSTM",
+    "GRU",
+    "RNNCellBase",
+    "RNNCell",
+    "LSTMCell",
+    "GRUCell",
+]
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/functional.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/functional.py
new file mode 100644
index 0000000000000000000000000000000000000000..d763e171fdb432c8ba2059cc2332e7ac6424854a
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/functional.py
@@ -0,0 +1,10 @@
+r"""nn.quantized.functional.
+
+Quantized equivalents of the `nn.functional`.
+
+Note::
+    This location is in the process of being deprecated.
+    Please, use the `torch.ao.nn.quantized.functional` instead.
+"""
+
+from torch.ao.nn.quantized.functional import *  # noqa: F401,F403
diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae76d1968b0faaf30f861ab009b9011ce2960cc5
--- /dev/null
+++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/__init__.py
@@ -0,0 +1,97 @@
+r"""Quantized Modules.
+
+Note::
+    The `torch.nn.quantized` namespace is in the process of being deprecated.
+    Please, use `torch.ao.nn.quantized` instead.
+"""
+
+# The following imports are needed in case the user decides
+# to import the files directly,
+# s.a. `from torch.nn.quantized.modules.conv import ...`.
+# No need to add them to the `__all__`.
+from torch.ao.nn.quantized.modules import (
+    activation,
+    batchnorm,
+    conv,
+    DeQuantize,
+    dropout,
+    embedding_ops,
+    functional_modules,
+    linear,
+    MaxPool2d,
+    normalization,
+    Quantize,
+    rnn,
+    utils,
+)
+from torch.ao.nn.quantized.modules.activation import (
+    ELU,
+    Hardswish,
+    LeakyReLU,
+    MultiheadAttention,
+    PReLU,
+    ReLU6,
+    Sigmoid,
+    Softmax,
+)
+from torch.ao.nn.quantized.modules.batchnorm import BatchNorm2d, BatchNorm3d
+from torch.ao.nn.quantized.modules.conv import (
+    Conv1d,
+    Conv2d,
+    Conv3d,
+    ConvTranspose1d,
+    ConvTranspose2d,
+    ConvTranspose3d,
+)
+from torch.ao.nn.quantized.modules.dropout import Dropout
+from torch.ao.nn.quantized.modules.embedding_ops import Embedding, EmbeddingBag
+from torch.ao.nn.quantized.modules.functional_modules import (
+    FloatFunctional,
+    FXFloatFunctional,
+    QFunctional,
+)
+from torch.ao.nn.quantized.modules.linear import Linear
+from torch.ao.nn.quantized.modules.normalization import (
+    GroupNorm,
+    InstanceNorm1d,
+    InstanceNorm2d,
+    InstanceNorm3d,
+    LayerNorm,
+)
+from torch.ao.nn.quantized.modules.rnn import LSTM
+
+
+__all__ = [
+    "BatchNorm2d",
+    "BatchNorm3d",
+    "Conv1d",
+    "Conv2d",
+    "Conv3d",
+    "ConvTranspose1d",
+    "ConvTranspose2d",
+    "ConvTranspose3d",
+    "DeQuantize",
+    "ELU",
+    "Embedding",
+    "EmbeddingBag",
+    "GroupNorm",
+    "Hardswish",
+    "InstanceNorm1d",
+    "InstanceNorm2d",
+    "InstanceNorm3d",
+    "LayerNorm",
+    "LeakyReLU",
+    "Linear",
+    "LSTM",
+    "MultiheadAttention",
+    "Quantize",
+    "ReLU6",
+    "Sigmoid",
+    "Softmax",
+    "Dropout",
+    "PReLU",
+    # Wrapper modules
+    "FloatFunctional",
+    "FXFloatFunctional",
+    "QFunctional",
+]