From 5e1a84d5e2df18a434b20e7e9e4948ad6e616a22 Mon Sep 17 00:00:00 2001 From: Paul Fuchs Date: Wed, 8 Jul 2026 08:19:10 +0200 Subject: [PATCH 1/3] Fix backward pass for fctp --- openequivariance_extjax/src/libjax_tp_jit.cpp | 44 ++-- tests/conftest.py | 1 - tests/tp_adjoint_test.py | 224 ++++++++++++++++++ 3 files changed, 251 insertions(+), 18 deletions(-) create mode 100644 tests/tp_adjoint_test.py diff --git a/openequivariance_extjax/src/libjax_tp_jit.cpp b/openequivariance_extjax/src/libjax_tp_jit.cpp index ec567a9..47a1423 100644 --- a/openequivariance_extjax/src/libjax_tp_jit.cpp +++ b/openequivariance_extjax/src/libjax_tp_jit.cpp @@ -164,11 +164,13 @@ struct KernelProp { } }; -std::unordered_map>, - KernelProp - >> tp_cache; +struct CachedTP { + std::unique_ptr> impl; + KernelProp prop; + std::string json_payload; +}; + +std::unordered_map tp_cache; std::unordered_map*, KernelProp> const std::lock_guard lock(mut); auto it = tp_cache.find(hash); if (it == tp_cache.end()) { + std::string incoming_json(json_payload); std::string err; - json root = json::parse(std::string(json_payload), err); + json root = json::parse(incoming_json, err); if (!err.empty()) throw std::runtime_error("JSON Parse Error: " + err); std::string kernel_src = root["kernel"].string_value(); @@ -202,12 +205,18 @@ std::pair*, KernelProp> backward_cfg, dbackward_cfg, kernel_prop_map); - tp_cache.insert({hash, - std::make_pair(std::move(jit_tp_impl), - KernelProp(kernel_prop_map, is_convolution))}); + tp_cache.emplace(hash, + CachedTP{ + std::move(jit_tp_impl), + KernelProp(kernel_prop_map, is_convolution), + std::move(incoming_json), + }); it = tp_cache.find(hash); + } else if (it->second.json_payload != json_payload) { + throw std::runtime_error( + "OpenEquivariance JAX TP cache collision: same hash, different kernel JSON"); } - return {it->second.first.get(), it->second.second}; + return {it->second.impl.get(), it->second.prop}; } } @@ -340,9 +349,9 @@ ffi::Error tp_backward_impl( check_tensor(*W_grad, {num_batch, k.weight_numel}, k.weight_dtype, "W_grad"); } - if (k.shared_weights) { - zero_buffer(*W_grad, stream); - } + zero_buffer(*L1_grad, stream); + zero_buffer(*L2_grad, stream); + zero_buffer(*W_grad, stream); jit_kernel->backward( num_batch, @@ -391,9 +400,10 @@ ffi::Error tp_double_backward_impl( check_tensor(W_dgrad, {num_batch, k.weight_numel}, k.weight_dtype, "W_dgrad"); } - if (k.shared_weights) { - zero_buffer(*W_grad, stream); - } + zero_buffer(*L1_grad, stream); + zero_buffer(*L2_grad, stream); + zero_buffer(*W_grad, stream); + zero_buffer(*L3_dgrad, stream); jit_kernel->double_backward( num_batch, @@ -748,4 +758,4 @@ NB_MODULE(openequivariance_extjax, m) { .def("start", &GPUTimer::start) .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed) .def("clear_L2_cache", &GPUTimer::clear_L2_cache); -} \ No newline at end of file +} diff --git a/tests/conftest.py b/tests/conftest.py index 4a51566..595e6bb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,7 +2,6 @@ import pytest os.environ["JAX_ENABLE_X64"] = "True" -os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "False" os.environ["JAX_TRACEBACK_FILTERING"] = "off" diff --git a/tests/tp_adjoint_test.py b/tests/tp_adjoint_test.py new file mode 100644 index 0000000..81c6043 --- /dev/null +++ b/tests/tp_adjoint_test.py @@ -0,0 +1,224 @@ +import os + +import numpy as np +import pytest + + +# Regression coverage for JAX FFI backward outputs. The JAX custom-call result +# buffers are not guaranteed to be zero-initialized, while the OEQ backward +# kernels accumulate into those buffers. +CASES = { + "shared_uvu_first": { + "irreps": ( + "128x0e+128x1o+128x2e+128x3o", + "89x0e", + "128x0e+128x1o+128x2e+128x3o", + ), + "mode": "uvu", + "shared_weights": True, + }, + "shared_uvu_second": { + "irreps": ( + "128x0e+128x1o", + "89x0e", + "128x0e", + ), + "mode": "uvu", + "shared_weights": True, + }, + "unshared_uvw": { + "irreps": ( + "8x0e+8x1o", + "3x0e", + "8x0e+8x1o", + ), + "mode": "uvw", + "shared_weights": False, + }, +} + +ORDER_CASES = [ + ( + "shared-uvu-first-then-second", + ("shared_uvu_first", "shared_uvu_second"), + ), + ( + "shared-uvu-second-then-first", + ("shared_uvu_second", "shared_uvu_first"), + ), + ( + "shared-uvu-then-unshared-uvw", + ("shared_uvu_first", "unshared_uvw"), + ), + ( + "unshared-uvw-then-shared-uvu", + ("unshared_uvw", "shared_uvu_first"), + ), +] + + +@pytest.fixture(scope="module") +def ctx(with_jax): + if not with_jax: + pytest.skip("Skipping JAX tests") + + os.environ["OEQ_NOTORCH"] = "1" + + import jax + import jax.numpy as jnp + import openequivariance as oeq + + if not any(device.platform == "gpu" for device in jax.devices()): + pytest.skip("JAX GPU device is required") + + return {"jax": jax, "jnp": jnp, "oeq": oeq} + + +@pytest.fixture( + params=ORDER_CASES, + ids=lambda case: case[0], + scope="module", +) +def operator_order(request): + return request.param[1] + + +def make_problem(oeq, irreps_in1, irreps_in2, irreps_out, mode, shared_weights): + irreps1 = oeq.Irreps(irreps_in1) + irreps2 = oeq.Irreps(irreps_in2) + requested_out = oeq.Irreps(irreps_out) + + generated_out = [] + instructions = [] + for i_in1, (mul, ir_in1) in enumerate(irreps1): + for i_in2, (_, ir_in2) in enumerate(irreps2): + for ir_out in ir_in1 * ir_in2: + if ir_out in requested_out: + i_out = len(generated_out) + generated_out.append((mul, ir_out)) + instructions.append((i_in1, i_in2, i_out, mode, True)) + + generated_out = oeq.Irreps(generated_out) + generated_out, perm, _ = generated_out.sort() + instructions = [ + (i_in1, i_in2, perm[i_out], mode, train) + for i_in1, i_in2, i_out, mode, train in instructions + ] + instructions = sorted(instructions, key=lambda x: x[2]) + + return oeq.TPProblem( + irreps1, + irreps2, + generated_out, + instructions, + shared_weights=shared_weights, + internal_weights=False, + irrep_dtype=np.float32, + weight_dtype=np.float32, + ) + + +def make_tensor_product(oeq, case): + spec = CASES[case] + return oeq.jax.TensorProduct( + make_problem(oeq, *spec["irreps"], spec["mode"], spec["shared_weights"]) + ) + + +def assert_forward_adjoint(jax, jnp, tp, seed): + keys = jax.random.split(jax.random.PRNGKey(seed), 8) + batch = 17 + + x1 = jax.random.normal( + keys[0], (batch, tp.config.irreps_in1.dim), dtype=jnp.float32 + ) + species = jax.random.randint(keys[1], (batch,), 0, tp.config.irreps_in2.dim) + x2 = jax.nn.one_hot(species, tp.config.irreps_in2.dim, dtype=jnp.float32) + weight_shape = ( + (tp.weight_numel,) if tp.config.shared_weights else (batch, tp.weight_numel) + ) + weights = jax.random.normal(keys[2], weight_shape, dtype=jnp.float32) + + dx1 = jax.random.normal(keys[3], x1.shape, dtype=jnp.float32) + dspecies = jax.random.randint(keys[4], (batch,), 0, tp.config.irreps_in2.dim) + dx2 = jax.nn.one_hot(dspecies, tp.config.irreps_in2.dim, dtype=jnp.float32) + dweights = jax.random.normal(keys[5], weights.shape, dtype=jnp.float32) + cotangent = jax.random.normal(keys[6], (batch, tp.L3_dim), dtype=jnp.float32) + + def fn(a, b, c): + return tp(a, b, c) + + _, jvp_out = jax.jvp(fn, (x1, x2, weights), (dx1, dx2, dweights)) + _, vjp_fn = jax.vjp(fn, x1, x2, weights) + grad1, grad2, gradw = vjp_fn(cotangent) + + lhs = jnp.vdot(cotangent, jvp_out) + rhs = jnp.vdot(grad1, dx1) + jnp.vdot(grad2, dx2) + jnp.vdot(gradw, dweights) + err = jnp.abs(lhs - rhs) + scale = jnp.maximum(jnp.maximum(jnp.abs(lhs), jnp.abs(rhs)), jnp.array(1.0)) + + relative_error = float(np.asarray(err / scale)) + assert relative_error <= 5e-4, f"relative adjoint error={relative_error:.5f}" + + +def assert_backward_adjoint(jax, jnp, tp, seed): + keys = jax.random.split(jax.random.PRNGKey(seed), 11) + batch = 17 + + x1 = jax.random.normal( + keys[0], (batch, tp.config.irreps_in1.dim), dtype=jnp.float32 + ) + species = jax.random.randint(keys[1], (batch,), 0, tp.config.irreps_in2.dim) + x2 = jax.nn.one_hot(species, tp.config.irreps_in2.dim, dtype=jnp.float32) + weight_shape = ( + (tp.weight_numel,) if tp.config.shared_weights else (batch, tp.weight_numel) + ) + weights = jax.random.normal(keys[2], weight_shape, dtype=jnp.float32) + cotangent = jax.random.normal(keys[3], (batch, tp.L3_dim), dtype=jnp.float32) + + dx1 = jax.random.normal(keys[4], x1.shape, dtype=jnp.float32) + dspecies = jax.random.randint(keys[5], (batch,), 0, tp.config.irreps_in2.dim) + dx2 = jax.nn.one_hot(dspecies, tp.config.irreps_in2.dim, dtype=jnp.float32) + dweights = jax.random.normal(keys[6], weights.shape, dtype=jnp.float32) + dcotangent = jax.random.normal(keys[7], cotangent.shape, dtype=jnp.float32) + + cgrad1 = jax.random.normal(keys[8], x1.shape, dtype=jnp.float32) + cgrad2 = jax.random.normal(keys[9], x2.shape, dtype=jnp.float32) + cgradw = jax.random.normal(keys[10], weights.shape, dtype=jnp.float32) + + def backward_fn(a, b, c, d): + _, vjp_fn = jax.vjp(lambda x, y, w: tp(x, y, w), a, b, c) + return vjp_fn(d) + + _, jvp_out = jax.jvp( + backward_fn, + (x1, x2, weights, cotangent), + (dx1, dx2, dweights, dcotangent), + ) + _, vjp_fn = jax.vjp(backward_fn, x1, x2, weights, cotangent) + input_grads = vjp_fn((cgrad1, cgrad2, cgradw)) + + lhs = ( + jnp.vdot(cgrad1, jvp_out[0]) + + jnp.vdot(cgrad2, jvp_out[1]) + + jnp.vdot(cgradw, jvp_out[2]) + ) + rhs = ( + jnp.vdot(input_grads[0], dx1) + + jnp.vdot(input_grads[1], dx2) + + jnp.vdot(input_grads[2], dweights) + + jnp.vdot(input_grads[3], dcotangent) + ) + err = jnp.abs(lhs - rhs) + scale = jnp.maximum(jnp.maximum(jnp.abs(lhs), jnp.abs(rhs)), jnp.array(1.0)) + + relative_error = float(np.asarray(err / scale)) + assert relative_error <= 5e-4, f"relative adjoint error={relative_error:.5f}" + + +def test_tensor_product_adjoint_after_multi_operator_construction(ctx, operator_order): + jax, jnp, oeq = ctx["jax"], ctx["jnp"], ctx["oeq"] + operators = {case: make_tensor_product(oeq, case) for case in operator_order} + for i, case in enumerate(operator_order): + assert_forward_adjoint(jax, jnp, operators[case], seed=1234 + i) + assert_backward_adjoint(jax, jnp, operators[case], seed=4321 + i) From 19f50c886783d9aa4a89fa9030934865c8b9da22 Mon Sep 17 00:00:00 2001 From: Paul Fuchs Date: Thu, 9 Jul 2026 16:35:51 +0200 Subject: [PATCH 2/3] Preserve conftest.py --- tests/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/conftest.py b/tests/conftest.py index 595e6bb..4a51566 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ import pytest os.environ["JAX_ENABLE_X64"] = "True" +os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "False" os.environ["JAX_TRACEBACK_FILTERING"] = "off" From 3f9530b0e535c8e5ced14c12632429b9b5dbca82 Mon Sep 17 00:00:00 2001 From: Paul Fuchs Date: Fri, 10 Jul 2026 09:35:10 +0200 Subject: [PATCH 3/3] Reverted the kernel cache collision check --- openequivariance_extjax/src/libjax_tp_jit.cpp | 29 +++++++------------ 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/openequivariance_extjax/src/libjax_tp_jit.cpp b/openequivariance_extjax/src/libjax_tp_jit.cpp index 47a1423..8889005 100644 --- a/openequivariance_extjax/src/libjax_tp_jit.cpp +++ b/openequivariance_extjax/src/libjax_tp_jit.cpp @@ -164,13 +164,11 @@ struct KernelProp { } }; -struct CachedTP { - std::unique_ptr> impl; - KernelProp prop; - std::string json_payload; -}; - -std::unordered_map tp_cache; +std::unordered_map>, + KernelProp + >> tp_cache; std::unordered_map*, KernelProp> const std::lock_guard lock(mut); auto it = tp_cache.find(hash); if (it == tp_cache.end()) { - std::string incoming_json(json_payload); std::string err; - json root = json::parse(incoming_json, err); + json root = json::parse(std::string(json_payload), err); if (!err.empty()) throw std::runtime_error("JSON Parse Error: " + err); std::string kernel_src = root["kernel"].string_value(); @@ -205,18 +202,12 @@ std::pair*, KernelProp> backward_cfg, dbackward_cfg, kernel_prop_map); - tp_cache.emplace(hash, - CachedTP{ - std::move(jit_tp_impl), - KernelProp(kernel_prop_map, is_convolution), - std::move(incoming_json), - }); + tp_cache.insert({hash, + std::make_pair(std::move(jit_tp_impl), + KernelProp(kernel_prop_map, is_convolution))}); it = tp_cache.find(hash); - } else if (it->second.json_payload != json_payload) { - throw std::runtime_error( - "OpenEquivariance JAX TP cache collision: same hash, different kernel JSON"); } - return {it->second.impl.get(), it->second.prop}; + return {it->second.first.get(), it->second.second}; } }