[WIP] call backend JAX bindings

[mariogeiger]

import jax
import jax.numpy as jnp
import numpy as np
from jax.experimental.mosaic.gpu import profiler

import cuequivariance as cue
import cuequivariance_jax as cuex
import logging

jax.config.update("jax_enable_x64", True)
logging.basicConfig(level=logging.INFO)

e = (
    cue.descriptors.channelwise_tensor_product(
        128 * cue.Irreps(cue.O3, "0e + 1o + 2e"),
        cue.Irreps(cue.O3, "0e + 1o + 2e + 3o"),
        128 * cue.Irreps(cue.O3, "0e + 1o + 2e"),
    )
    .squeeze_modes()
    .flatten_coefficient_modes()
)

dtype = jnp.float32
dtype_math = jnp.float32
batch_size = 20_000

inputs = [
    cuex.randn(jax.random.key(i), rep, (batch_size,), dtype)
    for i, rep in enumerate(e.inputs)
]


def foo(use_custom_kernels: bool):
    return cuex.equivariant_tensor_product(
        e, use_custom_kernels=use_custom_kernels, dtype_math=dtype_math
    )


output_custom, runtime_custom = profiler.measure(foo(True), mode="cupti")(*inputs)
print(f"Runtime: {runtime_custom:.3f} ms")

output_xla, runtime_xla = profiler.measure(foo(False), mode="cupti")(*inputs)
print(f"Runtime: {runtime_xla:.3f} ms")

Runtime: 0.759 ms
Runtime: 1.164 ms