From 8f6b4d17566a13f034615b259e984111bcc8e77e Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Fri, 7 Aug 2026 17:05:22 -0700 Subject: [PATCH 1/4] fix: replace `torch.tensor` copy-constructs with `as_tensor` clears up some warnings, and removes the double copy if you do torch.tensor(x).clone() and stops the autograd graph from detaching with torch.tensor() Signed-off-by: Sricharan Reddy Varra --- tests/test_stokes.py | 47 ++++++++++++++++++++++++++-- waveorder/optim/losses.py | 4 +-- waveorder/stokes.py | 10 +++--- waveorder/waveorder_reconstructor.py | 14 ++++----- waveorder/waveorder_simulator.py | 2 +- 5 files changed, 59 insertions(+), 18 deletions(-) diff --git a/tests/test_stokes.py b/tests/test_stokes.py index 1792e5c2..036e966e 100644 --- a/tests/test_stokes.py +++ b/tests/test_stokes.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np import pytest import torch @@ -53,7 +55,7 @@ def test_stokes_recon(device): s012 = stokes.stokes012_after_ar(*ar) ar1 = stokes.estimate_ar_from_stokes012(*s012) for i in range(3): - tt.assert_close(torch.tensor(ar[i]), ar1[i]) + tt.assert_close(torch.as_tensor(ar[i]), ar1[i]) # Test attenuating depolarizing retarder (adr) functions for depolarization in torch.arange(1e-3, 1, 0.1, device=device): @@ -67,7 +69,7 @@ def test_stokes_recon(device): adr1 = stokes.estimate_adr_from_stokes(*s0123) for i in range(4): - tt.assert_close(torch.tensor(adr[i]), adr1[i]) + tt.assert_close(torch.as_tensor(adr[i]), adr1[i]) def test_stokes_after_adr_usage(): @@ -130,6 +132,47 @@ def test_copying(device): assert a[0] == 1 +@pytest.mark.parametrize(*_DEVICE) +def test_estimate_copying(device): + s = torch.tensor([1.0, 1.0], device=device) + + _, _, transmittance, _ = stokes.estimate_adr_from_stokes(s, s, s, s) + transmittance[0] = 2 # modify the output + assert s[0] == 1 # check that the input hasn't changed + + _, _, transmittance012 = stokes.estimate_ar_from_stokes012(s, s, s) + transmittance012[0] = 2 + assert s[0] == 1 + + +def test_no_copy_construct_warning(): + t = torch.ones((2, 2)) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + stokes.stokes_after_adr(t, t, t, t) + stokes.stokes012_after_ar(t, t, t) + stokes.estimate_adr_from_stokes(t, t, t, t) + stokes.estimate_ar_from_stokes012(t, t, t) + stokes.mueller_from_stokes(t, t, t, t, direction="forward") + stokes.mueller_from_stokes(t, t, t, t, direction="inverse") + assert [w for w in caught if "copy construct" in str(w.message)] == [] + + +def test_gradients_reach_copied_outputs(): + """s0 and transmittance are copies of an input, so they must not detach.""" + ones = torch.ones((2, 2)) + + transmittance = torch.ones((2, 2), requires_grad=True) + s0, _, _, _ = stokes.stokes_after_adr(ones, ones, transmittance, ones) + s0.sum().backward() + assert transmittance.grad is not None + + s0_input = torch.ones((2, 2), requires_grad=True) + _, _, estimated, _ = stokes.estimate_adr_from_stokes(s0_input, ones, ones, ones) + estimated.sum().backward() + assert s0_input.grad is not None + + @pytest.mark.parametrize(*_DEVICE) def test_orientation_offset(device): ori = torch.tensor( diff --git a/waveorder/optim/losses.py b/waveorder/optim/losses.py index 57650996..14731cf0 100644 --- a/waveorder/optim/losses.py +++ b/waveorder/optim/losses.py @@ -243,8 +243,6 @@ def loss_fn(recon: Tensor) -> Tensor: def _make_spectral_flatness_loss(NA_det, wavelength, pixel_size, midband_fractions): - import numpy as np - from waveorder import util def _flatness_2d(img: Tensor) -> Tensor: @@ -252,7 +250,7 @@ def _flatness_2d(img: Tensor) -> Tensor: device = img.device _, _, fxx, fyy = util.gen_coordinate((Y, X), pixel_size) - frr = torch.tensor(np.sqrt(fxx**2 + fyy**2), device=device) + frr = torch.sqrt(fxx**2 + fyy**2).to(device) cutoff = 2 * NA_det / wavelength mask = torch.logical_and( frr > cutoff * midband_fractions[0], diff --git a/waveorder/stokes.py b/waveorder/stokes.py index bd276c18..ac57ff7f 100644 --- a/waveorder/stokes.py +++ b/waveorder/stokes.py @@ -170,7 +170,7 @@ def stokes_after_adr(retardance, orientation, transmittance, depolarization, inp raise NotImplementedError("input != cpl") # without copying transmittance, downstream changes to s0 will affect transmittance - s0 = torch.tensor(transmittance).clone() + s0 = torch.as_tensor(transmittance).clone() s1 = transmittance * depolarization * torch.sin(retardance) * torch.sin(2 * orientation) s2 = transmittance * depolarization * -torch.sin(retardance) * torch.cos(2 * orientation) s3 = transmittance * depolarization * torch.cos(retardance) @@ -208,7 +208,7 @@ def stokes012_after_ar(retardance, orientation, transmittance, input="cpl"): raise NotImplementedError("input != cpl") # without copying transmittance, downstream changes to s0 will affect transmittance - s0 = torch.tensor(transmittance).clone() + s0 = torch.as_tensor(transmittance).clone() s1 = transmittance * torch.sin(retardance) * torch.sin(2 * orientation) s2 = transmittance * -torch.sin(retardance) * torch.cos(2 * orientation) return s0, s1, s2 @@ -273,7 +273,7 @@ def estimate_adr_from_stokes(s0, s1, s2, s3, input="cpl"): retardance = torch.arcsin(((s1**2 + s2**2) ** 0.5) / len_pol) orientation = _s12_to_orientation(s1, s2) # without copying s0, downstream changes to transmittance will affect s0 - transmittance = torch.tensor(s0).clone() + transmittance = torch.as_tensor(s0).clone() depolarization = len_pol / s0 return retardance, orientation, transmittance, depolarization @@ -306,7 +306,7 @@ def estimate_ar_from_stokes012(s0, s1, s2, input="cpl"): retardance = torch.arcsin(((s1**2 + s2**2) ** 0.5) / s0) orientation = _s12_to_orientation(s1, s2) # without copying s0, downstream changes to transmittance will affect s0 - transmittance = torch.tensor(s0).clone() + transmittance = torch.as_tensor(s0).clone() return retardance, orientation, transmittance @@ -355,7 +355,7 @@ def mueller_from_stokes( raise NotImplementedError("direction must be `forward` or `inverse`") if direction == "forward": - M = torch.zeros((4, 4) + torch.tensor(s0).shape, device=s0.device) + M = torch.zeros((4, 4) + s0.shape, device=s0.device) denom = s1**2 + s2**2 M[0, 0] = s0 M[1, 1] = (s0 * s2**2 + s1**2 * s3) / denom diff --git a/waveorder/waveorder_reconstructor.py b/waveorder/waveorder_reconstructor.py index b5c89b33..3a18ea62 100644 --- a/waveorder/waveorder_reconstructor.py +++ b/waveorder/waveorder_reconstructor.py @@ -540,7 +540,7 @@ def Hz_det_setup(self, phase_deconv, ph_deconv_layer, bire_in_plane_deconv, inc_ # generate defocus kernel based on Pupil function and z_defocus self.Hz_det_2D = ( generate_propagation_kernel( - torch.tensor(self.frr), + torch.as_tensor(self.frr), torch.tensor(self.Pupil_support), self.lambda_illu, torch.tensor(self.z_defocus), @@ -580,7 +580,7 @@ def Hz_det_setup(self, phase_deconv, ph_deconv_layer, bire_in_plane_deconv, inc_ z = ifftshift((np.r_[0 : self.N_defocus_3D] - self.N_defocus_3D // 2) * self.psz) self.Hz_det_3D = ( generate_propagation_kernel( - torch.tensor(self.frr), + torch.as_tensor(self.frr), torch.tensor(self.Pupil_support), self.lambda_illu, torch.tensor(z), @@ -590,7 +590,7 @@ def Hz_det_setup(self, phase_deconv, ph_deconv_layer, bire_in_plane_deconv, inc_ ) self.G_fun_z_3D = ( generate_greens_function_z( - torch.tensor(self.frr), + torch.as_tensor(self.frr), torch.tensor(self.Pupil_support), self.lambda_illu, torch.tensor(z), @@ -772,7 +772,7 @@ def gen_WOTF(self): if self.N_pattern == 1: for i in range(self.N_defocus): Hu_temp, Hp_temp = compute_weak_object_transfer_function_2d( - torch.tensor(self.Source), + torch.as_tensor(self.Source), torch.tensor(self.Pupil_obj * self.Hz_det_2D[:, :, i]), ) self.Hu[:, :, i] = Hu_temp.numpy() @@ -781,7 +781,7 @@ def gen_WOTF(self): for i, j in itertools.product(range(self.N_defocus), range(self.N_pattern)): idx = i * self.N_pattern + j Hu_temp, Hp_temp = compute_weak_object_transfer_function_2d( - torch.tensor(self.Source[j]), + torch.as_tensor(self.Source[j]), torch.tensor(self.Pupil_obj * self.Hz_det_2D[idx, :, :]), ) self.Hu[:, :, idx] = Hu_temp.numpy() @@ -903,7 +903,7 @@ def gen_2D_vec_WOTF(self, inc_option=False): # generate dyadic Green's tensor G_fun_z = ( generate_greens_function_z( - torch.tensor(self.frr), + torch.as_tensor(self.frr), torch.tensor(self.Pupil_support), self.lambda_illu, torch.tensor(self.z_defocus), @@ -1241,7 +1241,7 @@ def gen_3D_vec_WOTF(self, inc_option): z = ifftshift((np.r_[0:N_defocus] - N_defocus // 2) * psz) G_fun_z = ( generate_greens_function_z( - torch.tensor(self.frr), + torch.as_tensor(self.frr), torch.tensor(self.Pupil_support), self.lambda_illu, torch.tensor(z), diff --git a/waveorder/waveorder_simulator.py b/waveorder/waveorder_simulator.py index 5d99034c..712cedec 100644 --- a/waveorder/waveorder_simulator.py +++ b/waveorder/waveorder_simulator.py @@ -109,7 +109,7 @@ def __init__( self.Hz_det = ( generate_propagation_kernel( - torch.tensor(self.frr), + torch.as_tensor(self.frr), torch.tensor(self.Pupil_support), self.lambda_illu, torch.tensor(self.z_defocus), From 7d6dde07a067c331a82d90fdbaf2cff4573989a6 Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Tue, 11 Aug 2026 23:41:03 +0000 Subject: [PATCH 2/4] fix(util): stack meshgrid arrays before tensor convert torch.tensor(np.meshgrid(...)) get list of 2 ndarrays. torch walk it as generic python sequence, one scalar read+store per element. warn "Creating a tensor from a list of numpy.ndarrays is extremely slow". np.stack first = 2 bulk memcpy, then as_tensor wrap zero-copy. 256x256: 10104 us -> 88 us 2048x2048: 682144 us -> 60850 us output bit-identical, dtype stay int64. last remaining copy warning in repo, so all paths now clean. Signed-off-by: Sricharan Reddy Varra --- waveorder/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/waveorder/util.py b/waveorder/util.py index 42a22902..df97d7c4 100644 --- a/waveorder/util.py +++ b/waveorder/util.py @@ -114,7 +114,7 @@ def generate_star_target(yx_shape, blur_px=2, margin=60): x = np.arange(X) - X // 2 y = np.arange(Y) - Y // 2 - xx, yy = torch.tensor(np.meshgrid(x, y)) + xx, yy = torch.as_tensor(np.stack(np.meshgrid(x, y))) rho = torch.sqrt(xx**2 + yy**2) theta = torch.arctan2(yy, xx) From 1ef2278106b22fbae20a9971477ec71e176d0603 Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Wed, 12 Aug 2026 00:25:03 +0000 Subject: [PATCH 3/4] test(stokes): drop no-copy-construct-warning test assert on torch warning substring. couple suite to vendor message text, which already drift across versions ("clone().detach()" in torch 2.0 vs "detach().clone()" in 2.10). also redundant. only way to bring warning back is torch.tensor(tensor), which detach autograd too, so test_gradients_reach_copied_outputs catch same regression. keep test_estimate_copying (copy contract, sibling of existing test_copying) and test_gradients_reach_copied_outputs (catch silent detach; warning itself invisible in prod since cli/main.py set PYTHONWARNINGS=ignore::UserWarning). Signed-off-by: Sricharan Reddy Varra --- tests/test_stokes.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/tests/test_stokes.py b/tests/test_stokes.py index 036e966e..d3cf86b1 100644 --- a/tests/test_stokes.py +++ b/tests/test_stokes.py @@ -1,5 +1,3 @@ -import warnings - import numpy as np import pytest import torch @@ -145,19 +143,6 @@ def test_estimate_copying(device): assert s[0] == 1 -def test_no_copy_construct_warning(): - t = torch.ones((2, 2)) - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - stokes.stokes_after_adr(t, t, t, t) - stokes.stokes012_after_ar(t, t, t) - stokes.estimate_adr_from_stokes(t, t, t, t) - stokes.estimate_ar_from_stokes012(t, t, t) - stokes.mueller_from_stokes(t, t, t, t, direction="forward") - stokes.mueller_from_stokes(t, t, t, t, direction="inverse") - assert [w for w in caught if "copy construct" in str(w.message)] == [] - - def test_gradients_reach_copied_outputs(): """s0 and transmittance are copies of an input, so they must not detach.""" ones = torch.ones((2, 2)) From c1df7bc2b5fe8d69de9afd65b23ce195691fe313 Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Wed, 12 Aug 2026 00:26:12 +0000 Subject: [PATCH 4/4] test(stokes): drop estimate copying test existing test_copying already cover copy contract for stokes_after_adr and mueller_from_stokes. estimate_* variant add no new coverage. Signed-off-by: Sricharan Reddy Varra --- tests/test_stokes.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/tests/test_stokes.py b/tests/test_stokes.py index d3cf86b1..7beeba98 100644 --- a/tests/test_stokes.py +++ b/tests/test_stokes.py @@ -130,19 +130,6 @@ def test_copying(device): assert a[0] == 1 -@pytest.mark.parametrize(*_DEVICE) -def test_estimate_copying(device): - s = torch.tensor([1.0, 1.0], device=device) - - _, _, transmittance, _ = stokes.estimate_adr_from_stokes(s, s, s, s) - transmittance[0] = 2 # modify the output - assert s[0] == 1 # check that the input hasn't changed - - _, _, transmittance012 = stokes.estimate_ar_from_stokes012(s, s, s) - transmittance012[0] = 2 - assert s[0] == 1 - - def test_gradients_reach_copied_outputs(): """s0 and transmittance are copies of an input, so they must not detach.""" ones = torch.ones((2, 2))