From 9046315b8cc566e44425141d17e6f7db388c92b9 Mon Sep 17 00:00:00 2001 From: HGSilveri Date: Mon, 16 Dec 2024 18:41:19 +0100 Subject: [PATCH 1/3] Define AbstractArray.is_differentiable --- pulser-core/pulser/math/abstract_array.py | 10 ++++++---- tests/test_channels.py | 6 ++---- tests/test_eom.py | 6 ++---- tests/test_math.py | 10 ++++------ tests/test_parametrized.py | 5 ++--- tests/test_pulse.py | 6 +++--- tests/test_register.py | 4 ++-- tests/test_sequence.py | 12 ++++++------ tests/test_sequence_sampler.py | 10 +++++----- tests/test_waveforms.py | 10 ++++------ 10 files changed, 36 insertions(+), 43 deletions(-) diff --git a/pulser-core/pulser/math/abstract_array.py b/pulser-core/pulser/math/abstract_array.py index c74805a69..bb7ed9821 100644 --- a/pulser-core/pulser/math/abstract_array.py +++ b/pulser-core/pulser/math/abstract_array.py @@ -71,6 +71,11 @@ def is_tensor(self) -> bool: """Whether the stored array is a tensor.""" return self.has_torch() and isinstance(self._array, torch.Tensor) + @property + def is_differentiable(self) -> bool: + """Whether the stored array is a differentiable tensor.""" + return self.is_tensor and cast(torch.Tensor, self._array).requires_grad + def astype(self, dtype: DTypeLike) -> AbstractArray: """Casts the data type of the array contents.""" if self.is_tensor: @@ -271,10 +276,7 @@ def __setitem__(self, indices: Any, values: AbstractArrayLike) -> None: self._process_indices(indices) ] = values # type: ignore[assignment] except RuntimeError as e: - if ( - self.is_tensor - and cast(torch.Tensor, self._array).requires_grad - ): + if self.is_differentiable: raise RuntimeError( "Failed to modify a tensor that requires grad in place." ) from e diff --git a/tests/test_channels.py b/tests/test_channels.py index 5f2e2cfe0..d61cfca05 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -290,8 +290,7 @@ def test_modulation(channel, tr, eom, side_buffer_len, requires_grad): tr, tr, ) - if requires_grad: - assert out_.as_tensor().requires_grad + assert out_.is_differentiable == requires_grad wf2 = BlackmanWaveform(800, wf_vals[1]) out_ = channel.modulate(wf2.samples, eom=eom) @@ -300,8 +299,7 @@ def test_modulation(channel, tr, eom, side_buffer_len, requires_grad): side_buffer_len, side_buffer_len, ) - if requires_grad: - assert out_.as_tensor().requires_grad + assert out_.is_differentiable == requires_grad @pytest.mark.parametrize( diff --git a/tests/test_eom.py b/tests/test_eom.py index ea63a4b2d..defebe11b 100644 --- a/tests/test_eom.py +++ b/tests/test_eom.py @@ -190,8 +190,7 @@ def calc_offset(amp): ] ) assert calculated_det_off == min(det_off_options, key=abs) - if requires_grad: - assert calculated_det_off.as_tensor().requires_grad + assert calculated_det_off.is_differentiable == requires_grad # Case where the EOM pulses are off-resonant detuning_on = detuning_on + 1.0 @@ -210,5 +209,4 @@ def calc_offset(amp): assert off_options[0] == eom_.calculate_detuning_off( amp, detuning_on, optimal_detuning_off=0.0 ) - if requires_grad: - assert off_options.as_tensor().requires_grad + assert off_options.is_differentiable == requires_grad diff --git a/tests/test_math.py b/tests/test_math.py index 75aa0d50a..ed945fd0c 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -39,8 +39,7 @@ def test_pad(cast_to, requires_grad): arr = torch.tensor(arr, requires_grad=requires_grad) def check_match(arr1: pm.AbstractArray, arr2): - if requires_grad: - assert arr1.as_tensor().requires_grad + assert arr1.is_differentiable == requires_grad np.testing.assert_array_equal( arr1.as_array(detach=requires_grad), arr2 ) @@ -260,8 +259,7 @@ def test_items(self, use_tensor, requires_grad, indices): assert item == val[i] assert isinstance(item, pm.AbstractArray) assert item.is_tensor == use_tensor - if use_tensor: - assert item.as_tensor().requires_grad == requires_grad + assert item.is_differentiable == requires_grad # setitem if not requires_grad: @@ -292,8 +290,8 @@ def test_items(self, use_tensor, requires_grad, indices): new_val[indices] = 0.0 assert np.all(arr_np == new_val) assert arr_np.is_tensor - # The resulting tensor requires grad if the assing one did - assert arr_np.as_tensor().requires_grad == requires_grad + # The resulting tensor requires grad if the assigned one did + assert arr_np.is_differentiable == requires_grad @pytest.mark.parametrize("scalar", [False, True]) @pytest.mark.parametrize( diff --git a/tests/test_parametrized.py b/tests/test_parametrized.py index 7d0c4ccc8..367d476b6 100644 --- a/tests/test_parametrized.py +++ b/tests/test_parametrized.py @@ -105,8 +105,7 @@ def test_var_diff(a, b, requires_grad): for var in [a, b]: assert ( - a.value is not None - and a.value.as_tensor().requires_grad == requires_grad + a.value is not None and a.value.is_differentiable == requires_grad ) @@ -167,7 +166,7 @@ def test_paramobj(bwf, t, a, b): def test_opsupport(a, b, with_diff_tensor): def check_var_grad(var): if with_diff_tensor: - assert var.build().as_tensor().requires_grad + assert var.build().is_differentiable a._assign(-2.0) if with_diff_tensor: diff --git a/tests/test_pulse.py b/tests/test_pulse.py index fe51866a6..cbe0b6770 100644 --- a/tests/test_pulse.py +++ b/tests/test_pulse.py @@ -234,9 +234,9 @@ def test_eq(): def _assert_pulse_requires_grad(pulse: Pulse, invert: bool = False) -> None: - assert pulse.amplitude.samples.as_tensor().requires_grad == (not invert) - assert pulse.detuning.samples.as_tensor().requires_grad == (not invert) - assert pulse.phase.as_tensor().requires_grad == (not invert) + assert pulse.amplitude.samples.is_differentiable == (not invert) + assert pulse.detuning.samples.is_differentiable == (not invert) + assert pulse.phase.is_differentiable == (not invert) @pytest.mark.parametrize("requires_grad", [True, False]) diff --git a/tests/test_register.py b/tests/test_register.py index 5cbacf0ae..51debfbda 100644 --- a/tests/test_register.py +++ b/tests/test_register.py @@ -508,9 +508,9 @@ def _assert_reg_requires_grad( ) -> None: for coords in reg.qubits.values(): if invert: - assert not coords.as_tensor().requires_grad + assert not coords.is_differentiable else: - assert coords.is_tensor and coords.as_tensor().requires_grad + assert coords.is_tensor and coords.is_differentiable @pytest.mark.parametrize( diff --git a/tests/test_sequence.py b/tests/test_sequence.py index d13d5bc3d..d45b2e573 100644 --- a/tests/test_sequence.py +++ b/tests/test_sequence.py @@ -2870,12 +2870,12 @@ def test_sequence_diff(device, parametrized, with_modulation, with_eom): seq_samples = sample(seq, modulation=with_modulation) ryd_ch_samples = seq_samples.channel_samples["ryd_global"] - assert ryd_ch_samples.amp.as_tensor().requires_grad - assert ryd_ch_samples.det.as_tensor().requires_grad - assert ryd_ch_samples.phase.as_tensor().requires_grad + assert ryd_ch_samples.amp.is_differentiable + assert ryd_ch_samples.det.is_differentiable + assert ryd_ch_samples.phase.is_differentiable if "dmm_0" in seq_samples.channel_samples: dmm_ch_samples = seq_samples.channel_samples["dmm_0"] # Only detuning is modulated - assert not dmm_ch_samples.amp.as_tensor().requires_grad - assert dmm_ch_samples.det.as_tensor().requires_grad - assert not dmm_ch_samples.phase.as_tensor().requires_grad + assert not dmm_ch_samples.amp.is_differentiable + assert dmm_ch_samples.det.is_differentiable + assert not dmm_ch_samples.phase.is_differentiable diff --git a/tests/test_sequence_sampler.py b/tests/test_sequence_sampler.py index 6b27d5be8..9594ce782 100644 --- a/tests/test_sequence_sampler.py +++ b/tests/test_sequence_sampler.py @@ -503,11 +503,11 @@ def test_phase_modulation(off_center, with_diff): seq_samples = sample(seq).channel_samples["rydberg_global"] if with_diff: - assert full_phase.samples.as_tensor().requires_grad - assert not seq_samples.amp.as_tensor().requires_grad - assert seq_samples.det.as_tensor().requires_grad - assert seq_samples.phase.as_tensor().requires_grad - assert seq_samples.phase_modulation.as_tensor().requires_grad + assert full_phase.samples.is_differentiable + assert not seq_samples.amp.is_differentiable + assert seq_samples.det.is_differentiable + assert seq_samples.phase.is_differentiable + assert seq_samples.phase_modulation.is_differentiable np.testing.assert_allclose( seq_samples.phase_modulation.as_array(detach=with_diff) diff --git a/tests/test_waveforms.py b/tests/test_waveforms.py index 59648cfbe..c7bc5aad7 100644 --- a/tests/test_waveforms.py +++ b/tests/test_waveforms.py @@ -491,8 +491,7 @@ def test_waveform_diff( samples_tensor = wf.samples.as_tensor() assert samples_tensor.requires_grad == requires_grad assert ( - wf.modulated_samples(rydberg_global).as_tensor().requires_grad - == requires_grad + wf.modulated_samples(rydberg_global).is_differentiable == requires_grad ) wfx2_tensor = (-wf * 2).samples.as_tensor() assert torch.equal(wfx2_tensor, samples_tensor * -2.0) @@ -501,14 +500,13 @@ def test_waveform_diff( wfdiv2 = wf / torch.tensor(2.0, requires_grad=True) assert torch.equal(wfdiv2.samples.as_tensor(), samples_tensor / 2.0) # Should always be true because it was divided by diff tensor - assert wfdiv2.samples.as_tensor().requires_grad + assert wfdiv2.samples.is_differentiable - assert wf[-1].as_tensor().requires_grad == requires_grad + assert wf[-1].is_differentiable == requires_grad try: assert ( - wf.change_duration(1000).samples.as_tensor().requires_grad - == requires_grad + wf.change_duration(1000).samples.is_differentiable == requires_grad ) except NotImplementedError: pass From c22d33c9d75363a5f745315144b5cd5c0e27d826 Mon Sep 17 00:00:00 2001 From: HGSilveri Date: Tue, 17 Dec 2024 12:02:47 +0100 Subject: [PATCH 2/3] Remove warnings.simplefilter that was overriding pytest filterwarnings --- pulser-core/pulser/register/base_register.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pulser-core/pulser/register/base_register.py b/pulser-core/pulser/register/base_register.py index 9b30c33bd..ef73e43ac 100644 --- a/pulser-core/pulser/register/base_register.py +++ b/pulser-core/pulser/register/base_register.py @@ -79,7 +79,6 @@ def __init__( ) self._ids: tuple[QubitId, ...] = tuple(qubits.keys()) if any(not isinstance(id, str) for id in self._ids): - warnings.simplefilter("always") warnings.warn( "Usage of `int`s or any non-`str`types as `QubitId`s will be " "deprecated. Define your `QubitId`s as `str`s, prefer setting " From 8d5a5d2a1150ff16854f396a898befa1fc62f9d7 Mon Sep 17 00:00:00 2001 From: HGSilveri Date: Wed, 18 Dec 2024 17:39:13 +0100 Subject: [PATCH 3/3] Rename is_differentiable -> requires_grad --- pulser-core/pulser/math/abstract_array.py | 6 +++--- tests/test_channels.py | 4 ++-- tests/test_eom.py | 4 ++-- tests/test_math.py | 6 +++--- tests/test_parametrized.py | 6 ++---- tests/test_pulse.py | 6 +++--- tests/test_register.py | 4 ++-- tests/test_sequence.py | 12 ++++++------ tests/test_sequence_sampler.py | 10 +++++----- tests/test_waveforms.py | 12 ++++-------- 10 files changed, 32 insertions(+), 38 deletions(-) diff --git a/pulser-core/pulser/math/abstract_array.py b/pulser-core/pulser/math/abstract_array.py index bb7ed9821..eb38bcddd 100644 --- a/pulser-core/pulser/math/abstract_array.py +++ b/pulser-core/pulser/math/abstract_array.py @@ -72,8 +72,8 @@ def is_tensor(self) -> bool: return self.has_torch() and isinstance(self._array, torch.Tensor) @property - def is_differentiable(self) -> bool: - """Whether the stored array is a differentiable tensor.""" + def requires_grad(self) -> bool: + """Whether the stored array is a tensor that needs a gradient.""" return self.is_tensor and cast(torch.Tensor, self._array).requires_grad def astype(self, dtype: DTypeLike) -> AbstractArray: @@ -276,7 +276,7 @@ def __setitem__(self, indices: Any, values: AbstractArrayLike) -> None: self._process_indices(indices) ] = values # type: ignore[assignment] except RuntimeError as e: - if self.is_differentiable: + if self.requires_grad: raise RuntimeError( "Failed to modify a tensor that requires grad in place." ) from e diff --git a/tests/test_channels.py b/tests/test_channels.py index d61cfca05..1f067b5ac 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -290,7 +290,7 @@ def test_modulation(channel, tr, eom, side_buffer_len, requires_grad): tr, tr, ) - assert out_.is_differentiable == requires_grad + assert out_.requires_grad == requires_grad wf2 = BlackmanWaveform(800, wf_vals[1]) out_ = channel.modulate(wf2.samples, eom=eom) @@ -299,7 +299,7 @@ def test_modulation(channel, tr, eom, side_buffer_len, requires_grad): side_buffer_len, side_buffer_len, ) - assert out_.is_differentiable == requires_grad + assert out_.requires_grad == requires_grad @pytest.mark.parametrize( diff --git a/tests/test_eom.py b/tests/test_eom.py index defebe11b..e10d508b1 100644 --- a/tests/test_eom.py +++ b/tests/test_eom.py @@ -190,7 +190,7 @@ def calc_offset(amp): ] ) assert calculated_det_off == min(det_off_options, key=abs) - assert calculated_det_off.is_differentiable == requires_grad + assert calculated_det_off.requires_grad == requires_grad # Case where the EOM pulses are off-resonant detuning_on = detuning_on + 1.0 @@ -209,4 +209,4 @@ def calc_offset(amp): assert off_options[0] == eom_.calculate_detuning_off( amp, detuning_on, optimal_detuning_off=0.0 ) - assert off_options.is_differentiable == requires_grad + assert off_options.requires_grad == requires_grad diff --git a/tests/test_math.py b/tests/test_math.py index ed945fd0c..51b8abb38 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -39,7 +39,7 @@ def test_pad(cast_to, requires_grad): arr = torch.tensor(arr, requires_grad=requires_grad) def check_match(arr1: pm.AbstractArray, arr2): - assert arr1.is_differentiable == requires_grad + assert arr1.requires_grad == requires_grad np.testing.assert_array_equal( arr1.as_array(detach=requires_grad), arr2 ) @@ -259,7 +259,7 @@ def test_items(self, use_tensor, requires_grad, indices): assert item == val[i] assert isinstance(item, pm.AbstractArray) assert item.is_tensor == use_tensor - assert item.is_differentiable == requires_grad + assert item.requires_grad == requires_grad # setitem if not requires_grad: @@ -291,7 +291,7 @@ def test_items(self, use_tensor, requires_grad, indices): assert np.all(arr_np == new_val) assert arr_np.is_tensor # The resulting tensor requires grad if the assigned one did - assert arr_np.is_differentiable == requires_grad + assert arr_np.requires_grad == requires_grad @pytest.mark.parametrize("scalar", [False, True]) @pytest.mark.parametrize( diff --git a/tests/test_parametrized.py b/tests/test_parametrized.py index 367d476b6..87e555843 100644 --- a/tests/test_parametrized.py +++ b/tests/test_parametrized.py @@ -104,9 +104,7 @@ def test_var_diff(a, b, requires_grad): b._assign(torch.tensor([-1.0, 1.0], requires_grad=requires_grad)) for var in [a, b]: - assert ( - a.value is not None and a.value.is_differentiable == requires_grad - ) + assert a.value is not None and a.value.requires_grad == requires_grad def test_varitem(a, b, d): @@ -166,7 +164,7 @@ def test_paramobj(bwf, t, a, b): def test_opsupport(a, b, with_diff_tensor): def check_var_grad(var): if with_diff_tensor: - assert var.build().is_differentiable + assert var.build().requires_grad a._assign(-2.0) if with_diff_tensor: diff --git a/tests/test_pulse.py b/tests/test_pulse.py index cbe0b6770..e5a265660 100644 --- a/tests/test_pulse.py +++ b/tests/test_pulse.py @@ -234,9 +234,9 @@ def test_eq(): def _assert_pulse_requires_grad(pulse: Pulse, invert: bool = False) -> None: - assert pulse.amplitude.samples.is_differentiable == (not invert) - assert pulse.detuning.samples.is_differentiable == (not invert) - assert pulse.phase.is_differentiable == (not invert) + assert pulse.amplitude.samples.requires_grad == (not invert) + assert pulse.detuning.samples.requires_grad == (not invert) + assert pulse.phase.requires_grad == (not invert) @pytest.mark.parametrize("requires_grad", [True, False]) diff --git a/tests/test_register.py b/tests/test_register.py index 51debfbda..c7c387a7f 100644 --- a/tests/test_register.py +++ b/tests/test_register.py @@ -508,9 +508,9 @@ def _assert_reg_requires_grad( ) -> None: for coords in reg.qubits.values(): if invert: - assert not coords.is_differentiable + assert not coords.requires_grad else: - assert coords.is_tensor and coords.is_differentiable + assert coords.is_tensor and coords.requires_grad @pytest.mark.parametrize( diff --git a/tests/test_sequence.py b/tests/test_sequence.py index d45b2e573..3352f7e41 100644 --- a/tests/test_sequence.py +++ b/tests/test_sequence.py @@ -2870,12 +2870,12 @@ def test_sequence_diff(device, parametrized, with_modulation, with_eom): seq_samples = sample(seq, modulation=with_modulation) ryd_ch_samples = seq_samples.channel_samples["ryd_global"] - assert ryd_ch_samples.amp.is_differentiable - assert ryd_ch_samples.det.is_differentiable - assert ryd_ch_samples.phase.is_differentiable + assert ryd_ch_samples.amp.requires_grad + assert ryd_ch_samples.det.requires_grad + assert ryd_ch_samples.phase.requires_grad if "dmm_0" in seq_samples.channel_samples: dmm_ch_samples = seq_samples.channel_samples["dmm_0"] # Only detuning is modulated - assert not dmm_ch_samples.amp.is_differentiable - assert dmm_ch_samples.det.is_differentiable - assert not dmm_ch_samples.phase.is_differentiable + assert not dmm_ch_samples.amp.requires_grad + assert dmm_ch_samples.det.requires_grad + assert not dmm_ch_samples.phase.requires_grad diff --git a/tests/test_sequence_sampler.py b/tests/test_sequence_sampler.py index 9594ce782..708688782 100644 --- a/tests/test_sequence_sampler.py +++ b/tests/test_sequence_sampler.py @@ -503,11 +503,11 @@ def test_phase_modulation(off_center, with_diff): seq_samples = sample(seq).channel_samples["rydberg_global"] if with_diff: - assert full_phase.samples.is_differentiable - assert not seq_samples.amp.is_differentiable - assert seq_samples.det.is_differentiable - assert seq_samples.phase.is_differentiable - assert seq_samples.phase_modulation.is_differentiable + assert full_phase.samples.requires_grad + assert not seq_samples.amp.requires_grad + assert seq_samples.det.requires_grad + assert seq_samples.phase.requires_grad + assert seq_samples.phase_modulation.requires_grad np.testing.assert_allclose( seq_samples.phase_modulation.as_array(detach=with_diff) diff --git a/tests/test_waveforms.py b/tests/test_waveforms.py index c7bc5aad7..5d46de56b 100644 --- a/tests/test_waveforms.py +++ b/tests/test_waveforms.py @@ -490,9 +490,7 @@ def test_waveform_diff( samples_tensor = wf.samples.as_tensor() assert samples_tensor.requires_grad == requires_grad - assert ( - wf.modulated_samples(rydberg_global).is_differentiable == requires_grad - ) + assert wf.modulated_samples(rydberg_global).requires_grad == requires_grad wfx2_tensor = (-wf * 2).samples.as_tensor() assert torch.equal(wfx2_tensor, samples_tensor * -2.0) assert wfx2_tensor.requires_grad == requires_grad @@ -500,14 +498,12 @@ def test_waveform_diff( wfdiv2 = wf / torch.tensor(2.0, requires_grad=True) assert torch.equal(wfdiv2.samples.as_tensor(), samples_tensor / 2.0) # Should always be true because it was divided by diff tensor - assert wfdiv2.samples.is_differentiable + assert wfdiv2.samples.requires_grad - assert wf[-1].is_differentiable == requires_grad + assert wf[-1].requires_grad == requires_grad try: - assert ( - wf.change_duration(1000).samples.is_differentiable == requires_grad - ) + assert wf.change_duration(1000).samples.requires_grad == requires_grad except NotImplementedError: pass