pyxlog/

lib.rs

//! Python bindings for XLOG via PyO3.

use std::collections::{HashMap, HashSet};
use std::os::raw::{c_char, c_void};
use std::sync::Arc;

use pyo3::exceptions::{PyBufferError, PyMemoryError, PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBool, PyDict, PyInt, PyList, PyTuple};

use xlog_core::{MemoryBudget, Schema};
#[cfg(feature = "arrow-device-import")]
use xlog_cuda::{ArrowDeviceArray, ArrowDeviceArrayOwned};
use xlog_cuda::{CudaBuffer, CudaKernelProvider, CudaProviderBuilder, DlpackManagedTensor};
use xlog_gpu::logic as gpu_logic;
use xlog_logic::ast::ProbEngine;
use xlog_neural::{NetworkRegistry, TensorSourceRegistry};
use xlog_prob::exact::GpuConfig;

use xlog_core::RelId;
use xlog_ir::ExecutionPlan;
use xlog_logic::ast::Program as AstProgram;
use xlog_runtime::{Executor, RelationStore};

mod neural_registry;
use neural_registry::NeuralPredicateRegistry;
mod diagnostic_serialization;
mod dlpack;
mod epistemic;
mod ilp;
mod ilp_exact;
mod ilp_gpu;
mod joint_carrier;
mod logic;
mod neural;
mod program;
mod relation_metadata;
mod semantic_transition;
mod training;
mod types;
pub(crate) use diagnostic_serialization::{pack_query_proof_traces, pack_rule_provenance};
pub(crate) use program::{
    CachedCircuit, CompiledProbProgram, HardFilter, InputSource, JoinPlan, NeuralGroup,
    QuerySignature,
};
use relation_metadata::RelationMetadataStore;

const DLPACK_CAPSULE_NAME: &[u8] = b"dltensor\0";
const USED_DLPACK_CAPSULE_NAME: &[u8] = b"used_dltensor\0";
// Ordinary callers use the legacy default; prepared callers supply their
// actual stream. The Python Array API reserves integer 1 for the legacy stream.
const DLPACK_CUDA_LEGACY_DEFAULT_STREAM: i64 = 1;

#[cfg(feature = "arrow-device-import")]
const ARROW_DEVICE_ARRAY_CAPSULE_NAME: &[u8] = b"arrow_device_array\0";
#[cfg(feature = "arrow-device-import")]
const USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME: &[u8] = b"used_arrow_device_array\0";

unsafe extern "C" fn dlpack_capsule_destructor(capsule: *mut pyo3::ffi::PyObject) {
    if capsule.is_null() {
        return;
    }

    let valid =
        pyo3::ffi::PyCapsule_IsValid(capsule, DLPACK_CAPSULE_NAME.as_ptr() as *const c_char);
    if valid == 0 {
        return;
    }

    let ptr =
        pyo3::ffi::PyCapsule_GetPointer(capsule, DLPACK_CAPSULE_NAME.as_ptr() as *const c_char);
    if ptr.is_null() {
        pyo3::ffi::PyErr_Clear();
        return;
    }

    let managed = ptr as *mut xlog_cuda::DLManagedTensor;
    drop(DlpackManagedTensor::from_raw(managed));
}

pub(crate) fn dlpack_capsule_from_tensor(
    py: Python<'_>,
    tensor: DlpackManagedTensor,
) -> PyResult<Py<PyAny>> {
    let raw = tensor.into_raw();
    let ptr = raw as *mut c_void;
    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let capsule = unsafe {
        pyo3::ffi::PyCapsule_New(
            ptr,
            DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
            Some(dlpack_capsule_destructor),
        )
    };
    if capsule.is_null() {
        // SAFETY: the pointer is a valid owned Python object pointer returned by the C API
        unsafe {
            drop(DlpackManagedTensor::from_raw(raw));
        }
        return Err(PyRuntimeError::new_err("Failed to create DLPack capsule"));
    }
    // SAFETY: capsule is a non-null owned pointer returned by PyCapsule_New; PyO3 takes ownership
    let obj: Py<PyAny> = unsafe { Bound::from_owned_ptr(py, capsule) }.unbind();
    Ok(obj)
}

#[cfg(feature = "arrow-device-import")]
unsafe extern "C" fn arrow_device_array_capsule_destructor(capsule: *mut pyo3::ffi::PyObject) {
    if capsule.is_null() {
        return;
    }

    let valid = pyo3::ffi::PyCapsule_IsValid(
        capsule,
        ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
    );
    if valid == 0 {
        return;
    }

    let ptr = pyo3::ffi::PyCapsule_GetPointer(
        capsule,
        ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
    );
    if ptr.is_null() {
        pyo3::ffi::PyErr_Clear();
        return;
    }

    drop(ArrowDeviceArrayOwned::from_raw(
        ptr as *mut ArrowDeviceArray,
    ));
}

#[cfg(feature = "arrow-device-import")]
pub(crate) fn arrow_device_capsule_from_device_array(
    py: Python<'_>,
    device_array: ArrowDeviceArrayOwned,
) -> PyResult<Py<PyAny>> {
    let raw = device_array.into_raw();
    let ptr = raw as *mut c_void;
    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let capsule = unsafe {
        pyo3::ffi::PyCapsule_New(
            ptr,
            ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
            Some(arrow_device_array_capsule_destructor),
        )
    };
    if capsule.is_null() {
        // SAFETY: the pointer is a valid owned Python object pointer returned by the C API
        unsafe {
            drop(ArrowDeviceArrayOwned::from_raw(raw));
        }
        return Err(PyRuntimeError::new_err(
            "Failed to create Arrow device array capsule",
        ));
    }
    // SAFETY: capsule is a non-null owned pointer returned by PyCapsule_New; PyO3 takes ownership
    let obj: Py<PyAny> = unsafe { Bound::from_owned_ptr(py, capsule) }.unbind();
    Ok(obj)
}

