warp-lang 1.7.2rc1__py3-none-win_amd64.whl → 1.8.1__py3-none-win_amd64.whl

warp/jax_experimental/xla_ffi.py

@@ -130,14 +130,14 @@ class XLA_FFI_DataType(enum.IntEnum):
 #   int64_t* dims;  // length == rank
 # };
 class XLA_FFI_Buffer(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("dtype", ctypes.c_int),  # XLA_FFI_DataType
         ("data", ctypes.c_void_p),
         ("rank", ctypes.c_int64),
         ("dims", ctypes.POINTER(ctypes.c_int64)),
-    ]
+    )
 
 
 # typedef enum {
@@ -162,13 +162,13 @@ class XLA_FFI_RetType(enum.IntEnum):
 #   void** args;             // length == size
 # };
 class XLA_FFI_Args(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("size", ctypes.c_int64),
         ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_ArgType*
         ("args", ctypes.POINTER(ctypes.c_void_p)),
-    ]
+    )
 
 
 # struct XLA_FFI_Rets {
@@ -179,13 +179,13 @@ class XLA_FFI_Args(ctypes.Structure):
 #   void** rets;             // length == size
 # };
 class XLA_FFI_Rets(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("size", ctypes.c_int64),
         ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_RetType*
         ("rets", ctypes.POINTER(ctypes.c_void_p)),
-    ]
+    )
 
 
 # typedef struct XLA_FFI_ByteSpan {
@@ -193,7 +193,10 @@ class XLA_FFI_Rets(ctypes.Structure):
 #   size_t len;
 # } XLA_FFI_ByteSpan;
 class XLA_FFI_ByteSpan(ctypes.Structure):
-    _fields_ = [("ptr", ctypes.POINTER(ctypes.c_char)), ("len", ctypes.c_size_t)]
+    _fields_ = (
+        ("ptr", ctypes.POINTER(ctypes.c_char)),
+        ("len", ctypes.c_size_t),
+    )
 
 
 # typedef struct XLA_FFI_Scalar {
@@ -201,7 +204,10 @@ class XLA_FFI_ByteSpan(ctypes.Structure):
 #   void* value;
 # } XLA_FFI_Scalar;
 class XLA_FFI_Scalar(ctypes.Structure):
-    _fields_ = [("dtype", ctypes.c_int), ("value", ctypes.c_void_p)]
+    _fields_ = (
+        ("dtype", ctypes.c_int),
+        ("value", ctypes.c_void_p),
+    )
 
 
 # typedef struct XLA_FFI_Array {
@@ -210,7 +216,11 @@ class XLA_FFI_Scalar(ctypes.Structure):
 #   void* data;
 # } XLA_FFI_Array;
 class XLA_FFI_Array(ctypes.Structure):
-    _fields_ = [("dtype", ctypes.c_int), ("size", ctypes.c_size_t), ("data", ctypes.c_void_p)]
+    _fields_ = (
+        ("dtype", ctypes.c_int),
+        ("size", ctypes.c_size_t),
+        ("data", ctypes.c_void_p),
+    )
 
 
 # typedef enum {
@@ -235,14 +245,14 @@ class XLA_FFI_AttrType(enum.IntEnum):
 #   void** attrs;              // length == size
 # };
 class XLA_FFI_Attrs(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("size", ctypes.c_int64),
         ("types", ctypes.POINTER(ctypes.c_int)),  # XLA_FFI_AttrType*
         ("names", ctypes.POINTER(ctypes.POINTER(XLA_FFI_ByteSpan))),
         ("attrs", ctypes.POINTER(ctypes.c_void_p)),
-    ]
+    )
 
 
 # struct XLA_FFI_Api_Version {
@@ -252,12 +262,12 @@ class XLA_FFI_Attrs(ctypes.Structure):
 #   int minor_version;  // out
 # };
 class XLA_FFI_Api_Version(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("major_version", ctypes.c_int),
         ("minor_version", ctypes.c_int),
-    ]
+    )
 
 
 # enum XLA_FFI_Handler_TraitsBits {
@@ -276,11 +286,11 @@ class XLA_FFI_Handler_TraitsBits(enum.IntEnum):
 #   XLA_FFI_Handler_Traits traits;
 # };
 class XLA_FFI_Metadata(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("api_version", XLA_FFI_Api_Version),  # XLA_FFI_Extension_Type
         ("traits", ctypes.c_uint32),  # XLA_FFI_Handler_Traits
-    ]
+    )
 
 
 # struct XLA_FFI_Metadata_Extension {
