Chapter 5: Interoperability and Custom Operations

# Shared setup: imports, version check, C++ directory, timing and aligned-allocation helpers
import sys, time, ctypes, subprocess, warnings, shutil
from pathlib import Path
from functools import partial
import numpy as np
import jax
import jax.numpy as jnp

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| devices", jax.devices())
assert jax.__version__.startswith("0.12")
assert jax.default_backend() == "cpu", "The examples in this chapter are written for CPU (on Colab, use a CPU runtime)"

NB_DIR = Path.cwd()                 # directory of this notebook
CPP_DIR = NB_DIR / "cpp"            # where C++ sources and .so files go
CPP_DIR.mkdir(exist_ok=True)
def short_path(s):
    """Shorten local absolute paths when printing (notebook dir → relative path, virtualenv → <venv>) so no personal directory shows up in the output"""
    return str(s).replace(str(NB_DIR) + "/", "").replace(sys.prefix, "<venv>")

print("C++ directory:", short_path(CPP_DIR))

def bench(f, *args, n=20):
    """Average time (ms): warm up once, loop n times, then block_until_ready"""
    jax.block_until_ready(f(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    jax.block_until_ready(out)
    return (time.perf_counter() - t0) / n * 1e3

def aligned_zeros(n, dtype=np.float32, align=64):
    """Allocate a 1-D NumPy array whose start address is aligned to `align` bytes (required for zero-copy DLPack import into JAX)"""
    itemsize = np.dtype(dtype).itemsize
    buf = np.zeros(n * itemsize + align, dtype=np.uint8)
    off = (-buf.ctypes.data) % align
    out = buf[off:off + n * itemsize].view(dtype)
    assert out.ctypes.data % align == 0
    return out

def shares_memory(j, n):
    """Whether JAX array j and NumPy array n point to the same memory (compares start pointers)"""
    return j.unsafe_buffer_pointer() == n.ctypes.data

# ✅ Recommended: NumPy → JAX (course code block 1; use the .device attribute for the device)
numpy_arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
print("Original NumPy type:", type(numpy_arr), "(implicitly on the host CPU)")

jax_arr = jnp.array(numpy_arr)
print("Converted type:", type(jax_arr), "| isinstance(jax.Array):", isinstance(jax_arr, jax.Array))
print("Device:", jax_arr.device)

assert isinstance(jax_arr, jax.Array) and jax_arr.dtype == jnp.float32
assert jax_arr.device == jax.devices()[0]

# ❌ Old style (outdated/wrong): jax_arr.device() — in 0.12, .device is an attribute (a Device object) and can't be called
try:
    print(jax_arr.device())
except TypeError as e:
    print("TypeError:", e)

# ✅ Recommended: the .device attribute (single-device arrays); .devices() returns the set of devices (works for sharded arrays too)
print(jax_arr.device, jax_arr.devices())
assert jax_arr.devices() == {jax_arr.device}

# Course code block 2: explicit placement with jax.device_put
available_devices = jax.devices()
print("Available devices:", available_devices)
target_device = available_devices[1] if len(available_devices) > 1 else available_devices[0]
jax_arr_explicit = jax.device_put(numpy_arr, device=target_device)
print("Device after explicit placement:", jax_arr_explicit.device)
assert jax_arr_explicit.device == target_device

# ❌ Old style (outdated/wrong): jax.random.PRNGKey(0) returns a uint32[2] "legacy raw key"; 0.12 recommends typed keys
try:
    old_key = jax.random.PRNGKey(0)
    print("Legacy key:", old_key, "| dtype:", old_key.dtype, "| shape:", old_key.shape)
except Exception as e:
    print(type(e).__name__, e)

# ✅ Recommended: jax.random.key creates a typed key; JAX → NumPy (course code block 3)
key = jax.random.key(0)
print("Typed key:", key, "| dtype:", key.dtype)
assert jax.dtypes.issubdtype(key.dtype, jax.dtypes.prng_key)

jax_arr = jax.random.normal(key, (3,)) * 2.0 + 1.0
print("JAX array:", jax_arr, "| device:", jax_arr.device)
numpy_result = np.array(jax_arr)     # sync point: waits for the computation, then copies back to the host
print("NumPy type:", type(numpy_result), "| value:", numpy_result)
assert isinstance(numpy_result, np.ndarray)
np.testing.assert_array_equal(numpy_result, jax_arr)

# Try it: which conversions actually copy (decided by comparing pointers)
big = aligned_zeros(1 << 20); big[:] = 1.0           # 4 MB, 64-byte aligned
a_copy = jnp.array(big)
a_asarray = jnp.asarray(big)
a_put = jax.device_put(big)
a_put_noalias = jax.device_put(big, may_alias=False)
for name, arr in [("jnp.array", a_copy), ("jnp.asarray", a_asarray),
                  ("jax.device_put", a_put), ("device_put(may_alias=False)", a_put_noalias)]:
    print(f"{name:28s} shares memory with NumPy: {shares_memory(arr, big)}")

big[0] = 42.0                                           # modify the NumPy source array
print("After modifying the source: jnp.array[0] =", float(a_copy[0]), "| device_put[0] =", float(a_put[0]))
assert not shares_memory(a_copy, big) and float(a_copy[0]) == 1.0
assert not shares_memory(a_asarray, big)
assert shares_memory(a_put, big) and float(a_put[0]) == 42.0   # ⚠️ the "immutable" JAX array changed
assert shares_memory(a_put_noalias, big)                        # measured on 0.12.0.1: may_alias=False aliases too

# JAX → NumPy
y = jax.jit(lambda v: v * 2)(jnp.arange(1024.0))
n_view, n_copy = np.asarray(y), np.array(y)
print("np.asarray zero-copy:", n_view.ctypes.data == y.unsafe_buffer_pointer(), "| writeable:", n_view.flags.writeable)
print("np.array   zero-copy:", n_copy.ctypes.data == y.unsafe_buffer_pointer(), "| writeable:", n_copy.flags.writeable)
assert n_view.ctypes.data == y.unsafe_buffer_pointer() and not n_view.flags.writeable
assert n_copy.ctypes.data != y.unsafe_buffer_pointer() and n_copy.flags.writeable

# The cost of sync points: convert back to NumPy every step vs. only once at the end (timings are printed, not asserted)
@jax.jit
def step(x):
    return jnp.sin(x) * 0.999 + 0.001

x0 = jnp.ones((512, 512))

def loop_convert_each_step(x, n=100):
    hist = []
    for _ in range(n):
        x = step(x)
        hist.append(np.array(x)[0, 0])        # sync and copy every step
    return x

def loop_convert_at_end(x, n=100):
    for _ in range(n):
        x = step(x)                           # queued asynchronously, doesn't block
    return x

t_each = bench(loop_convert_each_step, x0, n=3)
t_end = bench(loop_convert_at_end, x0, n=3)
print(f"convert every step: {t_each:.2f} ms | convert only at the end: {t_end:.2f} ms | ratio {t_each / t_end:.2f}x")
print("Reading: " + ("syncing every step is slower, as expected." if t_each > t_end * 1.1 else
      "CPU only here: device memory is host memory, so \"copying back\" is nearly free and async queuing gains little. The gap is small and may even flip. "
      "The cost of sync points mostly shows up on GPU/TPU, where every sync waits for the device and goes over PCIe."))
np.testing.assert_allclose(loop_convert_each_step(x0), loop_convert_at_end(x0))

# Check: calling a NumPy function on a tracer inside jit
@jax.jit
def bad(x):
    return np.sin(x)      # during tracing x is a tracer, and NumPy can't turn it into a concrete array

try:
    bad(jnp.ones(3))
    err = None
except Exception as e:
    err = e
    print(type(e).__name__, ":", str(e).splitlines()[0][:120])
assert isinstance(err, jax.errors.TracerArrayConversionError)

host = jax.device_get({"w": jnp.ones(2), "b": jnp.zeros(1)})
assert all(isinstance(v, np.ndarray) for v in host.values())

# ❌ Old style (outdated/wrong): to_dlpack / device_buffer / from_dlpack(capsule) / random.normal(device=) all fail in 0.12
x_src = jax.random.normal(jax.random.key(0), (1024, 1024))
attempts = {
    "jax.dlpack.to_dlpack(x)": lambda: jax.dlpack.to_dlpack(x_src),
    "x.device_buffer.unsafe_buffer_pointer()": lambda: x_src.device_buffer.unsafe_buffer_pointer(),
    "jax.dlpack.from_dlpack(x.__dlpack__())": lambda: jax.dlpack.from_dlpack(x_src.__dlpack__()),
    "jax.random.normal(key, shape, device=...)": lambda: jax.random.normal(jax.random.key(0), (2,), device=jax.devices()[0]),
}
for desc, fn in attempts.items():
    try:
        fn(); print(f"{desc:45s} → still works")
    except Exception as e:
        print(f"{desc:45s} → {type(e).__name__}: {str(e).splitlines()[0][:90]}")
print("Public members of jax.dlpack:", [a for a in dir(jax.dlpack) if not a.startswith("_")])
print("(torch / cupy aren't installed here, so those two examples can't run in this environment)")

# ✅ Recommended: the DLPack protocol — zero-copy JAX → NumPy
x = jax.random.normal(jax.random.key(0), (1024, 1024))
x.block_until_ready()                       # sync before exporting (caveat 2)
print("__dlpack_device__:", x.__dlpack_device__())
print("__dlpack__() returns:", type(x.__dlpack__()).__name__)

n = np.from_dlpack(x)                       # NumPy calls x.__dlpack__() itself
print("JAX pointer  :", x.unsafe_buffer_pointer())
print("NumPy pointer:", n.ctypes.data, "| writeable:", n.flags.writeable)
assert not hasattr(jax.dlpack, "to_dlpack")
assert n.ctypes.data == x.unsafe_buffer_pointer() == x.addressable_data(0).unsafe_buffer_pointer()
assert not n.flags.writeable                # the export is a read-only view, which protects JAX's immutability
np.testing.assert_array_equal(n, x)

# ✅ Recommended: zero-copy NumPy → JAX (jax.dlpack.from_dlpack / jnp.from_dlpack)
src = aligned_zeros(1 << 20); src[:] = np.arange(src.size, dtype=np.float32)
j1 = jax.dlpack.from_dlpack(src)
j2 = jnp.from_dlpack(src)                   # equivalent entry point
print("from_dlpack shares memory:", shares_memory(j1, src), "| jnp.from_dlpack shares memory:", shares_memory(j2, src))
assert shares_memory(j1, src) and shares_memory(j2, src)

# Alignment requirement: if the start address isn't 64-byte aligned, the default (copy=None) silently copies; copy=False raises
unaligned = aligned_zeros(1025)[1:]         # offset by 4 bytes
assert unaligned.ctypes.data % 64 != 0
j_u = jax.dlpack.from_dlpack(unaligned)
print("unaligned array, from_dlpack shares memory:", shares_memory(j_u, unaligned))
assert not shares_memory(j_u, unaligned)
try:
    jax.dlpack.from_dlpack(unaligned, copy=False)
    raised = False
except ValueError as e:
    raised = True
    print("copy=False →", str(e)[:110])
assert raised

# Check: after a zero-copy import, modifying the source array changes the JAX array; lifetime is guaranteed by the consumer holding a reference
src = aligned_zeros(1024); src[:] = 1.0
j = jax.dlpack.from_dlpack(src)
before = float(j[0])
src[0] = 99.0
print("before: j[0] =", before, "| after modifying the source: j[0] =", float(j[0]))
assert before == 1.0 and float(j[0]) == 99.0

# Lifetime: after deleting the JAX array, the NumPy view is still valid (DLPack's deleter mechanism makes the consumer hold a reference to the producer)
x = jnp.full((1000,), 3.0)
view = np.from_dlpack(x)
del x
import gc; gc.collect()
print("After deleting the original JAX array, NumPy view sum =", view.sum())
assert view.sum() == 3000.0

# Timing: zero-copy vs. copy (64 MB; only the comparison is printed)
big = aligned_zeros(16 << 20); big[:] = 1.0
t_copy = bench(lambda a: jnp.array(a), big, n=10)
t_dl = bench(lambda a: jax.dlpack.from_dlpack(a), big, n=10)
print(f"jnp.array (copy): {t_copy:.3f} ms | from_dlpack (zero-copy): {t_dl:.4f} ms | {t_copy / t_dl:.0f}x faster")

# ❌ Old style (outdated/wrong): from jax.experimental import host_callback — the module was removed in 0.12
import importlib.util
print("find_spec:", importlib.util.find_spec("jax.experimental.host_callback"))
try:
    from jax.experimental import host_callback
    host_callback.id_tap(print, jnp.ones(2))
except ImportError as e:
    print("ImportError:", short_path(e))
assert importlib.util.find_spec("jax.experimental.host_callback") is None

# ✅ Recommended: id_tap → jax.debug.callback (course code block 1)
records = []
def log_intermediate_data(arg):
    records.append(type(arg))
    arg = np.asarray(arg)                    # the callback receives a jax.Array; convert to NumPy first
    print(f"[host callback] shape={arg.shape} first 5={np.round(arg.ravel()[:5], 4)}")

@jax.jit
def process_data(x):
    y = jnp.sin(x) * 2.0
    jax.debug.callback(log_intermediate_data, y)   # no return value, doesn't affect the data flow
    z = jnp.mean(jnp.cos(y))
    return z

input_data = jax.random.uniform(jax.random.key(42), (8, 8))
result = process_data(input_data)
result.block_until_ready()
jax.effects_barrier()                        # wait for all callbacks so prints don't leak into other cells
print("Final result:", result)
assert len(records) == 1 and issubclass(records[0], jax.Array)
ref = np.mean(np.cos(np.sin(np.asarray(input_data, np.float64)) * 2.0))
np.testing.assert_allclose(result, ref, rtol=1e-5)

# ✅ Recommended: id_print → jax.debug.print (format string, formatted lazily)
@jax.jit
def simple_print_example(x):
    y = x + 5.0
    jax.debug.print("intermediate y = {y}", y=y)
    return y * 2.0

output = simple_print_example(jnp.arange(3.0))
output.block_until_ready(); jax.effects_barrier()
np.testing.assert_array_equal(output, [10.0, 12.0, 14.0])

# ✅ Recommended: host_callback.call (pure computation) → jax.pure_callback (course code block 2)
def external_cpu_calculation(data, parameter):
    data = np.asarray(data)                  # host side: ordinary NumPy computation
    return (np.tanh(data) + np.asarray(parameter)).astype(data.dtype)

@jax.jit
def jax_workflow(x, ext_param):
    intermediate = jnp.log1p(jnp.abs(x))
    external_result = jax.pure_callback(
        external_cpu_calculation,
        jax.ShapeDtypeStruct(intermediate.shape, intermediate.dtype),   # output shape and dtype
        intermediate, ext_param)
    return external_result / (1.0 + jnp.mean(intermediate))

input_array = jnp.linspace(-2.0, 2.0, 5, dtype=jnp.float32)
final_result = jax_workflow(input_array, 0.5)
print("Final result:", final_result)
xi = np.log1p(np.abs(np.linspace(-2, 2, 5)))
np.testing.assert_allclose(final_result, (np.tanh(xi) + 0.5) / (1 + xi.mean()), rtol=1e-5)

# ✅ Recommended: calls with side effects → io_callback; ordered=True keeps execution order the same as program order
from jax.experimental import io_callback

order = []
def write_log(i):
    order.append(int(i))                     # side effect: writes external state

def body(i, carry):
    io_callback(write_log, None, i, ordered=True)   # when the callback returns None, pass None for result_shape_dtypes
    return carry + i

total = jax.jit(lambda: jax.lax.fori_loop(0, 5, body, 0))()
jax.effects_barrier()
print("Write order:", order, "| total =", total)
assert order == [0, 1, 2, 3, 4] and int(total) == 10

# Test each entry of the table above
def run(fn):
    try:
        jax.block_until_ready(fn()); jax.effects_barrier()
        return "OK"
    except Exception as e:
        return type(e).__name__

calls = []
def host_log(a):  calls.append(np.shape(a))
def host_double(a):
    calls.append(np.shape(a)); return np.asarray(a) * 2

sds = lambda v: jax.ShapeDtypeStruct(v.shape, v.dtype)
variants = {
    "debug.callback":           lambda v: (jax.debug.callback(host_log, v), v * 2)[1],
    "io_callback":              lambda v: io_callback(host_double, sds(v), v),
    "io_callback(ordered)":     lambda v: io_callback(host_double, sds(v), v, ordered=True),
    "pure_callback(default)":      lambda v: jax.pure_callback(host_double, sds(v), v),
    "pure_callback(expand_dims)": lambda v: jax.pure_callback(host_double, sds(v), v, vmap_method="expand_dims"),
}
x3, xb = jnp.arange(3.0), jnp.ones((4, 3))
table = {}
for name, f in variants.items():
    row = {}
    row["jit"] = run(lambda: jax.jit(f)(x3))
    calls.clear(); row["vmap"] = run(lambda: jax.vmap(f)(xb)); row["vmap calls"] = len(calls)
    row["grad"] = run(lambda: jax.grad(lambda v: f(v).sum())(x3))
    calls.clear(); jax.jit(lambda v: (f(v), v + 1)[1])(x3).block_until_ready(); jax.effects_barrier()
    row["runs when unused"] = len(calls)
    table[name] = row

print(f"{'callback':28s}{'jit':6s}{'vmap':22s}{'vmap calls    ':10s}{'grad':12s}{'runs when unused'}")
for name, r in table.items():
    print(f"{name:28s}{r['jit']:6s}{r['vmap']:22s}{r['vmap calls']!s:14s}{r['grad']:12s}{r['runs when unused']}")

assert all(r["jit"] == "OK" for r in table.values())
assert table["debug.callback"]["vmap"] == "OK" and table["debug.callback"]["vmap calls"] == 4
assert table["debug.callback"]["grad"] == "OK"
assert table["io_callback"]["vmap"] == "OK" and table["io_callback"]["grad"] == "ValueError"
assert table["io_callback(ordered)"]["vmap"] == "ValueError"
assert table["pure_callback(default)"]["vmap"] == "NotImplementedError"
assert table["pure_callback(default)"]["grad"] == "ValueError"
assert table["pure_callback(expand_dims)"]["vmap"] == "OK" and table["pure_callback(expand_dims)"]["vmap calls"] == 1
assert table["debug.callback"]["runs when unused"] == 1 and table["io_callback"]["runs when unused"] == 1
assert table["pure_callback(default)"]["runs when unused"] == 0      # the pure callback is removed as dead code

# ✅ Recommended: the course's main example (jit part), with the key created by jax.random.key
def external_pure_python_computation(x, y):
    """Plain Python/NumPy function: deterministic, no side effects"""
    x = np.asarray(x)                     # the callback receives a jax.Array; convert to NumPy
    return (np.sin(x) * y + 1.0).astype(x.dtype)

def jax_function_using_callback(a, b):
    output_shape_dtype = jax.ShapeDtypeStruct(a.shape, a.dtype)
    result = jax.pure_callback(
        external_pure_python_computation, output_shape_dtype, a, b,
        vmap_method="expand_dims")        # declare the vmap semantics explicitly (see the comparison below)
    return result + a

key = jax.random.key(0)
input_array = jax.random.normal(key, (3, 3))
input_scalar = 2.0
output = jax.jit(jax_function_using_callback)(input_array, input_scalar)
expected = np.sin(np.asarray(input_array, np.float64)) * input_scalar + 1.0 + np.asarray(input_array, np.float64)
print("Output:\n", output)
np.testing.assert_allclose(output, expected, rtol=1e-5)

# vmap: with expand_dims the callback is called only once and receives the whole (10, 3, 3) batch
batched_input = jax.random.normal(key, (10, 3, 3))
batched_output = jax.vmap(jax_function_using_callback, in_axes=(0, None))(batched_input, input_scalar)
print("Batched output shape:", batched_output.shape)
assert batched_output.shape == (10, 3, 3)
np.testing.assert_allclose(batched_output, np.sin(np.asarray(batched_input)) * 2 + 1 + np.asarray(batched_input), rtol=1e-5)

# ❌ Old style (outdated/wrong): vmap without vmap_method; grad straight through pure_callback
def course_version(a, b):
    out = jax.pure_callback(external_pure_python_computation, jax.ShapeDtypeStruct(a.shape, a.dtype), a, b)
    return out + a

try:
    jax.vmap(course_version, in_axes=(0, None))(batched_input, input_scalar)
except NotImplementedError as e:
    print("vmap → NotImplementedError:", str(e)[:150])
try:
    jax.grad(lambda x: jnp.sum(course_version(x, input_scalar)))(input_array)
except ValueError as e:
    print("grad → ValueError:", str(e)[:150])

# ✅ Recommended: give pure_callback a derivative with custom_jvp (d/dx [sin(x)·y + 1] = cos(x)·y)
@jax.custom_jvp
def ext_fn(a, b):
    return jax.pure_callback(external_pure_python_computation, jax.ShapeDtypeStruct(a.shape, a.dtype),
                             a, b, vmap_method="expand_dims")

@ext_fn.defjvp
def ext_fn_jvp(primals, tangents):
    a, b = primals
    a_dot, b_dot = tangents
    y = ext_fn(a, b)
    # write the derivative in jnp, so it is itself differentiable, jittable and vmappable
    return y, jnp.cos(a) * b * a_dot + jnp.sin(a) * b_dot

grad_fn = jax.jit(jax.grad(lambda x: jnp.sum(ext_fn(x, input_scalar) + x)))
gradients = grad_fn(input_array)
print("Gradients:\n", gradients)
np.testing.assert_allclose(gradients, np.cos(np.asarray(input_array)) * 2.0 + 1.0, rtol=1e-5)
# derivative with respect to the scalar argument b
gb = jax.grad(lambda b: jnp.sum(ext_fn(input_array, b)))(2.0)
np.testing.assert_allclose(gb, np.sin(np.asarray(input_array)).sum(), rtol=1e-5)
# vmap ∘ grad composes too
vg = jax.vmap(jax.grad(lambda x: jnp.sum(ext_fn(x, 2.0))))(batched_input)
np.testing.assert_allclose(vg, np.cos(np.asarray(batched_input)) * 2.0, rtol=1e-5)

# The course's multi-output example: result_shape_dtypes described by a tuple (pytree)
def multi_output_pure_function(x):
    x = np.asarray(x)
    return np.sum(x).astype(x.dtype), np.mean(x).astype(x.dtype)

def jax_multi_output_callback(arr):
    output_structure = (jax.ShapeDtypeStruct((), arr.dtype),   # scalar shape is ()
                        jax.ShapeDtypeStruct((), arr.dtype))
    sum_val, mean_val = jax.pure_callback(multi_output_pure_function, output_structure, arr)
    return sum_val * 2, mean_val * 3

res1, res2 = jax.jit(jax_multi_output_callback)(jnp.arange(5.0))
print(f"Multiple outputs: {res1=}, {res2=}")
assert float(res1) == 20.0 and float(res2) == 6.0

# Try each vmap_method: number of calls and the shapes the callback receives
seen = []
def h2(a, b):
    seen.append((np.shape(a), np.shape(b)))
    return (np.asarray(a) * np.asarray(b)).astype(np.float32)

def call_with(method):
    def f(a):
        return jax.pure_callback(h2, jax.ShapeDtypeStruct(a.shape, a.dtype), a, jnp.ones(3), vmap_method=method)
    return f

results = {}
for m in ["sequential", "sequential_unrolled", "expand_dims", "broadcast_all", "legacy_vectorized"]:
    seen.clear()
    out = jax.vmap(call_with(m))(jnp.ones((4, 3)))
    results[m] = list(seen)
    print(f"{m:20s} called {len(seen)} time(s), (a, b) shapes received = {seen[0]}  output {out.shape}")

assert len(results["sequential"]) == 4 and results["sequential"][0] == ((3,), (3,))
assert results["expand_dims"] == [((4, 3), (1, 3))]       # the unbatched b just gets one extra size-1 dim
assert results["broadcast_all"] == [((4, 3), (4, 3))]     # b is broadcast to the batch size
assert results["legacy_vectorized"] == [((4, 3), (3,))]   # b keeps its original shape
try:
    jax.vmap(call_with("bogus"))(jnp.ones((4, 3)))
except ValueError as e:
    print("Invalid value →", str(e)[:120])

# Check: when the derivative also uses pure_callback, hessian fails; wrapping the derivative in custom_jvp too fixes it
np_cos_raw = lambda x: jax.pure_callback(lambda a: np.cos(np.asarray(a)).astype(a.dtype),
                                          jax.ShapeDtypeStruct(x.shape, x.dtype), x, vmap_method="expand_dims")
@jax.custom_jvp
def np_sin_v1(x):
    return jax.pure_callback(lambda a: np.sin(np.asarray(a)).astype(a.dtype),
                             jax.ShapeDtypeStruct(x.shape, x.dtype), x, vmap_method="expand_dims")
np_sin_v1.defjvp(lambda p, t: (np_sin_v1(p[0]), np_cos_raw(p[0]) * t[0]))

xs = jnp.linspace(0.0, 1.0, 4)
np.testing.assert_allclose(jax.grad(lambda v: np_sin_v1(v).sum())(xs), np.cos(np.asarray(xs)), rtol=1e-6)
try:
    jax.hessian(lambda v: np_sin_v1(v).sum())(xs)
    ok = True
except ValueError as e:
    ok = False
    print("hessian(v1) → ValueError:", str(e)[:80])
assert not ok

# Fix: sin and cos are each other's derivatives, so give both a custom_jvp
@jax.custom_jvp
def np_sin(x):
    return jax.pure_callback(lambda a: np.sin(np.asarray(a)).astype(a.dtype), jax.ShapeDtypeStruct(x.shape, x.dtype), x, vmap_method="expand_dims")
@jax.custom_jvp
def np_cos(x):
    return np_cos_raw(x)
np_sin.defjvp(lambda p, t: (np_sin(p[0]), np_cos(p[0]) * t[0]))
np_cos.defjvp(lambda p, t: (np_cos(p[0]), -np_sin(p[0]) * t[0]))
H = jax.hessian(lambda v: np_sin(v).sum())(xs)
np.testing.assert_allclose(np.diag(H), -np.sin(np.asarray(xs)), rtol=1e-5)
print("Hessian diagonal after the fix:", np.diag(H))

# Check: a shape mismatch errors at run time; the callback doesn't run when its output is unused
try:
    jax.jit(lambda v: jax.pure_callback(lambda a: np.asarray(a)[:2], jax.ShapeDtypeStruct(v.shape, v.dtype), v))(jnp.arange(3.0)).block_until_ready()
    shape_err = None
except Exception as e:
    shape_err = type(e).__name__
print("Wrong return shape →", shape_err)
assert shape_err == "JaxRuntimeError"

counter = []
def impure(a):
    counter.append(1)                         # ❌ a side effect inside pure_callback (counterexample)
    return np.asarray(a)
out = jax.jit(lambda v: (jax.pure_callback(impure, jax.ShapeDtypeStruct(v.shape, v.dtype), v), v + 1)[1])(jnp.arange(3.0))
out.block_until_ready()
print("Callback runs when the output is unused:", len(counter))
assert len(counter) == 0

# Timing: pure_callback (NumPy sin) vs. native jnp.sin (printed only)
xbig = jax.random.normal(jax.random.key(1), (1000, 1000))
f_cb = jax.jit(lambda v: jax.pure_callback(lambda a: np.sin(np.asarray(a)), jax.ShapeDtypeStruct(v.shape, v.dtype), v) * 2 + 1)
f_native = jax.jit(lambda v: jnp.sin(v) * 2 + 1)
np.testing.assert_allclose(f_cb(xbig), f_native(xbig), rtol=1e-5, atol=1e-6)
t_cb, t_nat = bench(f_cb, xbig), bench(f_native, xbig)
print(f"pure_callback: {t_cb:.3f} ms | native jnp: {t_nat:.3f} ms | ratio {t_cb / t_nat:.1f}x")

# See which primitives a function breaks down into
def f(x, y):
    return jnp.sum(jnp.sin(x) * y + jnp.dot(x, y))

closed = jax.make_jaxpr(f)(jnp.ones(3), jnp.ones(3))
print(closed)
names = [eqn.primitive.name for eqn in closed.jaxpr.eqns]
print("Primitive sequence:", names)
assert {"sin", "mul", "dot_general", "add", "reduce_sum"} <= set(names)

# Names of the control-flow primitives
def g(x):
    y = jax.lax.cond(x.sum() > 0, jnp.sin, jnp.cos, x)
    return jax.lax.scan(lambda c, v: (c + v, c), 0.0, y)[0]
cf_names = {e.primitive.name for e in jax.make_jaxpr(g)(jnp.ones(3)).jaxpr.eqns}
print("Control-flow primitives:", cf_names & {"cond", "scan"})
assert {"cond", "scan"} <= cf_names

import jax.extend.core.primitives as P          # semi-public entry point for built-in primitive objects
print("P.scan_p.name =", P.scan_p.name, "| P.cond_p.name =", P.cond_p.name, "| jax.lax.add_p.name =", jax.lax.add_p.name)
assert (P.scan_p.name, P.cond_p.name, jax.lax.add_p.name) == ("scan", "cond", "add")

# Primitives carry various rule tables: JVP, transpose, batching, lowering
from jax.interpreters import ad, batching, mlir
print("sin has a JVP rule       :", P.sin_p in ad.primitive_jvps)
print("mul has a transpose rule :", P.mul_p in ad.primitive_transposes)
# Note: in 0.12, batching.primitive_batchers is a "write-only proxy" (only [p] = rule works); you can't query it with `in`.
# It wraps the rule and stores it in batching.fancy_primitive_batchers
try:
    P.sin_p in batching.primitive_batchers
except TypeError as e:
    print("primitive_batchers doesn't support `in`:", e)
print("sin has a batching rule  :", P.sin_p in batching.fancy_primitive_batchers)
assert P.sin_p in ad.primitive_jvps and P.mul_p in ad.primitive_transposes and P.sin_p in batching.fancy_primitive_batchers

# The compilation pipeline: jaxpr → StableHLO (MLIR text)
hlo_text = jax.jit(f).lower(jnp.ones(3), jnp.ones(3)).as_text()
print(hlo_text)
assert "stablehlo.sine" in hlo_text and "stablehlo.dot_general" in hlo_text

# ❌ Old style (outdated/wrong): from jax.core import Primitive — jax.core no longer exports Primitive in 0.12
try:
    from jax.core import Primitive as _P
except ImportError as e:
    print("ImportError:", short_path(e))
import jax.core, jax.extend.core
print("jax.core has Primitive:", hasattr(jax.core, "Primitive"), "| has ShapedArray:", hasattr(jax.core, "ShapedArray"))
print("jax.extend.core has Primitive:", hasattr(jax.extend.core, "Primitive"), "| has ShapedArray:", hasattr(jax.extend.core, "ShapedArray"))

# ✅ Recommended: steps 1–3 — the primitive object, the user function, impl and abstract evaluation
from jax.extend.core import Primitive          # semi-internal API: path verified on 0.12
from jax.core import ShapedArray               # abstract value type
from jax.interpreters import mlir, ad, batching

my_custom_op_p = Primitive("my_custom_op")      # step 1: a unique name

def my_custom_op(x, *, some_parameter):         # step 2: the user-facing function
    return my_custom_op_p.bind(x, some_parameter=some_parameter)

@my_custom_op_p.def_impl                         # eager implementation: receives concrete arrays
def _my_custom_op_impl(x, *, some_parameter):
    return x * x + some_parameter

@my_custom_op_p.def_abstract_eval                # step 3: looks only at shape/dtype, computes no values
def _my_custom_op_abstract_eval(abstract_x, *, some_parameter):
    if some_parameter < 0:
        raise ValueError("some_parameter must be non-negative")
    return ShapedArray(abstract_x.shape, abstract_x.dtype)

x = jnp.array([1.0, 2.0, 3.0])
print("eager:", my_custom_op(x, some_parameter=1.0))
print(jax.make_jaxpr(partial(my_custom_op, some_parameter=1.0))(x))
np.testing.assert_array_equal(my_custom_op(x, some_parameter=1.0), [2.0, 5.0, 10.0])

# No lowering yet → jit fails; no JVP → grad fails; no batching → vmap fails
for name, fn in [("jit", lambda: jax.jit(partial(my_custom_op, some_parameter=1.0))(x)),
                 ("grad", lambda: jax.grad(lambda v: my_custom_op(v, some_parameter=1.0).sum())(x)),
                 ("vmap", lambda: jax.vmap(partial(my_custom_op, some_parameter=1.0))(jnp.ones((2, 3))))]:
    try:
        fn(); print(name, "→ unexpectedly succeeded")
    except NotImplementedError as e:
        print(f"{name:5s}→ NotImplementedError: {str(e).splitlines()[0][:80]}")

# ✅ Recommended: step 4 — lowering (details in 5.8; here mlir.lower_fun lowers a jnp implementation to StableHLO)
mlir.register_lowering(
    my_custom_op_p,
    mlir.lower_fun(lambda x, *, some_parameter: x * x + some_parameter, multiple_results=False))

y_jit = jax.jit(partial(my_custom_op, some_parameter=1.0))(x)
print("jit:", y_jit)
np.testing.assert_array_equal(y_jit, [2.0, 5.0, 10.0])

# ❌ Old style (outdated/wrong): building a zero tangent with ad.Zero.from_value(primal_out) in a JVP rule — this method doesn't exist in 0.12
print("ad.Zero has from_value:", hasattr(ad.Zero, "from_value"))
try:
    ad.Zero.from_value(jnp.ones(3))
except AttributeError as e:
    print("AttributeError:", e)

# ✅ Recommended: step 5 — the JVP rule (ẏ = 2x·ẋ); build zero tangents with ad.Zero(tangent aval)
def _my_custom_op_jvp(primals, tangents, *, some_parameter):
    x, = primals
    x_dot, = tangents
    primal_out = my_custom_op_p.bind(x, some_parameter=some_parameter)
    if type(x_dot) is ad.Zero:                                  # the input is treated as a constant
        tangent_out = ad.Zero(jax.typeof(primal_out).to_tangent_aval())
    else:
        tangent_out = 2 * x * x_dot                             # expressed with existing primitives, linear in x_dot
    return primal_out, tangent_out

ad.primitive_jvps[my_custom_op_p] = _my_custom_op_jvp

g = jax.grad(lambda v: my_custom_op(v, some_parameter=1.0).sum())(x)
gj = jax.jit(jax.grad(lambda v: my_custom_op(v, some_parameter=1.0).sum()))(x)
print("grad:", g, "| jit(grad):", gj)
np.testing.assert_allclose(g, 2 * np.asarray(x)); np.testing.assert_allclose(gj, 2 * np.asarray(x))

# Batching for an elementwise primitive: batching.defvectorized does it in one line (the batch dim is passed straight through)
batching.defvectorized(my_custom_op_p)
xb = jnp.arange(6.0).reshape(2, 3)
yb = jax.vmap(partial(my_custom_op, some_parameter=1.0))(xb)
np.testing.assert_array_equal(yb, np.asarray(xb) ** 2 + 1)
vg = jax.vmap(jax.grad(lambda v: my_custom_op(v, some_parameter=0.0).sum()))(xb)
np.testing.assert_array_equal(vg, 2 * np.asarray(xb))
print("vmap / vmap(grad) passed")

# ✅ Course example 1: abstract evaluation for custom_add_one (with input checks)
custom_add_one_p = Primitive("custom_add_one")

@custom_add_one_p.def_abstract_eval
def custom_add_one_abstract_eval(*avals, **params):
    if len(avals) != 1:
        raise ValueError("custom_add_one needs exactly one input operand.")
    input_aval = avals[0]
    if not isinstance(input_aval, ShapedArray):
        raise TypeError(f"Expected a ShapedArray, got {type(input_aval)}")
    if not jnp.issubdtype(input_aval.dtype, jnp.floating):
        raise TypeError(f"custom_add_one only supports floating-point inputs, got {input_aval.dtype}")
    return ShapedArray(input_aval.shape, input_aval.dtype)

impl_calls = []
@custom_add_one_p.def_impl
def _add_one_impl(x):
    impl_calls.append(1)
    return x + 1.0

# jax.eval_shape only does abstract evaluation: you get the shape, and the impl never runs
out_sds = jax.eval_shape(custom_add_one_p.bind, jax.ShapeDtypeStruct((4, 5), jnp.float32))
print("eval_shape:", out_sds, "| impl calls:", len(impl_calls))
assert out_sds.shape == (4, 5) and out_sds.dtype == jnp.float32 and impl_calls == []

# Type errors are caught at tracing time (before compiling or running anything)
try:
    jax.eval_shape(custom_add_one_p.bind, jax.ShapeDtypeStruct((3,), jnp.int32))
    caught = False
except TypeError as e:
    caught = True
    print("TypeError at tracing time:", e)
assert caught

# ✅ Course example 2: custom_reduce_sum — the output shape is decided by the axis parameter
custom_reduce_sum_p = Primitive("custom_reduce_sum")

def custom_reduce_sum(x, axis):
    axis = tuple(sorted(a % x.ndim for a in (axis if isinstance(axis, tuple) else (axis,))))
    return custom_reduce_sum_p.bind(x, axis=axis)         # axis must be hashable (a tuple)

@custom_reduce_sum_p.def_abstract_eval
def custom_reduce_sum_abstract_eval(aval, *, axis):
    if not isinstance(aval, ShapedArray):
        raise TypeError("The input must be a ShapedArray")
    output_shape = tuple(d for i, d in enumerate(aval.shape) if i not in axis)
    return ShapedArray(output_shape, aval.dtype)

custom_reduce_sum_p.def_impl(lambda x, *, axis: jnp.sum(x, axis=axis))
mlir.register_lowering(custom_reduce_sum_p,
                       mlir.lower_fun(lambda x, *, axis: jnp.sum(x, axis=axis), multiple_results=False))

my_array = jnp.arange(24.0).reshape(2, 3, 4)
print(jax.eval_shape(lambda a: custom_reduce_sum(a, (0,)), my_array))
assert jax.eval_shape(lambda a: custom_reduce_sum(a, (0, 2)), my_array).shape == (3,)
r = jax.jit(lambda a: custom_reduce_sum(a, (0,)))(my_array)
np.testing.assert_array_equal(r, np.asarray(my_array).sum(0))
print(jax.make_jaxpr(lambda a: custom_reduce_sum(a, -1))(my_array))

# A multi-output primitive: minmax(x) -> (min, max)
minmax_bad = Primitive("minmax_bad")                    # forgot to set multiple_results
minmax_bad.def_abstract_eval(lambda a: (ShapedArray((), a.dtype), ShapedArray((), a.dtype)))
try:
    jax.eval_shape(minmax_bad.bind, jnp.arange(5.0))
    err = None
except ValueError as e:
    err = e; print("ValueError:", str(e)[:110])
assert err is not None

minmax_p = Primitive("minmax")
minmax_p.multiple_results = True
minmax_p.def_impl(lambda x: (jnp.min(x), jnp.max(x)))
minmax_p.def_abstract_eval(lambda a: (ShapedArray((), a.dtype), ShapedArray((), a.dtype)))
mlir.register_lowering(minmax_p, mlir.lower_fun(lambda x: (jnp.min(x), jnp.max(x)), multiple_results=True))
lo, hi = jax.jit(minmax_p.bind)(jnp.array([3.0, -1.0, 7.0]))
print("min, max =", lo, hi)
assert (float(lo), float(hi)) == (-1.0, 7.0)