tests

src/beignet/polynomial/__init__.py

@@ -10,7 +10,6 @@
 from .__fit import _fit
 from .__normed_hermite_e_n import _normed_hermite_e_n
 from .__normed_hermite_n import _normed_hermite_n
-from .__nth_slice import _nth_slice
 from ._add_chebyshev_polynomial import add_chebyshev_polynomial
 from ._add_laguerre_polynomial import add_laguerre_polynomial
 from ._add_legendre_polynomial import add_legendre_polynomial

src/beignet/polynomial/_chebyshev_polynomial_vandermonde_2d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._chebyshev_polynomial_vandermonde import chebyshev_polynomial_vandermonde
 
 
@@ -36,7 +35,7 @@ def chebyshev_polynomial_vandermonde_2d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_chebyshev_polynomial_vandermonde_3d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._chebyshev_polynomial_vandermonde import chebyshev_polynomial_vandermonde
 
 
@@ -38,7 +37,7 @@ def chebyshev_polynomial_vandermonde_3d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_laguerre_polynomial_vandermonde_2d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._laguerre_polynomial_vandermonde import laguerre_polynomial_vandermonde
 
 
@@ -36,7 +35,7 @@ def laguerre_polynomial_vandermonde_2d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_laguerre_polynomial_vandermonde_3d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._laguerre_polynomial_vandermonde import laguerre_polynomial_vandermonde
 
 
@@ -38,7 +37,7 @@ def laguerre_polynomial_vandermonde_3d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_legendre_polynomial_vandermonde_2d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._legendre_polynomial_vandermonde import legendre_polynomial_vandermonde
 
 
@@ -36,7 +35,7 @@ def legendre_polynomial_vandermonde_2d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_legendre_polynomial_vandermonde_3d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._legendre_polynomial_vandermonde import legendre_polynomial_vandermonde
 
 
@@ -38,7 +37,7 @@ def legendre_polynomial_vandermonde_3d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_physicists_hermite_polynomial_vandermonde_2d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._physicists_hermite_polynomial_vandermonde import (
     physicists_hermite_polynomial_vandermonde,
 )
@@ -38,7 +37,7 @@ def physicists_hermite_polynomial_vandermonde_2d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_physicists_hermite_polynomial_vandermonde_3d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._physicists_hermite_polynomial_vandermonde import (
     physicists_hermite_polynomial_vandermonde,
 )
@@ -40,7 +39,7 @@ def physicists_hermite_polynomial_vandermonde_3d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_polynomial_vandermonde_2d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._polynomial_vandermonde import polynomial_vandermonde
 
 
@@ -45,7 +44,7 @@ def polynomial_vandermonde_2d(x: Tensor, y: Tensor, degree: Tensor) -> Tensor:
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_polynomial_vandermonde_3d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._polynomial_vandermonde import polynomial_vandermonde
 
 
@@ -50,7 +49,7 @@ def polynomial_vandermonde_3d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_probabilists_hermite_polynomial_vandermonde_2d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._probabilists_hermite_polynomial_vandermonde import (
     probabilists_hermite_polynomial_vandermonde,
 )
@@ -38,7 +37,7 @@ def probabilists_hermite_polynomial_vandermonde_2d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(

src/beignet/polynomial/_probabilists_hermite_polynomial_vandermonde_3d.py

@@ -4,7 +4,6 @@
 import torch
 from torch import Tensor
 
-from . import _nth_slice
 from ._probabilists_hermite_polynomial_vandermonde import (
     probabilists_hermite_polynomial_vandermonde,
 )
@@ -40,7 +39,7 @@ def probabilists_hermite_polynomial_vandermonde_3d(
 
         matrices = [
             *matrices,
-            matrix[(..., *_nth_slice(i, n))],
+            matrix[(..., *tuple(slice(None) if j == i else None for j in range(n)))],
         ]
 
     vandermonde = functools.reduce(