Fix backend and scalar parameter handling

src/pyrecest/_backend/pytorch/__init__.py

@@ -690,7 +690,7 @@ def _is_scalar_index(index):
 def _assignment_index_length(indices, zip_indices):
     if zip_indices:
         return len(indices)
-    if isinstance(indices, tuple) and all(
+    if isinstance(indices, tuple) and _builtins.all(
         _is_scalar_index(index) for index in indices
     ):
         return 1
@@ -784,7 +784,11 @@ def assignment(x, values, indices, axis=0):
         indices = tuple(zip(*indices))
     if not use_vectorization:
         len_values = _assignment_value_length(values)
-        if len_values > 1 and len_values != len_indices:
+        if (
+            not _contains_slice(indices)
+            and len_values > 1
+            and len_values != len_indices
+        ):
             raise ValueError("Either one value or as many values as indices")
         _apply_assignment(x_new, indices, values, accumulate=False)
     else:
@@ -837,7 +841,11 @@ def assignment_by_sum(x, values, indices, axis=0):
         indices = tuple(zip(*indices))
     if not use_vectorization:
         len_values = _assignment_value_length(values)
-        if len_values > 1 and len_values != len_indices:
+        if (
+            not _contains_slice(indices)
+            and len_values > 1
+            and len_values != len_indices
+        ):
             raise ValueError("Either one value or as many values as indices")
         _apply_assignment(x_new, indices, values, accumulate=True)
     else:

src/pyrecest/distributions/hypertorus/toroidal_vm_rivest_distribution.py

@@ -1,5 +1,5 @@
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
-from pyrecest.backend import all, array, cos, exp, mod, pi, sin, sqrt
+from pyrecest.backend import all, array, asarray, cos, exp, mod, pi, sin, sqrt
 from scipy.special import iv
 
 from .abstract_toroidal_distribution import AbstractToroidalDistribution
@@ -19,6 +19,8 @@ def __init__(self, mu, kappa, alpha, beta):
         AbstractToroidalDistribution.__init__(self)
         assert mu.shape == (2,)
         assert kappa.shape == (2,)
+        alpha = asarray(alpha)
+        beta = asarray(beta)
         assert alpha.shape == ()
         assert beta.shape == ()
         assert all(kappa >= 0.0)

src/pyrecest/distributions/hypertorus/toroidal_von_mises_cosine_distribution.py

@@ -2,7 +2,7 @@
 import math
 
 import numpy as np
-from pyrecest.backend import cos, mod, pi
+from pyrecest.backend import asarray, cos, mod, pi
 from scipy.special import iv
 
 from .abstract_toroidal_bivar_vm_distribution import AbstractToroidalBivarVMDistribution
@@ -60,6 +60,7 @@ class ToroidalVonMisesCosineDistribution(AbstractToroidalBivarVMDistribution):
 
     def __init__(self, mu, kappa, kappa3):
         AbstractToroidalBivarVMDistribution.__init__(self, mu, kappa)
+        kappa3 = asarray(kappa3)
         assert kappa3.shape == ()
         self.kappa3 = kappa3
         self.C = 1.0 / self.norm_const

src/pyrecest/distributions/hypertorus/toroidal_von_mises_sine_distribution.py

@@ -2,7 +2,7 @@
 # pylint: disable=no-name-in-module,no-member
 import math
 
-from pyrecest.backend import sin
+from pyrecest.backend import asarray, sin
 from scipy.special import gammaln, iv
 
 from .abstract_toroidal_bivar_vm_distribution import AbstractToroidalBivarVMDistribution
@@ -50,6 +50,7 @@ class ToroidalVonMisesSineDistribution(AbstractToroidalBivarVMDistribution):
 
     def __init__(self, mu, kappa, lambda_):
         AbstractToroidalBivarVMDistribution.__init__(self, mu, kappa)
+        lambda_ = asarray(lambda_)
         assert lambda_.shape == ()
         self.lambda_ = lambda_
         self.C = 1.0 / self.norm_const

tests/backend_support/test_pytorch_assignment_contract.py

@@ -56,6 +56,8 @@ def test_pytorch_assignment_accepts_list_and_boolean_indices():
 
 vector = backend.zeros(3)
 by_list = backend.assignment(vector, [4.0, 5.0], [0, 2])
+by_slice = backend.assignment(vector, [1.0, 2.0], slice(1, 3))
+by_slice_sum = backend.assignment_by_sum(vector, [1.0, 2.0], slice(1, 3))
 
 matrix = backend.zeros((3, 3))
 by_mask = backend.assignment(
@@ -71,6 +73,8 @@ def test_pytorch_assignment_accepts_list_and_boolean_indices():
 logical = backend.logical_and(backend.asarray([1, 0]), backend.asarray([2, 3]))
 
 assert backend.to_numpy(by_list).tolist() == [4.0, 0.0, 5.0]
+assert backend.to_numpy(by_slice).tolist() == [0.0, 1.0, 2.0]
+assert backend.to_numpy(by_slice_sum).tolist() == [0.0, 1.0, 2.0]
 assert backend.to_numpy(by_mask).tolist() == [
     [1.0, 0.0, 0.0],
     [0.0, 2.0, 0.0],

tests/distributions/test_toroidal_vm_rivest_distribution.py

@@ -36,6 +36,13 @@ def test_mu_kappa_alpha_beta(self):
         self.assertEqual(self.dist.alpha, self.alpha)
         self.assertEqual(self.dist.beta, self.beta)
 
+    def test_accepts_python_scalar_correlation_parameters(self):
+        dist = ToroidalVMRivestDistribution(self.mu, self.kappa, 0.5, 0.3)
+        x = array([1.3, 2.4])
+
+        npt.assert_allclose(dist.pdf(x), self.dist.pdf(x), rtol=5e-7)
+        npt.assert_allclose(dist.C, self.dist.C, rtol=5e-7)
+
     @unittest.skipIf(
         pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
         reason="Not supported on this backend",

tests/distributions/test_toroidal_von_mises_cosine_distribution.py

@@ -80,6 +80,15 @@ def test_mu_kappa_kappa3(self):
         npt.assert_allclose(self.tvm.kappa, self.kappa)
         self.assertEqual(self.tvm.kappa3, self.kappa3)
 
+    def test_accepts_python_scalar_coupling_parameter(self):
+        tvm = ToroidalVonMisesCosineDistribution(self.mu, self.kappa, 0.5)
+        x = array([1.3, 2.4])
+
+        npt.assert_allclose(tvm.pdf(x), self.tvm.pdf(x), rtol=5e-7)
+        npt.assert_allclose(
+            tvm.trigonometric_moment(1), self.tvm.trigonometric_moment(1), rtol=5e-7
+        )
+
     def _unnormalized_pdf(self, xs):
         return exp(
             self.kappa[0] * cos(xs[..., 0] - self.mu[0])

tests/distributions/test_toroidal_von_mises_sine_distribution.py

@@ -85,6 +85,13 @@ def test_mu_kappa_lambda(self):
         npt.assert_allclose(self.tvm.kappa, self.kappa)
         self.assertEqual(self.tvm.lambda_, self.lambda_)
 
+    def test_accepts_python_scalar_coupling_parameter(self):
+        tvm = ToroidalVonMisesSineDistribution(self.mu, self.kappa, 0.5)
+        x = array([1.3, 2.4])
+
+        npt.assert_allclose(tvm.pdf(x), self.tvm.pdf(x), rtol=5e-7)
+        npt.assert_allclose(tvm.norm_const, self.tvm.norm_const, rtol=5e-7)
+
     @unittest.skipIf(
         pyrecest.backend.__backend_name__ != "numpy",
         reason="Regression test uses NumPy/scipy scalar semantics",