Skip to content

Interop: copy and device semantics, explicit dtypes and a Julia adapter - #127

Merged
max-models merged 1 commit into
develfrom
interop-semantics
Oct 31, 2026
Merged

max-models merged 1 commit into
develfrom
interop-semantics

Conversation

@max-models

Copy link
Copy Markdown
Owner

Copy and device semantics

  • copy=None | True | False on every interop conversion (as in the array API): False raises when the conversion cannot share memory, True always copies.
  • Memory is read from __dlpack_device__ (Apple's unified memory counts as host). from_framework copies between host and device explicitly instead of falling back to asarray silently; a framework that cannot take an array's device raises a ValueError naming the device.
  • Transfer profiler: a conversion that shares memory records nothing (so assert_no_transfers checks zero copy), host↔device records to_host / to_device, a copy within one memory device_copy.
  • from_jax raises a clear TypeError for traced values.

Explicit dtypes

  • No silent float64 → float32: JAX without jax_enable_x64, PyTorch on MPS and MLX for the GPU raise a TypeError with advice. MLX keeps float64 for device=mlx.core.cpu.
  • supports_float64(device) in the JAX, PyTorch and MLX adapters, interop.jax.enable_x64().
  • Fixes to_mlx(device=...), which passed device to mlx.core.array (it takes none).

Julia adapter (host only)

  • cunumpy.interop.julia (extra cunumpy[julia], juliacall): to_julia / from_julia between NumPy arrays and Julia Arrays sharing memory, kept alive across both GCs. C order arrives with reversed dimensions (a component-major (3, n) array is a Julia n×3 matrix), F order with the same; from_julia(order="C" | "F"). Strided or read-only arrays are copied; non-Array Julia arrays are collected.
  • CuPy ↔ CUDA.jl is left for later (needs a GPU).

Tests

Shared rules tested with fake device arrays; the framework and Julia tests run with the real JAX, PyTorch, MLX and juliacall and are skipped without them (CI installs none of them yet, except MLX on macOS). Locally: 810 passed, interop coverage 98%.

@max-models
max-models merged commit 66788f6 into devel Oct 31, 2026
23 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant