cuequivariance-ops-jax-cu13 0.7.0__py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

cuequivariance_ops_jax/VERSION

@@ -0,0 +1 @@
+0.7.0

cuequivariance_ops_jax/__init__.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+from ._version import __version__, __git_commit__
+from ._tensor_product_uniform_1d_jit import (
+    tensor_product_uniform_1d_jit,
+    Operation,
+    Path,
+)
+from ._indexed_linear import indexed_linear
+from ._triangle_attention import (
+    triangle_attention_cuda_fwd,
+    triangle_attention_cuda_bwd,
+    triangle_attention_jax_fwd,
+)
+from ._gpu_utilities import noop, sleep, synchronize, event_record, event_elapsed
+
+__all__ = [
+    "__version__",
+    "__git_commit__",
+    "tensor_product_uniform_1d_jit",
+    "Operation",
+    "Path",
+    "indexed_linear",
+    "triangle_attention_cuda_fwd",
+    "triangle_attention_cuda_bwd",
+    "triangle_attention_jax_fwd",
+    "noop",
+    "sleep",
+    "synchronize",
+    "event_record",
+    "event_elapsed",
+]

cuequivariance_ops_jax/_common.py

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+import ctypes
+import importlib.metadata
+import os
+from enum import IntEnum
+
+import jax.numpy as jnp
+from jax import ffi
+
+import cuequivariance_ops  # noqa: F401
+
+# Load libcue_ops_jax.so
+try:
+    dist = importlib.metadata.distribution("cuequivariance_ops_jax")
+    root = dist.locate_file("cuequivariance_ops_jax")
+except Exception:
+    # last resort, will fail with writeable install
+    root = os.path.dirname(__file__)
+
+path = os.path.join(root, "lib/libcue_ops_jax.so")
+library = ctypes.cdll.LoadLibrary(path)
+
+# Register the c++ functions with JAX
+CUSTOM_FUNCS = [
+    (
+        "tensor_product_uniform_1d_jit",
+        "tensor_product_uniform_1d_jit",
+        "tensor_product_uniform_1d_cpu",
+    ),
+    ("indexed_linear_B", "indexed_linear_B", None),
+    ("indexed_linear_C", "indexed_linear_C", None),
+    ("triangle_attention_cuda_fwd", "triangle_attention_cuda_fwd", None),
+    ("triangle_attention_cuda_bwd", "triangle_attention_cuda_bwd", None),
+    ("noop", "noop_gpu", "noop_cpu"),
+    ("sleep", "sleep_gpu", "sleep_cpu"),
+    ("synchronize", "synchronize_gpu", "synchronize_cpu"),
+    ("event_record", "event_record_gpu", "event_record_cpu"),
+    ("event_elapsed", "event_elapsed_gpu", "event_elapsed_cpu"),
+]
+
+for name, cuda_fn, cpu_fn in CUSTOM_FUNCS:
+    if cuda_fn is not None:
+        ffi.register_ffi_target(
+            name=name, fn=ffi.pycapsule(getattr(library, cuda_fn)), platform="CUDA"
+        )
+    if cpu_fn is not None:
+        ffi.register_ffi_target(
+            name=name, fn=ffi.pycapsule(getattr(library, cpu_fn)), platform="cpu"
+        )
+
+
+class DataType(IntEnum):
+    FLOAT32 = 0
+    FLOAT64 = 1
+    FLOAT16 = 2
+    BFLOAT16 = 3
+    INT32 = 4
+    INT64 = 5
+
+
+def _dtype(jax_dtype: jnp.dtype) -> DataType:
+    try:
+        return {
+            jnp.float32: DataType.FLOAT32,
+            jnp.float64: DataType.FLOAT64,
+            jnp.float16: DataType.FLOAT16,
+            jnp.bfloat16: DataType.BFLOAT16,
+            jnp.int32: DataType.INT32,
+            jnp.int64: DataType.INT64,
+        }[jnp.dtype(jax_dtype).type]
+    except KeyError:
+        raise ValueError(f"Unsupported dtype: {jax_dtype}")

cuequivariance_ops_jax/_cublas_enums.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+# CUBLAS Compute Types
+# Converted from cublas_api.h cublasComputeType_t enum
+
+
+class CublasComputeTypes:
+    """Class to provide getattr access to CUBLAS compute types."""
+
+    CUBLAS_COMPUTE_16F = 64
+    CUBLAS_COMPUTE_16F_PEDANTIC = 65
+    CUBLAS_COMPUTE_32F = 68
+    CUBLAS_COMPUTE_32F_PEDANTIC = 69
+    CUBLAS_COMPUTE_32F_FAST_16F = 74
+    CUBLAS_COMPUTE_32F_FAST_16BF = 75
+    CUBLAS_COMPUTE_32F_FAST_TF32 = 77
+    CUBLAS_COMPUTE_64F = 70
+    CUBLAS_COMPUTE_64F_PEDANTIC = 71
+    CUBLAS_COMPUTE_32I = 72
+    CUBLAS_COMPUTE_32I_PEDANTIC = 73
+
+
+# Create instance for getattr access
+cublas_compute_types = CublasComputeTypes()

cuequivariance_ops_jax/_gpu_utilities.py

@@ -0,0 +1,157 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+"""GPU utility functions including no-op, sleep, synchronize, and GPU events.
+
+GPU Events Implementation Attribution:
+The GPU events functionality (event_record, event_elapsed) has been adapted from
+JAX's GPU events implementation that was removed in version 0.7.2.
+
+Original source: https://github.com/jax-ml/jax/
+License: Apache License 2.0
+
+JAX Copyright 2018 The JAX Authors.
+Licensed under the Apache License, Version 2.0.
+"""
+
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from jax import ffi
+
+
+def _flatten(pytree: Any):
+    """Helper to apply FFI function to JAX arrays in a pytree."""
+    leaves, treedef = jax.tree.flatten(pytree)
+    arrays = [(i, leaf) for i, leaf in enumerate(leaves) if isinstance(leaf, jax.Array)]
+    _, values = zip(*arrays)
+
+    def unflatten(outputs):
+        for idx, (leaf_idx, _) in enumerate(arrays):
+            leaves[leaf_idx] = outputs[idx]
+        return jax.tree.unflatten(treedef, leaves)
+
+    return values, unflatten
+
+
+def noop(pytree: Any) -> Any:
+    """
+    No-op function that returns input pytree unchanged through FFI.
+
+    Args:
+        pytree: Any pytree structure containing JAX arrays
+
+    Returns:
+        The same pytree structure with arrays passed through FFI
+    """
+    vals, unflatten = _flatten(pytree)
+
+    vals = ffi.ffi_call(
+        "noop",
+        [jax.ShapeDtypeStruct(arr.shape, arr.dtype) for arr in vals],
+        input_output_aliases={i: i for i in range(len(vals))},
+    )(*vals)
+
+    return unflatten(vals)
+
+
+def sleep(seconds: jax.Array, pytree: Any) -> tuple[jax.Array, Any]:
+    """
+    Sleep for the specified number of seconds and return input pytree unchanged.
+
+    Args:
+        seconds: Number of seconds to sleep (as a JAX array)
+        pytree: Any pytree structure containing JAX arrays
+
+    Returns:
+        A tuple of (elapsed_ticks, pytree) where elapsed_ticks is the number of
+        clock ticks that elapsed during the sleep operation
+    """
+    seconds = jnp.asarray(seconds, dtype=jnp.float32)
+    vals, unflatten = _flatten(pytree)
+
+    outputs = ffi.ffi_call(
+        "sleep",
+        [jax.ShapeDtypeStruct((), jnp.int64)]
+        + [jax.ShapeDtypeStruct(arr.shape, arr.dtype) for arr in vals],
+        input_output_aliases={i: i for i in range(1, len(vals) + 1)},
+    )(seconds, *vals)
+    elapsed_ticks, vals = outputs[0], outputs[1:]
+
+    return elapsed_ticks, unflatten(vals)
+
+
+def synchronize(pytree: Any) -> tuple[jax.Array, Any]:
+    """
+    Synchronize the current CUDA stream and return input pytree unchanged.
+
+    Args:
+        pytree: Any pytree structure containing JAX arrays
+
+    Returns:
+        A tuple of (elapsed_seconds, pytree) where elapsed_seconds is the time
+        in seconds it took to synchronize the CUDA stream
+    """
+    vals, unflatten = _flatten(pytree)
+
+    outputs = ffi.ffi_call(
+        "synchronize",
+        [jax.ShapeDtypeStruct((), jnp.float32)]
+        + [jax.ShapeDtypeStruct(arr.shape, arr.dtype) for arr in vals],
+        input_output_aliases={i: i + 1 for i in range(len(vals))},
+    )(*vals)
+    elapsed_seconds, vals = outputs[0], outputs[1:]
+
+    return elapsed_seconds, unflatten(vals)
+
+
+def event_record(pytree: Any, *, copy_before: bool = False) -> tuple[jax.Array, Any]:
+    """
+    Record a GPU event on the current CUDA stream and return the event handle.
+
+    Args:
+        pytree: Any pytree structure containing JAX arrays
+        copy_before: If True, copy event handle to device before recording.
+                    If False, copy after recording (default).
+
+    Returns:
+        A tuple of (event_handle, pytree) where event_handle is a uint64
+        representing the CUDA event, and pytree is passed through unchanged.
+    """
+    vals, unflatten = _flatten(pytree)
+
+    outputs = ffi.ffi_call(
+        "event_record",
+        [jax.ShapeDtypeStruct((), jnp.uint64)]  # event_handle
+        + [jax.ShapeDtypeStruct(arr.shape, arr.dtype) for arr in vals],
+        input_output_aliases={i: i + 1 for i in range(len(vals))},
+    )(*vals, copy_before=copy_before)
+    event_handle, vals = outputs[0], outputs[1:]
+
+    return event_handle, unflatten(vals)
+
+
+def event_elapsed(start_event: jax.Array, end_event: jax.Array) -> jax.Array:
+    """
+    Calculate elapsed time between two GPU events.
+
+    Args:
+        start_event: uint64 event handle from event_record()
+        end_event: uint64 event handle from event_record()
+
+    Returns:
+        Elapsed time in milliseconds as a float32 scalar.
+    """
+    elapsed_ms = ffi.ffi_call(
+        "event_elapsed",
+        jax.ShapeDtypeStruct((), jnp.float32),
+    )(start_event, end_event)
+
+    return elapsed_ms

cuequivariance_ops_jax/_indexed_linear.py

@@ -0,0 +1,148 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+from jax import ffi
+
+from ._common import _dtype
+from ._cublas_enums import cublas_compute_types
+
+
+def indexed_linear(
+    A: jax.Array,
+    B: jax.Array,
+    D: jax.ShapeDtypeStruct,
+    counts: jax.Array,
+    u: int,
+    v: int,
+    C: int,
+    Z: int,
+    subscripts: tuple[str, str, str],
+    coefficient: float,
+    math_dtype: str | None = None,
+) -> jax.Array:
+    """
+    Performance benchmarks for compatible dtype + math_dtype combinations
+    (Problem size: u=512, v=512, Z=10,000):
+
+    | Rank | Dtype + Math Dtype                    | Time (ms) | Speedup vs float64 |
+    |------|---------------------------------------|-----------|-------------------|
+    | 🥇   | bfloat16 + CUBLAS_COMPUTE_32F         | 0.0891    | 83.9x            |
+    | 🥈   | float16 + CUBLAS_COMPUTE_32F          | 0.0911    | 82.1x            |
+    | 🥉   | float32 + CUBLAS_COMPUTE_32F_FAST_TF32| 0.1403    | 53.3x            |
+    | 4th  | float32 + CUBLAS_COMPUTE_32F_PEDANTIC | 0.1987    | 37.6x            |
+    | 5th  | float32 + CUBLAS_COMPUTE_32F          | 0.2038    | 36.7x            |
+    | 6th  | float64 + CUBLAS_COMPUTE_64F          | 7.4742    | 1.0x (baseline)  |
+    """
+    subscripts = tuple(subscripts)
+    original_subscripts = subscripts
+    assert len(subscripts) == 3
+    swap_u_v = False
+    swap_A_B = False
+
+    dtype = jnp.dtype(A.dtype)
+    assert dtype == B.dtype
+    assert dtype == D.dtype
+
+    if math_dtype is None:
+        dtype_to_compute_type = {
+            jnp.bfloat16: cublas_compute_types.CUBLAS_COMPUTE_32F,  # Use 32F compute type for bf16
+            jnp.float16: cublas_compute_types.CUBLAS_COMPUTE_32F,  # Use 32F compute type for fp16
+            jnp.float32: cublas_compute_types.CUBLAS_COMPUTE_32F,
+            jnp.float64: cublas_compute_types.CUBLAS_COMPUTE_64F,
+        }
+        if dtype.type not in dtype_to_compute_type:
+            raise ValueError(
+                f"For dtype '{dtype}', please specify math_dtype manually and check CUBLAS documentation "
+                f"for compatible compute types."
+            )
+        compute_type = dtype_to_compute_type[dtype.type]
+    else:
+        if not hasattr(cublas_compute_types, math_dtype):
+            supported_types = [
+                attr
+                for attr in dir(cublas_compute_types)
+                if attr.startswith("CUBLAS_COMPUTE_")
+            ]
+            raise ValueError(
+                f"Unsupported math_dtype '{math_dtype}'. "
+                f"The supported compute types are: {supported_types}. "
+                f"Be careful as they are not compatible with all I/O dtype combinations. "
+                f"Have a look at the CUBLAS documentation for compatibility details."
+            )
+        compute_type = getattr(cublas_compute_types, math_dtype)
+
+    if subscripts in [("u", "v", "vu"), ("uv", "v", "u"), ("vu", "v", "u")]:
+        swap_A_B = True
+        swap_u_v = True
+    if subscripts in [("v", "uv", "u"), ("v", "vu", "u"), ("v", "u", "vu")]:
+        swap_u_v = True
+    if subscripts in [("v", "u", "uv"), ("uv", "u", "v"), ("vu", "u", "v")]:
+        swap_A_B = True
+
+    if swap_u_v:
+        subscripts = tuple(
+            x.replace("u", "q").replace("v", "u").replace("q", "v") for x in subscripts
+        )
+        u, v = v, u
+
+    if swap_A_B:
+        subscripts = (subscripts[1], subscripts[0], subscripts[2])
+        A, B = B, A
+
+    temp_storage_bytes_cub_ExclusiveSum = 1024  # TODO this seems to be sufficient but we never know if it's enough for all use cases and GPUs
+    workspace_size = (
+        counts.size * (3 + 1) * jnp.dtype(jnp.int64).itemsize
+        + temp_storage_bytes_cub_ExclusiveSum
+    )
+    workspace = jnp.empty((workspace_size,), dtype=jnp.int8)
+
+    if subscripts == ("u", "v", "uv"):
+        (D, _) = ffi.ffi_call("indexed_linear_C", (D, workspace))(
+            A,
+            B,
+            counts,
+            compute_type=compute_type,
+            u=u,
+            v=v,
+            C=C,
+            Z=Z,
+            coefficient=coefficient,
+            dtype_A=_dtype(A.dtype),
+            dtype_B=_dtype(B.dtype),
+            dtype_D=_dtype(D.dtype),
+        )
+        return D
+
+    if subscripts == ("u", "uv", "v"):
+        transpose_B = False
+    elif subscripts == ("u", "vu", "v"):
+        transpose_B = True
+    else:
+        raise ValueError(f"Invalid subscripts: {original_subscripts}.")
+
+    (D, _) = ffi.ffi_call("indexed_linear_B", (D, workspace))(
+        A,
+        B,
+        counts,
+        compute_type=compute_type,
+        u=u,
+        v=v,
+        C=C,
+        Z=Z,
+        transpose_B=transpose_B,
+        coefficient=coefficient,
+        dtype_A=_dtype(A.dtype),
+        dtype_B=_dtype(B.dtype),
+        dtype_D=_dtype(D.dtype),
+    )
+    return D

cuequivariance_ops_jax/_tensor_product_uniform_1d_jit.py

@@ -0,0 +1,210 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+from __future__ import annotations
+
+from enum import IntEnum
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import ffi
+
+from ._common import _dtype
+
+
+class BatchDimension(IntEnum):
+    BATCHED = 0
+    SHARED = 1
+    INDEXED = 2
+
+
+class SegmentDimension(IntEnum):
+    SCALAR = 0
+    VECTOR = 1
+
+
+def _batch_dim(idx: int, size: int, batch_size: int) -> BatchDimension:
+    if idx < 0:
+        if size == batch_size:
+            return BatchDimension.BATCHED
+        else:
+            return BatchDimension.SHARED
+    else:
+        return BatchDimension.INDEXED
+
+
+def _seg_dim(buffer: jax.ShapeDtypeStruct, operand_extent: int) -> SegmentDimension:
+    if buffer.shape[-1] == operand_extent:
+        return SegmentDimension.VECTOR
+    else:
+        return SegmentDimension.SCALAR
+
+
+class Operation(NamedTuple):
+    buffers: list[int]
+    start_path: int
+    num_paths: int
+
+
+class Path(NamedTuple):
+    indices: list[int]
+    coefficient: float
+
+
+def _batch_size(sizes: list[int]) -> int:
+    batch_size = 1
+    for size in sizes:
+        if size != 1:
+            assert batch_size in {1, size}
+            batch_size = size
+    return batch_size
+
+
+def _operand_extent(
+    buffers: list[jax.ShapeDtypeStruct],
+):
+    operand_extent = max(x.shape[-1] for x in buffers)
+    for x in buffers:
+        assert x.shape[-1] in {1, operand_extent}, x.shape[-1]
+    return operand_extent
+
+
+def _operation_start_indices(
+    paths: list[Path], operation_start_paths: np.ndarray
+) -> np.ndarray:
+    path_num_operands = np.array([len(path.indices) for path in paths], dtype=np.int32)
+    start_indices = np.append(0, np.cumsum(path_num_operands))
+    return start_indices[operation_start_paths].astype(np.int64)
+
+
+def tensor_product_uniform_1d_jit(
+    input_buffers: list[jax.Array],  # ndim = num_batch_axes + 2
+    output_buffers_shape_dtype: list[jax.ShapeDtypeStruct],  # ndim = num_batch_axes + 2
+    index_buffers: list[jax.Array],  # ndim = num_batch_axes
+    buffer_index: list[list[int]],  # -1 if not indexed
+    *,
+    operations: list[Operation],
+    paths: list[Path],
+    math_dtype: jnp.dtype,
+    name: str = "untitled",
+) -> list[jax.Array]:
+    """JIT-compiled CUDA implementation of tensor_product_uniform_1d."""
+    input_buffers = list(input_buffers)
+    output_buffers_shape_dtype = list(output_buffers_shape_dtype)
+    index_buffers = list(index_buffers)
+    buffer_index = np.array(buffer_index, dtype=np.int32)
+    operations, paths = list(operations), list(paths)
+
+    io_buffers = input_buffers + output_buffers_shape_dtype
+    buffers = io_buffers + index_buffers
+    assert buffer_index.shape[0] == len(buffers)
+
+    # trick: ensure all outputs are "used" by adding dummy operations for unused outputs
+    # this ensures that the kernel writes zeros to unused outputs (as expected by the XLA bindings)
+    for i in range(len(input_buffers), len(io_buffers)):
+        if not any(i in op.buffers for op in operations):
+            operations.append(Operation([i], 0, 0))
+
+    num_batch_axes = buffer_index.shape[1]
+    for x in io_buffers:
+        assert x.ndim == num_batch_axes + 2
+    for i in index_buffers:
+        assert i.ndim == num_batch_axes
+        assert i.dtype.type in {jnp.int32, jnp.int64}
+
+    batch_sizes = [
+        _batch_size(
+            [x.shape[i] for x, idx in zip(buffers, buffer_index[:, i]) if idx < 0],
+        )
+        for i in range(num_batch_axes)
+    ]
+
+    buffer_batch_dim = np.array(
+        [
+            [
+                _batch_dim(idx, size, batch_size)
+                for idx, size, batch_size in zip(idxs, x.shape, batch_sizes)
+            ]
+            for x, idxs in zip(buffers, buffer_index)
+        ]
+    )
+
+    index_extent = [set() for _ in range(len(index_buffers))]
+    for x, idxs in zip(buffers, buffer_index):
+        for size, idx in zip(x.shape, idxs):
+            if idx >= 0:
+                index_extent[idx].add(size)
+    assert all(len(x) == 1 for x in index_extent), index_extent
+    index_extent = [x.pop() for x in index_extent]
+
+    operand_extent = _operand_extent(io_buffers)
+
+    math_dtype = jnp.dtype(math_dtype)
+    assert math_dtype.type in {jnp.float32, jnp.float64}
+
+    def ii(items):
+        return np.array([i for i in items], dtype=np.int64)
+
+    buffer_batch_dim = ii(buffer_batch_dim.flatten())
+    buffer_num_segments = ii(x.shape[-2] for x in io_buffers)
+    buffer_segments_dim = ii(_seg_dim(x, operand_extent) for x in io_buffers)
+    buffer_index = ii(buffer_index.flatten())
+    index_extent = ii(index_extent)
+    buffer_dtype = ii(_dtype(x.dtype) for x in io_buffers + index_buffers)
+    operation_num_operands = ii(len(op.buffers) for op in operations)
+    operation_buffers = ii(b for op in operations for b in op.buffers)
+    operation_num_paths = ii(op.num_paths for op in operations)
+    operation_start_coeffs = ii(op.start_path for op in operations)
+    operation_start_indices = _operation_start_indices(paths, operation_start_coeffs)
+    path_indices = ii(i for path in paths for i in path.indices)
+    path_coefficients = np.array([path.coefficient for path in paths], dtype=np.float64)
+    batch_sizes = ii(batch_sizes)
+
+    # print(f"{operand_extent=}")
+    # print(f"num_indices={len(index_buffers)}")
+    # print(f"{buffer_batch_dim=}")
+    # print(f"{buffer_num_segments=}")
+    # print(f"{buffer_segments_dim=}")
+    # print(f"{buffer_index=}")
+    # print(f"{index_extent=}")
+    # print(f"{buffer_dtype=}")
+    # print(f"{operation_num_operands=}")
+    # print(f"{operation_buffers=}")
+    # print(f"{operation_num_paths=}")
+    # print(f"{operation_start_indices=}")
+    # print(f"{operation_start_coeffs=}")
+    # print(f"{path_indices=}")
+    # print(f"{path_coefficients=}")
+    # print(f"{batch_sizes=}", flush=True)
+
+    call = ffi.ffi_call("tensor_product_uniform_1d_jit", output_buffers_shape_dtype)
+    return call(
+        *input_buffers,
+        *index_buffers,
+        name=name,
+        math_dtype=_dtype(math_dtype),
+        operand_extent=operand_extent,
+        num_indices=len(index_buffers),
+        buffer_batch_dim=buffer_batch_dim,
+        buffer_num_segments=buffer_num_segments,
+        buffer_segments_dim=buffer_segments_dim,
+        buffer_index=buffer_index,
+        index_extent=index_extent,
+        buffer_dtype=buffer_dtype,
+        operation_num_operands=operation_num_operands,
+        operation_buffers=operation_buffers,
+        operation_num_paths=operation_num_paths,
+        operation_start_indices=operation_start_indices,
+        operation_start_coeffs=operation_start_coeffs,
+        path_indices=path_indices,
+        path_coefficients=path_coefficients.view(np.int64),
+        batch_sizes=batch_sizes,
+    )