@@ -288,7 +298,10 @@ class XLA_FFI_Metadata(ctypes.Structure):
 #   XLA_FFI_Metadata* metadata;
 # };
 class XLA_FFI_Metadata_Extension(ctypes.Structure):
-    _fields_ = [("extension_base", XLA_FFI_Extension_Base), ("metadata", ctypes.POINTER(XLA_FFI_Metadata))]
+    _fields_ = (
+        ("extension_base", XLA_FFI_Extension_Base),
+        ("metadata", ctypes.POINTER(XLA_FFI_Metadata)),
+    )
 
 
 # typedef enum {
@@ -337,12 +350,12 @@ class XLA_FFI_Error_Code(enum.IntEnum):
 #   XLA_FFI_Error_Code errc;
 # };
 class XLA_FFI_Error_Create_Args(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("message", ctypes.c_char_p),
         ("errc", ctypes.c_int),
-    ]  # XLA_FFI_Error_Code
+    )  # XLA_FFI_Error_Code
 
 
 XLA_FFI_Error_Create = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_Error_Create_Args))
@@ -355,12 +368,12 @@ XLA_FFI_Error_Create = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_
 #   void* stream;  // out
 # };
 class XLA_FFI_Stream_Get_Args(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("ctx", ctypes.c_void_p),  # XLA_FFI_ExecutionContext*
         ("stream", ctypes.c_void_p),
-    ]  # // out
+    )  # // out
 
 
 XLA_FFI_Stream_Get = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_Stream_Get_Args))
@@ -391,7 +404,7 @@ XLA_FFI_Stream_Get = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_St
 #   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Future_SetError);
 # };
 class XLA_FFI_Api(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("api_version", XLA_FFI_Api_Version),
@@ -412,7 +425,7 @@ class XLA_FFI_Api(ctypes.Structure):
         ("XLA_FFI_Future_Create", ctypes.c_void_p),  # XLA_FFI_Future_Create
         ("XLA_FFI_Future_SetAvailable", ctypes.c_void_p),  # XLA_FFI_Future_SetAvailable
         ("XLA_FFI_Future_SetError", ctypes.c_void_p),  # XLA_FFI_Future_SetError
-    ]
+    )
 
 
 # struct XLA_FFI_CallFrame {
@@ -431,7 +444,7 @@ class XLA_FFI_Api(ctypes.Structure):
 #   XLA_FFI_Future* future;  // out
 # };
 class XLA_FFI_CallFrame(ctypes.Structure):
-    _fields_ = [
+    _fields_ = (
         ("struct_size", ctypes.c_size_t),
         ("extension_start", ctypes.POINTER(XLA_FFI_Extension_Base)),
         ("api", ctypes.POINTER(XLA_FFI_Api)),
@@ -441,7 +454,7 @@ class XLA_FFI_CallFrame(ctypes.Structure):
         ("rets", XLA_FFI_Rets),
         ("attrs", XLA_FFI_Attrs),
         ("future", ctypes.c_void_p),  # XLA_FFI_Future* // out
-    ]
+    )
 
 
 _xla_data_type_to_constructor = {

warp/math.py

@@ -22,11 +22,13 @@ Vector norm functions
 """
 
 __all__ = [
+    "norm_huber",
     "norm_l1",
     "norm_l2",
-    "norm_huber",
     "norm_pseudo_huber",
     "smooth_normalize",
+    "transform_compose",
+    "transform_decompose",
     "transform_from_matrix",
     "transform_to_matrix",
 ]
@@ -142,6 +144,19 @@ def create_transform_from_matrix_func(dtype):
         """
         Construct a transformation from a 4x4 matrix.
 
+        .. math::
+            M = \\begin{bmatrix}
+            R_{00} & R_{01} & R_{02} & p_x \\\\
+            R_{10} & R_{11} & R_{12} & p_y \\\\
+            R_{20} & R_{21} & R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+
         Args:
             mat (Matrix[4, 4, Float]): Matrix to convert.
 
@@ -177,6 +192,19 @@ def create_transform_to_matrix_func(dtype):
         """
         Convert a transformation to a 4x4 matrix.
 
+        .. math::
+            M = \\begin{bmatrix}
+            R_{00} & R_{01} & R_{02} & p_x \\\\
+            R_{10} & R_{11} & R_{12} & p_y \\\\
+            R_{20} & R_{21} & R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+
         Args:
             xform (Transformation[Float]): Transformation to convert.
 
@@ -212,6 +240,140 @@ wp.func(
 )
 
 
+def create_transform_compose_func(dtype):
+    mat44 = wp.types.matrix((4, 4), dtype)
+    quat = wp.types.quaternion(dtype)
+    vec3 = wp.types.vector(3, dtype)
+
+    def transform_compose(position: vec3, rotation: quat, scale: vec3):
+        """
+        Compose a 4x4 transformation matrix from a 3D position, quaternion orientation, and 3D scale.
+
+        .. math::
+            M = \\begin{bmatrix}
+            s_x R_{00} & s_y R_{01} & s_z R_{02} & p_x \\\\
+            s_x R_{10} & s_y R_{11} & s_z R_{12} & p_y \\\\
+            s_x R_{20} & s_y R_{21} & s_z R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+        * :math:`s` is the 3D scale vector :math:`[s_x, s_y, s_z]` of the input transform.
+
+        Args:
+            position (Vector[3, Float]): The 3D position vector.
+            rotation (Quaternion[Float]): The quaternion orientation.
+            scale (Vector[3, Float]): The 3D scale vector.
+
+        Returns:
+            Matrix[4, 4, Float]: The transformation matrix.
+        """
+        R = wp.quat_to_matrix(rotation)
+        # fmt: off
+        return mat44(
+            scale[0] * R[0,0], scale[1] * R[0,1], scale[2] * R[0,2], position[0],
+            scale[0] * R[1,0], scale[1] * R[1,1], scale[2] * R[1,2], position[1],
+            scale[0] * R[2,0], scale[1] * R[2,1], scale[2] * R[2,2], position[2],
+            dtype(0.0), dtype(0.0), dtype(0.0), dtype(1.0),
+        )
+        # fmt: on
+
+    return transform_compose
+
+
+transform_compose = wp.func(
+    create_transform_compose_func(wp.float32),
+    name="transform_compose",
+)
+wp.func(
+    create_transform_compose_func(wp.float16),
+    name="transform_compose",
+)
+wp.func(
+    create_transform_compose_func(wp.float64),
+    name="transform_compose",
+)
+
+
+def create_transform_decompose_func(dtype):
+    mat44 = wp.types.matrix((4, 4), dtype)
+    vec3 = wp.types.vector(3, dtype)
+    mat33 = wp.types.matrix((3, 3), dtype)
+    zero = dtype(0.0)
+
+    def transform_decompose(m: mat44):
+        """
+        Decompose a 4x4 transformation matrix into 3D position, quaternion orientation, and 3D scale.
+
+        .. math::
+            M = \\begin{bmatrix}
+            s_x R_{00} & s_y R_{01} & s_z R_{02} & p_x \\\\
+            s_x R_{10} & s_y R_{11} & s_z R_{12} & p_y \\\\
+            s_x R_{20} & s_y R_{21} & s_z R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+        * :math:`s` is the 3D scale vector :math:`[s_x, s_y, s_z]` of the input transform.
+
+        Args:
+            m (Matrix[4, 4, Float]): The matrix to decompose.
+
+        Returns:
+            Tuple[Vector[3, Float], Quaternion[Float], Vector[3, Float]]: A tuple containing the position vector, quaternion orientation, and scale vector.
+        """
+        # extract position
+        position = vec3(m[0, 3], m[1, 3], m[2, 3])
+        # extract rotation matrix components
+        r00, r01, r02 = m[0, 0], m[0, 1], m[0, 2]
+        r10, r11, r12 = m[1, 0], m[1, 1], m[1, 2]
+        r20, r21, r22 = m[2, 0], m[2, 1], m[2, 2]
+        # get scale magnitudes
+        sx = wp.sqrt(r00 * r00 + r10 * r10 + r20 * r20)
+        sy = wp.sqrt(r01 * r01 + r11 * r11 + r21 * r21)
+        sz = wp.sqrt(r02 * r02 + r12 * r12 + r22 * r22)
+        # normalize rotation matrix components
+        if sx != zero:
+            r00 /= sx
+            r10 /= sx
+            r20 /= sx
+        if sy != zero:
+            r01 /= sy
+            r11 /= sy
+            r21 /= sy
+        if sz != zero:
+            r02 /= sz
+            r12 /= sz
+            r22 /= sz
+        # extract rotation (quaternion)
+        rotation = wp.quat_from_matrix(mat33(r00, r01, r02, r10, r11, r12, r20, r21, r22))
+        # extract scale
+        scale = vec3(sx, sy, sz)
+        return position, rotation, scale
+
+    return transform_decompose
+
+
+transform_decompose = wp.func(
+    create_transform_decompose_func(wp.float32),
+    name="transform_decompose",
+)
+wp.func(
+    create_transform_decompose_func(wp.float16),
+    name="transform_decompose",
+)
+wp.func(
+    create_transform_decompose_func(wp.float64),
+    name="transform_decompose",
+)
+
+
 # register API functions so they appear in the documentation
 
 wp.context.register_api_function(
@@ -242,3 +404,11 @@ wp.context.register_api_function(
     transform_to_matrix,
     group="Transformations",
 )
+wp.context.register_api_function(
+    transform_compose,
+    group="Transformations",
+)
+wp.context.register_api_function(
+    transform_decompose,
+    group="Transformations",
+)

warp/native/array.h

@@ -161,7 +161,7 @@ inline CUDA_CALLABLE void print(shape_t s)
     // should probably store ndim with shape
     printf("(%d, %d, %d, %d)\n", s.dims[0], s.dims[1], s.dims[2], s.dims[3]);
 }
-inline CUDA_CALLABLE void adj_print(shape_t s, shape_t& shape_t) {}
+inline CUDA_CALLABLE void adj_print(shape_t s, shape_t& adj_s) {}
 
 
 template <typename T>
@@ -665,11 +665,11 @@ CUDA_CALLABLE inline indexedarray_t<T> view(indexedarray_t<T>& src, int i, int j
 }
 
 template<template<typename> class A1, template<typename> class A2, template<typename> class A3, typename T>
-inline CUDA_CALLABLE void adj_view(A1<T>& src, int i, A2<T>& adj_src, int adj_i, A3<T> adj_ret) {}
+inline CUDA_CALLABLE void adj_view(A1<T>& src, int i, A2<T>& adj_src, int adj_i, A3<T>& adj_ret) {}
 template<template<typename> class A1, template<typename> class A2, template<typename> class A3, typename T>
-inline CUDA_CALLABLE void adj_view(A1<T>& src, int i, int j, A2<T>& adj_src, int adj_i, int adj_j, A3<T> adj_ret) {}
+inline CUDA_CALLABLE void adj_view(A1<T>& src, int i, int j, A2<T>& adj_src, int adj_i, int adj_j, A3<T>& adj_ret) {}
 template<template<typename> class A1, template<typename> class A2, template<typename> class A3, typename T>
-inline CUDA_CALLABLE void adj_view(A1<T>& src, int i, int j, int k, A2<T>& adj_src, int adj_i, int adj_j, int adj_k, A3<T> adj_ret) {}
+inline CUDA_CALLABLE void adj_view(A1<T>& src, int i, int j, int k, A2<T>& adj_src, int adj_i, int adj_j, int adj_k, A3<T>& adj_ret) {}
 
 // TODO: lower_bound() for indexed arrays?
 
@@ -743,6 +743,24 @@ inline CUDA_CALLABLE T atomic_max(const A<T>& buf, int i, int j, int k, T value)
 template<template<typename> class A, typename T>
 inline CUDA_CALLABLE T atomic_max(const A<T>& buf, int i, int j, int k, int l, T value) { return atomic_max(&index(buf, i, j, k, l), value); }
 
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_cas(const A<T>& buf, int i, T old_value, T new_value) { return atomic_cas(&index(buf, i), old_value, new_value); }
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_cas(const A<T>& buf, int i, int j, T old_value, T new_value) { return atomic_cas(&index(buf, i, j), old_value, new_value); }
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_cas(const A<T>& buf, int i, int j, int k, T old_value, T new_value) { return atomic_cas(&index(buf, i, j, k), old_value, new_value); }
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_cas(const A<T>& buf, int i, int j, int k, int l, T old_value, T new_value) { return atomic_cas(&index(buf, i, j, k, l), old_value, new_value); }
+
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_exch(const A<T>& buf, int i, T value) { return atomic_exch(&index(buf, i), value); }
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_exch(const A<T>& buf, int i, int j, T value) { return atomic_exch(&index(buf, i, j), value); }
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_exch(const A<T>& buf, int i, int j, int k, T value) { return atomic_exch(&index(buf, i, j, k), value); }
+template<template<typename> class A, typename T>
+inline CUDA_CALLABLE T atomic_exch(const A<T>& buf, int i, int j, int k, int l, T value) { return atomic_exch(&index(buf, i, j, k, l), value); }
+
 template<template<typename> class A, typename T>
 inline CUDA_CALLABLE T* address(const A<T>& buf, int i) { return &index(buf, i); }
 template<template<typename> class A, typename T>
@@ -1128,6 +1146,87 @@ inline CUDA_CALLABLE void adj_atomic_max(const A1<T>& buf, int i, int j, int k,
     FP_VERIFY_ADJ_4(value, adj_value)
 }
 
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_cas(const A1<T>& buf, int i, T compare, T value, const A2<T>& adj_buf, int adj_i, T& adj_compare, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_cas(&index(buf, i), compare, value, &index(adj_buf, i), adj_compare, adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_cas(&index(buf, i), compare, value, &index_grad(buf, i), adj_compare, adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_1(value, adj_value)
+}
+
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_cas(const A1<T>& buf, int i, int j, T compare, T value, const A2<T>& adj_buf, int adj_i, int adj_j, T& adj_compare, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_cas(&index(buf, i, j), compare, value, &index(adj_buf, i, j), adj_compare, adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_cas(&index(buf, i, j), compare, value, &index_grad(buf, i, j), adj_compare, adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_2(value, adj_value)
+}
+
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_cas(const A1<T>& buf, int i, int j, int k, T compare, T value, const A2<T>& adj_buf, int adj_i, int adj_j, int adj_k, T& adj_compare, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_cas(&index(buf, i, j, k), compare, value, &index(adj_buf, i, j, k), adj_compare, adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_cas(&index(buf, i, j, k), compare, value, &index_grad(buf, i, j, k), adj_compare, adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_3(value, adj_value)
+}
+
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_cas(const A1<T>& buf, int i, int j, int k, int l, T compare, T value, const A2<T>& adj_buf, int adj_i, int adj_j, int adj_k, int adj_l, T& adj_compare, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_cas(&index(buf, i, j, k, l), compare, value, &index(adj_buf, i, j, k, l), adj_compare, adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_cas(&index(buf, i, j, k, l), compare, value, &index_grad(buf, i, j, k, l), adj_compare, adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_4(value, adj_value)
+}
+
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_exch(const A1<T>& buf, int i, T value, const A2<T>& adj_buf, int adj_i, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_exch(&index(buf, i), value, &index(adj_buf, i), adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_exch(&index(buf, i), value, &index_grad(buf, i), adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_1(value, adj_value)
+}
+
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_exch(const A1<T>& buf, int i, int j, T value, const A2<T>& adj_buf, int adj_i, int adj_j, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_exch(&index(buf, i, j), value, &index(adj_buf, i, j), adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_exch(&index(buf, i, j), value, &index_grad(buf, i, j), adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_2(value, adj_value)
+}
+
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_exch(const A1<T>& buf, int i, int j, int k, T value, const A2<T>& adj_buf, int adj_i, int adj_j, int adj_k, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_exch(&index(buf, i, j, k), value, &index(adj_buf, i, j, k), adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_exch(&index(buf, i, j, k), value, &index_grad(buf, i, j, k), adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_3(value, adj_value)
+}
+
+template<template<typename> class A1, template<typename> class A2, typename T>
+inline CUDA_CALLABLE void adj_atomic_exch(const A1<T>& buf, int i, int j, int k, int l, T value, const A2<T>& adj_buf, int adj_i, int adj_j, int adj_k, int adj_l, T& adj_value, const T& adj_ret) {
+    if (adj_buf.data)
+        adj_atomic_exch(&index(buf, i, j, k, l), value, &index(adj_buf, i, j, k, l), adj_value, adj_ret);
+    else if (buf.grad)
+        adj_atomic_exch(&index(buf, i, j, k, l), value, &index_grad(buf, i, j, k, l), adj_value, adj_ret);
+
+    FP_VERIFY_ADJ_4(value, adj_value)
+}
+
+
 template<template<typename> class A, typename T>
 CUDA_CALLABLE inline int len(const A<T>& a)
 {