#[cfg(feature = "arrow-device-import")]
pub(crate) fn arrow_device_from_py(obj: &Bound<'_, PyAny>) -> PyResult<ArrowDeviceArrayOwned> {
    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    if unsafe {
        pyo3::ffi::PyCapsule_IsValid(
            obj.as_ptr(),
            ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    } == 0
    {
        return Err(PyValueError::new_err(
            "Expected an Arrow device array capsule (arrow_device_array)",
        ));
    }

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let ptr = unsafe {
        pyo3::ffi::PyCapsule_GetPointer(
            obj.as_ptr(),
            ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if ptr.is_null() {
        return Err(PyRuntimeError::new_err(
            "Failed to get Arrow device array pointer",
        ));
    }

    // Mark consumed so the capsule destructor doesn't free the pointer we now own.
    // SAFETY: capsule is valid (checked above); renaming marks it consumed so the destructor skips cleanup
    let rc = unsafe {
        pyo3::ffi::PyCapsule_SetName(
            obj.as_ptr(),
            USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if rc != 0 {
        return Err(PyRuntimeError::new_err(
            "Failed to mark Arrow device array capsule as consumed",
        ));
    }

    // SAFETY: ptr is non-null (checked above) and points to an ArrowDeviceArray matching the Arrow C Data Interface layout
    Ok(unsafe { ArrowDeviceArrayOwned::from_raw(ptr as *mut ArrowDeviceArray) })
}

pub(crate) fn provider_from_config(config: GpuConfig) -> xlog_core::Result<CudaKernelProvider> {
    CudaProviderBuilder::new(
        config.device_ordinal,
        MemoryBudget::with_limit(config.memory_bytes),
    )
    .build()
}

pub(crate) fn enforce_call_memory_limit(
    provider: &Arc<CudaKernelProvider>,
    memory_mb: Option<u64>,
) -> PyResult<()> {
    let Some(memory_mb) = memory_mb else {
        return Ok(());
    };
    if memory_mb == 0 {
        return Err(PyValueError::new_err("memory_mb must be > 0"));
    }
    let memory_limit_bytes = memory_mb.saturating_mul(1024 * 1024);
    let allocated_bytes = provider.memory().allocated_bytes();
    if allocated_bytes > memory_limit_bytes {
        return Err(PyMemoryError::new_err(format!(
            "per-call memory limit exceeded before evaluation: allocated_bytes={} memory_limit_bytes={}",
            allocated_bytes, memory_limit_bytes
        )));
    }
    Ok(())
}

pub(crate) fn provider_memory_stats(
    py: Python<'_>,
    provider: &Arc<CudaKernelProvider>,
) -> PyResult<Py<PyAny>> {
    let dict = PyDict::new(py);
    let memory = provider.memory();
    dict.set_item("allocated_bytes", memory.allocated_bytes())?;
    dict.set_item("memory_limit_bytes", memory.budget().device_bytes)?;
    dict.set_item("peak_memory_bytes", memory.peak_bytes())?;
    dict.set_item("status", "available")?;
    Ok(dict.into())
}

pub(crate) fn parse_prob_engine_override(s: &str) -> PyResult<ProbEngine> {
    let v = s.trim().to_ascii_lowercase();
    match v.as_str() {
        "exact_ddnnf" | "exact" | "ddnnf" => Ok(ProbEngine::ExactDdnnf),
        "mc" => Ok(ProbEngine::Mc),
        other => Err(PyValueError::new_err(format!(
            "Unknown prob_engine '{}'; expected 'exact_ddnnf' or 'mc'",
            other
        ))),
    }
}

/// No native owner lock crosses a Python callback. Recheck even when a producer
/// raises, so a caught exception cannot hide a revoked original authority.
pub(crate) fn guarded_python_callback<T>(
    check: impl Fn() -> PyResult<()>,
    callback: impl FnOnce() -> PyResult<T>,
) -> PyResult<T> {
    check()?;
    let result = callback();
    let after = check();
    match result {
        Err(error) => Err(error),
        Ok(value) => after.map(|()| value),
    }
}

/// Export storage after the caller selected the consumer's DLPack stream.
#[cfg(test)]
pub(crate) fn dlpack_export_for_stream(
    obj: &Bound<'_, PyAny>,
    stream: Option<i64>,
) -> PyResult<Py<PyAny>> {
    dlpack_export_for_stream_guarded(obj, stream, &|| Ok(()))
}

fn dlpack_export_for_stream_guarded(
    obj: &Bound<'_, PyAny>,
    stream: Option<i64>,
    check: &dyn Fn() -> PyResult<()>,
) -> PyResult<Py<PyAny>> {
    check()?;
    let kwargs = PyDict::new(obj.py());
    if let Some(stream) = stream {
        kwargs.set_item("stream", stream)?;
    }
    // PyTorch intentionally rejects exporting autograd semantics via DLPack.
    // Export only a shared-storage alias; the caller retains the original
    // primal and applies returned adjoints to its original graph. Using the
    // alias's standard protocol preserves producer/consumer stream ordering,
    // which a raw torch.utils.dlpack.to_dlpack capsule cannot negotiate.
    // Do not import an optional framework for other DLPack producers.
    let sys = guarded_python_callback(check, || obj.py().import("sys"))?;
    let modules = guarded_python_callback(check, || sys.getattr("modules"))?;
    let mut producer = obj.clone();
    if let Some(torch) =
        guarded_python_callback(check, || modules.cast::<PyDict>()?.get_item("torch"))?
    {
        let tensor_type = guarded_python_callback(check, || torch.getattr("Tensor"))?;
        let is_tensor = guarded_python_callback(check, || obj.is_instance(&tensor_type))?;
        let requires_grad = if is_tensor {
            let value = guarded_python_callback(check, || obj.getattr("requires_grad"))?;
            guarded_python_callback(check, || value.extract::<bool>())?
        } else {
            false
        };
        if requires_grad {
            let detach = guarded_python_callback(check, || tensor_type.getattr("detach"))?;
            let alias = guarded_python_callback(check, || detach.call1((obj,)))?;
            for method in ["data_ptr", "storage_offset", "stride"] {
                let alias_method = guarded_python_callback(check, || alias.getattr(method))?;
                let alias_value = guarded_python_callback(check, || alias_method.call0())?;
                let original_method = guarded_python_callback(check, || obj.getattr(method))?;
                let original_value = guarded_python_callback(check, || original_method.call0())?;
                let equal = guarded_python_callback(check, || {
                    alias_value.rich_compare(&original_value, pyo3::basic::CompareOp::Eq)
                })?;
                if !guarded_python_callback(check, || equal.is_truthy())? {
                    return Err(PyBufferError::new_err(
                        "DLPack transport alias changed the original tensor storage",
                    ));
                }
            }
            for attribute in ["shape", "dtype", "device"] {
                let alias_value = guarded_python_callback(check, || alias.getattr(attribute))?;
                let original_value = guarded_python_callback(check, || obj.getattr(attribute))?;
                let equal = guarded_python_callback(check, || {
                    alias_value.rich_compare(&original_value, pyo3::basic::CompareOp::Eq)
                })?;
                if !guarded_python_callback(check, || equal.is_truthy())? {
                    return Err(PyBufferError::new_err(
                        "DLPack transport alias changed the original tensor layout",
                    ));
                }
            }
            producer = alias;
        }
    }
    let export = guarded_python_callback(check, || producer.getattr("__dlpack__"))?;
    Ok(guarded_python_callback(check, || export.call((), Some(&kwargs)))?.unbind())
}

pub(crate) fn dlpack_from_py(obj: &Bound<'_, PyAny>) -> PyResult<DlpackManagedTensor> {
    dlpack_from_py_for_stream(obj, DLPACK_CUDA_LEGACY_DEFAULT_STREAM)
}

/// Consume the standard producer protocol on the actual selected consumer stream.
pub(crate) fn dlpack_from_py_for_stream(
    obj: &Bound<'_, PyAny>,
    consumer_stream: i64,
) -> PyResult<DlpackManagedTensor> {
    dlpack_from_py_for_stream_guarded(obj, consumer_stream, &|| Ok(()))
}

/// Read the producer's DLPack device declaration without Python conversions.
/// Device kinds may be integer enums; the tuple and ordinal remain exact builtins.
pub(crate) fn dlpack_device_pair(device: &Bound<'_, PyAny>) -> PyResult<(i32, i32)> {
    let invalid_pair = || {
        PyValueError::new_err(
            "DLPack device must be an exact tuple (integer kind, exact integer id)",
        )
    };
    if !device.is_exact_instance_of::<PyTuple>() {
        return Err(invalid_pair());
    }
    let fields = device.cast::<PyTuple>()?;
    if fields.len() != 2 {
        return Err(invalid_pair());
    }
    let kind = fields.get_item(0)?;
    let ordinal = fields.get_item(1)?;
    if kind.is_instance_of::<PyBool>()
        || !kind.is_instance_of::<PyInt>()
        || !ordinal.is_exact_instance_of::<PyInt>()
    {
        return Err(invalid_pair());
    }
    // PyO3's signed integer extraction reads the checked PyLong payload, not
    // overridden __int__/__index__. Never extract an arbitrary numeric object.
    let device_type = kind.extract::<i32>()?;
    let device_id = ordinal.extract::<i32>()?;
    if device_type != xlog_cuda::dlpack::K_DLCUDA {
        return Err(PyBufferError::new_err(format!(
            "Unsupported DLPack producer device type {device_type} (device {device_id}); \
             XLOG requires CUDA device memory (kDLCUDA=2)"
        )));
    }
    if device_id < 0 {
        return Err(PyValueError::new_err(
            "DLPack device id must be nonnegative",
        ));
    }
    Ok((device_type, device_id))
}

pub(crate) fn dlpack_from_py_for_stream_guarded(
    obj: &Bound<'_, PyAny>,
    consumer_stream: i64,
    check: &dyn Fn() -> PyResult<()>,
) -> PyResult<DlpackManagedTensor> {
    if consumer_stream <= 0 || consumer_stream == 2 {
        return Err(PyValueError::new_err(
            "DLPack requires an explicit supported consumer stream",
        ));
    }
    check()?;
    let py = obj.py();

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let capsule_obj: Bound<'_, PyAny> = if unsafe {
        pyo3::ffi::PyCapsule_IsValid(obj.as_ptr(), DLPACK_CAPSULE_NAME.as_ptr() as *const c_char)
    } != 0
    {
        obj.clone()
    } else if guarded_python_callback(check, || obj.hasattr("__dlpack__"))? {
        let method = guarded_python_callback(check, || obj.getattr("__dlpack_device__"))?;
        let device = guarded_python_callback(check, || method.call0())?;
        dlpack_device_pair(&device)?;
        // Passing the consumer stream makes the producer order any pending
        // non-default-stream writes before XLOG reads the tensor. A raw
        // capsule cannot negotiate synchronization and must already be ready
        // for the selected consumer stream when supplied by the caller.
        dlpack_export_for_stream_guarded(obj, Some(consumer_stream), check)?.into_bound(py)
    } else {
        return Err(PyValueError::new_err(
            "Expected a DLPack capsule or an object with __dlpack__",
        ));
    };
    check()?;

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    if unsafe {
        pyo3::ffi::PyCapsule_IsValid(
            capsule_obj.as_ptr(),
            DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    } == 0
    {
        return Err(PyValueError::new_err("Invalid DLPack capsule"));
    }

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let ptr = unsafe {
        pyo3::ffi::PyCapsule_GetPointer(
            capsule_obj.as_ptr(),
            DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if ptr.is_null() {
        return Err(PyRuntimeError::new_err("Failed to get DLPack pointer"));
    }

    // SAFETY: capsule is valid (checked above); renaming marks it consumed so the destructor skips cleanup
    let rc = unsafe {
        pyo3::ffi::PyCapsule_SetName(
            capsule_obj.as_ptr(),
            USED_DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if rc != 0 {
        return Err(PyRuntimeError::new_err(
            "Failed to mark DLPack capsule as consumed",
        ));
    }

    // SAFETY: ptr is non-null (checked above) and points to a DLManagedTensor matching the DLPack specification layout
    Ok(unsafe { DlpackManagedTensor::from_raw(ptr as *mut xlog_cuda::DLManagedTensor) })
}

#[cfg(test)]
mod dlpack_guard_tests {
    use super::*;

    #[test]
    fn dlpack_device_callback_refusal_stops_all_later_producer_callbacks() {
        Python::initialize();
        Python::attach(|py| {
            let globals = PyDict::new(py);
            py.run(
                cr#"
from enum import IntEnum
class DeviceKind(IntEnum):
    CUDA = 2
refused = False
calls = []
class Producer:
    def __getattribute__(self, name):
        if refused:
            calls.append(('late_lookup', name))
        return object.__getattribute__(self, name)
    def __dlpack_device__(self):
        global refused
        calls.append('device')
        refused = True
        return (DeviceKind.CUDA, 0)
    def __dlpack__(self, *, stream):
        calls.append('export')
        raise AssertionError('export ran after authority refusal')
producer = Producer()
"#,
                Some(&globals),
                None,
            )
            .unwrap();
            let producer = globals.get_item("producer").unwrap().unwrap();
            let check = || {
                if globals.get_item("refused")?.unwrap().extract::<bool>()? {
                    Err(PyValueError::new_err(
                        "original producer authority was refused",
                    ))
                } else {
                    Ok(())
                }
            };
            // The ordinary caller has no authority guard: the same production
            // importer reaches export. This controls the guarded-path check.
            assert!(super::dlpack_from_py_for_stream(&producer, 19).is_err());
            py.run(
                c"assert 'export' in calls\ncalls.clear()\nrefused = False",
                Some(&globals),
                None,
            )
            .unwrap();
            let error = match super::dlpack_from_py_for_stream_guarded(&producer, 19, &check) {
                Ok(_) => panic!("revoked producer handed off a native tensor"),
                Err(error) => error,
            };
            assert!(error
                .to_string()
                .contains("original producer authority was refused"));
            assert_eq!(
                globals
                    .get_item("calls")
                    .unwrap()
                    .unwrap()
                    .extract::<Vec<String>>()
                    .unwrap(),
                ["device"]
            );
        });
    }

    #[test]
    fn dlpack_integer_device_kinds_reach_original_export_and_capsule_validation() {
        Python::initialize();
        Python::attach(|py| {
            let globals = PyDict::new(py);
            py.run(
                cr#"
from enum import IntEnum
class DeviceKind(IntEnum):
    CUDA = 2
    def __int__(self):
        raise AssertionError('device kind conversion invoked')
    def __index__(self):
        raise AssertionError('device kind index invoked')
class IntegerKind(int):
    def __int__(self):
        raise AssertionError('integer subclass conversion invoked')
    def __index__(self):
        raise AssertionError('integer subclass index invoked')
calls = []
class Producer:
    def __dlpack_device__(self):
        calls.append('device')
        return device
    def __dlpack__(self, *, stream):
        assert stream == 19
        calls.append('export')
        return None
producer = Producer()
devices = [(2, 0), (DeviceKind.CUDA, 0), (IntegerKind(2), 0)]
"#,
                Some(&globals),
                None,
            )
            .unwrap();
            let producer = globals.get_item("producer").unwrap().unwrap();
            let devices = globals.get_item("devices").unwrap().unwrap();
            for device in devices.cast::<PyList>().unwrap().iter() {
                globals.set_item("device", device).unwrap();
                py.run(c"calls.clear()", Some(&globals), None).unwrap();
                let error = super::dlpack_from_py_for_stream_guarded(&producer, 19, &|| Ok(()))
                    .err()
                    .expect("invalid capsule must not be consumed");
                assert!(error.to_string().contains("Invalid DLPack capsule"));
                assert_eq!(
                    globals
                        .get_item("calls")
                        .unwrap()
                        .unwrap()
                        .extract::<Vec<String>>()
                        .unwrap(),
                    ["device", "export"]
                );
            }
        });
    }

    #[test]
    fn dlpack_device_pair_rejects_invalid_fields_before_export() {
        Python::initialize();
        Python::attach(|py| {
            let globals = PyDict::new(py);
            py.run(
                cr#"
class Pair(tuple):
    def __iter__(self):
        raise AssertionError('tuple subclass iterated')
    def __getitem__(self, index):
        raise AssertionError('tuple subclass indexed')
class IntegerKind(int):
    def __int__(self):
        raise AssertionError('kind converted')
    def __index__(self):
        raise AssertionError('kind index invoked')
calls = []
class Producer:
    def __dlpack_device__(self):
        calls.append('device')
        return device
    def __dlpack__(self, *, stream):
        calls.append('export')
        return None
producer = Producer()
devices = [None, [2, 0], Pair((2, 0)), (), (2,), (2, 0, 0),
           (True, 0), (2, False), (2.0, 0), (2, 0.0), ('2', 0),
           (1, 0), (-1, 0), (2, -1), (2, 1 << 31), (1 << 31, 0),
           (IntegerKind(1), 0), (IntegerKind(-1), 0),
           (IntegerKind(1 << 80), 0), (2, IntegerKind(0))]
"#,
                Some(&globals),
                None,
            )
            .unwrap();
            let producer = globals.get_item("producer").unwrap().unwrap();
            let devices = globals.get_item("devices").unwrap().unwrap();
            for device in devices.cast::<PyList>().unwrap().iter() {
                globals.set_item("device", device).unwrap();
                py.run(c"calls.clear()", Some(&globals), None).unwrap();
                assert!(
                    super::dlpack_from_py_for_stream_guarded(&producer, 19, &|| Ok(())).is_err()
                );
                assert_eq!(
                    globals
                        .get_item("calls")
                        .unwrap()
                        .unwrap()
                        .extract::<Vec<String>>()
                        .unwrap(),
                    ["device"]
                );
            }
        });
    }

    #[test]
    fn dlpack_device_metadata_never_invokes_custom_integer_conversion() {
        Python::initialize();
        Python::attach(|py| {
            let globals = PyDict::new(py);
            py.run(
                cr#"
calls = []
class DeviceIndex:
    def __int__(self):
        calls.append('int')
        return 2
    def __index__(self):
        calls.append('index')
        return 2
class Producer:
    def __dlpack_device__(self):
        calls.append('device')
        return device
    def __dlpack__(self, *, stream):
        calls.append('export')
        raise AssertionError('export ran after invalid device metadata')
producer = Producer()
devices = [(2, DeviceIndex()), (DeviceIndex(), 0)]
"#,
                Some(&globals),
                None,
            )
            .unwrap();
            let producer = globals.get_item("producer").unwrap().unwrap();
            let devices = globals.get_item("devices").unwrap().unwrap();
            for device in devices.cast::<PyList>().unwrap().iter() {
                globals.set_item("device", device).unwrap();
                py.run(c"calls.clear()", Some(&globals), None).unwrap();
                let error = super::dlpack_from_py_for_stream_guarded(&producer, 19, &|| Ok(()))
                    .err()
                    .expect("custom device conversion must not hand off a native tensor");
                assert!(error
                    .to_string()
                    .contains("exact tuple (integer kind, exact integer id)"));
                assert_eq!(
                    globals
                        .get_item("calls")
                        .unwrap()
                        .unwrap()
                        .extract::<Vec<String>>()
                        .unwrap(),
                    ["device"]
                );
            }
        });
    }
}

#[pyfunction]
fn dlpack_is_cuda(obj: &Bound<'_, PyAny>) -> PyResult<bool> {
    // SAFETY: capsule validity is checked before reading the DLPack header. This
    // does not consume the capsule; ownership remains with its destructor.
    if unsafe {
        pyo3::ffi::PyCapsule_IsValid(obj.as_ptr(), DLPACK_CAPSULE_NAME.as_ptr() as *const c_char)
    } == 0
    {
        return Err(PyValueError::new_err(
            "Expected a DLPack capsule (dltensor)",
        ));
    }

    // SAFETY: capsule validity was checked immediately before this call.
    let ptr = unsafe {
        pyo3::ffi::PyCapsule_GetPointer(obj.as_ptr(), DLPACK_CAPSULE_NAME.as_ptr() as *const c_char)
    };
    if ptr.is_null() {
        return Err(PyRuntimeError::new_err("Failed to get DLPack pointer"));
    }

    // SAFETY: ptr is non-null and points to a DLManagedTensor owned by the capsule.
    let managed = unsafe { &*(ptr as *const xlog_cuda::DLManagedTensor) };
    Ok(managed.dl_tensor.device.device_type == xlog_cuda::dlpack::K_DLCUDA)
}

#[pyfunction]
fn intern_symbols(symbols: Vec<String>) -> Vec<u32> {
    symbols
        .iter()
        .map(|symbol| xlog_core::symbol::intern(symbol))
        .collect()
}

#[pyfunction]
fn resolve_symbols(symbol_ids: Vec<u32>) -> PyResult<Vec<String>> {
    symbol_ids
        .into_iter()
        .enumerate()
        .map(|(index, symbol_id)| {
            xlog_core::symbol::resolve_checked(symbol_id).ok_or_else(|| {
                PyValueError::new_err(format!("unknown symbol ID {symbol_id} at index {index}"))
            })
        })
        .collect()
}

#[pyclass(name = "DifferentiableProofTraceMap")]
pub struct PyDifferentiableProofTraceMap {
    inner: xlog_logic::DifferentiableProofTraceMap,
}

fn pack_differentiable_proof_trace(
    py: Python<'_>,
    trace: &xlog_logic::ProofTrace,
) -> PyResult<Py<PyAny>> {
    let dict = PyDict::new(py);
    dict.set_item("proof_id", trace.proof_id)?;
    dict.set_item("answer_key", &trace.answer_key)?;
    dict.set_item("clause_id", &trace.clause_id)?;
    dict.set_item("support_atoms", &trace.support_atoms)?;
    dict.set_item("weight", trace.weight)?;
    dict.set_item("gradient", trace.gradient)?;
    Ok(dict.into())
}

#[pymethods]
impl PyDifferentiableProofTraceMap {
    #[new]
    fn new() -> Self {
        Self {
            inner: xlog_logic::DifferentiableProofTraceMap::new(),
        }
    }

    fn insert(
        &mut self,
        answer_key: String,
        clause_id: String,
        support_atoms: Vec<String>,
        initial_weight: f64,
    ) -> PyResult<u64> {
        if !initial_weight.is_finite() {
            return Err(PyValueError::new_err(
                "initial_weight must be a finite float",
            ));
        }
        Ok(self.inner.insert(xlog_logic::ProofTraceSpec {
            answer_key,
            clause_id,
            support_atoms,
            initial_weight,
        }))
    }

    fn trace(&self, py: Python<'_>, proof_id: u64) -> PyResult<Option<Py<PyAny>>> {
        self.inner
            .trace(proof_id)
            .map(|trace| pack_differentiable_proof_trace(py, trace))
            .transpose()
    }

    fn traces(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let list = PyList::empty(py);
        for trace in self.inner.traces() {
            list.append(pack_differentiable_proof_trace(py, trace)?)?;
        }
        Ok(list.into())
    }

    fn accumulate_binary_logistic_gradients(
        &mut self,
        targets: Vec<(String, f64)>,
    ) -> PyResult<f64> {
        if targets.iter().any(|(_, target)| !target.is_finite()) {
            return Err(PyValueError::new_err("targets must be finite floats"));
        }
        Ok(self.inner.accumulate_binary_logistic_gradients(&targets))
    }

    fn apply_gradients(&mut self, learning_rate: f64) -> PyResult<()> {
        if !learning_rate.is_finite() || learning_rate < 0.0 {
            return Err(PyValueError::new_err(
                "learning_rate must be a finite non-negative float",
            ));
        }
        self.inner.apply_gradients(learning_rate);
        Ok(())
    }
}

#[pyclass]
pub struct Program;

#[pyclass]
pub struct CompiledProgram {
    pub(crate) program: CompiledProbProgram,
    pub(crate) output_provider: Arc<CudaKernelProvider>,
    /// Registry for neural networks
    pub(crate) network_registry: NetworkRegistry,
    /// Registry for neural predicate metadata (predicate -> network/labels)
    pub(crate) neural_registry: NeuralPredicateRegistry,
    /// Names of neural networks declared in the program (from nn() declarations)
    pub(crate) declared_networks: HashSet<String>,
    /// Map from network name to form: true = embedding, false = classification
    pub(crate) declared_network_forms: HashMap<String, bool>,
    /// Registry for tensor data sources (images, embeddings, etc.)
    pub(crate) tensor_sources: TensorSourceRegistry,
    /// Name of the Stage-B existential-join domain tensor source, as supplied by
    /// the Python driver (single source of truth). `None` until a domain source is
    /// registered; read by the join forward to resolve `DomainRow`/`ConstDummy`
    /// groups instead of an engine-side hardcoded name.
    pub(crate) domain_source: Option<String>,
    /// Which ROW of the join-domain tensor holds which domain CONSTANT, as stated by
    /// the Python driver: `domain_ids[j]` is the constant whose feature vector is row
    /// `j`. This is the SINGLE source of truth for "constant -> feature row" — the
    /// torch-side mixture and this circuit look the row up in the same list, so they
    /// cannot drift apart.
    ///
    /// The ids arrive WITH the tensor (`register_domain_tensor_source` requires them), so
    /// there is no "no ids registered" state to fall back from. The fallback that used to
    /// live here — row = the constant's position in the relation's materialized domain —
    /// is the defect this map exists to close: it reads a different row from the torch
    /// path for any domain that is not exactly `0..D-1`, silently, and no test that
    /// stays inside that coincidence can see it. Keeping it as a reachable branch would
    /// leave the same silent-wrong mode one call away.
    ///
    /// Empty until a domain source is registered — and that state IS reachable from a
    /// caller mistake, not only from an internal bug: a join signature is compiled for
    /// every rule defining the train head (`joint_candidate_eligibility`), so a program
    /// whose driver forgot `domain_inputs` reaches this map while it is still empty. The
    /// lookup then fails with "constant N is not in domain_ids", which is true but points
    /// at the wrong parameter, so the Python driver checks `domain_inputs` for every join
    /// candidate BEFORE it asks for eligibility. A lookup miss here is therefore a real
    /// error to surface (a joined constant with no id), never something to paper over.
    pub(crate) domain_ids: Vec<i64>,
    /// Original program source (for dynamic query compilation)
    pub(crate) _source: String,
    /// Parsed program AST (for signature analysis)
    pub(crate) ast: xlog_logic::ast::Program,
    /// GPU configuration
    pub(crate) _gpu_config: GpuConfig,
    /// Probabilistic inference engine
    pub(crate) _prob_engine: ProbEngine,
    /// Cache of analyzed query signatures.
    pub(crate) query_signature_cache: HashMap<String, QuerySignature>,
    /// Cache of compiled circuits by template signature
    pub(crate) circuit_cache: HashMap<String, CachedCircuit>,
    /// Number of circuit-template cache hits observed by neural training paths.
    pub(crate) circuit_cache_hits: usize,
    /// Number of circuit-template cache misses observed by neural training paths.
    pub(crate) circuit_cache_misses: usize,
    /// Number of times the template compilation path executed.
    pub(crate) template_compile_count: usize,
    /// When true, batch queries sharing the same circuit template in training.
    pub(crate) batch_queries: bool,
    /// Latest circuit compilation profile (populated on cache miss when profiling).
    pub(crate) last_compile_profile: Option<xlog_prob::compilation::CircuitCompileProfile>,
}

#[pyclass]
pub struct LogicProgram;

#[pyclass]
pub struct CompiledLogicProgram {
    pub(crate) program: Arc<gpu_logic::LogicProgram>,
    pub(crate) provider: Arc<CudaKernelProvider>,
}

/// A fixed accepted-evidence exact circuit with mutable independent fact priors.
///
/// Clones share one serialized native state. Updating or evaluating one clone
/// excludes concurrent work on the same circuit and is visible to every clone.
#[pyclass(skip_from_py_object)]
#[derive(Clone)]
pub struct CompiledConditionedProgram {
    #[cfg(feature = "host-io")]
    pub(crate) program: xlog_prob::epistemic_production::PreparedConditionedProgram,
    #[cfg(feature = "host-io")]
    pub(crate) result_provider: Arc<CudaKernelProvider>,
}

#[pyclass]
pub struct LogicRelationSession {
    pub(crate) program: Arc<gpu_logic::LogicProgram>,
    pub(crate) provider: Arc<CudaKernelProvider>,
    pub(crate) relation_store: RelationStore,
    pub(crate) evaluation_store: Option<gpu_logic::LogicMaterializedStore>,
    pub(crate) session_runtime: Option<gpu_logic::LogicSessionRuntime>,
    pub(crate) last_delta_stats: Option<LogicDeltaStats>,
    pub(crate) relation_callbacks: Vec<RelationChangeCallback>,
    pub(crate) next_relation_callback_id: u64,
    pub(crate) relation_generations: HashMap<String, u64>,
    pub(crate) relation_metadata: RelationMetadataStore,
}

pub(crate) struct RelationChangeCallback {
    pub id: u64,
    pub callback: Py<PyAny>,
}

#[derive(Clone, Debug)]
pub(crate) struct LogicDeltaStats {
    pub input_delta_count: usize,
    pub changed_relations: usize,
    pub changed_relation_names: Vec<String>,
    pub insert_rows: u64,
    pub delete_rows: u64,
    pub has_deletes: bool,
    pub affected_sccs: usize,
    pub recomputed_sccs: usize,
    pub incremental_sccs: usize,
    pub coalesced_insert_rows: u64,
    pub coalesced_delete_rows: u64,
    pub canceled_rows: u64,
    pub equivalent_to_full_recompute: Option<bool>,
    pub planner_telemetry: gpu_logic::DeltaPlannerTelemetry,
    pub debug_trace: Vec<String>,
}

#[pyclass]
pub struct LogicQueryResult {
    #[pyo3(get)]
    pub relation_name: String,
    #[pyo3(get)]
    pub columns: Vec<String>,
    #[pyo3(get)]
    pub sort_labels: Vec<String>,
    #[pyo3(get)]
    pub tensors: Vec<Py<PyAny>>,
    #[pyo3(get)]
    pub num_rows: usize,
    #[pyo3(get)]
    pub is_true: bool,
}

#[pyclass]
pub struct LogicEvalResult {
    #[pyo3(get)]
    pub queries: Vec<Py<LogicQueryResult>>,
}

#[pyclass]
pub struct IlpTaggedCreditDeviceResult {
    #[pyo3(get)]
    pub fact_row_offsets: Py<PyAny>,
    #[pyo3(get)]
    pub entry_indices: Py<PyAny>,
    #[pyo3(get)]
    pub entry_i: Py<PyAny>,
    #[pyo3(get)]
    pub entry_j: Py<PyAny>,
    #[pyo3(get)]
    pub entry_k: Py<PyAny>,
}

#[pyclass]
pub struct McDeviceEvalResult {
    /// Per-query satisfying-sample counts. DLPack int32 tensor on CUDA.
    #[pyo3(get)]
    pub query_counts: Py<PyAny>,
    /// Evidence satisfying-sample count. DLPack int32 tensor with shape [1] on CUDA.
    #[pyo3(get)]
    pub evidence_count: Py<PyAny>,
    #[pyo3(get)]
    pub total_samples: usize,
    #[pyo3(get)]
    pub seed: u64,
    #[pyo3(get)]
    pub confidence: f64,
    #[pyo3(get)]
    pub nonmonotone_semantics: String,
    #[pyo3(get)]
    pub nonmonotone_sccs: usize,
    #[pyo3(get)]
    pub nonmonotone_cycles: usize,
    #[pyo3(get)]
    pub nonmonotone_iteration_limit_hits: usize,
    #[pyo3(get)]
    pub sampling_method: String,
    #[pyo3(get)]
    pub resident_no_host_certified: bool,
    #[pyo3(get)]
    pub resident_no_host_policy_result: String,
    #[pyo3(get)]
    pub resident_no_host_tracked_dtoh_calls: u64,
    #[pyo3(get)]
    pub resident_no_host_tracked_htod_calls: u64,
    #[pyo3(get)]
    pub resident_no_host_host_loop_iterations: u64,
    #[pyo3(get)]
    pub resident_no_host_per_sample_host_launches: u64,
    #[pyo3(get)]
    pub resident_no_host_untracked_metadata_reads: u64,
    #[pyo3(get)]
    pub resident_no_host_engine_launches: u64,
    #[pyo3(get)]
    pub resident_no_host_host_fixpoint_iterations: u64,
    #[pyo3(get)]
    pub resident_no_host_per_operator_host_allocations: u64,
}

#[pyclass]
pub struct EvalResult {
    #[pyo3(get)]
    pub atoms: Vec<String>,
    #[pyo3(get)]
    pub prob: Py<PyAny>,
    #[pyo3(get)]
    pub log_prob: Py<PyAny>,
    #[pyo3(get)]
    pub num_vars: usize,
    /// Exact log-evidence `log Z_E` (natural log). `None` for Monte Carlo results.
    #[pyo3(get)]
    pub log_z_e: Option<f64>,
    #[pyo3(get)]
    pub grad_true: Option<Vec<Py<PyAny>>>,
    #[pyo3(get)]
    pub grad_false: Option<Vec<Py<PyAny>>>,
    #[pyo3(get)]
    pub approx: bool,
    #[pyo3(get)]
    pub stderr: Option<Py<PyAny>>,
    #[pyo3(get)]
    pub ci_low: Option<Py<PyAny>>,
    #[pyo3(get)]
    pub ci_high: Option<Py<PyAny>>,
    #[pyo3(get)]
    pub samples: Option<usize>,
    #[pyo3(get)]
    pub evidence_samples: Option<usize>,
    #[pyo3(get)]
    pub seed: Option<u64>,
    #[pyo3(get)]
    pub confidence: Option<f64>,
    #[pyo3(get)]
    pub nonmonotone_semantics: Option<String>,
    #[pyo3(get)]
    pub nonmonotone_sccs: Option<usize>,
    #[pyo3(get)]
    pub nonmonotone_cycles: Option<usize>,
    #[pyo3(get)]
    pub nonmonotone_iteration_limit_hits: Option<usize>,
    #[pyo3(get)]
    pub sampling_method: Option<String>,
    /// MC only: which engine produced the result — `"gpu-resident"` for the
    /// resident megakernel engine, `"cpu-oracle"` for the explicitly opted-in
    /// CPU oracle. `None` for exact inference.
    #[pyo3(get)]
    pub mc_engine: Option<String>,
}

/// Exact probabilities conditioned on an accepted epistemic world view.
///
/// `prob`/`log_prob` are DLPack capsules over device memory, like `EvalResult`.
/// `trace` carries the production-path counters of the epistemic->probability
/// adapter: they are the evidence that conditioning actually happened on the GPU.
///
/// `log_z_e` is log P(evidence): the exact log-probability of the conditioned
/// evidence under the probabilistic program's distribution, computed by weighted
/// model counting over the compiled circuit. Query probabilities are
/// `exp(log_z_eq - log_z_e)`. When the conditioned atoms are independent root
/// facts it coincides with the log of the product of their priors, but that is a
/// special case, not the definition: evidence on a derived atom, on atoms sharing
/// an ancestor, or negated evidence all diverge from the product form.
#[pyclass]
pub struct EpistemicEvalResult {
    #[pyo3(get)]
    pub atoms: Vec<String>,
    #[pyo3(get)]
    pub prob: Py<PyAny>,
    #[pyo3(get)]
    pub log_prob: Py<PyAny>,
    #[pyo3(get)]
    pub log_z_e: f64,
    #[pyo3(get)]
    pub trace: Py<PyAny>,
}

/// Summary of one accepted epistemic GPU execution.
///
/// `accepted_world_views == 0` means the program ran but nothing was accepted.
/// `evaluate_conditioned` on that same program RAISES `RuntimeError` rather than
/// returning an unconditioned result — this method is the non-raising way to detect
/// the state. `know_operator_count`/`possible_operator_count` are plan-level
/// censuses and stay non-zero even when nothing is accepted.
#[pyclass]
pub struct EpistemicEvidence {
    #[pyo3(get)]
    pub epistemic_mode: String,
    #[pyo3(get)]
    pub know_operator_count: usize,
    #[pyo3(get)]
    pub possible_operator_count: usize,
    #[pyo3(get)]
    pub accepted_candidates: usize,
    #[pyo3(get)]
    pub rejected_candidates: usize,
    #[pyo3(get)]
    pub accepted_world_views: usize,
    #[pyo3(get)]
    pub final_output_rows: usize,
}

// =========================================================================
// Training Infrastructure
// =========================================================================

/// Statistics for a single training epoch.
// Result-only class: nothing extracts it back from Python, so the
// Clone-derived automatic `FromPyObject` is explicitly skipped.
#[pyclass(skip_from_py_object)]
#[derive(Clone)]
pub struct EpochStats {
    /// Average loss across all batches in the epoch
    #[pyo3(get)]
    pub avg_loss: f64,
    /// Number of batches processed
    #[pyo3(get)]
    pub num_batches: usize,
    /// Total number of queries processed
    #[pyo3(get)]
    pub total_queries: usize,
}

/// Training history tracking loss over epochs and batches.
// Result-only class: nothing extracts it back from Python, so the
// Clone-derived automatic `FromPyObject` is explicitly skipped.
#[pyclass(skip_from_py_object)]
#[derive(Clone)]
pub struct TrainingHistory {
    /// Loss at the end of each epoch
    #[pyo3(get)]
    pub epoch_losses: Vec<f64>,
    /// Wall-clock time (seconds) for each epoch
    #[pyo3(get)]
    pub epoch_times: Vec<f64>,
    /// Loss for each batch across all epochs
    #[pyo3(get)]
    pub batch_losses: Vec<f64>,
    /// True if training was stopped early due to validation loss plateau.
    #[pyo3(get)]
    pub stopped_early: bool,
}

#[pyclass]
pub struct IlpProgramFactory;

#[pyclass]
pub struct CompiledIlpProgram {
    pub(crate) base_source: String,
    pub(crate) _learnable_source: String,
    pub(crate) ast: AstProgram,
    pub(crate) executor: Executor,
    pub(crate) provider: Arc<CudaKernelProvider>,
    pub(crate) plan: ExecutionPlan,
    pub(crate) rel_index: Vec<(RelId, String)>,
    pub(crate) schemas: HashMap<String, Schema>,
    pub(crate) left_keys: Vec<usize>,
    pub(crate) right_keys: Vec<usize>,
    pub(crate) head_projection: Vec<usize>,
    pub(crate) compiled_schema_size: usize,
    pub(crate) head_rel_name: String,
    pub(crate) max_active_rules: usize,
    pub(crate) candidate_map: Option<HashMap<(u32, u32, u32), u32>>,
    pub(crate) candidate_order: Option<Vec<(u32, u32, u32)>>,
    pub(crate) relation_overrides: HashMap<String, CudaBuffer>,
    /// Maximum bytes for per-chunk temp allocations (masks, prefix sums,
    /// chunk-local COO scratch). The final merged COO buffer is exact-NNZ
    /// sized and may exceed this budget. Default: 16 MB.
    pub(crate) coo_chunk_budget: u64,
    /// When true, raise instead of falling back to chunked COO path.
    /// Use in zero-D2H benchmarks and CI gates. Default: false.
    pub(crate) strict_zero_dtoh: bool,
    /// Per-phase wall-clock of the compile that built this program (ms).
    pub(crate) compile_timing: Vec<(&'static str, f64)>,
}

#[pymodule]
#[pyo3(name = "_native")]
fn pyxlog(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add_class::<Program>()?;
    m.add_class::<CompiledProgram>()?;
    m.add_class::<LogicProgram>()?;
    m.add_class::<CompiledLogicProgram>()?;
    m.add_class::<CompiledConditionedProgram>()?;
    m.add_class::<LogicRelationSession>()?;
    m.add_class::<semantic_transition::PySemanticTransitionSession>()?;
    m.add_class::<semantic_transition::PySemanticTransitionController>()?;
    m.add_class::<semantic_transition::cold_task::PySemanticTransitionColdTask>()?;
    m.add_class::<semantic_transition::cold_task::PySemanticTransitionFreshParent>()?;
    m.add_class::<semantic_transition::PySemanticTransitionTaskUse>()?;
    m.add_class::<semantic_transition::PySemanticInitialPrefillStage>()?;
    m.add_class::<semantic_transition::PySemanticInitialPrefillContentWitness>()?;
    m.add_class::<semantic_transition::PySemanticInitialPrefillReceipt>()?;
    m.add_class::<semantic_transition::PySemanticPublishedParent>()?;
    m.add_class::<semantic_transition::PySemanticTransitionRestoredCheckpoint>()?;
    m.add_class::<semantic_transition::learning_phase::PySemanticLearningPhaseRecipe>()?;
    m.add_class::<semantic_transition::learning_phase::PySemanticLearningPhaseTransition>()?;
    m.add_class::<semantic_transition::PyNativeTensorAllocation>()?;
    m.add_class::<semantic_transition::PySemanticPreparedStep>()?;
    #[cfg(feature = "semantic-policy")]
    {
        m.add(
            "SemanticCompletedSegmentPending",
            m.py()
                .get_type::<semantic_transition::SemanticCompletedSegmentPending>(),
        )?;
        m.add(
            "SemanticModelEvaluationPending",
            m.py()
                .get_type::<semantic_transition::model_evaluation::SemanticModelEvaluationPending>(
                ),
        )?;
        m.add_class::<semantic_transition::model_evaluation::PySemanticEvaluationCohort>()?;
        m.add_class::<semantic_transition::model_evaluation::PySemanticModelEvaluation>()?;
        m.add_class::<semantic_transition::PySemanticCompletedExecutionObservation>()?;
        m.add_class::<semantic_transition::model_evaluation::PySemanticCompletedModelEvaluation>()?;
    }
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedModelCarrier>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedMaterial>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedActionComponent>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedTextAction>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedDecodedEdit>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedActionLane>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedTheoryDelta>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedTaskFacts>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedLaneOutcome>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedTaskGround>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedEditSolution>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticCompletedActionProjection>()?;
    m.add_class::<semantic_transition::PySemanticGradientDelivery>()?;
    m.add_class::<semantic_transition::PySemanticTensorContentWitness>()?;
    m.add_class::<semantic_transition::PySemanticModelForwardWitness>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticPolicyInvocation>()?;
    #[cfg(feature = "semantic-policy")]
    m.add_class::<semantic_transition::PySemanticRetainedReplayMember>()?;
    m.add_class::<relation_metadata::RelationEvidence>()?;
    m.add_class::<LogicQueryResult>()?;
    m.add_class::<LogicEvalResult>()?;
    m.add_class::<McDeviceEvalResult>()?;
    m.add_class::<EvalResult>()?;
    m.add_class::<EpistemicEvalResult>()?;
    m.add_class::<EpistemicEvidence>()?;
    // Training infrastructure
    m.add_class::<PyDifferentiableProofTraceMap>()?;
    m.add_class::<EpochStats>()?;
    m.add_class::<TrainingHistory>()?;
    // ILP bindings
    m.add_class::<IlpProgramFactory>()?;
    m.add_class::<CompiledIlpProgram>()?;
    m.add_class::<IlpTaggedCreditDeviceResult>()?;
    m.add_function(wrap_pyfunction!(training::train_model, m)?)?;
    m.add_function(wrap_pyfunction!(training::train_model_tensor, m)?)?;
    m.add_function(wrap_pyfunction!(dlpack::dlpack_roundtrip, m)?)?;
    m.add_function(wrap_pyfunction!(dlpack_is_cuda, m)?)?;
    m.add_function(wrap_pyfunction!(intern_symbols, m)?)?;
    m.add_function(wrap_pyfunction!(resolve_symbols, m)?)?;
    #[cfg(feature = "arrow-device-import")]
    m.add_function(wrap_pyfunction!(dlpack::export_arrow_device, m)?)?;
    #[cfg(feature = "arrow-device-import")]
    m.add_function(wrap_pyfunction!(dlpack::import_arrow_device, m)?)?;
    // Joint constraint carrier bindings
    m.add_class::<joint_carrier::JointConstraintCarrier>()?;
    m.add(
        "CarrierRefused",
        _py.get_type::<joint_carrier::CarrierRefused>(),
    )?;
    m.add(
        "SolverResourceExhausted",
        _py.get_type::<joint_carrier::SolverResourceExhausted>(),
    )?;
    m.add(
        "RelationMetadataError",
        _py.get_type::<relation_metadata::RelationMetadataError>(),
    )?;
    m.add("SOLVER_ABI_IDENTITY", xlog_cuda::SOLVER_ABI_IDENTITY)?;
    Ok(())
}