Skip to content

Add a scoped_env utility for tests #4474

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zcbenz merged 1 commit into from
Sep 8, 2026
Merged

Add a scoped_env utility for tests #4474

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions python/tests/mlx_tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# Copyright © 2023 Apple Inc.

import contextlib
import os
import platform
import sys
Expand All @@ -10,6 +11,26 @@
import numpy as np


@contextlib.contextmanager
def scoped_env(**environ):
"""
Temporarily set the process environment variables.

Passing a value of None removes the variable for the duration of the context.
"""
old_environ = dict(os.environ)
for key, value in environ.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
try:
yield
finally:
os.environ.clear()
os.environ.update(old_environ)


class MLXTestRunner(unittest.TestProgram):
def __init__(self, *args, **kwargs):
# Do not exit in runTests
Expand Down
134 changes: 60 additions & 74 deletions python/tests/test_conv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1220,7 +1220,6 @@ def test_conv2d_winograd_batch_tiling(self):
# Use envs to test tiling without allocating large buffers.
tile_key = "MLX_CONV_WINOGRAD_TILE_BATCH"
ws_key = "MLX_CONV_WINOGRAD_WORKING_SET"
prev = {k: os.environ.get(k) for k in (tile_key, ws_key)}

# Winograd needs 3x3 stride-1, channels in multiples of 32,
# C + O >= 256 and N * iH * iW >= 4096.
Expand All @@ -1231,81 +1230,68 @@ def test_conv2d_winograd_batch_tiling(self):
)

def run(x, w, env={}):
for k in (tile_key, ws_key):
os.environ.pop(k, None)
os.environ.update(env)
y = mx.conv2d(x, w, padding=1)
mx.eval(y)
return np.array(y)

try:
for in_shape, wt_shape in cases:
np.random.seed(0)
x = mx.array(np.random.normal(size=in_shape).astype(np.float32))
# Small weights keep the output near unit scale.
w = mx.array(
(np.random.normal(size=wt_shape) * 0.05).astype(np.float32)
)
b = mx.zeros((wt_shape[0],))
mx.eval(x, w, b)
cpu_ref = np.array(mx.conv2d(x, w, padding=1, stream=mx.cpu))

untiled = run(x, w)
self.assertGreater(np.abs(untiled).max(), 0)
self.assertTrue(np.allclose(untiled, cpu_ref, atol=1e-3))

# Tiled winograd keeps the same per-element reduction order,
# so it is bit-identical to untiled; the implicit gemm
# fallback never is. Exact equality pins each run to its path.
# 3 divides none of the batches, so it also covers a short
# final tile.
for tile in (1, 3):
with self.subTest(in_shape=in_shape, tile=tile):
tiled = run(x, w, {tile_key: str(tile)})
self.assertTrue(np.array_equal(untiled, tiled))

# A consumer op checks the output is fenced across
# command encoders.
os.environ[tile_key] = str(tile)
fused = mx.conv2d(x, w, padding=1) + b
mx.eval(fused)
self.assertTrue(np.allclose(untiled, fused, atol=1e-4))
os.environ.pop(tile_key, None)

