File size: 2,660 Bytes
5ccb4fd | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import contextvars
import jax
import jax.numpy as jnp
from jax.extend import core
from jax.interpreters import batching
_custom_lap_active_var: contextvars.ContextVar[bool] = contextvars.ContextVar(
"spin_custom_lap_active",
default=False,
)
def custom_lap_active() -> bool:
return _custom_lap_active_var.get()
def enter_custom_lap():
return _custom_lap_active_var.set(True)
def restore_custom_lap(token):
_custom_lap_active_var.reset(token)
quadrilinear_merge_p = core.Primitive("quadrilinear_merge_lap")
quadrilinear_merge_p.multiple_results = False
def _quadrilinear_merge_impl(T, u_a, u_b):
groups, rank, _, _ = T.shape
if u_a.ndim == 1:
u_a_2d = u_a.reshape(groups, rank)
u_b_2d = u_b.reshape(groups, rank)
contracted = jnp.einsum("ijkl,ik->ijl", T, u_a_2d)
out_2d = jnp.einsum("ijl,il->ij", contracted, u_b_2d)
return out_2d.reshape(-1)
batch = u_a.shape[0]
u_a_3d = u_a.reshape(batch, groups, rank)
u_b_3d = u_b.reshape(batch, groups, rank)
out_3d = jnp.einsum("ijkl,Bik,Bil->Bij", T, u_a_3d, u_b_3d)
return out_3d.reshape(batch, -1)
def _quadrilinear_merge_abstract_eval(T_aval, u_a_aval, u_b_aval):
del u_b_aval
return jax.core.ShapedArray(u_a_aval.shape, T_aval.dtype)
quadrilinear_merge_p.def_impl(_quadrilinear_merge_impl)
quadrilinear_merge_p.def_abstract_eval(_quadrilinear_merge_abstract_eval)
def _quadrilinear_merge_batched(args, dims):
T, u_a, u_b = args
T_axis, u_a_axis, u_b_axis = dims
if T_axis is not None:
raise ValueError("quadrilinear merge parameters cannot be batched")
batch = None
if u_a_axis is not None:
u_a = jnp.moveaxis(u_a, u_a_axis, 0)
batch = u_a.shape[0]
if u_b_axis is not None:
u_b = jnp.moveaxis(u_b, u_b_axis, 0)
batch = u_b.shape[0] if batch is None else batch
if batch is None:
return quadrilinear_merge_p.bind(T, u_a, u_b), None
if u_a_axis is None:
u_a = jnp.broadcast_to(u_a[None], (batch,) + u_a.shape)
if u_b_axis is None:
u_b = jnp.broadcast_to(u_b[None], (batch,) + u_b.shape)
u_a_flat = u_a.reshape((-1, u_a.shape[-1]))
u_b_flat = u_b.reshape((-1, u_b.shape[-1]))
out_flat = quadrilinear_merge_p.bind(T, u_a_flat, u_b_flat)
out = out_flat.reshape(u_a.shape[:-1] + (out_flat.shape[-1],))
return out, 0
batching.primitive_batchers[quadrilinear_merge_p] = _quadrilinear_merge_batched
__all__ = [
"custom_lap_active",
"quadrilinear_merge_p",
]
|