From d6a3d7749ed9d500f0927dd0e747f42a13ecff90 Mon Sep 17 00:00:00 2001 From: RoomWithOutRoof Date: Wed, 15 Apr 2026 20:22:51 +0800 Subject: [PATCH 1/3] Fix NaN issue in mean_kernel for empty dimension Handle edge case when reduction dimension size is 0 to avoid division by zero leading to NaN values. This fix adds explicit checks in both mean_kernel (Triton kernel) and mean_dim (Python wrapper) to properly handle empty tensors. --- batch_invariant_ops/batch_invariant_ops.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/batch_invariant_ops/batch_invariant_ops.py b/batch_invariant_ops/batch_invariant_ops.py index b9021bb..793eced 100644 --- a/batch_invariant_ops/batch_invariant_ops.py +++ b/batch_invariant_ops/batch_invariant_ops.py @@ -340,6 +340,18 @@ def mean_kernel( Kernel for computing mean along a single dimension. Input is viewed as (M, N, K) where N is the dimension being reduced. """ + # Handle edge case: empty reduction dimension returns NaN + if N == 0: + pid = tl.program_id(0) + m_idx = pid // K + k_idx = pid % K + if m_idx >= M or k_idx >= K: + return + mean_val = float("nan") + output_idx = m_idx * output_stride0 + k_idx * output_stride1 + tl.store(output_ptr + output_idx, mean_val) + return + # Program ID gives us which output element we're computing pid = tl.program_id(0) @@ -390,6 +402,16 @@ def mean_dim( assert -input.ndim <= dim < input.ndim, ( f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" ) + + # Handle empty dimension case explicitly to avoid NaN + if input.shape[dim] == 0: + # Return a tensor of NaN values with the appropriate shape + shape = list(input.shape) + if keepdim: + shape[dim] = 1 + else: + shape.pop(dim) + return torch.full(shape, float("nan"), dtype=dtype or input.dtype, device=input.device) # Handle negative dim if dim < 0: From 20de7fff1b9f22e41953deb208298ef8dbc5bc1d Mon Sep 17 00:00:00 2001 From: RoomWithRoof Date: Wed, 15 Apr 2026 22:15:27 +0800 Subject: [PATCH 2/3] Fix: Replace tl.range flatten with Python range for Triton 3.x compatibility The flatten parameter was removed in Triton 3.x public release. This changes tl.range(..., flatten=True) to Python's built-in range(), which works inside @triton.jit functions and is compatible with all versions. Fixes: TypeError: range.__init__() got an unexpected keyword argument 'flatten' --- batch_invariant_ops/batch_invariant_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/batch_invariant_ops/batch_invariant_ops.py b/batch_invariant_ops/batch_invariant_ops.py index 793eced..2741ab6 100644 --- a/batch_invariant_ops/batch_invariant_ops.py +++ b/batch_invariant_ops/batch_invariant_ops.py @@ -75,7 +75,7 @@ def matmul_kernel_persistent( offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) num_pid_in_group = GROUP_SIZE_M * num_pid_n - for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): + for tile_id in range(start_pid, num_tiles, NUM_SMS): pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) start_m = pid_m * BLOCK_SIZE_M start_n = pid_n * BLOCK_SIZE_N From e58b36a4f3807e1e94c718088a6cf672a96e0cf7 Mon Sep 17 00:00:00 2001 From: RoomWithOutRoof Date: Wed, 15 Apr 2026 22:31:25 +0800 Subject: [PATCH 3/3] Add test suite for batch-invariant operations - Add pytest tests for mm, addmm, log_softmax, mean operations - Test batch-invariant property: op(x[:1], y) == op(x, y)[:1] - Include tests for multiple dtypes (float32, float16, bfloat16) - Add tests for batch invariant mode context manager --- .../conftest.cpython-312-pytest-7.4.4.pyc | Bin 0 -> 637 bytes ...invariant_ops.cpython-312-pytest-7.4.4.pyc | Bin 0 -> 15610 bytes tests/conftest.py | 9 + tests/test_batch_invariant_ops.py | 255 ++++++++++++++++++ 4 files changed, 264 insertions(+) create mode 100644 tests/__pycache__/conftest.cpython-312-pytest-7.4.4.pyc create mode 100644 tests/__pycache__/test_batch_invariant_ops.cpython-312-pytest-7.4.4.pyc create mode 100644 tests/conftest.py create mode 100644 tests/test_batch_invariant_ops.py diff --git a/tests/__pycache__/conftest.cpython-312-pytest-7.4.4.pyc b/tests/__pycache__/conftest.cpython-312-pytest-7.4.4.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e695439cda71e47c047765a2902e7384835fec0e GIT binary patch literal 637 zcmZvZv5V9|6vk(gm}Ena9C(NdLW;Fm&Pxih6S1(c5V<KYTU!Cp>(=`p^sA0P&8A=O zKXLw1i+7*^87N?3l`ylG(XO1_f}pyTxo}nSGEZ0R(ys>D00Ik;0SS)b%4c?^H{Qhm z<+6}?v+J9lUd?XwY7+~n84(G4#<8`kLC{$v=ipAT-$^25oU?jDpR^oH%IbK+cs#~3 zpYGB68Rit%GG>j4O{++PU&iC{2r6eE%jsxac?cCOM@aeOmX?y%LWTF64y||BP_yJD zSO~%mBjIGusU&K!PARX?1^TdIAmpj_FpezJ_LL1Ia&tb93j5N9mMRByGoSE!#9<4XN`I&(k^MPep@svL`V;py_gW zv>qPm;nU^tsrgYxB~5qX->q2DmgF8cFK1p+JkX!tHiYoZvY_+F1rWXluY&L3>X{3+ vcNP!dyf}@9C(-aU+C7POUxwd+edHhdN6}%lfNyUvo`1eP`~q(B9c}a%v8k;x literal 0 HcmV?d00001 diff --git a/tests/__pycache__/test_batch_invariant_ops.cpython-312-pytest-7.4.4.pyc b/tests/__pycache__/test_batch_invariant_ops.cpython-312-pytest-7.4.4.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fce3f22b4aa585973bf3142075676c20839cdbad GIT binary patch literal 15610 zcmeHOYit`=cAgmy$st8awk0{XY$xSrv_4>E(-suv}!BB zDB7NL?+iI2M^C$(P1hsw@Xnq4nsd)R=ey?~{(W_|pM$IEPriphG!qLd z&vCaol?!t!uL?0fEYNRPT#N`j7k7s}JbU5|d+8G?EWx)c=8OBo{&-cmDjo<2;??2m zculwl%8EC*aIL^;wRg-e!ke^Bs{6{z97pzQ(jbha{DRzl<#`+jxy_Zi{jNm7xs`rC zTrY5EIMs85Q@wWtRyW*0^CZag74x>zJU`@B74x2>c>&0)F6M2cc{Pw%Tg=<8HL9Dg z9DuihmNoC-;mR>;JJrqo!X4l&UObl-vO;K!@)czGFKT)^sAr;SE!dYN!CoaD8EAt$@lBLO>1Be+L{JIpkR3C74Qeqht^q^!fn+A82K%GeG|DkuiEBZk>F~OK z#ok^Nn4h3U`ZBQ?<1BQdVix2pSFlsc%qW2=G}~|WKf*cg3Av4Jx|Wvde9Fc7l;cTN zGd2T{fxj|LjULt0Ff5~ia<{y=QDtMLA>k{_YBV0fF09K92^V}85Jm+o}WKgI?Q>nY30YuIoG4CR_zB~HZ$2wP+I5X1@O=TKI3(uwp3#`(MKXjS zf0alCSOGE%!mMi7q7kjt5YyR|X1F4$Oh_=qNJdo**AU!@hPPjXsZSHb4do*PMwP9) zQD?R(+cHDEQN@~-u~8l0Hu(5^Z3D@;)|TnbB+{9-R*?BbB&J+%yPhOh^pq0OTGC21 z7i}w_jbuu1!}Y{yIrn9F+?vW7ei)?8Iwf_`-%j}GM}Z7;54UXpNz-iO!KucBvyI26 z8jnvmp8WFg&$_3#oF6`W-@D~A?JMt&`?dA=ef6WYQ@;K8Ya1t}#`{&97N9yuP#-Uu zOlzVR9>9~^oW_NDP{p8v!!Au6r0N?hmaf6-Qrb$Xf>u;_s_O=4O>0=%U>l&qi*5Ke z*oNEIhCk~Gd5wBheWy_0y=|$xbLVo0F_&qoWwlMnT#eSULK$>~u*2$o~j>*Hkkay+9JfF5Egu5%;!IcJA!Lv$TG4C#`X|_`E z+G=`OYj&FVs4i7(Z9)Uo4q_Fw;N zV*aahz3(2M0EGa$B4&sL=!A#VU|BZ-MVl}@Z=SyN_8X@SAs#opxK2Uk8UaPu!KBHN zSW?#}{D!D&u|C!Xf6Cy`8T?s;?=koxgNLP?%%o*z2Mh_*sLOR+M`&f(MYdp$cr~<<={s zCr7V-{`%OSv5v{QrV;6(RQ+l7&FWcc&y=)h?8x}gU#LH;pOp6eTDVU0AdTU60gp2XfFcN@(9C-y&WEYSg@&YDd z>l0MkcfUc2Qhu{OiY2T->?!tK*f|wrHP11@_I0uFeMl8 zq?XWOhXFFmBoTIIOm!2Q+JKDiSxzys0mv|b+{-wuHEYyT#kZGA>+Ww_T6;oOh9t`g z*hk2+;g=EOXJVKR$nw>U5;LE8Wm!!|AlsA5B6|!mu8=E6RZ1aH74}>?4GiAJ_uoBz zfpp;v96^F>bAqGVo7yR|6?2hL1u+DP8{99%6aVC{`sDIQmuKBOr`$V7zdP19?QWSD z0`7PC1sCUOoukRQot!7&P|RN{hRL~Ys}!3TuoxQ@4#gS5j>3~)PY9cJhdjn}mGXB| zRRMv%fe#o0UE`Lw!UnWbJmPF_6TgmN)CmHYgIH@rxEdstY%tcU<68yBLLC6IW@!`1 z`ZoY%eR;nlkZs2~g=tzA$g-I(VOg>rq}S0rn;4ZeVF7eV8E#scfD4EtAs`b~qag!lm%dpec3v$h4DX(CI77$T9<{E7S~Q zr+*CO2{7t|+WJpFy!qj12Su&5ua1ZewO+n;d8~fy>iCxN-D8dG$I%>5s(C4%lvf;O0(1+spDWhiwc0!wQSVy8a_ z@&s75QeY9iI*BXd*A0nlPsmS>A(LK5avaGCBqx!aLh=TZH<6q}@(z;INN}UZO#K;5 zA^9$R^z%T_)Sq)U_0=UCj>);_oo#rPRcv0sBbb}yx! zbA@Rp__{X3ER`8%gTG+#mzX^+8DZS0TSk~!;N~>I2iZ2vW$ZA8j!XA6C0Sy|c-2Po z(7Ub!p&G7k=S?qj8#z0TsHvT8Xr5|lwoL84Ju?m6)YQ(_?VGCG2d4JaOkLYO;oj!S zxQUQOynyS$qc?Ad?>y8*g7 zNpKTtfr>>>J)oU$0z6CTNdi6|feP%S6Tm8)tdfg({~aVjB-F=)Dm~;a6Gwt@tXvw= zi)~6H6UM!(3S+05JAotKxFmL}zX#pwaUfI-H#OW8Q4DQDSoBeX@pKeGFrH6A0P*#y z0RG0S(U+Ws(C%*3wW}F8`a*xBi#c!5%NXTg&8|ft$$L2I8&Q7@)jgs5qmBF#4&f>g zrv9E^S^YhA*#mFRnt|J87$x}t$q$f}p*O^91W6V?y5*?pa7J%MST;vt*<7PDdh=Q6 zZBD@C+}2f!&C_BGumiHE!^{5`sqV0LSnOS7p5+$=F({QqRGGeC|ce zH%~0#zH}f~VK7#53wqoEePcsOkf-e5<8^n`k;X+=3O4+A0P3=aPc zhN~xmtR2_g4&Mro@p!y9<{s~uX=q^%lq0u}z;WRDnY#Vs!Z={MgX_d}&rB*u&$oTc zo#0_07L_(7mV_W5*kc6SE9DJ0Oy#gn#k^&#=<~}97WB}b>&6AJF1f;LU!YhOaoTFz z1$1;&3KhCJo|;hoHw@@kzTIUT0K|!Ktwoj54TD+!36+`ChM~h$W}rXdq>TmhMOcI% z6)Xzcu4q~crtsKBlF=%}>muwXF2TzfUbC2O^RSYX?OtBW%!YCf&bL(*m|62n7>^B` zt+F+(khOd2OvBL;zfIcii8FQG_j>PLoveHHTcyNorf@CNRx-0sC~HoJlj4sGtid8O zDWoQ)9jdu2_&!WzMF7H5`$<+3vK4G3+X<}SM$&r!fC5%hX1i7e|4&bu^@IXz2L64> zz5ygxkl=vJ(SI-IA^B7I=syL5zT2)}Df-{QCILshWRn0}19~l>18(VQt|~xeH+XA+ zy0an#ygC772zO=3jsawfsEQgf9mH^0h~Z77wIR4s5()&(v_`B>v1bgrLp{Cl1E6SE zASj)S^94RG6oj~|;NoF-o98X8UJ&zQfGfE3Zas@@E$zcb@N$Ztv}#@B^oDdbD+eVp z2CeLZ9?Ve@dr1s+E4xZ16Fl@^YSpQ%Xtk=ke*cg>vfS{$*ykuwDmxZ+x#E6=0`qH?7iLoQlE>5Z|d zdKK`mZr(X9-&TS^IOH;3i0TzPA6(xuZA?>*9W3Ek$k@qx(Hx&X4zZxmU?Q0gHXTT1 zJ43qvLo*RPMzh*O`hk$YK(4}QAga`r%C;lv?0MW|R$jRxgz?nT_p@#IAOM`UM+h`Z zbRb6HL@F!yW}@)xh=guq?)7 z*+$Jt8pR4>SEmU{5~EIrO$S5^%GyvW26_;G@)0tREAWkMBdEd=uvh9ZY&VD44Qr>% zesjpdr+0u3xBP`|^94UNz-iZV#}Uh(S{kPf)5`X{YiQsB3}63iAj90EDEj>K9FWBt z&euIFb*<}vQtF!baK3#rQq$K``|NuDr=@nR+B73If4$CLVTe;w7j>d7u}pY)K`m3h z%rcc}Tp=%jZgJ&ZrbR-FMz29^q3p8s6SYVy+OCqF0gF^B87ONIZc8PpMOwLR@n=d! zizKFNtu}4_tZ2>bdKWAb9v2m@8vGQ?w4~jz4PRorfm6UE32f81jYsY^^T=(zz4c9P zlb`~!SBr+mzg_d#-a?hig6a=p)8=R>|+E9$6nAW6^XsDdg( zoTka*nQq5(LYB%vs^=m>Up>Tcpx>V5=P>ml5ep`5XlW7I<9wil*dl;Ld)x56BElOCU!c{JGZhDE^Fc6Z~%V?WSILG>C7VN!i}h3 zIBu?K9#f~Q56*MEulwsj*X(-!rvqJ!E~vLCZJ>Kl!izqrJcm6$))83$&5l#n*a z+v$GmsJ93B?-Xj2eD~lyC z!Z)r&6SABWTCRn5(P$n+vTy)F2{E?I3 zQs!XPaWsLi(-*}cFFz`{@nR#f$tQdV_Q^#_!>~ u=Q$vAzr}N6U?lgIYx{k%cKEGN&VF?Evu&eYlbeD+7kAIOII)pV;r{~B1pvta literal 0 HcmV?d00001 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..73bfe13 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,9 @@ +""" +Pytest configuration for batch-invariant ops tests. +""" +import sys +from pathlib import Path + +# Add the project root to the path +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) diff --git a/tests/test_batch_invariant_ops.py b/tests/test_batch_invariant_ops.py new file mode 100644 index 0000000..917033f --- /dev/null +++ b/tests/test_batch_invariant_ops.py @@ -0,0 +1,255 @@ +""" +Test suite for batch-invariant operations. + +This test suite verifies that the batch-invariant property holds for the following operations: +- mm (matrix multiplication) +- addmm (matrix multiplication with bias) +- log_softmax +- mean + +The batch-invariant property states that: + op(x[:1], y) == op(x, y)[:1] + +This means computing an operation on a single batch element should give the same result +as computing it on the full batch and then taking the first element. +""" + +import pytest +import torch +from batch_invariant_ops import ( + set_batch_invariant_mode, + matmul_persistent, + addmm_batch_invariant, + log_softmax, + mean_dim, +) + + +def get_device(): + """Get the current accelerator device.""" + device_type = getattr(torch.accelerator.current_accelerator(), "type", "cpu") + if device_type == "cpu": + return "cpu" + elif device_type in ("cuda", "xpu"): + return device_type + return "cpu" + + +DEVICE = get_device() +DTYPES = [torch.float32, torch.float16, torch.bfloat16] if DEVICE != "cpu" else [torch.float32] + + +class TestBatchInvariantMM: + """Tests for mm (matrix multiplication) batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mm_batch_invariant_basic(self, dtype): + """Test basic batch invariance for mm operation.""" + B, M, K, N = 4, 16, 32, 16 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.mm(x[:1], y) + # Full batch, take first + out_full = torch.mm(x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mm_batch_invariant_large(self, dtype): + """Test batch invariance with larger matrices.""" + B, M, K, N = 8, 128, 256, 128 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + + out_single = torch.mm(x[:1], y) + out_full = torch.mm(x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mm_batch_invariant_single_row(self, dtype): + """Test with single row in batch dimension.""" + B, M, K, N = 1, 8, 16, 8 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + + out_single = torch.mm(x[:1], y) + out_full = torch.mm(x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantAddMM: + """Tests for addmm (matrix multiplication with bias) batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_addmm_batch_invariant_basic(self, dtype): + """Test basic batch invariance for addmm operation.""" + B, M, K, N = 4, 16, 32, 16 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + bias = torch.randn(N, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.addmm(bias, x[:1], y) + # Full batch, take first + out_full = torch.addmm(bias, x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_addmm_batch_invariant_large(self, dtype): + """Test batch invariance with larger matrices.""" + B, M, K, N = 8, 128, 256, 128 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + bias = torch.randn(N, dtype=dtype, device=DEVICE) + + out_single = torch.addmm(bias, x[:1], y) + out_full = torch.addmm(bias, x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantLogSoftmax: + """Tests for log_softmax batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_log_softmax_batch_invariant_basic(self, dtype): + """Test basic batch invariance for log_softmax operation.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.log_softmax(x[:1], dim=-1) + # Full batch, take first + out_full = torch.log_softmax(x, dim=-1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_log_softmax_batch_invariant_2d(self, dtype): + """Test batch invariance for 2D tensors.""" + B, V = 4, 64 + + x = torch.randn(B, V, dtype=dtype, device=DEVICE) + + out_single = torch.log_softmax(x[:1], dim=-1) + out_full = torch.log_softmax(x, dim=-1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_log_softmax_batch_invariant_large(self, dtype): + """Test batch invariance with larger sequences.""" + B, S, V = 8, 128, 512 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.log_softmax(x[:1], dim=-1) + out_full = torch.log_softmax(x, dim=-1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantMean: + """Tests for mean operation batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_basic(self, dtype): + """Test basic batch invariance for mean operation.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.mean(x[:1], dim=1) + # Full batch, take first + out_full = torch.mean(x, dim=1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_dim0(self, dtype): + """Test batch invariance for mean along dim 0.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.mean(x[:1], dim=0) + out_full = torch.mean(x, dim=0)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_keepdim(self, dtype): + """Test batch invariance for mean with keepdim=True.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.mean(x[:1], dim=1, keepdim=True) + out_full = torch.mean(x, dim=1, keepdim=True)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_large(self, dtype): + """Test batch invariance with larger tensors.""" + B, S, V = 8, 64, 128 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.mean(x[:1], dim=1) + out_full = torch.mean(x, dim=1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantMode: + """Tests for batch invariant mode context manager.""" + + def test_batch_invariant_mode_enable_disable(self): + """Test enabling and disabling batch invariant mode.""" + from batch_invariant_ops import ( + is_batch_invariant_mode_enabled, + enable_batch_invariant_mode, + disable_batch_invariant_mode, + ) + + # Initially disabled + assert not is_batch_invariant_mode_enabled() + + # Enable + enable_batch_invariant_mode() + assert is_batch_invariant_mode_enabled() + + # Disable + disable_batch_invariant_mode() + assert not is_batch_invariant_mode_enabled() + + def test_batch_invariant_mode_context_manager(self): + """Test context manager for batch invariant mode.""" + from batch_invariant_ops import is_batch_invariant_mode_enabled + + # Initially disabled + assert not is_batch_invariant_mode_enabled() + + # Enable via context manager + with set_batch_invariant_mode(True): + assert is_batch_invariant_mode_enabled() + + # Should be disabled after context + assert not is_batch_invariant_mode_enabled() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])