From b15fb893317c6f497abbff632f3232b257b23150 Mon Sep 17 00:00:00 2001 From: Dangyi Liu Date: Tue, 7 Oct 2025 23:36:01 -0700 Subject: [PATCH] Mark `ragged_dot_qt_fwd` with `disable_interceptions` The same reason as dot_general_qt_fwd. PiperOrigin-RevId: 816555366 --- qwix/_src/core/ragged_dot_qt.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/qwix/_src/core/ragged_dot_qt.py b/qwix/_src/core/ragged_dot_qt.py index 98c06f6..4192086 100644 --- a/qwix/_src/core/ragged_dot_qt.py +++ b/qwix/_src/core/ragged_dot_qt.py @@ -18,6 +18,7 @@ import functools import jax from jax import numpy as jnp +from qwix._src import interception from qwix._src.core import qarray from qwix._src.core import ragged_dot @@ -90,6 +91,7 @@ def _ragged_dot_general( return out.astype(result_type) +@interception.disable_interceptions def ragged_dot_qt_fwd( lhs: jax.Array, rhs: jax.Array,