# Budget for ~2 batch elements so the selector itself must
# tile; mirrors the winograd_batch_step arithmetic.
n, iH, iW, C = in_shape
O = wt_shape[0]
pH = 6 * ((iH + 2 - 2 + 5) // 6) + 2
pW = 6 * ((iW + 2 - 2 + 5) // 6) + 2
per_n = (
pH * pW * C * 4
+ 64 * ((iH + 5) // 6) * ((iW + 5) // 6) * (C + O) * 4
)
used = (n * iH * iW * (C + O) + 64 * C * O) * 4
with self.subTest(in_shape=in_shape, budget="tiled"):
budget = str(int((used + 5 * per_n // 2) / 0.75))
tiled = run(x, w, {ws_key: budget})
with mlx_tests.scoped_env(**env):
y = mx.conv2d(x, w, padding=1)
mx.eval(y)
return np.array(y)

for in_shape, wt_shape in cases:
np.random.seed(0)
x = mx.array(np.random.normal(size=in_shape).astype(np.float32))
# Small weights keep the output near unit scale.
w = mx.array((np.random.normal(size=wt_shape) * 0.05).astype(np.float32))
b = mx.zeros((wt_shape[0],))
mx.eval(x, w, b)
cpu_ref = np.array(mx.conv2d(x, w, padding=1, stream=mx.cpu))

untiled = run(x, w)
self.assertGreater(np.abs(untiled).max(), 0)
self.assertTrue(np.allclose(untiled, cpu_ref, atol=1e-3))

# Tiled winograd keeps the same per-element reduction order,
# so it is bit-identical to untiled; the implicit gemm
# fallback never is. Exact equality pins each run to its path.
# 3 divides none of the batches, so it also covers a short
# final tile.
for tile in (1, 3):
with self.subTest(in_shape=in_shape, tile=tile):
tiled = run(x, w, {tile_key: str(tile)})
self.assertTrue(np.array_equal(untiled, tiled))

# Too small for even one batch element: must fall back.
with self.subTest(in_shape=in_shape, budget="infeasible"):
fallback = run(x, w, {ws_key: "1"})
self.assertFalse(np.array_equal(untiled, fallback))
self.assertTrue(np.allclose(fallback, cpu_ref, atol=1e-3))

# A forced tile is capped by the budget, so this must still
# fall back.
with self.subTest(in_shape=in_shape, budget="forced+infeasible"):
capped = run(x, w, {ws_key: "1", tile_key: "1"})
self.assertFalse(np.array_equal(untiled, capped))
self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3))
finally:
for k, v in prev.items():
if v is None:
os.environ.pop(k, None)
else:
os.environ[k] = v
# A consumer op checks the output is fenced across
# command encoders.
with mlx_tests.scoped_env(MLX_CONV_WINOGRAD_TILE_BATCH=str(tile)):
fused = mx.conv2d(x, w, padding=1) + b
mx.eval(fused)
self.assertTrue(np.allclose(untiled, fused, atol=1e-4))

# Budget for ~2 batch elements so the selector itself must
# tile; mirrors the winograd_batch_step arithmetic.
n, iH, iW, C = in_shape
O = wt_shape[0]
pH = 6 * ((iH + 2 - 2 + 5) // 6) + 2
pW = 6 * ((iW + 2 - 2 + 5) // 6) + 2
per_n = (
pH * pW * C * 4 + 64 * ((iH + 5) // 6) * ((iW + 5) // 6) * (C + O) * 4
)
used = (n * iH * iW * (C + O) + 64 * C * O) * 4
with self.subTest(in_shape=in_shape, budget="tiled"):
budget = str(int((used + 5 * per_n // 2) / 0.75))
tiled = run(x, w, {ws_key: budget})
self.assertTrue(np.array_equal(untiled, tiled))

# Too small for even one batch element: must fall back.
with self.subTest(in_shape=in_shape, budget="infeasible"):
fallback = run(x, w, {ws_key: "1"})
self.assertFalse(np.array_equal(untiled, fallback))
self.assertTrue(np.allclose(fallback, cpu_ref, atol=1e-3))

# A forced tile is capped by the budget, so this must still
# fall back.
with self.subTest(in_shape=in_shape, budget="forced+infeasible"):
capped = run(x, w, {ws_key: "1", tile_key: "1"})
self.assertFalse(np.array_equal(untiled, capped))
self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3))

def test_conv2d_large_filter_small_channels(self):
x = mx.random.normal(shape=(1, 181, 181, 1))
Expand Down
32 changes: 11 additions & 21 deletions python/tests/test_conv_transpose.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -813,37 +813,27 @@ def test_conv_transpose_unfold_tiling(self):
# The explicit-GEMM conv path unfolds into one buffer that can exceed
# maxBufferLength for large outputs; it unfolds and runs the gemm in row
# tiles instead (issue #3082).
key = "MLX_CONV_UNFOLD_TILE_ROWS"
prev = os.environ.get(key)
cases = (
(mx.conv_transpose1d, (2, 9, 4), (5, 3, 4), {"stride": 2}),
(mx.conv_transpose2d, (2, 5, 5, 4), (5, 3, 3, 4), {"stride": 2}),
(mx.conv_transpose3d, (1, 4, 4, 4, 2), (3, 2, 2, 2, 2), {"stride": 2}),
(mx.conv_transpose1d, (2, 9, 4), (6, 3, 2), {"stride": 2, "groups": 2}),
)
try:
# tile_rows=1 forces uniform tiles; 7 does not divide any output, so
# it exercises a partial final tile too.
for conv, in_shape, wt_shape, kwargs in cases:
for tile_rows in (1, 7):
with self.subTest(
conv=conv.__name__, tile_rows=tile_rows, **kwargs
):
np.random.seed(0)
x = mx.array(np.random.normal(size=in_shape).astype(np.float32))
w = mx.array(np.random.normal(size=wt_shape).astype(np.float32))
os.environ.pop(key, None)
# tile_rows=1 forces uniform tiles; 7 does not divide any output, so
# it exercises a partial final tile too.
for conv, in_shape, wt_shape, kwargs in cases:
for tile_rows in (1, 7):
with self.subTest(conv=conv.__name__, tile_rows=tile_rows, **kwargs):
np.random.seed(0)
x = mx.array(np.random.normal(size=in_shape).astype(np.float32))
w = mx.array(np.random.normal(size=wt_shape).astype(np.float32))
with mlx_tests.scoped_env(MLX_CONV_UNFOLD_TILE_ROWS=None):
untiled = conv(x, w, **kwargs)
mx.eval(untiled)
os.environ[key] = str(tile_rows)
with mlx_tests.scoped_env(MLX_CONV_UNFOLD_TILE_ROWS=str(tile_rows)):
tiled = conv(x, w, **kwargs)
mx.eval(tiled)
self.assertTrue(np.allclose(untiled, tiled, atol=1e-4))
finally:
if prev is None:
os.environ.pop(key, None)
else:
os.environ[key] = prev
self.assertTrue(np.allclose(untiled, tiled, atol=1e-4))


if __name__ == "__main__":
Expand Down
Loading
Loading